module Wingman.CodeGen.Utils where

import Data.String
import Data.List
import Development.IDE.GHC.Compat
import GHC.SourceGen (RdrNameStr (UnqualStr), recordConE, string)
import GHC.SourceGen.Overloaded as SourceGen
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' :: SourceGen.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
"$"