From 7da8a681a8dfc43e1d634481486ba7298632ca4d Mon Sep 17 00:00:00 2001
From: Rishikesh Vaishnav <rishhvaishnav@gmail.com>
Date: Fri, 17 May 2024 16:42:19 +0200
Subject: [PATCH 1/6] docs: document defeq-related functions

---
 Lean4Lean/Inductive/Reduce.lean |  31 +++++++
 Lean4Lean/Quot.lean             |  19 +++-
 Lean4Lean/TypeChecker.lean      | 151 +++++++++++++++++++++++++++++++-
 3 files changed, 198 insertions(+), 3 deletions(-)

diff --git a/Lean4Lean/Inductive/Reduce.lean b/Lean4Lean/Inductive/Reduce.lean
index da73b54..5925d5e 100644
--- a/Lean4Lean/Inductive/Reduce.lean
+++ b/Lean4Lean/Inductive/Reduce.lean
@@ -17,6 +17,15 @@ def mkNullaryCtor (type : Expr) (nparams : Nat) : Option Expr :=
   let name ← getFirstCtor env dName
   return mkAppRange (.const name ls) 0 nparams args
 
+/--
+When `e` has the type of a K-like inductive, converts it into a constructor
+application.
+
+For instance if we have `e : Eq a a`, it is converted into `Eq.refl a` (which
+it is definitionally equal to by proof irrelevance). Note that the indices of
+`e`'s type must match those of the constructor application (for instance,
+`e : Eq a b` cannot be converted if `a` and `b` are not defeq).
+-/
 def toCtorWhenK (rval : RecursorVal) (e : Expr) : m Expr := do
   assert! rval.k
   let appType ← whnf (← inferType e)
@@ -27,6 +36,7 @@ def toCtorWhenK (rval : RecursorVal) (e : Expr) : m Expr := do
     for h : i in [rval.numParams:appTypeArgs.size] do
       if (appTypeArgs[i]'h.2).hasExprMVar then return e
   let some newCtorApp := mkNullaryCtor env appType rval.numParams | return e
+  -- check that the indices of types of `e` and `newCtorApp` match
   unless ← isDefEq appType (← inferType newCtorApp) do return e
   return newCtorApp
 
@@ -40,6 +50,14 @@ def expandEtaStruct (eType e : Expr) : Expr :=
     result := .app result (.proj I i e)
   pure result
 
+/--
+When `e` is of struct type, converts it into a constructor application using
+projections.
+
+For instance if we have `e : String`, it is converted into
+`String.mk (String.data e)` (which is definitionally equal to `e` by struct
+eta).
+-/
 def toCtorWhenStruct (inductName : Name) (e : Expr) : m Expr := do
   if !isStructureLike env inductName || (e.isConstructorApp?' env).isSome then
     return e
@@ -52,6 +70,16 @@ def getRecRuleFor (rval : RecursorVal) (major : Expr) : Option RecursorRule := d
   let .const fn _ := major.getAppFn | none
   rval.rules.find? (·.ctor == fn)
 
+/--
+Performs recursor reduction on `e` (returning `none` if not applicable).
+
+For recursor reduction to occur, `e` must be a recursor application where the
+major premise is either a complete constructor application or of a K- or
+structure-like inductive type (in which case it is converted into an equivalent
+constructor application). The reduction is done by applying the
+`RecursorRule.rhs` associated with the constructor to the parameters from the
+recursor application and the fields of the constructor application.
+-/
 def inductiveReduceRec [Monad m] (env : Environment) (e : Expr)
     (whnf : Expr → m Expr) (inferType : Expr → m Expr) (isDefEq : Expr → Expr → m Bool) :
     m (Option Expr) := do
@@ -71,7 +99,10 @@ def inductiveReduceRec [Monad m] (env : Environment) (e : Expr)
   if rule.nfields > majorArgs.size then return none
   if ls.length != info.levelParams.length then return none
   let mut rhs := rule.rhs.instantiateLevelParams info.levelParams ls
+  -- get parameters from recursor application (recursor rules don't need the indices,
+  -- as these are determined by the constructor and its parameters/fields)
   rhs := mkAppRange rhs 0 info.getFirstIndexIdx recArgs
+  -- get fields from constructor application
   rhs := mkAppRange rhs (majorArgs.size - rule.nfields) majorArgs.size majorArgs
   if majorIdx + 1 < recArgs.size then
     rhs := mkAppRange rhs (majorIdx + 1) recArgs.size recArgs
diff --git a/Lean4Lean/Quot.lean b/Lean4Lean/Quot.lean
index 7d5fad5..a846dc4 100644
--- a/Lean4Lean/Quot.lean
+++ b/Lean4Lean/Quot.lean
@@ -71,7 +71,7 @@ def Environment.addQuot (env : Environment) : Except KernelException Environment
   let all_quot := (← read).mkForall #[a] <| .app β quotMk_a
   withLocalDecl `q quot_r .implicit fun q => do
   -- constant Quot.ind.{u} {α : Sort u} {r : α → α → Prop} {β : @Quot.{u} α r → Prop} :
-  --   (∀ a : α, β (@Quot.mk.{u} α r a)) → ∀ q : @Quot.{u} α r, β q */
+  --   (∀ a : α, β (@Quot.mk.{u} α r a)) → ∀ q : @Quot.{u} α r, β q
   let env := add env <| .quotInfo {
     name := ``Quot.ind, kind := .ind, levelParams := [`u]
     type := (← read).mkForall #[α, r, β] <|
@@ -79,6 +79,23 @@ def Environment.addQuot (env : Environment) : Except KernelException Environment
   }
   return markQuotInit env
 
+/--
+Reduces the head application of a quotient eliminator as follows:
+
+```
+Quot.lift.{u, v} {α : Sort u} {r : α → α → Prop} {β : Sort v} (f : α → β) :
+  (∀ a b : α, r a b → f a = f b) → @Quot.{u} α r → β
+
+Quot.lift f h (Quot.mk r a) ... ⟶ f a ...
+```
+
+```
+Quot.ind.{u} {α : Sort u} {r : α → α → Prop} {β : @Quot.{u} α r → Prop} :
+  (∀ a : α, β (@Quot.mk.{u} α r a)) → ∀ q : @Quot.{u} α r, β q
+
+Quot.ind p (Quot.mk r a) ... ⟶ p a ...
+```
+-/
 def quotReduceRec [Monad m] (e : Expr) (whnf : Expr → m Expr) : m (Option Expr) := do
   let .const fn _ := e.getAppFn | return none
   let cont mkPos argPos := do
