1module Wingman.CodeGen.Utils where
2
3import ConLike (ConLike(RealDataCon), conLikeName)
4import Data.List
5import DataCon
6import Development.IDE.GHC.Compat
7import GHC.Exts
8import GHC.SourceGen (RdrNameStr (UnqualStr), recordConE, string)
9import GHC.SourceGen.Overloaded
10import GhcPlugins (nilDataCon, charTy, eqType)
11import Name
12import Wingman.GHC (getRecordFields)
13
14
15------------------------------------------------------------------------------
16-- | Make a data constructor with the given arguments.
17mkCon :: ConLike -> [Type] -> [LHsExpr GhcPs] -> LHsExpr GhcPs
18mkCon con apps (fmap unLoc -> args)
19  | RealDataCon dcon <- con
20  , dcon == nilDataCon
21  , [ty] <- apps
22  , ty `eqType` charTy = noLoc $ string ""
23
24  | RealDataCon dcon <- con
25  , isTupleDataCon dcon =
26      noLoc $ tuple args
27
28  | RealDataCon dcon <- con
29  , dataConIsInfix dcon
30  , (lhs : rhs : args') <- args =
31      noLoc $ foldl' (@@) (op lhs (coerceName con_name) rhs) args'
32
33  | Just fields <- getRecordFields con
34  , length fields >= 2 =  --  record notation is unnatural on single field ctors
35      noLoc $ recordConE (coerceName con_name) $ do
36        (arg, (field, _)) <- zip args fields
37        pure (coerceName field, arg)
38
39  | otherwise =
40      noLoc $ foldl' (@@) (bvar' $ occName con_name) args
41  where
42    con_name = conLikeName con
43
44
45coerceName :: HasOccName a => a -> RdrNameStr
46coerceName = UnqualStr . fromString . occNameString . occName
47
48
49------------------------------------------------------------------------------
50-- | Like 'var', but works over standard GHC 'OccName's.
51var' :: Var a => OccName -> a
52var' = var . fromString . occNameString
53
54
55------------------------------------------------------------------------------
56-- | Like 'bvar', but works over standard GHC 'OccName's.
57bvar' :: BVar a => OccName -> a
58bvar' = bvar . fromString . occNameString
59
60
61------------------------------------------------------------------------------
62-- | Get an HsExpr corresponding to a function name.
63mkFunc :: String -> HsExpr GhcPs
64mkFunc = var' . mkVarOcc
65
66
67------------------------------------------------------------------------------
68-- | Get an HsExpr corresponding to a value name.
69mkVal :: String -> HsExpr GhcPs
70mkVal = var' . mkVarOcc
71
72
73------------------------------------------------------------------------------
74-- | Like 'op', but easier to call.
75infixCall :: String -> HsExpr GhcPs -> HsExpr GhcPs -> HsExpr GhcPs
76infixCall s = flip op (fromString s)
77
78
79------------------------------------------------------------------------------
80-- | Like '(@@)', but uses a dollar instead of parentheses.
81appDollar :: HsExpr GhcPs -> HsExpr GhcPs -> HsExpr GhcPs
82appDollar = infixCall "$"
83
84