Utils.hs

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, 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 con apps (fmap unLoc -> args)
  | RealDataCon dcon <- con
  , dcon == nilDataCon
  , [ty] <- apps
  , ty `eqType` charTy = noLoc $ string ""

  | RealDataCon dcon <- con
  , isTupleDataCon dcon =
      noLoc $ tuple args

  | RealDataCon dcon <- con
  , dataConIsInfix dcon
  , (lhs : rhs : args') <- args =
      noLoc $ foldl' (@@) (op lhs (coerceName con_name) rhs) args'

  | Just fields <- getRecordFields con
  , length fields >= 2 =  --  record notation is unnatural on single field ctors
      noLoc $ recordConE (coerceName con_name) $ do
        (arg, (field, _)) <- zip args fields
        pure (coerceName field, arg)

  | otherwise =
      noLoc $ foldl' (@@) (bvar' $ occName con_name) args
  where
    con_name = conLikeName con


coerceName :: HasOccName a => a -> RdrNameStr
coerceName = fromString . occNameString . occName


------------------------------------------------------------------------------
-- | Like 'var', but works over standard GHC 'OccName's.
var' :: Var a => OccName -> a
var' = var . fromString . occNameString


------------------------------------------------------------------------------
-- | Like 'bvar', but works over standard GHC 'OccName's.
bvar' :: BVar a => OccName -> a
bvar' = bvar . fromString . occNameString


------------------------------------------------------------------------------
-- | Get an HsExpr corresponding to a function name.
mkFunc :: String -> HsExpr GhcPs
mkFunc = var' . mkVarOcc


------------------------------------------------------------------------------
-- | Get an HsExpr corresponding to a value name.
mkVal :: String -> HsExpr GhcPs
mkVal = var' . mkVarOcc


------------------------------------------------------------------------------
-- | Like 'op', but easier to call.
infixCall :: String -> HsExpr GhcPs -> HsExpr GhcPs -> HsExpr GhcPs
infixCall s = flip op (fromString s)


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