chore: Sym cleanup (#11833)

This PR fixes a few typos, adds missing docstrings, and adds a (simple)
missing optimization.

src/Lean/Meta/Sym/InstantiateS.lean

@@ -79,6 +79,7 @@ def instantiateRangeS' (e : Expr) (beginIdx endIdx : Nat) (subst : Array Expr) :
       else
         return none
 
+/-- Internal variant of `instantiateS` that runs in `AlphaShareBuilderM`. -/
 def instantiateS' (e : Expr) (subst : Array Expr) : AlphaShareBuilderM Expr :=
   instantiateRangeS' e 0 subst.size subst
 
@@ -101,7 +102,8 @@ where
     loop revArgs start (← mkAppS b revArgs[i]!) i
 
 /--
-Similar to `betaRev`, but ensures maximally shared terms.
+Beta-reduces `f` applied to reversed arguments `revArgs`, ensuring maximally shared terms.
+`betaRevS f #[a₃, a₂, a₁]` computes the beta-normal form of `f a₁ a₂ a₃`.
 -/
 partial def betaRevS (f : Expr) (revArgs : Array Expr) : AlphaShareBuilderM Expr :=
   if revArgs.size == 0 then
@@ -121,12 +123,15 @@ partial def betaRevS (f : Expr) (revArgs : Array Expr) : AlphaShareBuilderM Expr
         mkAppRevRangeS (← instantiateRangeS' e n sz revArgs) 0 n revArgs
     go f 0
 
+/-- Monad for `instantiateRevBetaS'` with caching keyed by `(expression pointer, binder offset)`. -/
 abbrev M := StateT (Std.HashMap (ExprPtr × Nat) Expr) AlphaShareBuilderM
 
+/-- Caches the result `r` for `key` and returns `r`. -/
 def save (key : ExprPtr × Nat) (r : Expr) : M Expr := do
   modify fun cache => cache.insert key r
   return r
 
+/-- Internal variant of `instantiateRevBetaS` that runs in `AlphaShareBuilderM`. -/
 partial def instantiateRevBetaS' (e : Expr) (subst : Array Expr) : AlphaShareBuilderM Expr := do
   if subst.isEmpty || !e.hasLooseBVars then return e
   visit e 0 |>.run' {}

src/Lean/Meta/Sym/Intro.lean

@@ -120,7 +120,7 @@ Throws an error if the target type does not have a leading binder.
 public def intro (goal : Goal) (name : Name) : SymM (FVarId × Goal) := do
   let (fvarIds, goal') ← introCore goal 1 #[name]
   if h : 0 < fvarIds.size then
-    return (fvarIds[0], goal)
+    return (fvarIds[0], goal')
   else
     throwError "`intro` failed, binder expected"
 

src/Lean/Meta/Sym/LooseBVarsS.lean

@@ -9,7 +9,7 @@ public import Lean.Meta.Sym.SymM
 public import Lean.Meta.Sym.ReplaceS
 public section
 namespace Lean.Meta.Sym
-
+open Grind
 /--
 Lowers the loose bound variables `>= s` in `e` by `d`.
 That is, a loose bound variable `bvar i` with `i >= s` is mapped to `bvar (i-d)`.
@@ -24,7 +24,7 @@ def lowerLooseBVarsS' (e : Expr) (s d : Nat) : AlphaShareBuilderM Expr := do
     else match e with
     | .bvar idx =>
       if idx >= s₁ then
-        return some (← Grind.mkBVarS (idx - d))
+        return some (← mkBVarS (idx - d))
       else
         return none
     | _ => return none
@@ -45,7 +45,7 @@ def liftLooseBVarsS' (e : Expr) (s d : Nat) : AlphaShareBuilderM Expr := do
     else match e with
     | .bvar idx =>
       if idx >= s₁ then
-        return some (← Grind.mkBVarS (idx + d))
+        return some (← mkBVarS (idx + d))
       else
         return none
     | _ => return none

src/Lean/Meta/Sym/MaxFVar.lean

@@ -21,12 +21,16 @@ private def max (fvarId1? : Option FVarId) (fvarId2? : Option FVarId) : MetaM (O
     return fvarId2?
 
 private abbrev check (e : Expr) (k : SymM (Option FVarId)) : SymM (Option FVarId) := do
-  if let some fvarId? := (← get).maxFVar.find? { expr := e } then
-    return fvarId?
+  if e.hasFVar || e.hasMVar then
+    if let some fvarId? := (← get).maxFVar.find? { expr := e } then
+      return fvarId?
+    else
+      let fvarId? ← k
+      modify fun s => { s with maxFVar := s.maxFVar.insert { expr := e } fvarId? }
+      return fvarId?
   else
-    let fvarId? ← k
-    modify fun s => { s with maxFVar := s.maxFVar.insert { expr := e } fvarId? }
-    return fvarId?
+    -- `e` does not contain free variables or metavariables, then `maxFVar[e]` is definitely `none`.
+    return none
 
 /--
 Returns the maximal free variable occurring in `e`.

src/Lean/Meta/Sym/Pattern.lean

@@ -421,7 +421,7 @@ in `t`. We use this function when `t` is a metavariable, and we are trying to as
 -/
 def mayAssign (t s : Expr) : SymM Bool := do
   let some sMaxFVarId ← getMaxFVar? s
-    | return true -- `s` does not contain metavariables
+    | return true -- `s` does not contain free variables
   let some tMaxFVarId ← getMaxFVar? t
     | return false
   let sMaxFVarDecl ← sMaxFVarId.getDecl
@@ -495,7 +495,6 @@ def isDefEqAppWithInfo (t : Expr) (s : Expr) (i : Nat) (info : ProofInstInfo) :
       if (← tryAssignUnassigned a₁ a₂) then
         return true
       else
-        --
         isDefEqI a₁ a₂
     else if argInfo.isProof then
       discard <| tryAssignUnassigned a₁ a₂

src/Lean/Meta/Sym/SymM.lean

@@ -36,6 +36,24 @@ public structure ProofInstInfo where
   deriving Inhabited
 
 structure State where
+  /--
+  Maps expressions to their maximal free variable (by declaration index).
+
+  - `maxFVar[e] = some fvarId` means `fvarId` is the free variable with the largest declaration
+    index occurring in `e`.
+  - `maxFVar[e] = none` means `e` contains no free variables (but may contain metavariables).
+
+  Recall that if `e` contains a metavariable `?m`, then we assume `e` may contain any free variable
+  in the local context associated with `?m`.
+
+  This mapping enables O(1) local context compatibility checks during metavariable assignment.
+  Instead of traversing local contexts with `isSubPrefixOf`, we check if the maximal free variable
+  in the assigned value is in scope of the metavariable's local context.
+
+  **Note**: We considered using a mapping `PHashMap ExprPtr FVarId`. However, there is a corner
+  case that is not supported by this representation. `e` contains a metavariable (with an empty local context),
+  and no free variables.
+  -/
   maxFVar : PHashMap ExprPtr (Option FVarId) := {}
   proofInstInfo : PHashMap Name (Option ProofInstInfo) := {}