From f9b2f6b597ecd8ff257f08021d9972823bd31f93 Mon Sep 17 00:00:00 2001
From: Sebastian Graf <sgraf1337@gmail.com>
Date: Wed, 8 Apr 2026 18:34:51 +0200
Subject: [PATCH] fix: use `getDecLevel`/`isLevelDefEq` for `for` loop mut var
 universe constraints (#13332)

This PR fixes universe unification for `for` loops with `mut` variables
whose types span multiple implicit universes. The old approach used
`ensureHasType (mkSort mi.u.succ)` per variable, which generated
constraints like `max (?u+1) (?v+1) =?= ?u+1` that the universe solver
cannot decompose. The new approach uses `getDecLevel`/`isLevelDefEq` on
the decremented level, producing `max ?u ?v =?= ?u` which `solveSelfMax`
handles directly.

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/Lean/Elab/BuiltinDo/For.lean | 10 +++++--
 tests/elab/doForMutUniv.lean     | 45 ++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 tests/elab/doForMutUniv.lean

diff --git a/src/Lean/Elab/BuiltinDo/For.lean b/src/Lean/Elab/BuiltinDo/For.lean
index 8c3c130d3b..9b0f5cd644 100644
--- a/src/Lean/Elab/BuiltinDo/For.lean
+++ b/src/Lean/Elab/BuiltinDo/For.lean
@@ -111,8 +111,14 @@ open Lean.Meta
     for x in loopMutVars do
       let defn ← getLocalDeclFromUserName x.getId
       Term.addTermInfo' x defn.toExpr
-      -- ForIn forces all mut vars into the same universe: that of the do block result type.
-      discard <| Term.ensureHasType (mkSort (mi.u.succ)) defn.type
+      -- 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
+      discard <| isLevelDefEq u mi.u
       defs := defs.push defn.toExpr
     if info.returnsEarly && loopMutVars.isEmpty then
       defs := defs.push (mkConst ``Unit.unit)
diff --git a/tests/elab/doForMutUniv.lean b/tests/elab/doForMutUniv.lean
new file mode 100644
index 0000000000..fd42ae2bd3
--- /dev/null
+++ b/tests/elab/doForMutUniv.lean
@@ -0,0 +1,45 @@
+import Std.Data.HashSet
+
+/-!
+Test that `for` loops with `mut` variables whose type spans multiple implicit universes
+do not get stuck on `max (?u+1) (?v+1) =?= ?u+1`. This is solved by instead solving the decremented
+constraint `max ?u ?v =?= ?u` eagerly, which solves by `solveSelfMax`.
+-/
+
+-- Works with explicit universe variables
+example {α : Type (max u v)} {β : Type v} [Inhabited α] [Inhabited β] (as : Array α) : Array α := Id.run do
+  let mut m : α × β := default
+  for a in as do
+    m := (a, m.2)
+  return #[m.1]
+
+-- Works with legacy do elaborator
+set_option backward.do.legacy true in
+example [Inhabited α] [Inhabited β] (as : Array α) : Array α := Id.run do
+  let mut m : α × β := default
+  for a in as do
+    m := (a, m.2)
+  return #[m.1]
+
+-- Also works with implicit universe variables (the actual regression case)
+example [Inhabited α] [Inhabited β] (as : Array α) : Array α := Id.run do
+  let mut m : α × β := default
+  for a in as do
+    m := (a, m.2)
+  return #[m.1]
+
+-- Closer to the real world reproducer from aesop's UnionFind data structure:
+def cluster [BEq α] [Hashable α] [BEq β] [Hashable β] (f : α → Array β)
+    (as : Array α) : Array (Array α) := Id.run do
+  let mut clusters := Std.HashSet.ofArray as
+  let mut aOccs : Std.HashMap β (Array α) := {}
+  for a in as do
+    for b in f a do
+      match aOccs[b]? with
+      | some as' =>
+        for a' in as' do
+          clusters := clusters.insert a'
+        aOccs := aOccs.insert b (as'.push a)
+      | none =>
+        aOccs := aOccs.insert b #[a]
+  return #[clusters.toArray]