From 220a242f65448eb1007c1ccfd9f4ee1bb5027be7 Mon Sep 17 00:00:00 2001
From: Sebastian Graf <sgraf1337@gmail.com>
Date: Wed, 11 Mar 2026 16:04:22 +0800
Subject: [PATCH] feat: add `@[mvcgen_invariant_type]` attribute for extensible
 invariant classification (#12874)

This PR adds an `@[mvcgen_invariant_type]` tag attribute so that users
can mark
custom types as invariant types for the `mvcgen` tactic. Goals whose
type is an
application of a tagged type are classified as invariants rather than
verification
conditions. The hard-coded check for `Std.Do.Invariant` is kept as a
fallback
until a stage0 update allows applying the attribute directly.

A follow-up PR (after a stage0 update) will apply
`@[mvcgen_invariant_type]` to
`Std.Do.Invariant` and remove the hard-coded fallback.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/Lean/Elab/Tactic/Do/Attr.lean        | 19 +++++++++
 src/Lean/Elab/Tactic/Do/VCGen/Basic.lean |  2 +-
 stage0/src/stdlib_flags.h                |  2 +-
 tests/elab/mvcgenInvariantAttr.lean      | 50 ++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 tests/elab/mvcgenInvariantAttr.lean

diff --git a/src/Lean/Elab/Tactic/Do/Attr.lean b/src/Lean/Elab/Tactic/Do/Attr.lean
index 493eb84481..963d6a4c8e 100644
--- a/src/Lean/Elab/Tactic/Do/Attr.lean
+++ b/src/Lean/Elab/Tactic/Do/Attr.lean
@@ -250,3 +250,22 @@ def SpecExtension.getTheorems (ext : SpecExtension) : CoreM SpecTheorems :=
 
 def getSpecTheorems : CoreM SpecTheorems :=
   specAttr.getTheorems
+
+/--
+Marks a type as an invariant type for the `mvcgen` tactic.
+Goals whose type is an application of a tagged type will be classified
+as invariants rather than verification conditions.
+-/
+builtin_initialize mvcgenInvariantAttr : TagAttribute ←
+  registerTagAttribute `mvcgen_invariant_type
+    "marks a type as an invariant type for the `mvcgen` tactic"
+
+/--
+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
+  else
+    false
diff --git a/src/Lean/Elab/Tactic/Do/VCGen/Basic.lean b/src/Lean/Elab/Tactic/Do/VCGen/Basic.lean
index 8fb6f5d411..3b189f5e98 100644
--- a/src/Lean/Elab/Tactic/Do/VCGen/Basic.lean
+++ b/src/Lean/Elab/Tactic/Do/VCGen/Basic.lean
@@ -104,7 +104,7 @@ def addSubGoalAsVC (goal : MVarId) : VCGenM PUnit := do
   -- VC to the user as-is, without abstracting any variables in the local context.
   -- This only makes sense for synthetic opaque metavariables.
   goal.setKind .syntheticOpaque
-  if ty.isAppOf ``Std.Do.Invariant then
+  if isMVCGenInvariantType (← getEnv) ty then
     modify fun s => { s with invariants := s.invariants.push goal }
   else
     modify fun s => { s with vcs := s.vcs.push goal }
diff --git a/stage0/src/stdlib_flags.h b/stage0/src/stdlib_flags.h
index b625c5934f..ca321bef93 100644
--- a/stage0/src/stdlib_flags.h
+++ b/stage0/src/stdlib_flags.h
@@ -1,5 +1,5 @@
 #include "util/options.h"
-
+//update stage0 to make @[mvcgen_invariant_type] attribute available
 namespace lean {
 options get_default_options() {
     options opts;
diff --git a/tests/elab/mvcgenInvariantAttr.lean b/tests/elab/mvcgenInvariantAttr.lean
new file mode 100644
index 0000000000..5acfa31be2
--- /dev/null
+++ b/tests/elab/mvcgenInvariantAttr.lean
@@ -0,0 +1,50 @@
+import Std.Tactic.Do
+import Std
+set_option backward.do.legacy false
+
+open Std Do
+
+set_option grind.warning false
+set_option mvcgen.warning false
+
+-- Test that `@[mvcgen_invariant_type]` works end-to-end with a custom invariant type.
+-- We replicate the `StringInvariant` setup locally: define a type, tag it, provide a spec,
+-- and verify that `mvcgen` classifies the invariant goal as `inv1` rather than `vc1`.
+
+@[mvcgen_invariant_type]
+abbrev MyStringInvariant (s : String) (β : Type u) (ps : PostShape.{u}) :=
+  PostCond (s.Pos × β) ps
+
+section
+universe u₁ u₂ v
+variable {β : Type (max u₁ u₂)} {m : Type (max u₁ u₂) → Type v} {ps : PostShape.{max u₁ u₂}}
+variable [Monad m] [WPMonad m ps]
+
+-- This is the same as `Spec.forIn_string` from `Std.Do.Triple.SpecLemmas`, but with
+-- `MyStringInvariant` instead of `PostCond` so that `mvcgen` classifies the goal as an invariant.
+@[spec high]
+axiom mySpec_forIn_string
+    {s : String} {init : β} {f : Char → β → m (ForInStep β)}
+    (inv : MyStringInvariant s β ps)
+    (step : ∀ pos b (h : pos ≠ s.endPos),
+      Triple
+        (f (pos.get h) b)
+        (inv.1 (pos, b))
+        (fun r => match r with
+          | .yield b' => inv.1 (pos.next h, b')
+          | .done b' => inv.1 (s.endPos, b'), inv.2)) :
+    Triple (forIn s init f) (inv.1 (s.startPos, init)) (fun b => inv.1 (s.endPos, b), inv.2)
+
+end
+
+def stringLoop (s : String) : Id Unit := do
+  for _ in s do
+    pure ()
+
+theorem stringLoop_correct (s : String) :
+    ⦃⌜True⌝⦄ stringLoop s ⦃⇓ _ => ⌜True⌝⦄ := by
+  -- With `@[mvcgen_invariant_type]` on `MyStringInvariant`, mvcgen classifies the
+  -- loop invariant goal as `inv1` rather than `vc1`, enabling the `invariants` syntax.
+  mvcgen [stringLoop] invariants
+  · ⇓ _ => ⌜True⌝
+  with grind