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