{-# LANGUAGE ViewPatterns #-}

module Ide.Plugin.Tactic.CodeGen.Utils where

import Data.List
import DataCon
import Development.IDE.GHC.Compat
import GHC.Exts
import GHC.SourceGen (recordConE, RdrNameStr)
import GHC.SourceGen.Overloaded
import Ide.Plugin.Tactic.GHC (getRecordFields)
import Name


------------------------------------------------------------------------------
-- | Make a data constructor with the given arguments.
mkCon :: DataCon -> [LHsExpr GhcPs] -> LHsExpr GhcPs
mkCon :: DataCon -> [LHsExpr GhcPs] -> LHsExpr GhcPs
mkCon DataCon
dcon ((LHsExpr GhcPs -> HsExpr') -> [LHsExpr GhcPs] -> [HsExpr']
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LHsExpr GhcPs -> HsExpr'
forall a. HasSrcSpan a => a -> SrcSpanLess a
unLoc -> [HsExpr']
args)
  | DataCon -> Bool
isTupleDataCon DataCon
dcon =
      SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a. HasSrcSpan a => SrcSpanLess a -> a
noLoc (SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs)
-> SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a b. (a -> b) -> a -> b
$ [HsExpr'] -> HsExpr'
forall e. HasTuple e => [e] -> e
tuple [HsExpr']
args
  | DataCon -> Bool
dataConIsInfix DataCon
dcon
  , (HsExpr'
lhs : HsExpr'
rhs : [HsExpr']
args') <- [HsExpr']
args =
      SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a. HasSrcSpan a => SrcSpanLess a -> a
noLoc (SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs)
-> SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a b. (a -> b) -> a -> b
$ (HsExpr' -> HsExpr' -> HsExpr') -> HsExpr' -> [HsExpr'] -> HsExpr'
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' HsExpr' -> HsExpr' -> HsExpr'
forall e. App e => e -> e -> e
(@@) (HsExpr' -> RdrNameStr -> HsExpr' -> HsExpr'
forall e. App e => e -> RdrNameStr -> e -> e
op HsExpr'
lhs (Name -> RdrNameStr
forall a. HasOccName a => a -> RdrNameStr
coerceName Name
dcon_name) HsExpr'
rhs) [HsExpr']
args'
  | Just [(OccName, CType)]
fields <- DataCon -> Maybe [(OccName, CType)]
getRecordFields DataCon
dcon =
      SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a. HasSrcSpan a => SrcSpanLess a -> a
noLoc (SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs)
-> SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a b. (a -> b) -> a -> b
$ RdrNameStr -> [(RdrNameStr, HsExpr')] -> HsExpr'
recordConE (Name -> RdrNameStr
forall a. HasOccName a => a -> RdrNameStr
coerceName Name
dcon_name) ([(RdrNameStr, HsExpr')] -> HsExpr')
-> [(RdrNameStr, HsExpr')] -> HsExpr'
forall a b. (a -> b) -> a -> b
$ do
        (HsExpr'
arg, (OccName
field, CType
_)) <- [HsExpr'] -> [(OccName, CType)] -> [(HsExpr', (OccName, CType))]
forall a b. [a] -> [b] -> [(a, b)]
zip [HsExpr']
args [(OccName, CType)]
fields
        (RdrNameStr, HsExpr') -> [(RdrNameStr, HsExpr')]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (OccName -> RdrNameStr
forall a. HasOccName a => a -> RdrNameStr
coerceName OccName
field, HsExpr'
arg)
  | Bool
otherwise =
      SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a. HasSrcSpan a => SrcSpanLess a -> a
noLoc (SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs)
-> SrcSpanLess (LHsExpr GhcPs) -> LHsExpr GhcPs
forall a b. (a -> b) -> a -> b
$ (HsExpr' -> HsExpr' -> HsExpr') -> HsExpr' -> [HsExpr'] -> HsExpr'
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' HsExpr' -> HsExpr' -> HsExpr'
forall e. App e => e -> e -> e
(@@) (OccName -> HsExpr'
forall a. BVar a => OccName -> a
bvar' (OccName -> HsExpr') -> OccName -> HsExpr'
forall a b. (a -> b) -> a -> b
$ Name -> OccName
forall name. HasOccName name => name -> OccName
occName Name
dcon_name) [HsExpr']
args
  where
    dcon_name :: Name
dcon_name = DataCon -> Name
dataConName DataCon
dcon


coerceName :: HasOccName a => a -> RdrNameStr
coerceName :: a -> RdrNameStr
coerceName = String -> RdrNameStr
forall a. IsString a => String -> a
fromString (String -> RdrNameStr) -> (a -> String) -> a -> RdrNameStr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> String
occNameString (OccName -> String) -> (a -> OccName) -> a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> OccName
forall name. HasOccName name => name -> OccName
occName


------------------------------------------------------------------------------
-- | Like 'var', but works over standard GHC 'OccName's.
var' :: Var a => OccName -> a
var' :: OccName -> a
var' = RdrNameStr -> a
forall a. Var a => RdrNameStr -> a
var (RdrNameStr -> a) -> (OccName -> RdrNameStr) -> OccName -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> RdrNameStr
forall a. IsString a => String -> a
fromString (String -> RdrNameStr)
-> (OccName -> String) -> OccName -> RdrNameStr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> String
occNameString


------------------------------------------------------------------------------
-- | Like 'bvar', but works over standard GHC 'OccName's.
bvar' :: BVar a => OccName -> a
bvar' :: OccName -> a
bvar' = OccNameStr -> a
forall a. BVar a => OccNameStr -> a
bvar (OccNameStr -> a) -> (OccName -> OccNameStr) -> OccName -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> OccNameStr
forall a. IsString a => String -> a
fromString (String -> OccNameStr)
-> (OccName -> String) -> OccName -> OccNameStr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OccName -> String
occNameString


------------------------------------------------------------------------------
-- | Get an HsExpr corresponding to a function name.
mkFunc :: String -> HsExpr GhcPs
mkFunc :: String -> HsExpr'
mkFunc = OccName -> HsExpr'
forall a. Var a => OccName -> a
var' (OccName -> HsExpr') -> (String -> OccName) -> String -> HsExpr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> OccName
mkVarOcc


------------------------------------------------------------------------------
-- | Get an HsExpr corresponding to a value name.
mkVal :: String -> HsExpr GhcPs
mkVal :: String -> HsExpr'
mkVal = OccName -> HsExpr'
forall a. Var a => OccName -> a
var' (OccName -> HsExpr') -> (String -> OccName) -> String -> HsExpr'
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> OccName
mkVarOcc


------------------------------------------------------------------------------
-- | Like 'op', but easier to call.
infixCall :: String -> HsExpr GhcPs -> HsExpr GhcPs -> HsExpr GhcPs
infixCall :: String -> HsExpr' -> HsExpr' -> HsExpr'
infixCall String
s = (HsExpr' -> RdrNameStr -> HsExpr' -> HsExpr')
-> RdrNameStr -> HsExpr' -> HsExpr' -> HsExpr'
forall a b c. (a -> b -> c) -> b -> a -> c
flip HsExpr' -> RdrNameStr -> HsExpr' -> HsExpr'
forall e. App e => e -> RdrNameStr -> e -> e
op (String -> RdrNameStr
forall a. IsString a => String -> a
fromString String
s)


------------------------------------------------------------------------------
-- | Like '(@@)', but uses a dollar instead of parentheses.
appDollar :: HsExpr GhcPs -> HsExpr GhcPs -> HsExpr GhcPs
appDollar :: HsExpr' -> HsExpr' -> HsExpr'
appDollar = String -> HsExpr' -> HsExpr' -> HsExpr'
infixCall String
"$"