fix: Meta.Closure: topologically sort abstracted vars (#10926)

This PR topologically sorts abstracted vars in
`Meta.Closure.mkValueTypeClosure` if MVars are being abstracted.
Fixes #10705

---------

Co-authored-by: Eric Wieser <[email protected]>

Author: nomeata

src/Lean/Meta/Closure.lean

@@ -8,6 +8,7 @@ module
 prelude
 public import Lean.Meta.Check
 public import Lean.Meta.Tactic.AuxLemma
+import Lean.Util.ForEachExpr
 
 public section
 
@@ -349,9 +350,67 @@ def mkValueTypeClosureAux (type : Expr) (value : Expr) : ClosureM (Expr × Expr)
     process
     pure (type, value)
 
+private structure TopoSort where
+  tempMark  : FVarIdHashSet := {}
+  doneMark  : FVarIdHashSet := {}
+  newDecls : Array LocalDecl := #[]
+  newArgs : Array Expr := #[]
+
+/--
+By construction, the `newLocalDecls` for fvars are in dependency order, but those for MVars may not be,
+and need to be interleaved appropriately. This we do a “topological insertion sort” of these.
+We care about efficiency for the common case of many fvars and no mvars.
+-/
+private partial def sortDecls (sortedDecls : Array LocalDecl) (sortedArgs : Array Expr)
+  (toSortDecls : Array LocalDecl) (toSortArgs : Array Expr) : CoreM (Array LocalDecl × Array Expr):= do
+  assert! sortedDecls.size = sortedArgs.size
+  assert! toSortDecls.size = toSortArgs.size
+  if toSortDecls.isEmpty then
+    return (sortedDecls, sortedArgs)
+  trace[Meta.Closure] "MVars to abstract, topologically sorting the abstracted variables"
+  let mut m : Std.HashMap FVarId (LocalDecl × Expr) := {}
+  for decl in sortedDecls, arg in sortedArgs do
+    m := m.insert decl.fvarId (decl, arg)
+  for decl in toSortDecls, arg in toSortArgs do
+    m := m.insert decl.fvarId (decl, arg)
+
+  let rec visit (fvarId : FVarId) : StateT TopoSort CoreM Unit := do
+    let some (decl, arg) := m.get? fvarId | return
+    if (← get).doneMark.contains decl.fvarId then
+      return ()
+    trace[Meta.Closure] "Sorting decl {mkFVar decl.fvarId} : {decl.type}"
+    if (← get).tempMark.contains decl.fvarId then
+      throwError "cycle detected in sorting abstracted variables"
+    assert! !decl.isLet (allowNondep := true) -- should all be cdecls
+    modify fun s => { s with tempMark := s.tempMark.insert decl.fvarId }
+    let type := decl.type
+    type.forEach' fun e => do
+      if e.hasFVar then
+        if e.isFVar then
+          visit e.fvarId!
+        return true
+      else
+        return false
+    modify fun s => { s with
+      newDecls := s.newDecls.push decl
+      newArgs := s.newArgs.push arg
+      doneMark := s.doneMark.insert decl.fvarId
+    }
+
+  let s₀ := { newDecls := .emptyWithCapacity m.size, newArgs := .emptyWithCapacity m.size }
+  StateT.run' (s := s₀) do
+    for decl in sortedDecls do
+      visit decl.fvarId
+    for decl in toSortDecls do
+      visit decl.fvarId
+    let {newDecls, newArgs, .. } ← get
+    trace[Meta.Closure] "Sorted fvars: {newDecls.map (mkFVar ·.fvarId)}"
+    return (newDecls, newArgs)
+
 def mkValueTypeClosure (type : Expr) (value : Expr) (zetaDelta : Bool) : MetaM MkValueTypeClosureResult := do
   let ((type, value), s) ← ((mkValueTypeClosureAux type value).run { zetaDelta }).run {}
-  let newLocalDecls := s.newLocalDecls.reverse ++ s.newLocalDeclsForMVars
+  let (newLocalDecls, newArgs) ← sortDecls s.newLocalDecls.reverse s.exprFVarArgs.reverse
+                                           s.newLocalDeclsForMVars s.exprMVarArgs
   let newLetDecls   := s.newLetDecls.reverse
   let type  := mkForall newLocalDecls (mkForall newLetDecls type)
   let value := mkLambda newLocalDecls (mkLambda newLetDecls value)
@@ -360,7 +419,7 @@ def mkValueTypeClosure (type : Expr) (value : Expr) (zetaDelta : Bool) : MetaM M
     value       := value,
     levelParams := s.levelParams,
     levelArgs   := s.levelArgs,
-    exprArgs    := s.exprFVarArgs.reverse ++ s.exprMVarArgs
+    exprArgs    := newArgs
   }
 
 end Closure
@@ -396,4 +455,7 @@ def mkAuxTheorem (type : Expr) (value : Expr) (zetaDelta : Bool := false) (kind?
   let name ← mkAuxLemma (kind? := kind?) (cache := cache) result.levelParams.toList result.type result.value
   return mkAppN (mkConst name result.levelArgs.toList) result.exprArgs
 
+builtin_initialize
+  registerTraceClass `Meta.Closure
+
 end Lean.Meta

tests/lean/run/issue10705.lean

@@ -0,0 +1,21 @@
+import Lean
+
+open Lean Meta
+
+-- set_option trace.Meta.Closure true
+
+/--
+info: before: h.2
+---
+info: after: _proof_2 ?m.1 h
+-/
+#guard_msgs(pass trace, all) in
+run_meta do
+  have l := ← Lean.Meta.mkFreshExprMVar (mkConst ``True) (kind := .syntheticOpaque)
+  let ty := mkAnd (mkConst ``True) (.letE `foo (mkConst ``True) l (mkConst ``True) false)
+  withLocalDecl `h default ty fun x => do
+    let e := mkProj ``And 1 x
+    Lean.logInfo m!"before: {e}"
+    -- works fine without this line
+    let e ← Lean.Meta.mkAuxTheorem (mkConst ``True) e (zetaDelta := true) -- or false, not really relevant
+    Lean.logInfo m!"after: {e}"