From 17807e1cbea7d45a56ab4ebe21ae7f3662f695c3 Mon Sep 17 00:00:00 2001
From: Sebastian Graf <sgraf1337@gmail.com>
Date: Wed, 11 Mar 2026 18:00:24 +0800
Subject: [PATCH] feat: apply `@[mvcgen_invariant_type]` to `Invariant`,
 `StringInvariant`, `StringSliceInvariant` (#12880)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This PR applies `@[mvcgen_invariant_type]` to `Std.Do.Invariant` and
removes the hard-coded fallback in `isMVCGenInvariantType` that was
needed for bootstrapping (cf. #12874). It also extracts
`StringInvariant` and `StringSliceInvariant` as named abbreviations
tagged with `@[mvcgen_invariant_type]`, so that `mvcgen` classifies
string and string slice loop invariants correctly.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/Lean/Elab/Tactic/Do/Attr.lean |  3 +--
 src/Std/Do/Triple/SpecLemmas.lean | 33 +++++++++++++++++++++++++++----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/src/Lean/Elab/Tactic/Do/Attr.lean b/src/Lean/Elab/Tactic/Do/Attr.lean
index 963d6a4c8e..2d0040d874 100644
--- a/src/Lean/Elab/Tactic/Do/Attr.lean
+++ b/src/Lean/Elab/Tactic/Do/Attr.lean
@@ -262,10 +262,9 @@ builtin_initialize mvcgenInvariantAttr : TagAttribute ←
 
 /--
 Returns `true` if `ty` is an application of a type tagged with `@[mvcgen_invariant_type]`.
-Falls back to checking for `Std.Do.Invariant` for bootstrapping purposes.
 -/
 def isMVCGenInvariantType (env : Environment) (ty : Expr) : Bool :=
   if let .const name .. := ty.getAppFn then
-    mvcgenInvariantAttr.hasTag env name || name == ``Std.Do.Invariant
+    mvcgenInvariantAttr.hasTag env name
   else
     false
diff --git a/src/Std/Do/Triple/SpecLemmas.lean b/src/Std/Do/Triple/SpecLemmas.lean
index 58d3f525cd..4f1b88cd4a 100644
--- a/src/Std/Do/Triple/SpecLemmas.lean
+++ b/src/Std/Do/Triple/SpecLemmas.lean
@@ -688,6 +688,7 @@ After leaving the loop, the cursor's prefix is `xs` and the suffix is empty.
 During the induction step, the invariant holds for a suffix with head element `x`.
 After running the loop body, the invariant then holds after shifting `x` to the prefix.
 -/
+@[mvcgen_invariant_type]
 abbrev Invariant {α : Type u₁} (xs : List α) (β : Type u₂) (ps : PostShape.{max u₁ u₂}) :=
   PostCond (List.Cursor xs × β) ps
 
@@ -1997,6 +1998,18 @@ theorem Spec.foldlM_array {α β : Type u} {m : Type u → Type v} {ps : PostSha
   simp
   apply Spec.foldlM_list inv step
 
+/--
+The type of loop invariants used by the specifications of `for ... in ...` loops over strings.
+A loop invariant is a `PostCond` that takes as parameters
+
+* A `String.Pos` representing the current position in the string `s`.
+* A state tuple of type `β`, which will be a nesting of `MProd`s representing the elaboration of
+  `let mut` variables and early return.
+-/
+@[mvcgen_invariant_type]
+abbrev StringInvariant (s : String) (β : Type u) (ps : PostShape.{u}) :=
+  PostCond (s.Pos × β) ps
+
 /--
 Helper definition for specifying loop invariants for loops with early return.
 
@@ -2019,7 +2032,7 @@ abbrev StringInvariant.withEarlyReturn {s : String}
   (onContinue : s.Pos → β → Assertion ps)
   (onReturn : γ → β → Assertion ps)
   (onExcept : ExceptConds ps := ExceptConds.false) :
-    PostCond (s.Pos × MProd (Option γ) β) ps
+    StringInvariant s (MProd (Option γ) β) ps
     :=
   ⟨fun ⟨pos, x, b⟩ => spred(
         (⌜x = none⌝ ∧ onContinue pos b)
@@ -2029,7 +2042,7 @@ abbrev StringInvariant.withEarlyReturn {s : String}
 @[spec]
 theorem Spec.forIn_string
     {s : String} {init : β} {f : Char → β → m (ForInStep β)}
-    (inv : PostCond (s.Pos × β) ps)
+    (inv : StringInvariant s β ps)
     (step : ∀ pos b (h : pos ≠ s.endPos),
       Triple
         (f (pos.get h) b)
@@ -2057,6 +2070,18 @@ theorem Spec.forIn_string
       next b => simp [ih _ _ hsp.next]
   | endPos => simpa using Triple.pure _ (by simp)
 
+/--
+The type of loop invariants used by the specifications of `for ... in ...` loops over string slices.
+A loop invariant is a `PostCond` that takes as parameters
+
+* A `String.Slice.Pos` representing the current position in the string slice `s`.
+* A state tuple of type `β`, which will be a nesting of `MProd`s representing the elaboration of
+  `let mut` variables and early return.
+-/
+@[mvcgen_invariant_type]
+abbrev StringSliceInvariant (s : String.Slice) (β : Type u) (ps : PostShape.{u}) :=
+  PostCond (s.Pos × β) ps
+
 /--
 Helper definition for specifying loop invariants for loops with early return.
 
@@ -2079,7 +2104,7 @@ abbrev StringSliceInvariant.withEarlyReturn {s : String.Slice}
   (onContinue : s.Pos → β → Assertion ps)
   (onReturn : γ → β → Assertion ps)
   (onExcept : ExceptConds ps := ExceptConds.false) :
-    PostCond (s.Pos × MProd (Option γ) β) ps
+    StringSliceInvariant s (MProd (Option γ) β) ps
     :=
   ⟨fun ⟨pos, x, b⟩ => spred(
         (⌜x = none⌝ ∧ onContinue pos b)
@@ -2089,7 +2114,7 @@ abbrev StringSliceInvariant.withEarlyReturn {s : String.Slice}
 @[spec]
 theorem Spec.forIn_stringSlice
     {s : String.Slice} {init : β} {f : Char → β → m (ForInStep β)}
-    (inv : PostCond (s.Pos × β) ps)
+    (inv : StringSliceInvariant s β ps)
     (step : ∀ pos b (h : pos ≠ s.endPos),
       Triple
         (f (pos.get h) b)