diff --git a/Lean4Lean/TypeChecker.lean b/Lean4Lean/TypeChecker.lean
index f826216..681a323 100644
--- a/Lean4Lean/TypeChecker.lean
+++ b/Lean4Lean/TypeChecker.lean
@@ -62,6 +62,9 @@ inductive ReductionStatus where
 
 namespace Inner
 
+/--
+Reduces `e` to its weak-head normal form.
+-/
 def whnf (e : Expr) : RecM Expr := fun m => m.whnf e
 
 @[inline] def withLCtx [MonadWithReaderOf LocalContext m] (lctx : LocalContext) (x : m α) : m α :=
@@ -139,8 +142,13 @@ def inferForall (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e wher
     let s ← ensureSortCore r e
     return .sort <| us.foldr mkLevelIMax' s.sortLevel!
 
+/--
+Returns whether `t` and `s` are definitionally equal according to Lean's
+algorithmic definitional equality judgment.
+-/
 def isDefEqCore (t s : Expr) : RecM Bool := fun m => m.isDefEqCore t s
 
+@[inherit_doc isDefEqCore]
 def isDefEq (t s : Expr) : RecM Bool := do
   let r ← isDefEqCore t s
   if r then
@@ -276,6 +284,19 @@ def inferType' (e : Expr) (inferOnly : Bool) : RecM Expr := do
     { s with inferTypeC := s.inferTypeC.insert e r }
   return r
 
+/--
+Reduces `e` to its weak-head normal form, without unfolding definitions. This
+is a conservative version of `whnf` (which does unfold definitions), to be used
+for efficiency purposes.
+
+Setting `cheapRec` or `cheapProj` to `true` will cause the major premise/struct
+argument to be reduced "lazily" (using `whnfCore` rather than `whnf`) when
+reducing recursor applications/struct projections. This can be a useful
+optimization if we're checking the definitional equality of two recursor
+applications/struct projections of the same recursor/projections, where we
+might save some work by directly checking if the major premises/struct
+arguments are defeq (rather than eagerly applying a recursor rule/projection).
+-/
 def whnfCore (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr :=
   fun m => m.whnfCore e cheapRec cheapProj
 
@@ -289,11 +310,19 @@ def reduceRecursor (e : Expr) (cheapRec cheapProj : Bool) : RecM (Option Expr) :
     return r
   return none
 
+/--
+Gets the weak-head normal form of the free variable `e`,
+which is the weak-head normal form of its definition if `e` is a let variable and itself if it is a lambda variable.
+-/
 def whnfFVar (e : Expr) (cheapRec cheapProj : Bool) : RecM Expr := do
   if let some (.ldecl (value := v) ..) := (← getLCtx).find? e.fvarId! then
     return ← whnfCore v cheapRec cheapProj
   return e
 
+/--
+Reduces a projection of `struct` at index `idx` (when `struct` is reducible to a
+constructor application).
+-/
 def reduceProj (idx : Nat) (struct : Expr) (cheapRec cheapProj : Bool) : RecM (Option Expr) := do
   let mut c ← (if cheapProj then whnfCore struct cheapRec cheapProj else whnf struct)
   if let .lit (.strVal s) := c then
@@ -307,6 +336,7 @@ def reduceProj (idx : Nat) (struct : Expr) (cheapRec cheapProj : Bool) : RecM (O
 def isLetFVar (lctx : LocalContext) (fvar : FVarId) : Bool :=
   lctx.find? fvar matches some (.ldecl ..)
 
+@[inherit_doc whnfCore]
 def whnfCore' (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr := do
   match e with
   | .bvar .. | .sort .. | .mvar .. | .forallE .. | .const .. | .lam .. | .lit .. => return e
@@ -323,8 +353,9 @@ def whnfCore' (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr :=
   | .bvar .. | .sort .. | .mvar .. | .forallE .. | .const .. | .lam .. | .lit ..
   | .mdata .. => unreachable!
   | .fvar _ => return ← whnfFVar e cheapRec cheapProj
-  | .app .. =>
+  | .app .. => -- beta-reduce at the head as much as possible, apply any remaining `rargs` to the resulting expression, and re-run `whnfCore`
     e.withAppRev fun f0 rargs => do
+    -- the head may still be a let variable/binding, projection, or mdata-wrapped expression
     let f ← whnfCore f0 cheapRec cheapProj
     if let .lam _ _ body _ := f then
       let rec loop m (f : Expr) : RecM Expr :=
@@ -344,6 +375,7 @@ def whnfCore' (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr :=
         pure e
     else
       let r := f.mkAppRevRange 0 rargs.size rargs
+      -- FIXME replace with reduceRecursor? adding arguments can only result in further normalization if the head reduced to a partial recursor application
       save <|← whnfCore r cheapRec cheapProj
   | .letE _ _ val body _ =>
     save <|← whnfCore (body.instantiate1 val) cheapRec cheapProj
@@ -353,6 +385,10 @@ def whnfCore' (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr :=
     else
       save e
 
+/--
+Checks if `e` has a head constant that can be delta-reduced (that is, it is a
+theorem or definition), returning its `ConstantInfo` if so.
+-/
 def isDelta (env : Environment) (e : Expr) : Option ConstantInfo := do
   if let .const c _ := e.getAppFn then
     if let some ci := env.find? c then
@@ -367,6 +403,10 @@ def unfoldDefinitionCore (env : Environment) (e : Expr) : Option Expr := do
         return d.instantiateValueLevelParams! ls
   none
 
+/--
+Unfolds the definition at the head of the application `e` (or `e` itself if it
+is not an application).
+-/
 def unfoldDefinition (env : Environment) (e : Expr) : Option Expr := do
   if e.isApp then
     let f0 := e.getAppFn
@@ -387,16 +427,34 @@ def reduceNative (_env : Environment) (e : Expr) : Except KernelException (Optio
 
 def natLitExt? (e : Expr) : Option Nat := if e == .natZero then some 0 else e.natLit?
 
+/--
+Reduces the application `f a b` to a Nat literal if `a` and `b` can be reduced
+to Nat literals.
+
+Note: `f` should have an (efficient) external implementation.
+-/
 def reduceBinNatOp (f : Nat → Nat → Nat) (a b : Expr) : RecM (Option Expr) := do
   let some v1 := natLitExt? (← whnf a) | return none
   let some v2 := natLitExt? (← whnf b) | return none
   return some <| .lit <| .natVal <| f v1 v2
 
+/--
+Reduces the application `f a b` to a boolean expression if `a` and `b` can be
+reduced to Nat literals.
+
+Note: `f` should have an (efficient) external implementation.
+-/
 def reduceBinNatPred (f : Nat → Nat → Bool) (a b : Expr) : RecM (Option Expr) := do
   let some v1 := natLitExt? (← whnf a) | return none
   let some v2 := natLitExt? (← whnf b) | return none
   return toExpr <| f v1 v2
 
+/--
+Reduces `e` to a natural number literal if possible, where binary operations
+and predicates may be applied (provided they have an external implementation).
+These include: `Nat.add`, `Nat.sub`, `Nat.mul`, `Nat.pow`, `Nat.gcd`,
+`Nat.mod`, `Nat.div`, `Nat.beq`, `Nat.ble`.
+-/
 def reduceNat (e : Expr) : RecM (Option Expr) := do
   if e.hasFVar then return none
   let nargs := e.getAppNumArgs
@@ -418,7 +476,7 @@ def reduceNat (e : Expr) : RecM (Option Expr) := do
     if f == ``Nat.ble then return ← reduceBinNatPred Nat.ble a b
   return none
 
-
+@[inherit_doc whnf]
 def whnf' (e : Expr) : RecM Expr := do
   -- Do not cache easy cases
   match e with
@@ -444,6 +502,12 @@ def whnf' (e : Expr) : RecM Expr := do
   modify fun s => { s with whnfCache := s.whnfCache.insert e r }
   return r
 
+/--
+If `t` and `s` are lambda expressions, checks that their domains are defeq and
+recurses on the bodies, substituting in a new free variable for that binder
+(this substitution is delayed for efficiency purposes using the `subst`
+parameter). Otherwise, does a normal defeq check.
+-/
 def isDefEqLambda (t s : Expr) (subst : Array Expr := #[]) : RecM Bool :=
   match t, s with
   | .lam _ tDom tBody _, .lam name sDom sBody bi => do
@@ -462,6 +526,12 @@ def isDefEqLambda (t s : Expr) (subst : Array Expr := #[]) : RecM Bool :=
       isDefEqLambda tBody sBody (subst.push default)
   | t, s => isDefEq (t.instantiateRev subst) (s.instantiateRev subst)
 
+/--
+If `t` and `s` are for-all expressions, checks that their domains are defeq and
+recurses on the bodies, substituting in a new free variable for that binder
+(this substitution is delayed for efficiency purposes using the `subst`
+parameter). Otherwise, does a normal defeq check.
+-/
 def isDefEqForall (t s : Expr) (subst : Array Expr := #[]) : RecM Bool :=
   match t, s with
   | .forallE _ tDom tBody _, .forallE name sDom sBody bi => do
@@ -480,7 +550,13 @@ def isDefEqForall (t s : Expr) (subst : Array Expr := #[]) : RecM Bool :=
       isDefEqForall tBody sBody (subst.push default)
   | t, s => isDefEq (t.instantiateRev subst) (s.instantiateRev subst)
 
+/--
+If `t` and `s` have matching head constructors and are not projections or
+(non-α-equivalent) applications, checks that they are definitionally equal.
+Otherwise, defers to the calling function.
+-/
 def quickIsDefEq (t s : Expr) (useHash := false) : RecM LBool := do
+  -- optimization for terms that are already α-equivalent
   if ← modifyGet fun (.mk a1 a2 a3 a4 a5 a6 (eqvManager := m)) =>
     let (b, m) := m.isEquiv useHash t s
     (b, .mk a1 a2 a3 a4 a5 a6 (eqvManager := m))
@@ -494,6 +570,11 @@ def quickIsDefEq (t s : Expr) (useHash := false) : RecM LBool := do
   | .lit a1, .lit a2 => pure (a1 == a2).toLBool
   | _, _ => return .undef
 
+/--
+Assuming that `t` and `s` have the same constant heads, returns true if they
+are applications with definitionally equal arguments (in which case they are
+defeq), and false otherwise (deferring further defeq checking to caller).
+-/
 def isDefEqArgs (t s : Expr) : RecM Bool := do
   match t, s with
   | .app tf ta, .app sf sa =>
@@ -502,15 +583,30 @@ def isDefEqArgs (t s : Expr) : RecM Bool := do
   | .app .., _ | _, .app .. => return false
   | _, _ => return true
 
+/--
+Checks if `t` and `s` are defeq after applying η-expansion to `s`.
+
+Assuming that `s` has a function type `(x : A) -> B x`, it η-expands to
+`fun (x : A) => s x` (which it is definitionally equal to by the η rule).
+-/
 def tryEtaExpansionCore (t s : Expr) : RecM Bool := do
   if t.isLambda && !s.isLambda then
     let .forallE name ty _ bi ← whnf (← inferType s) | return false
     isDefEq t (.lam name ty (.app s (.bvar 0)) bi)
   else return false
 
+@[inherit_doc tryEtaExpansionCore]
 def tryEtaExpansion (t s : Expr) : RecM Bool :=
   tryEtaExpansionCore t s <||> tryEtaExpansionCore s t
 
+/--
+Checks if `t` and `s` are application-defeq (as in `isDefEqApp`)
+after applying struct-η-expansion to `t`.
+
+Assuming that `s` has a struct type `S` constructor `S.mk` and projection
+functions `pᵢ : S → Tᵢ`, it struct-η-expands to `S.mk (p₁ s) ... (pₙ s)` (which
+it is definitionally equal to by the struct-η rule).
+-/
 def tryEtaStructCore (t s : Expr) : RecM Bool := do
   let .const f _ := s.getAppFn | return false
   let env ← getEnv
@@ -523,9 +619,14 @@ def tryEtaStructCore (t s : Expr) : RecM Bool := do
     unless ← isDefEq (.proj fInfo.induct (i - fInfo.numParams) t) (args[i]'h.2) do return false
   return true
 
+@[inherit_doc tryEtaStructCore]
 def tryEtaStruct (t s : Expr) : RecM Bool :=
   tryEtaStructCore t s <||> tryEtaStructCore s t
 
+/--
+Checks if applications `t` and `s` (should be WHNF) are defeq on account of
+their function heads and arguments being defeq.
+-/
 def isDefEqApp (t s : Expr) : RecM Bool := do
   unless t.isApp && s.isApp do return false
   t.withApp fun tf tArgs =>
@@ -536,6 +637,10 @@ def isDefEqApp (t s : Expr) : RecM Bool := do
     unless ← isDefEq ta sa do return false
   return true
 
+/--
+Checks if `t` and `s` are definitionally equivalent according to proof irrelevance
+(that is, they are proofs of the same proposition).
+-/
 def isDefEqProofIrrel (t s : Expr) : RecM LBool := do
   let tType ← inferType t
   if !(← isProp tType) then return .undef
@@ -559,6 +664,18 @@ def tryUnfoldProjApp (e : Expr) : RecM (Option Expr) := do
   let e' ← whnfCore e
   return if e' != e then e' else none
 
+/--
+Performs a single step of δ-reduction on `tn`, `sn`, or both (according to
+optimizations) followed by weak-head normalization (without further
+δ-reduction). If the resulting terms have matching head constructors (excluding
+non-α-equivalent applications and projections) returns whether `tn` and `sn`
+are defeq. Otherwise, a return value indicates to the calling
+`lazyDeltaReduction` that δ-reduction is to be continued.
+
+If δ-reduction+weak-head-normalization cannot be continued (i.e. we have a
+weak-head normal form (with cheapProj := true)), defers further defeq-checking
+to `isDefEq`.
+-/
 def lazyDeltaReductionStep (tn sn : Expr) : RecM ReductionStatus := do
   let env ← getEnv
   let delta e := whnfCore (unfoldDefinition env e).get! (cheapProj := true)
@@ -570,6 +687,7 @@ def lazyDeltaReductionStep (tn sn : Expr) : RecM ReductionStatus := do
   match isDelta env tn, isDelta env sn with
   | none, none => return .unknown tn sn
   | some _, none =>
+    -- FIXME hasn't whnfCore already been called on sn? so when would this case arise?
     if let some sn' ← tryUnfoldProjApp sn then
       cont tn sn'
     else
@@ -604,6 +722,11 @@ def isNatSuccOf? : Expr → Option Expr
   | .app (.const ``Nat.succ _) e => return e
   | _ => none
 
+/--
+If `t` and `s` are both successors of natural numbers `t'` and `s'`, either as
+literals or `Nat.succ` applications, checks that `t'` and `s'` are
+definitionally equal. Otherwise, defers to the calling function.
+-/
 def isDefEqOffset (t s : Expr) : RecM LBool := do
   if isNatZero t && isNatZero s then
     return .true
@@ -611,6 +734,15 @@ def isDefEqOffset (t s : Expr) : RecM LBool := do
   | some t', some s' => toLBoolM <| isDefEqCore t' s'
   | _, _ => return .undef
 
+/--
+Returns whether the `cheapProj := true` weak-head normal forms of `tn` and
+`sn` are defeq if:
+- they have matching head constructors (excluding non-α-equivalent applications
+  and projections)
+- they're both natural number successors (as literals or `Nat.succ` applications)
+- one of them can be converted to a natural number/boolean literal.
+Otherwise, defers to the calling function with these normal forms.
+-/
 def lazyDeltaReduction (tn sn : Expr) : RecM ReductionStatus := loop tn sn 1000 where
   loop tn sn
   | 0 => throw .deterministicTimeout
@@ -631,17 +763,27 @@ def lazyDeltaReduction (tn sn : Expr) : RecM ReductionStatus := loop tn sn 1000
     | .continue tn sn => loop tn sn fuel
     | r => return r
 
+/--
+If `t` is a string literal and `s` is a string constructor application,
+checks that they are defeq after turning `t` into a constructor application.
+Otherwise, defers to the calling function.
+-/
 def tryStringLitExpansionCore (t s : Expr) : RecM LBool := do
   let .lit (.strVal st) := t | return .undef
   let .app sf _ := s | return .undef
   unless sf == .const ``String.mk [] do return .undef
   toLBoolM <| isDefEqCore (.strLitToConstructor st) s
 
+@[inherit_doc tryStringLitExpansionCore]
 def tryStringLitExpansion (t s : Expr) : RecM LBool := do
   match ← tryStringLitExpansionCore t s with
   | .undef => tryStringLitExpansionCore s t
   | r => return r
 
+/--
+Checks if `t` and `s` are defeq on account of both being of a unit type (a type
+with one constructor without any fields or indices).
+-/
 def isDefEqUnitLike (t s : Expr) : RecM Bool := do
   let tType ← whnf (← inferType t)
   let .const I _ := tType.getAppFn | return false
@@ -651,6 +793,7 @@ def isDefEqUnitLike (t s : Expr) : RecM Bool := do
   let .ctorInfo { numFields := 0, .. } ← env.get c | return false
   isDefEqCore tType (← inferType s)
 
+@[inherit_doc isDefEqCore]
 def isDefEqCore' (t s : Expr) : RecM Bool := do
   let r ← quickIsDefEq t s (useHash := true)
   if r != .undef then return r == .true
@@ -678,14 +821,18 @@ def isDefEqCore' (t s : Expr) : RecM Bool := do
     if tf == sf && Level.isEquivList tl sl then return true
   | .fvar tv, .fvar sv => if tv == sv then return true
   | .proj _ ti te, .proj _ si se =>
+    -- optimized by `lazyDeltaReductionStep` using `cheapProj = true`
     if ti == si then if ← isDefEq te se then return true
   | _, _ => pure ()
 
+  -- `lazyDeltaReductionStep` used `cheapProj = true`, so we may not have a complete WHNF
   let tnn ← whnfCore tn
   let snn ← whnfCore sn
   if !(unsafe ptrEq tnn tn && ptrEq snn sn) then
+    -- if projection reduced, need to re-run (as we may not have a WHNF)
     return ← isDefEqCore tnn snn
 
+  -- tn and sn are both in WHNF
   if ← isDefEqApp tn sn then return true
   if ← tryEtaExpansion tn sn then return true
   if ← tryEtaStruct tn sn then return true

From cbb7b90bca50aadfce77263e9a9999d3d31eac6d Mon Sep 17 00:00:00 2001
From: Rishikesh Vaishnav <rishhvaishnav@gmail.com>
Date: Tue, 21 May 2024 05:28:21 +0200
Subject: [PATCH 2/6] docs: document type-inference functions

---
 Lean4Lean/Instantiate.lean |  6 ++-
 Lean4Lean/Level.lean       |  3 ++
 Lean4Lean/TypeChecker.lean | 83 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/Lean4Lean/Instantiate.lean b/Lean4Lean/Instantiate.lean
index 039679c..1358117 100644
--- a/Lean4Lean/Instantiate.lean
+++ b/Lean4Lean/Instantiate.lean
@@ -2,6 +2,10 @@ import Lean.Expr
 
 namespace Lean
 
+/--
+Reduces an expression of the form (λ x₁ ... xₙ, x₁) a₁ ... aₙ aₙ₊₁ ... aₘ
+to aᵢ aₙ₊₁ ... aₘ.
+-/
 def Expr.cheapBetaReduce (e : Expr) : Expr := Id.run do
   if !e.isApp then return e
   let fn := e.getAppFn
@@ -18,7 +22,7 @@ def Expr.cheapBetaReduce (e : Expr) : Expr := Id.run do
   let rec loop i fn : Id Expr :=
     if i < args.size then
       match fn with
-      | .lam _ _ dom .. => loop (i + 1) dom
+      | .lam _ _ bod .. => loop (i + 1) bod
       | _ => cont i fn
     else cont i fn
   loop 0 fn
diff --git a/Lean4Lean/Level.lean b/Lean4Lean/Level.lean
index 4b7aa3e..f2e99c4 100644
--- a/Lean4Lean/Level.lean
+++ b/Lean4Lean/Level.lean
@@ -10,6 +10,9 @@ def forEach [Monad m] (l : Level) (f : Level → m Bool) : m Unit := do
   | .max l₁ l₂ | .imax l₁ l₂ => l₁.forEach f; l₂.forEach f
   | .zero | .param .. | .mvar .. => pure ()
 
+/--
+Returns `some n` if level parameter `n` appears in `l` and `n ∉ ps`.
+-/
 def getUndefParam (l : Level) (ps : List Name) : Option Name := Id.run do
   (·.2) <$> StateT.run (s := none) do
     l.forEach fun l => do
diff --git a/Lean4Lean/TypeChecker.lean b/Lean4Lean/TypeChecker.lean
index 681a323..c45c392 100644
--- a/Lean4Lean/TypeChecker.lean
+++ b/Lean4Lean/TypeChecker.lean
@@ -70,18 +70,30 @@ def whnf (e : Expr) : RecM Expr := fun m => m.whnf e
 @[inline] def withLCtx [MonadWithReaderOf LocalContext m] (lctx : LocalContext) (x : m α) : m α :=
   withReader (fun _ => lctx) x
 
+/--
+Ensures that `e` is defeq to some `e' := .sort ..`, returning `e'`. If not,
+throws an error with `s` (the expression required be a sort).
+-/
 def ensureSortCore (e s : Expr) : RecM Expr := do
   if e.isSort then return e
   let e ← whnf e
   if e.isSort then return e
   throw <| .typeExpected (← getEnv) (← getLCtx) s
 
+/--
+Ensures that `e` is defeq to some `e' := .forallE ..`, returning `e'`. If not,
+throws an error with `s := f a` (the application requiring `f` to be of
+function type).
+-/
 def ensureForallCore (e s : Expr) : RecM Expr := do
   if e.isForall then return e
   let e ← whnf e
   if e.isForall then return e
   throw <| .funExpected (← getEnv) (← getLCtx) s
 
+/--
+Checks that `l` does not contain any level parameters not found in the context `tc`.
+-/
 def checkLevel (tc : Context) (l : Level) : Except KernelException Unit := do
   if let some n2 := l.getUndefParam tc.lparams then
     throw <| .other s!"invalid reference to undefined universe level parameter '{n2}'"
@@ -91,6 +103,9 @@ def inferFVar (tc : Context) (name : FVarId) : Except KernelException Expr := do
     return decl.type
   throw <| .other "unknown free variable"
 
+/--
+Infers the type of `.const e ls`.
+-/
 def inferConstant (tc : Context) (name : Name) (ls : List Level) (inferOnly : Bool) :
     Except KernelException Expr := do
   let e := Expr.const name ls
@@ -110,8 +125,17 @@ def inferConstant (tc : Context) (name : Name) (ls : List Level) (inferOnly : Bo
       checkLevel tc l
   return info.instantiateTypeLevelParams ls
 
+/--
+Infers the type of expression `e`. If `inferOnly := false`, this function will
+throw an error if and only if `e` is not typeable according to Lean's
+algorithmic typing judgment. Setting `inferOnly := true` optimizes to avoid
+unnecessary checks in the case that `e` is already known to be well-typed.
+-/
 def inferType (e : Expr) (inferOnly := true) : RecM Expr := fun m => m.inferType e inferOnly
 
+/--
+Infers the type of lambda expression `e`.
+-/
 def inferLambda (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] e where
   loop fvars : Expr → RecM Expr
   | .lam name dom body bi => do
@@ -119,6 +143,7 @@ def inferLambda (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] e where
     let id := ⟨← mkFreshId⟩
     withLCtx ((← getLCtx).mkLocalDecl id name d bi) do
       let fvars := fvars.push (.fvar id)
+      -- FIXME this should happen before extending the local context (as it does in `inferForall`)
       if !inferOnly then
         _ ← ensureSortCore (← inferType d inferOnly) d
       loop fvars body
@@ -127,6 +152,9 @@ def inferLambda (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] e where
     let r := r.cheapBetaReduce
     return (← getLCtx).mkForall fvars r
 
+/--
+Infers the type of for-all expression `e`.
+-/
 def inferForall (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e where
   loop fvars us : Expr → RecM Expr
   | .forallE name dom body bi => do
@@ -145,6 +173,13 @@ def inferForall (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e wher
 /--
 Returns whether `t` and `s` are definitionally equal according to Lean's
 algorithmic definitional equality judgment.
+
+NOTE: This function does not do any typechecking of its own on `t` and `s`.
+So, when this is used as part of a typechecking routine, it is expected that
+they are already well-typed (that is, that `check t` and `check s` 
+did not/would not throw an error). This ensures in particular that any calls to
+`inferType e (inferOnly := false)` on subterms `e` would not fail, so we know
+that `e` types as the return value of `inferType e (inferOnly := true)`.
 -/
 def isDefEqCore (t s : Expr) : RecM Bool := fun m => m.isDefEqCore t s
 
@@ -155,6 +190,9 @@ def isDefEq (t s : Expr) : RecM Bool := do
     modify fun st => { st with eqvManager := st.eqvManager.addEquiv t s }
   pure r
 
+/--
+Infers the type of application `e`, assuming that `e` is already well-typed.
+-/
 def inferApp (e : Expr) : RecM Expr := do
   e.withApp fun f args => do
   let mut fType ← inferType f
@@ -181,6 +219,9 @@ def markUsed (n : Nat) (fvars : Array Expr) (b : Expr) (used : Array Bool) : Arr
             return false
       return true
 
+/--
+Infers the type of let-expression `e`.
+-/
 def inferLet (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e where
   loop fvars vals : Expr → RecM Expr
   | .letE name type val body _ => do
@@ -190,6 +231,7 @@ def inferLet (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e where
     withLCtx ((← getLCtx).mkLetDecl id name type val) do
       let fvars := fvars.push (.fvar id)
       let vals := vals.push val
+      -- FIXME this should happen before extending the local context
       if !inferOnly then
         _ ← ensureSortCore (← inferType type inferOnly) type
         let valType ← inferType val inferOnly
@@ -212,11 +254,18 @@ def inferLet (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e where
     for fvar in fvars, b in used do
       if b then
         usedFVars := usedFVars.push fvar
+    -- FIXME `usedFVars` is never used
     return (← getLCtx).mkForall fvars r
 
+/--
+Checks if `e` is definitionally equal to `Prop`.
+-/
 def isProp (e : Expr) : RecM Bool :=
   return (← whnf (← inferType e)) == .prop
 
+/--
+Infers the type of structure projection `e`.
+-/
 def inferProj (typeName : Name) (idx : Nat) (struct structType : Expr) : RecM Expr := do
   let e := Expr.proj typeName idx struct
   let type ← whnf structType
@@ -237,6 +286,7 @@ def inferProj (typeName : Name) (idx : Nat) (struct structType : Expr) : RecM Ex
   for i in [:idx] do
     let .forallE _ dom b _ ← whnf r | fail
     if b.hasLooseBVars then
+      -- prop structs cannot have non-prop dependent fields
       if isPropType then if !(← isProp dom) then fail
       r := b.instantiate1 (.proj I_name i struct)
     else
@@ -245,6 +295,7 @@ def inferProj (typeName : Name) (idx : Nat) (struct structType : Expr) : RecM Ex
   if isPropType then if !(← isProp dom) then fail
   return dom
 
+@[inherit_doc inferType]
 def inferType' (e : Expr) (inferOnly : Bool) : RecM Expr := do
   if e.isBVar then
     throw <| .other
@@ -275,6 +326,8 @@ def inferType' (e : Expr) (inferOnly : Bool) : RecM Expr := do
         let fType ← ensureForallCore (← inferType' f inferOnly) e
         let aType ← inferType' a inferOnly
         let dType := fType.bindingDomain!
+        -- it can be shown that if `e` is typeable as `T`, then `T` is typeable as `Sort l`
+        -- for some universe level `l`, so this use of `isDefEq` is valid
         if !(← isDefEq dType aType) then
           throw <| .appTypeMismatch (← getEnv) (← getLCtx) e fType aType
         pure <| fType.bindingBody!.instantiate1 a
@@ -293,7 +346,7 @@ Setting `cheapRec` or `cheapProj` to `true` will cause the major premise/struct
 argument to be reduced "lazily" (using `whnfCore` rather than `whnf`) when
 reducing recursor applications/struct projections. This can be a useful
 optimization if we're checking the definitional equality of two recursor
-applications/struct projections of the same recursor/projections, where we
+applications/struct projections of the same recursor/projection, where we
 might save some work by directly checking if the major premises/struct
 arguments are defeq (rather than eagerly applying a recursor rule/projection).
 -/
@@ -396,6 +449,11 @@ def isDelta (env : Environment) (e : Expr) : Option ConstantInfo := do
         return ci
   none
 
+/--
+Checks if `e` has a head constant that can be delta-reduced (that is, it is a
+theorem or definition), returning its value (instantiated by level parameters)
+if so.
+-/
 def unfoldDefinitionCore (env : Environment) (e : Expr) : Option Expr := do
   if let .const _ ls := e then
     if let some d := isDelta env e then
@@ -857,21 +915,44 @@ def Methods.withFuel : Nat → Methods
       whnf := fun e => whnf' e (withFuel n)
       inferType := fun e i => inferType' e i (withFuel n) }
 
+/--
+Runs `x` with a limit on the recursion depth.
+-/
 def RecM.run (x : RecM α) : M α := x (Methods.withFuel 1000)
 
+/--
+With the level context `lps`, infers the type of expression `e` and checks that
+`e` is well-typed according to Lean's typing judgment.
+
+Use `inferType` to infer type alone.
+-/
 def check (e : Expr) (lps : List Name) : M Expr :=
   withReader ({ · with lparams := lps }) (inferType e (inferOnly := false)).run
 
+@[inherit_doc whnf']
 def whnf (e : Expr) : M Expr := (Inner.whnf e).run
 
+/--
+Infers the type of expression `e`. Note that this uses the optimization
+`inferOnly := false`, and so should only be used for the purpose of type
+inference on terms that are known to be well-typed. To typecheck terms for the
+first time, use `check`.
+-/
 def inferType (e : Expr) : M Expr := (Inner.inferType e).run
 
+@[inherit_doc isDefEqCore]
 def isDefEq (t s : Expr) : M Bool := (Inner.isDefEq t s).run
 
+@[inherit_doc ensureSortCore]
 def ensureSort (t : Expr) (s := t) : M Expr := (ensureSortCore t s).run
 
+@[inherit_doc ensureForallCore]
 def ensureForall (t : Expr) (s := t) : M Expr := (ensureForallCore t s).run
 
+/--
+Ensures that `e` is defeq to some `e' := .sort (_ + 1)`, returning `e'`. If not,
+throws an error with `s` (the expression required be a type).
+-/
 def ensureType (e : Expr) : M Expr := do ensureSort (← inferType e) e
 
 def etaExpand (e : Expr) : M Expr :=

From 33ab65bcb6597a9404cdada8497e952eeb70e796 Mon Sep 17 00:00:00 2001
From: Rishikesh Vaishnav <rishhvaishnav@gmail.com>
Date: Mon, 1 Jul 2024 15:51:43 +0200
Subject: [PATCH 3/6] docs: some minor updates

---
 Lean4Lean/TypeChecker.lean | 35 +++++++++++++++++++++--------------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/Lean4Lean/TypeChecker.lean b/Lean4Lean/TypeChecker.lean
index c45c392..706d336 100644
--- a/Lean4Lean/TypeChecker.lean
+++ b/Lean4Lean/TypeChecker.lean
@@ -364,8 +364,9 @@ def reduceRecursor (e : Expr) (cheapRec cheapProj : Bool) : RecM (Option Expr) :
   return none
 
 /--
-Gets the weak-head normal form of the free variable `e`,
-which is the weak-head normal form of its definition if `e` is a let variable and itself if it is a lambda variable.
+Gets the weak-head normal form of the free variable `e`, which is the weak-head
+normal form of its definition if `e` is a let variable and itself if it is a
+lambda variable.
 -/
 def whnfFVar (e : Expr) (cheapRec cheapProj : Bool) : RecM Expr := do
   if let some (.ldecl (value := v) ..) := (← getLCtx).find? e.fvarId! then
@@ -609,8 +610,8 @@ def isDefEqForall (t s : Expr) (subst : Array Expr := #[]) : RecM Bool :=
   | t, s => isDefEq (t.instantiateRev subst) (s.instantiateRev subst)
 
 /--
-If `t` and `s` have matching head constructors and are not projections or
-(non-α-equivalent) applications, checks that they are definitionally equal.
+If `t` and `s` have matching head constructors and are not (non-α-equivalent)
+projections or applications, checks that they are definitionally equal.
 Otherwise, defers to the calling function.
 -/
 def quickIsDefEq (t s : Expr) (useHash := false) : RecM LBool := do
@@ -629,7 +630,7 @@ def quickIsDefEq (t s : Expr) (useHash := false) : RecM LBool := do
   | _, _ => return .undef
 
 /--
-Assuming that `t` and `s` have the same constant heads, returns true if they
+Assuming that `t` and `s` have the same function heads, returns true if they
 are applications with definitionally equal arguments (in which case they are
 defeq), and false otherwise (deferring further defeq checking to caller).
 -/
@@ -642,7 +643,8 @@ def isDefEqArgs (t s : Expr) : RecM Bool := do
   | _, _ => return true
 
 /--
-Checks if `t` and `s` are defeq after applying η-expansion to `s`.
+Assuming `t` and `s` are WHNF, checks if they are defeq on account of `t` being
+an η-expansion of `s`.
 
 Assuming that `s` has a function type `(x : A) -> B x`, it η-expands to
 `fun (x : A) => s x` (which it is definitionally equal to by the η rule).
@@ -658,11 +660,11 @@ def tryEtaExpansion (t s : Expr) : RecM Bool :=
   tryEtaExpansionCore t s <||> tryEtaExpansionCore s t
 
 /--
-Checks if `t` and `s` are application-defeq (as in `isDefEqApp`)
-after applying struct-η-expansion to `t`.
+Assuming `t` and `s` in WHNF, checks if they are defeq on account of `s` being
+defeq to the struct-η-expansion of `t`.
 
-Assuming that `s` has a struct type `S` constructor `S.mk` and projection
-functions `pᵢ : S → Tᵢ`, it struct-η-expands to `S.mk (p₁ s) ... (pₙ s)` (which
+Assuming that `t` has a struct type `S`, constructor `S.mk`, and projection
+functions `pᵢ : S → Tᵢ`, it struct-η-expands to `S.mk (p₁ t) ... (pₙ t)` (which
 it is definitionally equal to by the struct-η rule).
 -/
 def tryEtaStructCore (t s : Expr) : RecM Bool := do
@@ -674,11 +676,15 @@ def tryEtaStructCore (t s : Expr) : RecM Bool := do
   unless ← isDefEq (← inferType t) (← inferType s) do return false
   let args := s.getAppArgs
   for h : i in [fInfo.numParams:args.size] do
+    -- since `t` is in WHNF, and assuming it is not a constructor application, this projection cannot reduce
+    -- (so we are directly checking if `s` is defeq to the struct-η-expansion of `t`)
     unless ← isDefEq (.proj fInfo.induct (i - fInfo.numParams) t) (args[i]'h.2) do return false
   return true
 
 @[inherit_doc tryEtaStructCore]
 def tryEtaStruct (t s : Expr) : RecM Bool :=
+  -- FIXME can return false if `t` and `s` are both constructor applications
+  -- (we have already called `isDefEqApp` on them, which returned false)
   tryEtaStructCore t s <||> tryEtaStructCore s t
 
 /--
@@ -726,8 +732,9 @@ def tryUnfoldProjApp (e : Expr) : RecM (Option Expr) := do
 Performs a single step of δ-reduction on `tn`, `sn`, or both (according to
 optimizations) followed by weak-head normalization (without further
 δ-reduction). If the resulting terms have matching head constructors (excluding
-non-α-equivalent applications and projections) returns whether `tn` and `sn`
-are defeq. Otherwise, a return value indicates to the calling
+non-α-equivalent applications and projections), or are applications with the
+same defined constant function head and defeq args, returns whether `tn` and
+`sn` are defeq. Otherwise, a return value indicates to the calling
 `lazyDeltaReduction` that δ-reduction is to be continued.
 
 If δ-reduction+weak-head-normalization cannot be continued (i.e. we have a
@@ -879,11 +886,11 @@ def isDefEqCore' (t s : Expr) : RecM Bool := do
     if tf == sf && Level.isEquivList tl sl then return true
   | .fvar tv, .fvar sv => if tv == sv then return true
   | .proj _ ti te, .proj _ si se =>
-    -- optimized by `lazyDeltaReductionStep` using `cheapProj = true`
+    -- optimized by the previous reduction functions using `cheapProj = true`
     if ti == si then if ← isDefEq te se then return true
   | _, _ => pure ()
 
-  -- `lazyDeltaReductionStep` used `cheapProj = true`, so we may not have a complete WHNF
+  -- the previous reduction functions used `cheapProj = true`, so we may not have a complete WHNF
   let tnn ← whnfCore tn
   let snn ← whnfCore sn
   if !(unsafe ptrEq tnn tn && ptrEq snn sn) then

From 3e304eb2b37923b35157047350eccbc8b17d2ef8 Mon Sep 17 00:00:00 2001
From: Rishikesh Vaishnav <rishhvaishnav@gmail.com>
Date: Fri, 12 Jul 2024 17:46:02 +0200
Subject: [PATCH 4/6] docs: label FIXMEs

---
 Lean4Lean/TypeChecker.lean | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/Lean4Lean/TypeChecker.lean b/Lean4Lean/TypeChecker.lean
index 706d336..1ee8193 100644
--- a/Lean4Lean/TypeChecker.lean
+++ b/Lean4Lean/TypeChecker.lean
@@ -143,7 +143,7 @@ def inferLambda (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] e where
     let id := ⟨← mkFreshId⟩
     withLCtx ((← getLCtx).mkLocalDecl id name d bi) do
       let fvars := fvars.push (.fvar id)
-      -- FIXME this should happen before extending the local context (as it does in `inferForall`)
+      -- FIXME(kernel) this should happen before extending the local context (as it does in `inferForall`)
       if !inferOnly then
         _ ← ensureSortCore (← inferType d inferOnly) d
       loop fvars body
@@ -231,7 +231,7 @@ def inferLet (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e where
     withLCtx ((← getLCtx).mkLetDecl id name type val) do
       let fvars := fvars.push (.fvar id)
       let vals := vals.push val
-      -- FIXME this should happen before extending the local context
+      -- FIXME(kernel) this should happen before extending the local context
       if !inferOnly then
         _ ← ensureSortCore (← inferType type inferOnly) type
         let valType ← inferType val inferOnly
@@ -254,7 +254,7 @@ def inferLet (e : Expr) (inferOnly : Bool) : RecM Expr := loop #[] #[] e where
     for fvar in fvars, b in used do
       if b then
         usedFVars := usedFVars.push fvar
-    -- FIXME `usedFVars` is never used
+    -- FIXME(kernel) `usedFVars` is never used
     return (← getLCtx).mkForall fvars r
 
 /--
@@ -429,7 +429,7 @@ def whnfCore' (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr :=
         pure e
     else
       let r := f.mkAppRevRange 0 rargs.size rargs
-      -- FIXME replace with reduceRecursor? adding arguments can only result in further normalization if the head reduced to a partial recursor application
+      -- FIXME(kernel) replace with reduceRecursor? adding arguments can only result in further normalization if the head reduced to a partial recursor application
       save <|← whnfCore r cheapRec cheapProj
   | .letE _ _ val body _ =>
     save <|← whnfCore (body.instantiate1 val) cheapRec cheapProj
@@ -683,7 +683,7 @@ def tryEtaStructCore (t s : Expr) : RecM Bool := do
 
 @[inherit_doc tryEtaStructCore]
 def tryEtaStruct (t s : Expr) : RecM Bool :=
-  -- FIXME can return false if `t` and `s` are both constructor applications
+  -- FIXME(kernel) can return false if `t` and `s` are both constructor applications
   -- (we have already called `isDefEqApp` on them, which returned false)
   tryEtaStructCore t s <||> tryEtaStructCore s t
 
@@ -752,7 +752,7 @@ def lazyDeltaReductionStep (tn sn : Expr) : RecM ReductionStatus := do
   match isDelta env tn, isDelta env sn with
   | none, none => return .unknown tn sn
   | some _, none =>
-    -- FIXME hasn't whnfCore already been called on sn? so when would this case arise?
+    -- FIXME(kernel) hasn't whnfCore already been called on sn? so when would this case arise?
     if let some sn' ← tryUnfoldProjApp sn then
       cont tn sn'
     else

From 3d62e2d7294956cb553e28d3c013dc93025212c3 Mon Sep 17 00:00:00 2001
From: Rishikesh Vaishnav <rishhvaishnav@gmail.com>
Date: Mon, 15 Jul 2024 15:50:38 +0200
Subject: [PATCH 5/6] docs: small corrections

---
 Lean4Lean/TypeChecker.lean | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/Lean4Lean/TypeChecker.lean b/Lean4Lean/TypeChecker.lean
index 1ee8193..72a6dc3 100644
--- a/Lean4Lean/TypeChecker.lean
+++ b/Lean4Lean/TypeChecker.lean
@@ -941,7 +941,7 @@ def whnf (e : Expr) : M Expr := (Inner.whnf e).run
 
 /--
 Infers the type of expression `e`. Note that this uses the optimization
-`inferOnly := false`, and so should only be used for the purpose of type
+`inferOnly := true`, and so should only be used for the purpose of type
 inference on terms that are known to be well-typed. To typecheck terms for the
 first time, use `check`.
 -/
@@ -957,8 +957,7 @@ def ensureSort (t : Expr) (s := t) : M Expr := (ensureSortCore t s).run
 def ensureForall (t : Expr) (s := t) : M Expr := (ensureForallCore t s).run
 
 /--
-Ensures that `e` is defeq to some `e' := .sort (_ + 1)`, returning `e'`. If not,
-throws an error with `s` (the expression required be a type).
+Ensures that `e` is a type/proposition. If it is not, throws an error.
 -/
 def ensureType (e : Expr) : M Expr := do ensureSort (← inferType e) e
 

From 5d8a4e113e177b22615ccdfe690abe54b22babf8 Mon Sep 17 00:00:00 2001
From: Rishikesh Vaishnav <rishhvaishnav@gmail.com>
Date: Thu, 15 Aug 2024 11:30:52 +0200
Subject: [PATCH 6/6] docs: small corrections

---
 Lean4Lean/TypeChecker.lean | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/Lean4Lean/TypeChecker.lean b/Lean4Lean/TypeChecker.lean
index 72a6dc3..43ecfed 100644
--- a/Lean4Lean/TypeChecker.lean
+++ b/Lean4Lean/TypeChecker.lean
@@ -429,7 +429,7 @@ def whnfCore' (e : Expr) (cheapRec := false) (cheapProj := false) : RecM Expr :=
         pure e
     else
       let r := f.mkAppRevRange 0 rargs.size rargs
-      -- FIXME(kernel) replace with reduceRecursor? adding arguments can only result in further normalization if the head reduced to a partial recursor application
+      -- FIXME(kernel) guard with reduceRecursor (as in the previous case)? adding arguments can only result in further normalization if the head reduced to a partial recursor application
       save <|← whnfCore r cheapRec cheapProj
   | .letE _ _ val body _ =>
     save <|← whnfCore (body.instantiate1 val) cheapRec cheapProj
@@ -610,12 +610,15 @@ def isDefEqForall (t s : Expr) (subst : Array Expr := #[]) : RecM Bool :=
   | t, s => isDefEq (t.instantiateRev subst) (s.instantiateRev subst)
 
 /--
-If `t` and `s` have matching head constructors and are not (non-α-equivalent)
-projections or applications, checks that they are definitionally equal.
+Checks that `t` and `s` are definitionally equal if:
+- they are α-equivalent
+- they have matching head constructors and are not (non-α-equivalent)
+  projections or applications
+- they have previously been checked for definitional equality
 Otherwise, defers to the calling function.
 -/
 def quickIsDefEq (t s : Expr) (useHash := false) : RecM LBool := do
-  -- optimization for terms that are already α-equivalent
+  -- optimization for terms that are already α-equivalent or were previously checked
   if ← modifyGet fun (.mk a1 a2 a3 a4 a5 a6 (eqvManager := m)) =>
     let (b, m) := m.isEquiv useHash t s
     (b, .mk a1 a2 a3 a4 a5 a6 (eqvManager := m))