Synthesizing a function binding in Core

Question

I'm writing a GHC Core plugin that eventually adds new function declarations to the module it's called on.

Currently I'm struggling with combining the Core helpers to produce an Expr corresponding to a lambda function.

For example, say we want to synthesize the term \(x :: Double) -> x + x.

If we look at the pretty-printed Core of that :

f \ (x [Dmd=] :: Double) -> plusDouble x x

After copying the AST of that term into a Showable one we can see its structure of lambda abstractions and function applications , which I aim to reproduce by hand using Core combinators :

DLam "x::Double" 
  (DApp 
    (DApp (DVar "plusDouble::Double -> Double -> Double") (DVar "x::Double")) (DVar "x::Double"))

Long story short, I'm not able to synthesize that term L0L.

I produce a term that sort of looks right (?), GHC completes with no errors after splicing in my new term, the module with the new declaration loads correctly but GHCi crashes when I try using it.

My question(s) :

why do those empty (DEFAULT) cases appear? I suspect they are due to how I use mkCoreApps
how do I box the result such that it's a Double instead of a primitive Double# ?

GHCi signature of the new declaration :

PluginTest.f_new :: Double -> ghc-prim-0.7.0:GHC.Prim.Double#

pretty Core :

f_new \ (x :: Double) ->
  case x of x_ { D# x_ ->
  case x_ of wild_00 { __DEFAULT ->
  (case x_ of wild_00 { __DEFAULT -> +## wild_00 }) wild_00
  }
  }

Show instance of above :

DLam "x::Double" 
  (DCase (DVar "x::Double") "Double#" [
    ("D#",["x_"], DCase (DVar "x_::Double") "Double#" [
      ("__DEFAULT",[], DApp (DCase (DVar "x_::Double") "Double# -> Double#" [
        ("__DEFAULT",[], DApp (DVar "+##::Double# -> Double# -> Double#") (DVar "wild_00::Double#"))]) (DVar "wild_00::Double#"))])])

Full repro steps and code :

I reach into the ghc 9.0.1 API and after declaring imports and a couple helpers :

import GHC.Core.Make (mkCoreLams, mkSingleAltCase, mkCoreApps, mkCoreConApps) -- Core syntax combinators
import qualified GHC.Types.Name.Occurrence as ON (varName, mkOccName)
import GHC.Core.Opt.Monad (CoreM)
import GHC.Types.Unique.Supply (MonadUnique(..))
import GHC.Types.SrcLoc (noSrcSpan)
import GHC.Builtin.Types (manyDataConTy)
import GHC.Types.Id.Info (vanillaIdInfo)
import GHC.Types.Id.Make (mkPrimOpId)
import GHC.Builtin.PrimOps (PrimOp(..))
import GHC.Builtin.Types (doubleTy, floatTy, doubleDataCon, floatDataCon)
import GHC.Core.Utils (exprType)


-- | fresh name using the supply of unique symbols provided by MonadUnique
mkNameM :: String -> CoreM Name
mkNameM n = do
  u <- getUniqueM
  pure $ mkInternalName u (ON.mkOccName ON.varName n) noSrcSpan

-- | an external (= exported) name
mkExtNameM :: Module -- ^ module that will export this
           -> String -> CoreM Name
mkExtNameM modl n = do
  u <- getUniqueM
  pure $ mkExternalName u modl (ON.mkOccName ON.varName n) noSrcSpan

-- | variable identifier (of multiplicity "Many" which is the default)
mkId :: Name -> Type -> Id
mkId xname tyvar = mkLocalVar VanillaId xname manyDataConTy tyvar vanillaIdInfo

mkGlobalId :: Name -> Type -> Id
mkGlobalId xname tyvar = mkGlobalVar VanillaId xname tyvar vanillaIdInfo

I've also written helpers for declaring binary math operators, which uses the builtin primops and wraps them :

-- | apply an Expr corresponding to a binary operator to two argument expressions
appBin :: CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr
appBin f e1 e2 = mkCoreApps f [e1, e2]

-- | primop corrsponding to (+) :: Double -> Double -> Double
d_add :: CoreExpr
d_add = Var (mkPrimOpId DoubleAddOp)

Now we're ready to write a function that takes the current contents of the module Core (ModGuts), declares fresh variables and appends the new declaration to the module.

This function testAddLambda can be used as an additional Core plugin pass, as shown in the GHC user manual.

I've separated out the part that is supposed to declare the term \(x :: Double) -> x + x.

-- | add a single new binding corresponding to a lambda expression to ModGuts
testAddLambda :: ModGuts -> CoreM ModGuts
testAddLambda guts = do
  let
    modl = mg_module guts
    binds = mg_binds guts
    exports = mg_exports guts
  xn <- mkNameM "x"
  xn_ <- mkNameM "x_"
  fn <- mkExtNameM modl "f_new" -- new name to be exported
  let
    x = mkId xn doubleTy
    x_ = mkId xn_ doubleTy

    f = mkCoreLams [x] $
          mkSingleAltCase (Var x) x_ (DataAlt doubleDataCon) [x_] $
            appBin d_add (Var x_) (Var x_)
    fty = exprType f -- type of 'f'

    fv = mkGlobalId fn fty
    fbind = NonRec fv f
    fexp = Avail fn
    guts' = guts {
      mg_binds = binds ++ [fbind]
      , mg_exports = exports ++ [fexp]
      }
  pure guts'

Synthesizing a function binding in Core

0 Answers