module Wingman.CodeGen.Utils where

import ConLike (ConLike(RealDataCon), conLikeName)
import Data.List
import DataCon
import Development.IDE.GHC.Compat
import GHC.Exts
import GHC.SourceGen (RdrNameStr (UnqualStr), recordConE, string)
import GHC.SourceGen.Overloaded
import GhcPlugins (nilDataCon, charTy, eqType)
import Name
import Wingman.GHC (getRecordFields)


------------------------------------------------------------------------------
-- | Make a data constructor with the given arguments.
mkCon :: ConLike -> [Type] -> [LHsExpr GhcPs] -> LHsExpr GhcPs
mkCon :: ConLike -> [Type] -> [LHsExpr GhcPs] -> LHsExpr GhcPs
mkCon ConLike
con [Type]
apps ((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)
  | RealDataCon DataCon
dcon <- ConLike
con
  , DataCon
dcon DataCon -> DataCon -> Bool
forall a. Eq a => a -> a -> Bool
== DataCon
nilDataCon
  , [Type
ty] <- [Type]
apps
  , Type
ty Type -> Type -> Bool
`eqType` Type
charTy = 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
$ String -> HsExpr'
forall e. HasLit e => String -> e
string String
""

  | RealDataCon DataCon
dcon <- ConLike
con
  , 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

  | RealDataCon DataCon
dcon <- ConLike
con
  , 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
con_name) HsExpr'
rhs) [HsExpr']
args'

  | Just [(OccName, CType)]
fields <- ConLike -> Maybe [(OccName, CType)]
getRecordFields ConLike
con
  , [(OccName, CType)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(OccName, CType)]
fields Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
2 =  --  record notation is unnatural on single field ctors
      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
con_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
con_name) [HsExpr']
args
  where
    con_name :: Name
con_name = ConLike -> Name
conLikeName ConLike
con


coerceName :: HasOccName a => a -> RdrNameStr
coerceName :: a -> RdrNameStr
coerceName = OccNameStr -> RdrNameStr
UnqualStr (OccNameStr -> RdrNameStr) -> (a -> OccNameStr) -> a -> RdrNameStr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> OccNameStr
forall a. IsString a => String -> a
fromString (String -> OccNameStr) -> (a -> String) -> a -> OccNameStr
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
"$"