feat: add CompilerM.lean and LCNF.lean

src/Lean/Compiler.lean

@@ -14,3 +14,5 @@ import Lean.Compiler.IR
 import Lean.Compiler.CSimpAttr
 import Lean.Compiler.FFI
 import Lean.Compiler.NoncomputableAttr
+import Lean.Compiler.CompilerM
+import Lean.Compiler.LCNF

src/Lean/Compiler/CompilerM.lean

@@ -0,0 +1,89 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Meta.InferType
+
+namespace Lean.Compiler
+
+structure CompilerM.State where
+  lctx     : LocalContext
+  letFVars : Array Expr := #[]
+
+abbrev CompilerM := StateRefT CompilerM.State CoreM
+
+@[inline] def liftMetaM (x : MetaM α) : CompilerM α := do
+  x.run' { lctx := (← get).lctx }
+
+def inferType (e : Expr) : CompilerM Expr := liftMetaM <| Meta.inferType e
+
+def whnf (e : Expr) : CompilerM Expr := liftMetaM <| Meta.whnf e
+
+def isProp (e : Expr) : CompilerM Bool := liftMetaM <| Meta.isProp e
+
+def isTypeFormerType (e : Expr) : CompilerM Bool := liftMetaM <| Meta.isTypeFormerType e
+
+def mkLocalDecl (binderName : Name) (type : Expr) (bi := BinderInfo.default) : CompilerM Expr := do
+  let fvarId ← mkFreshFVarId
+  modify fun s => { s with lctx := s.lctx.mkLocalDecl fvarId binderName type bi }
+  return .fvar fvarId
+
+def mkLetDecl (binderName : Name) (type : Expr) (value : Expr) (nonDep : Bool) : CompilerM Expr := do
+  let fvarId ← mkFreshFVarId
+  let x := .fvar fvarId
+  modify fun s => { s with lctx := s.lctx.mkLetDecl fvarId binderName type value nonDep, letFVars := s.letFVars.push x }
+  return x
+
+def visitLambda (e : Expr) : CompilerM (Array Expr × Expr) :=
+  go e #[]
+where
+  go (e : Expr) (fvars : Array Expr) := do
+    if let .lam binderName type body binderInfo := e then
+      let type := type.instantiateRev fvars
+      let fvar ← mkLocalDecl binderName type binderInfo
+      go body (fvars.push fvar)
+    else
+      return (fvars, e.instantiateRev fvars)
+
+def withNewScopeImp (x : CompilerM α) : CompilerM α := do
+  let s ← get
+  modify fun s => { s with letFVars := #[] }
+  try x finally set s
+
+def withNewScope [MonadFunctorT CompilerM m] (x : m α) : m α :=
+  monadMap (m := CompilerM) withNewScopeImp x
+
+class VisitLet (m : Type → Type) where
+  visitLet : Expr → (Expr → m Expr) → m Expr
+
+export VisitLet (visitLet)
+
+def visitLetImp (e : Expr) (f : Expr → CompilerM Expr) : CompilerM Expr :=
+  go e #[]
+where
+  go (e : Expr) (fvars : Array Expr) : CompilerM Expr := do
+    if let .letE binderName type value body nonDep := e then
+      let type := type.instantiateRev fvars
+      let value := value.instantiateRev fvars
+      let value ← f value
+      let fvar ← mkLetDecl binderName type value nonDep
+      go body (fvars.push fvar)
+    else
+      return e.instantiateRev fvars
+
+instance : VisitLet CompilerM where
+  visitLet := visitLetImp
+
+instance [VisitLet m] : VisitLet (ReaderT ρ m) where
+  visitLet e f ctx := visitLet e (f · ctx)
+
+instance [VisitLet m] : VisitLet (StateRefT' ω σ m) := inferInstanceAs (VisitLet (ReaderT _ _))
+
+def mkLetUsingScope (e : Expr) : CompilerM Expr :=
+  return (← get).lctx.mkLambda (← get).letFVars e
+
+def mkLambda (xs : Array Expr) (e : Expr) : CompilerM Expr :=
+  return (← get).lctx.mkLambda xs e
+
+end Lean.Compiler

src/Lean/Compiler/LCNF.lean

@@ -0,0 +1,150 @@
+/-
+Copyright (c) 2022 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+import Lean.Meta.Match.MatcherInfo
+import Lean.Meta.Transform
+import Lean.Compiler.InlineAttrs
+import Lean.Compiler.CompilerM
+
+namespace Lean.Compiler
+/-!
+# Lean Compiler Normal Form (LCNF)
+
+It is based on the [A-normal form](https://en.wikipedia.org/wiki/A-normal_form),
+and the approach described in the paper
+[Compiling  without continuations](https://www.microsoft.com/en-us/research/wp-content/uploads/2016/11/compiling-without-continuations.pdf).
+-/
+
+/--
+Return `true` if `e` is a `lcProof` application.
+Recall that we use `lcProof` to erase all nested proofs.
+-/
+def isLCProof (e : Expr) : Bool :=
+  e.isAppOfArity ``lcProof 1
+
+/--
+Recursors, noConfusion, constructors, and matchers are not treated as regular
+functions by the code generator. We eta-expand all of them to make sure they are
+not partially applied.
+-/
+def shouldEtaExpand (declName : Name) : CoreM Bool := do
+  let env ← getEnv
+  if isCasesOnRecursor env declName then return true
+  if isNoConfusion env declName then return true
+  if (← isRec declName) then return true
+  if (← Meta.isMatcher declName) then return true
+  if env.isConstructor declName then return true
+  if declName == ``Quot.lift then return true
+  return false
+
+/--
+Return `true` if `mdata` should be preserved.
+Right now, we don't preserve any `MData`, but this may
+change in the future when we add support for debugging information
+-/
+def isCompilerRelevantMData (_mdata : MData) : Bool :=
+  false
+
+/--
+Inline constants tagged with the `[macroInline]` attribute occurring in `e`.
+-/
+def macroInline (e : Expr) : CoreM Expr :=
+  Core.transform e fun e => do
+    let .const declName us := e.getAppFn | return .visit e
+    unless hasMacroInlineAttribute (← getEnv) declName do return .visit e
+    let val ← Core.instantiateValueLevelParams (← getConstInfo declName) us
+    return .visit <| val.beta e.getAppArgs
+
+namespace ToLCNF
+
+structure Context where
+  root : Bool
+
+structure State where
+  cache     : ExprMap Expr := {}
+
+abbrev M := ReaderT Context $ StateRefT State CompilerM
+
+def mkFreshLetDecl (e : Expr) : M Expr := do
+  if (← read).root then
+    return e
+  else
+    mkLetDecl (← mkFreshUserName `x) (← inferType e) e (nonDep := false)
+
+def withNewRoot (x : M α) : M α := do
+  let s ← get
+  try
+    withReader (fun _ => { root := true }) <| Compiler.withNewScope x
+  finally
+    set s
+
+/--
+Put the given expression in `LCNF`.
+
+- Nested proofs are replaced with `lcProof`-applications.
+- Eta-expand applications of declarations that satisfy `shouldEtaExpand`.
+- Put computationally relevant expressions in A-normal form.
+-/
+partial def toLCNF (lctx : LocalContext) (e : Expr) : CoreM Expr := do
+  let ((e, _), s) ← visit e |>.run { root := true } |>.run {} |>.run { lctx }
+  return s.lctx.mkLambda s.letFVars e
+where
+  visitChild (e : Expr) : M Expr :=
+    withReader (fun _ => { root := false }) <| visit e
+
+  visitApp (f : Expr) (args : Array Expr) : M Expr := do
+    -- TODO: handle special cases
+    let f := f
+    let args ← args.mapM visitChild
+    mkFreshLetDecl <| mkAppN f args
+
+  visitLambda (e : Expr) : M Expr := do
+    let r ← withNewRoot do
+      let (as, e) ← Compiler.visitLambda e
+      let e ← mkLetUsingScope (← visit e)
+      mkLambda as e
+    mkFreshLetDecl r
+
+  visitMData (mdata : MData) (e : Expr) : M Expr := do
+    if isCompilerRelevantMData mdata then
+      mkFreshLetDecl <| .mdata mdata (← visitChild e)
+    else
+      visit e
+
+  visitProj (s : Name) (i : Nat) (e : Expr) : M Expr := do
+    mkFreshLetDecl <| .proj s i (← visitChild e)
+
+  visit (e : Expr) : M Expr := withIncRecDepth do
+    match e with
+    | .mvar .. => throwError "unexpected occurrence of metavariable in code generator{indentExpr e}"
+    | .bvar .. => unreachable!
+    | .fvar .. | .sort .. | .forallE .. => return e -- Do nothing
+    | _ =>
+    if isLCProof e then
+      return e
+    let type ← inferType e
+    if (← isProp type) then
+      /- We replace proofs with `lcProof` applications. -/
+      return mkApp (mkConst ``lcProof) type
+    if (← isTypeFormerType type) then
+      /- Types and Type formers are not put into A-normal form. -/
+      return e
+    if let some e := (← get).cache.find? e then
+      return e
+    let r ← match e with
+      | .app ..     => e.withApp fun f args => visitApp f args
+      | .const ..   => visitApp e #[]
+      | .proj s i e => visitProj s i e
+      | .mdata d e  => visitMData d e
+      | .lam ..     => visitLambda e
+      | .letE ..    => visit (← visitLet e visitChild)
+      | .lit ..     => mkFreshLetDecl e
+      | _           => pure e
+    modify fun s => { s with cache := s.cache.insert e r }
+    return r
+
+end ToLCNF
+
+end Lean.Compiler

src/Lean/Meta/InferType.lean

@@ -363,6 +363,9 @@ partial def isTypeQuick : Expr → MetaM LBool
   | .mvar mvarId      => do let mvarType  ← inferMVarType mvarId;  isArrowType mvarType 0
   | .app f ..         => isTypeQuickApp f 1
 
+/--
+Return `true` iff the type of `e` is a `Sort _`.
+-/
 def isType (e : Expr) : MetaM Bool := do
   match (← isTypeQuick e) with
   | .true  => return true
@@ -374,6 +377,9 @@ def isType (e : Expr) : MetaM Bool := do
     | .sort .. => return true
     | _        => return false
 
+/--
+Return true iff `type` is `Sort _` or `As → Sort _`.
+-/
 partial def isTypeFormerType (type : Expr) : MetaM Bool := do
   let type ← whnfD type
   match type with
@@ -383,8 +389,9 @@ partial def isTypeFormerType (type : Expr) : MetaM Bool := do
   | _ => return false
 
 /--
-  Return true iff `e : Sort _` or `e : (forall As, Sort _)`.
-  Remark: it subsumes `isType` -/
+Return true iff `e : Sort _` or `e : (forall As, Sort _)`.
+Remark: it subsumes `isType`
+-/
 def isTypeFormer (e : Expr) : MetaM Bool := do
   isTypeFormerType (← inferType e)
 

tests/lean/run/lcnf1.lean

@@ -0,0 +1,10 @@
+import Lean
+
+open Lean Compiler Meta
+
+def test (declName : Name) : MetaM Unit := do
+  let .defnInfo info ← getConstInfo declName | throwError "definition expected"
+  let val ← ToLCNF.toLCNF {} (← macroInline info.value)
+  IO.println (← ppExpr val)
+
+#eval test ``Lean.Elab.Term.synthesizeInstMVarCore