feat: erased declarations in do notation

src/Init/Data.lean

@@ -38,6 +38,7 @@ public import Init.Data.Queue
 public import Init.Data.Sum
 public import Init.Data.BEq
 public import Init.Data.Subtype
+public import Init.Data.Erased
 public import Init.Data.ULift
 public import Init.Data.PLift
 public import Init.Data.Zero

src/Init/Data/Erased.lean

@@ -0,0 +1,43 @@
+/-
+Copyright (c) 2018 Mario Carneiro. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Mario Carneiro, Sebastian Graf
+-/
+module
+
+prelude
+public import Init.Classical
+public import Init.Ext
+import Init.Grind.Attr
+
+public section
+
+/-- A value hidden from compiled code. `Erased.mk 42` erases to a dummy at runtime, and
+proofs recover the `42` as `(Erased.mk 42).out`. -/
+@[expose] def Erased (α : Sort u) : Sort (max 1 u) :=
+  { s : α → Prop // ∃ a, (a = ·) = s }
+
+namespace Erased
+
+/-- Hides `a` in an `Erased α`. Compiled code drops the argument. -/
+@[expose, macro_inline] def mk {α : Sort u} (a : α) : Erased α :=
+  ⟨fun b => a = b, a, rfl⟩
+
+/-- The value hidden in `e`, available to proofs only. -/
+noncomputable def out {α : Sort u} (e : Erased α) : α :=
+  Classical.choose e.property
+
+@[simp, grind =] theorem out_mk {α : Sort u} (a : α) : (mk a).out = a :=
+  cast (congrFun (Classical.choose_spec (mk a).property) a).symm rfl
+
+@[simp, grind =] theorem mk_out {α : Sort u} (e : Erased α) : mk e.out = e := by
+  cases e with
+  | mk s h => exact Subtype.ext (Classical.choose_spec h)
+
+@[ext] theorem out_inj {α : Sort u} {a b : Erased α} (h : a.out = b.out) : a = b := by
+  rw [← mk_out a, ← mk_out b, h]
+
+@[simp, grind =] theorem mk_inj {α : Sort u} {a b : α} : mk a = mk b ↔ a = b :=
+  ⟨fun h => by have := congrArg out h; rwa [out_mk, out_mk] at this, fun h => h ▸ rfl⟩
+
+end Erased

src/Lean/Compiler/LCNF/ToLCNF.lean

@@ -447,6 +447,8 @@ private def checkComputable (ref : Name) : M Unit := do
   -- `noncomputable section`, where the failure to compile the `_unsafe_rec` version is tolerated and
   -- only that auxiliary is marked `noncomputable`, leaving `ref` itself unmarked.
   if isNoncomputable (← getEnv) ref || isNoncomputable (← getEnv) (mkUnsafeRecName ref) then
+    if ref == `Erased.out then
+      throwNamedError lean.dependsOnNoncomputable m!"failed to compile definition: it depends on 'Erased.out', which recovers the value of an erased variable. An erased variable's value is available in specifications such as `invariant` clauses and `assert`s, but not in compiled code. Consider marking the definition as 'noncomputable'."
     throwNamedError lean.dependsOnNoncomputable m!"failed to compile definition, consider marking it as 'noncomputable' because it depends on '{.ofConstName ref}', which is 'noncomputable'"
   else if getOriginalConstKind? (← getEnv) ref matches some .axiom | some .quot | some .induct | some .thm then
     throwNamedError lean.dependsOnNoncomputable f!"`{ref}` not supported by code generator; consider marking definition as `noncomputable`"

src/Lean/Elab/BuiltinDo/For.lean

@@ -6,6 +6,7 @@ Authors: Sebastian Graf
 module
 
 prelude
+meta import Init.Data.Erased
 public import Lean.Elab.BuiltinDo.Basic
 meta import Lean.Parser.Do
 meta import Std.WP.Gadget.ForIn
@@ -169,10 +170,22 @@ structure ForInApp where
   σ : Expr
   /-- The pattern naming the loop's mutable variables in the state tuple. -/
   statePat : Term
+  /-- The erased variables among the loop's mutable variables; annotations bind their `.out`
+  projections over the state tuple. -/
+  erasedMutVars : Array MutVar := #[]
+
+/-- Bind the `.out` projection of each erased variable over `e`, so that an annotation names erased
+variables at their underlying type. The `+zeta` substitutes the binding away at elaboration, so
+annotation goals carry the projection inline like the compiled body does. -/
+private def ForInApp.wrapErasedProjs (g : ForInApp) (e : Term) : DoElabM Term := do
+  let mut e := e
+  for mv in g.erasedMutVars do
+    e ← `(let +zeta $(mv.ident):ident := Erased.out $(⟨mv.ident.raw⟩); $e)
+  return e
 
 /-- Abstract `e` over the loop's state tuple, so that `e` may name the loop's mutable variables. -/
-private def ForInApp.mkStateFun (g : ForInApp) (e : Term) : DoElabM Term :=
-  `(fun $(g.statePat) => $e)
+private def ForInApp.mkStateFun (g : ForInApp) (e : Term) : DoElabM Term := do
+  `(fun $(g.statePat) => $(← g.wrapErasedProjs e))
 
 /-- Elaborate the gadget application that replaces the loop. The gadgets live downstream of this
 module, so `gadget` is an unresolved name that resolves in the user's context. -/
@@ -242,7 +255,7 @@ private def mkForInLoopGadget (g : ForInApp)
     -- unfolded type, and a specification's instance arguments are synthesized before the check that
     -- would unfold it.
     return ((invClause : Syntax), ← `($(mkIdent ``Std.WP.WhileInvariant.mk)
-      fun $exitVar:ident $(g.statePat) => $invBody))
+      fun $exitVar:ident $(g.statePat) => $(← g.wrapErasedProjs invBody)))
   let varArg? ← dec?.mapM fun decClause => do
     let (binders, body) ← match decClause with
       | `(doLoopDecreasing| decreasing $binders* => $body) => pure (binders, body)
@@ -296,17 +309,17 @@ private def mkForInLoopGadget (g : ForInApp)
         | some e => mkSome oldReturnCont.resultType e
       defs := defs.push returnVar
     for x in loopMutVars do
-      let defn ← getLocalDeclFromUserName x.getId
-      Term.addTermInfo' x.ident defn.toExpr
+      Term.addTermInfo' x.ident (← getFVarFromUserName x.getId)
+      let v ← x.stateValue
       -- ForIn forces the mut tuple into the universe mi.u: that of the do block result type.
       -- If we don't do this, then we are stuck on solving constraints such as
       --   `max ?u.46 ?u.47 =?= max (max ?u.22 ?u.46) ?u.47`
       -- It's important we do this as a separate isLevelDefEq check on the decremented level because
       -- otherwise (`ensureHasType (mkSort mi.u.succ)`) we are stuck on constraints like
       --   `max (?u+1) (?v+1) =?= ?u+1`
-      let u ← getDecLevel defn.type
+      let u ← getDecLevel (← inferType v)
       discard <| isLevelDefEq u mi.u
-      defs := defs.push defn.toExpr
+      defs := defs.push v
     if info.returnsEarly && loopMutVars.isEmpty then
       defs := defs.push (mkConst ``Unit.unit)
     return defs
@@ -367,7 +380,8 @@ private def mkForInLoopGadget (g : ForInApp)
   let mut forIn := mkApp app body
   unless inv?.isNone && dec?.isNone do
     let g : ForInApp :=
-      { xs, init := preS, body, σ, statePat := ← mkStatePat loopMutVars info.returnsEarly }
+      { xs, init := preS, body, σ, statePat := ← mkStatePat loopMutVars info.returnsEarly,
+        erasedMutVars := loopMutVars.filter (·.erased) }
     if (← instantiateMVars ρ).isConstOf ``Lean.Loop then
       if let some e ← mkForInLoopGadget g inv? dec? then forIn := e
     else if let some decClause := dec? then

src/Lean/Elab/BuiltinDo/Let.lean

@@ -6,52 +6,69 @@ Authors: Sebastian Graf
 module
 
 prelude
+meta import Init.Data.Erased
 public import Lean.Elab.Do.Basic
 meta import Lean.Parser.Do
 import Lean.Elab.BuiltinDo.Basic
 import Lean.Elab.Do.PatternVar
 
 public section
 
+-- The `erased` doElem quotations below need the current stage's parser until stage0 catches up.
+set_option internal.parseQuotWithCurrentStage true
+
 namespace Lean.Elab.Do
 
 open Lean.Parser.Term
 open Lean.Meta
 
 inductive LetOrReassign
-  | let (mutTk? : Option Syntax)
+  | let (mutTk? : Option Syntax) (erased : Bool)
   | have
   | reassign
 
 def LetOrReassign.getLetMutTk? (letOrReassign : LetOrReassign) : Option Syntax :=
   match letOrReassign with
-  | .let mutTk? => mutTk?
-  | _           => none
+  | .let mutTk? _ => mutTk?
+  | _             => none
+
+def LetOrReassign.isErasedDecl (letOrReassign : LetOrReassign) : Bool :=
+  match letOrReassign with
+  | .let _ erased => erased
+  | _            => false
+
+def isErased (letOrReassign : LetOrReassign) (vars : Array Ident) : DoElabM Bool := do
+  match letOrReassign with
+  | .let _ erased => return erased
+  | .reassign    =>
+    let some v := vars[0]? | return false
+    let some mv ← findMutVar? v.getId | return false
+    return mv.erased
+  | _            => return false
 
 def LetOrReassign.checkMutVars (letOrReassign : LetOrReassign) (vars : Array Ident) : DoElabM Unit :=
   match letOrReassign with
   | .reassign => do
     throwUnlessMutVarsDeclared vars
+    -- Reassigning an erased variable wraps its value, which only the single-variable form can do.
+    unless vars.size == 1 do
+      for v in vars do
+        if ((← findMutVar? v.getId).map (·.erased)).getD false then
+          throwErrorAt v "an erased variable takes a plain reassignment, as in `{v.getId} := e`"
   | _         => checkMutVarsForShadowing vars
 
 def LetOrReassign.registerReassignAliasInfo (letOrReassign : LetOrReassign) (vars : Array Ident) : DoElabM Unit := do
   if letOrReassign matches .reassign then
     for var in vars do
       registerMutVarAlias var.getId
 
-def elabDoLetOrReassignWith (hint : MessageData) (letOrReassign : LetOrReassign) (vars : Array Ident)
-    (k : DoElabM Expr) (elabBody : (body : Term) → TermElabM Expr) : DoElabM Expr := do
-  -- letOrReassign.checkMutVars vars -- Should be done by the caller!
-  let elabCont : DoElabM Expr := do
-    declareMutVars? letOrReassign.getLetMutTk? vars do
-      letOrReassign.registerReassignAliasInfo vars
-      k
-  doElabToSyntax hint elabCont fun body => elabBody body
-
 def elabWithReassignments (letOrReassign : LetOrReassign) (vars : Array Ident) (k : DoElabM Expr) : DoElabM Expr := do
-  declareMutVars? letOrReassign.getLetMutTk? vars do
+  declareMutVars? letOrReassign.getLetMutTk? vars letOrReassign.isErasedDecl do
     letOrReassign.registerReassignAliasInfo vars
-    k
+    if ← isErased letOrReassign vars then
+      vars.foldr (init := k) withErasedProj
+    else
+      k
 
 private def pushTypeIntoReassignment (letOrReassign : LetOrReassign) (decl : TSyntax ``letDecl) : TermElabM (TSyntax ``letDecl) := do
   if letOrReassign matches .reassign then
@@ -83,14 +100,25 @@ private def checkLetConfigInDo (config : Term.LetConfig) : DoElabM Unit := do
   if config.generalize then
     throwError "`+generalize` is not supported in `do` blocks"
 
+/-- Wrap an erased binding `x : t := e` as `x : Erased t := Erased.mk e`. For a reassignment,
+`pushTypeIntoReassignment` has already checked the ascription and pinned `t`. -/
+private def wrapErasedDecl (decl : TSyntax ``letDecl) : DoElabM (TSyntax ``letDecl) := do
+  let `(letDecl| $x:ident $[: $t?]? := $e) := decl
+    | throwUnsupportedSyntax
+  match t? with
+  | some t => `(letDecl| $x:ident : Erased $t := Erased.mk ($e : $t))
+  | none   => `(letDecl| $x:ident := Erased.mk $e)
+
 partial def elabDoLetOrReassign (config : Term.LetConfig) (letOrReassign : LetOrReassign) (decl : TSyntax ``letDecl)
     (tk : Syntax) (dec : DoElemCont) : DoElabM Expr := do
   checkLetConfigInDo config
   let vars ← getLetDeclVars decl
   letOrReassign.checkMutVars vars
   let dec ← dec.ensureUnitAt tk
+  let isErased ← isErased letOrReassign vars
   -- Some decl preprocessing on the patterns and expected types:
   let decl ← pushTypeIntoReassignment letOrReassign decl
+  let decl ← if isErased then wrapErasedDecl decl else pure decl
   let mγ ← mkMonadApp (← read).doBlockResultType
   match decl with
   | `(letDecl| $decl:letEqnsDecl) =>
@@ -137,10 +165,9 @@ partial def elabDoLetOrReassign (config : Term.LetConfig) (letOrReassign : LetOr
     trace[Elab.let.decl] "{id.getId} : {type} := {val}"
     withLetDecl id.getId (kind := kind) type val (nondep := nondep) fun x => do
       Term.addLocalVarInfo id x
-      elabWithReassignments letOrReassign vars do
       match config.eq? with
       | none =>
-        let body ← dec.continueWithUnit
+        let body ← elabWithReassignments letOrReassign vars dec.continueWithUnit
         if config.zeta then
           pure <| (← body.abstractM #[x]).instantiate1 val
         else
@@ -149,7 +176,7 @@ partial def elabDoLetOrReassign (config : Term.LetConfig) (letOrReassign : LetOr
         let hTy ← mkEq x val
         withLetDecl h.getId hTy (← mkEqRefl x) (nondep := true) fun h' => do
           Term.addLocalVarInfo h h'
-          let body ← dec.continueWithUnit
+          let body ← elabWithReassignments letOrReassign vars dec.continueWithUnit
           if config.zeta then
             pure <| (← body.abstractM #[x, h']).instantiateRev #[val, ← mkEqRefl val]
           else if nondep then
@@ -159,44 +186,6 @@ partial def elabDoLetOrReassign (config : Term.LetConfig) (letOrReassign : LetOr
             mkLetFVars #[x, h'] body (usedLetOnly := config.usedOnly) (generalizeNondepLet := false)
   | _ => throwUnsupportedSyntax
 
-def elabDoArrow (letOrReassign : LetOrReassign) (stx : TSyntax [``doIdDecl, ``doPatDecl]) (tk : Syntax) (dec : DoElemCont) : DoElabM Expr := do
-  match stx with
-  | `(doIdDecl| $x:ident $[: $xType?]? ← $rhs) =>
-    letOrReassign.checkMutVars #[x]
-    let dec ← dec.ensureUnitAt tk
-    -- For plain variable reassignment, we know the expected type of the reassigned variable and
-    -- propagate it eagerly via type ascription if the user hasn't provided one themselves:
-    let xType? ← match letOrReassign, xType? with
-      | .reassign, none =>
-        let decl ← getLocalDeclFromUserName x.getId
-        some <$> Term.exprToSyntax decl.type
-      | _, _ => pure xType?
-    elabDoIdDecl x xType? rhs (declareMutVar? letOrReassign.getLetMutTk? x <| dec.continueWithUnit)
-      (kind := dec.kind)
-  | `(doPatDecl| _%$pattern $[: $patType?]? ← $rhs) =>
-    let x := mkIdentFrom pattern (← mkFreshUserName `__x)
-    let dec ← dec.ensureUnitAt tk
-    elabDoIdDecl x patType? rhs dec.continueWithUnit (kind := dec.kind)
-  | `(doPatDecl| $pattern:term $[: $patType?]? ← $rhs $[| $otherwise? $(rest?)?]?) =>
-    let rest? := rest?.join
-    let x := mkIdentFrom pattern (← mkFreshUserName `__x)
-    elabDoIdDecl x patType? rhs do
-      match letOrReassign, otherwise? with
-      | .let mutTk?, some otherwise =>
-        elabDoElem (← `(doElem| let $[mut%$mutTk?]? $pattern:term := $x | $otherwise $(rest?)?)) dec
-      | .let mutTk?, _ =>
-        elabDoElem (← `(doElem| let $[mut%$mutTk?]? $pattern:term := $x)) dec
-      | .have, some _otherwise =>
-        throwUnsupportedSyntax
-      | .have, _ =>
-        elabDoElem (← `(doElem| have $pattern:term := $x)) dec
-      | .reassign, _ =>
-        -- otherwise? is always `none`, because there is no `doReassignElse`
-        unless rest?.isNone do
-          throwError "reassignment with `|` (i.e., \"else clause\") is not supported"
-        elabDoElem (← `(doElem| $pattern:term := $x)) dec
-  | _ => throwUnsupportedSyntax
-
 private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letConfig)
     (mutTk? : Option Syntax) (initConfig : Term.LetConfig := {}) : DoElabM Term.LetConfig := do
   if mutTk?.isSome && !letConfigStx.raw[0].getArgs.isEmpty then
@@ -206,7 +195,20 @@ private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letCon
 @[builtin_doElem_elab Lean.Parser.Term.doLet] def elabDoLet : DoElab := fun stx dec => do
   let `(doLet| let%$tk $[mut%$mutTk?]? $config:letConfig $decl:letDecl) := stx | throwUnsupportedSyntax
   let config ← getLetConfigAndCheckMut config mutTk?
-  elabDoLetOrReassign config (.let mutTk?) decl tk dec
+  elabDoLetOrReassign config (.let mutTk? false) decl tk dec
+
+@[builtin_doElem_elab Lean.Parser.Term.doErased] def elabDoErased : DoElab := fun stx dec => do
+  let `(doErased| erased%$tk $[mut%$mutTk?]? $x:ident $[: $t?]? := $e) := stx | throwUnsupportedSyntax
+  elabDoLetOrReassign {} (.let mutTk? true) (← `(letDecl| $x:ident $[: $t?]? := $e)) tk dec
+
+@[builtin_macro Lean.Parser.Term.doErasedArrow] def expandDoErasedArrow : Macro := fun stx => do
+  match stx with
+  | `(doErasedArrow| erased%$tk $[mut%$mutTk?]? $x:ident $[: $t?]? ← $rhs) =>
+    let y := mkIdentFrom x (← MonadQuotation.addMacroScope `__x)
+    let letElem ← `(doElem| let $y:ident $[: $t?]? ← $rhs)
+    let erasedElem : TSyntax `doElem := ⟨(← `(doErased| erased%$tk $[mut%$mutTk?]? $x:ident := $y)).raw⟩
+    `(doElem| do $letElem:doElem; $erasedElem:doElem)
+  | _ => Macro.throwUnsupported
 
 @[builtin_doElem_elab Lean.Parser.Term.doHave] def elabDoHave : DoElab := fun stx dec => do
   let `(doHave| have%$tk $config:letConfig $decl:letDecl) := stx | throwUnsupportedSyntax
@@ -241,7 +243,7 @@ private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letCon
     | throwUnsupportedSyntax
   let config ← getLetConfigAndCheckMut cfg mutTk?
   checkLetConfigInDo config
-  let letOrReassign := LetOrReassign.let mutTk?
+  let letOrReassign := LetOrReassign.let mutTk? false
   let vars ← getPatternVarsEx pattern
   letOrReassign.checkMutVars vars
   let mut body ← body?.getDM `(doSeqIndent|pure PUnit.unit)
@@ -260,12 +262,42 @@ private def getLetConfigAndCheckMut (letConfigStx : TSyntax ``Parser.Term.letCon
   checkLetConfigInDo config
   if config.nondep || config.usedOnly || config.zeta || config.eq?.isSome then
     throwErrorAt cfg "configuration options are not supported with `←`"
-  elabDoArrow (.let mutTk?) decl tk dec
+  match decl with
+  | `(doIdDecl| $x:ident $[: $xType?]? ← $rhs) =>
+    checkMutVarsForShadowing #[x]
+    let dec ← dec.ensureUnitAt tk
+    elabDoIdDecl x xType? rhs (declareMutVar? mutTk? x false <| dec.continueWithUnit)
+      (kind := dec.kind)
+  | `(doPatDecl| _%$pattern $[: $patType?]? ← $rhs) =>
+    let x := mkIdentFrom pattern (← mkFreshUserName `__x)
+    let dec ← dec.ensureUnitAt tk
+    elabDoIdDecl x patType? rhs dec.continueWithUnit (kind := dec.kind)
+  -- No `ensureUnitAt` here: the else form swallows the rest of the block into `rest?`, so the
+  -- element keeps `dec`'s result type.
+  | `(doPatDecl| $pattern:term $[: $patType?]? ← $rhs $[| $otherwise? $(rest?)?]?) =>
+    let rest? := rest?.join
+    let x := mkIdentFrom pattern (← mkFreshUserName `__x)
+    elabDoIdDecl x patType? rhs do
+      match otherwise? with
+      | some otherwise =>
+        elabDoElem (← `(doElem| let $[mut%$mutTk?]? $pattern:term := $x | $otherwise $(rest?)?)) dec
+      | none =>
+        elabDoElem (← `(doElem| let $[mut%$mutTk?]? $pattern:term := $x)) dec
+  | _ => throwUnsupportedSyntax
 
 @[builtin_doElem_elab Lean.Parser.Term.doReassignArrow] def elabDoReassignArrow : DoElab := fun stx dec => do
   match stx with
-  | `(doReassignArrow| $decl:doIdDecl) =>
-    elabDoArrow .reassign decl decl dec
-  | `(doReassignArrow| $decl:doPatDecl) =>
-    elabDoArrow .reassign decl decl dec
+  | `(doReassignArrow| $x:ident $[: $t?]? ← $rhs) =>
+    throwUnlessMutVarDeclared x
+    -- Pin the variable's declared type on the bind, so a type error blames the action.
+    let t ← match t? with
+      | some t => pure t
+      | none   => Term.exprToSyntax (← getLocalDeclFromUserName x.getId).type
+    let y := mkIdentFrom x (← mkFreshUserName `__x)
+    elabDoIdDecl y (some t) rhs (elabDoElem (← `(doElem| $x:ident := $y)) dec) (kind := dec.kind)
+  | `(doReassignArrow| $pat:term $[: $t?]? ← $rhs $[| $otherwise? $(rest?)?]?) =>
+    unless otherwise?.isNone && rest?.join.isNone do
+      throwError "reassignment with `|` (i.e., \"else clause\") is not supported"
+    let y := mkIdentFrom pat (← mkFreshUserName `__x)
+    elabDoIdDecl y t? rhs (elabDoElem (← `(doElem| $pat:term := $y)) dec) (kind := dec.kind)
   | _ => throwUnsupportedSyntax