fix: fix vcgen frames feature for monads with multiple StateT transformers

src/Lean/Elab/Tactic/Do/Internal/VCGen/Context.lean

@@ -46,6 +46,14 @@ public def Deferred.force {α : Type} {n : Type → Type} [Monad n] [MonadLiftT
     writeback (.elaborated a)
     return a
 
+/-- Update an already-elaborated value in place; a not-yet-elaborated `d` is returned unchanged. -/
+public def Deferred.modifyElaborated {α : Type} (d : Deferred α) (f : α → α) : Deferred α :=
+  match d with
+  | .elaborated a => .elaborated (f a)
+  | d => d
+
+public instance {α : Type} [Inhabited α] : Inhabited (Deferred α) := ⟨.elaborated default⟩
+
 /-- A single elaborated `frames` alternative: its program pattern, the binder name of each pattern
 variable (`none` for `_`, index-aligned with `pat.varTypes`), the raw frame term (elaborated in the
 matched goal's context), the source position (a precedence tiebreak among matches, and its index
@@ -61,11 +69,13 @@ public structure FrameEntry where
 
 /-- The materialized frame database: program patterns keyed in a discrimination tree to the `srcIdx`
 of the matching alternative, alongside the alternatives themselves in source order. Materialized once
-in `State.frameDB?`; thereafter only the `retired` flags of `entries` change. -/
+in `State.frameDB`; thereafter only the `retired` flags of `entries` change. -/
 public structure FrameDB where
   tree : DiscrTree Nat := .empty
   entries : Array FrameEntry := #[]
 
+public instance : Inhabited FrameDB := ⟨{}⟩
+
 /--
 Common metadata for a goal whose right-hand side is a weakest-precondition application
 `pre ⊑ wp Prog Value Pred EPred instAL instEAL instWP prog post epost s₁ ... sₙ`.
@@ -142,9 +152,70 @@ public def VCGen.mkBackwardRules : MetaM VCGen.BackwardRules := do
     meetTop := ← mkBackwardRuleFromDecl ``Std.Internal.Do.CompleteLattice.meet_top_le_of_le
   }
 
+private opaque FrameInferenceProcRefPointed : NonemptyType.{0}
+
+/-- Opaque handle to a `FrameInferenceProc` stored in `Context`. Breaks the type-level cycle
+between the procedure, which runs in `VCGenM`, and `VCGenM`, a reader over `Context`. -/
+public def VCGen.FrameInferenceProcRef : Type := FrameInferenceProcRefPointed.type
+
+public instance : Nonempty VCGen.FrameInferenceProcRef := FrameInferenceProcRefPointed.property
+
+/-- A decomposition of a lattice logic connective on the RHS of an entailment. Custom frame operators
+register their own split here via `Context.customLatticeSplits`.
+
+A split with `applyLemma := some _` decomposes pointwise through the excess (state) arguments via
+`LatticeSplit.mkBackwardRuleForLattice`, chaining the `_apply` distribution lemma with the `⊑`-form
+split lemma; the remaining fields supply the structure it needs. A split with `applyLemma := none`
+applies `relLemma` directly as a backward rule, leaving the residual to the rules that follow; the
+pointwise-only fields are then unused. -/
+public structure VCGen.LatticeSplit where
+  /-- The `⊑`-form split lemma decomposing `pre ⊑ connective`. Applied directly as a backward rule
+  when `applyLemma` is `none`. -/
+  relLemma : Name
+  /-- The pointwise `_apply` lemma distributing the connective through function application, or `none`
+  for a split that applies `relLemma` directly without pointwise distribution. -/
+  applyLemma : Option Name := none
+  /-- Rebuild the connective from its fixed parameters `params` (e.g. the frame operator of
+  `PreservesSup.upperAdjoint`), its operands `as`, and the optional lattice carrier type. Unused when
+  `applyLemma` is `none`. -/
+  mkLattice : Array Expr → Array Expr → Option Expr → MetaM Expr := fun _ _ _ =>
+    throwError "LatticeSplit.mkLattice is unavailable for a direct split (applyLemma := none)"
+  /-- Whether the operands are functions of the excess (state) arguments, and so must be applied to
+  each excess argument when descending one lattice level during `mkApplyEq`.
+
+  For `⊓`/`⇨` the operands are themselves elements of the function lattice (`(a ⊓ b) s = a s ⊓ b s`),
+  so each operand `a` becomes `a s`. For `⌜·⌝`/`⊤` the operand is reused unchanged
+  (`(⌜p⌝ : σ→β) s = (⌜p⌝ : β)`, `(⊤ : σ→β) s = (⊤ : β)`), so it must not be applied to `s`. -/
+  needApplyArgs : Bool := false
+  /-- The number of fixed parameters before the lattice operands: `1` for the frame operator of
+  `PreservesSup.upperAdjoint`, `0` for `⊓`/`⌜·⌝`/`⊤`. -/
+  numParams : Nat := 0
+  /-- The number of explicit lattice operands the connective takes after its carrier type,
+  instance, and parameters: `2` for `⊓`/`⇨`, `1` for `⌜·⌝`, `0` for `⊤`. -/
+  numOperands : Nat := 0
+
 public structure VCGen.Context where
   /-- Pre-built backward rules used by `solve`. -/
   backwardRules : VCGen.BackwardRules
+  /-- The frame-inference procedure run on each spec-ready program to optionally produce a frame.
+  Initialized to the default `frames`-clause matcher; overriders wrap or replace it. -/
+  frameInferenceProc : VCGen.FrameInferenceProcRef
+  /-- Lattice splits for custom frame operators, keyed by head constant. Consulted by
+  `splitLatticeOp?` before the built-in connectives, so a custom frame proc can decompose
+  `pre ⊑ conj F rest` for its own `conj`. -/
+  customLatticeSplits : Std.HashMap Name VCGen.LatticeSplit := {}
+  /-- Lattice splits for the residual wands `PreservesSup.upperAdjoint conj F rest` of custom frame operators,
+  keyed by the `conj` head constant. Consulted by `splitLatticeOp?` (dispatching on the inner
+  operator) so a custom frame's magic wand decomposes instead of surfacing in a VC. -/
+  customImpSplits : Std.HashMap Name VCGen.LatticeSplit := {}
+  /-- Conj-reduction equations (`FrameProc.conjReduce`) for custom frame operators, keyed by the `conj`
+  head constant. Consulted by `reduceFrameConj?` to reduce `conj F rest` to its built-in connective over
+  a nested-base lattice. -/
+  customConjReduces : Std.HashMap Name Name := {}
+  /-- Wand-reduction equations (`FrameProc.impReduce`) for custom frame operators, keyed by the `conj`
+  head constant. Consulted by `reduceFrameImp?` to peel a residual
+  `PreservesSup.upperAdjoint conj F rest` over a nested-base lattice. -/
+  customImpReduces : Std.HashMap Name Name := {}
   /-- User-customizable simp methods used to pre-simplify hypotheses. -/
   hypSimpMethods : Option Sym.Simp.Methods := none
   /-- The `trivial` config option: when `true` (default), `Driver.emitVC` runs
@@ -224,7 +295,7 @@ public structure VCGen.State where
   `FrameEntry.retired`, so each applies at most once (first match wins) and a program whose
   alternative is retired is framed no further and reaches the normal `applySpec` path.
   -/
-  frameDB? : Option (Deferred FrameDB) := none
+  frameDB : Deferred FrameDB := default
   /--
   Holes of type `Invariant` that have been generated so far.
   -/
@@ -256,6 +327,25 @@ public structure VCGen.State where
 
 public abbrev VCGenM := ReaderT VCGen.Context (StateRefT VCGen.State Grind.GrindM)
 
+/-- A frame-inference procedure: given the resource type `R` of the applicable frame operator
+`op : R → Pred → Pred`, the goal's precondition, and the `wp` metadata of a spec-ready program,
+optionally produce a frame `F : R` to apply. -/
+public abbrev VCGen.FrameInferenceProc := Expr → Expr → VCGen.WPInfo → VCGenM (Option Expr)
+
+instance : Nonempty VCGen.FrameInferenceProc := ⟨fun _ _ _ => pure none⟩
+
+unsafe abbrev VCGen.FrameInferenceProc.toRefImpl (p : VCGen.FrameInferenceProc) :
+    VCGen.FrameInferenceProcRef := unsafeCast p
+@[implemented_by VCGen.FrameInferenceProc.toRefImpl]
+public opaque VCGen.FrameInferenceProc.toRef (p : VCGen.FrameInferenceProc) :
+    VCGen.FrameInferenceProcRef
+
+unsafe abbrev VCGen.FrameInferenceProcRef.toProcImpl (r : VCGen.FrameInferenceProcRef) :
+    VCGen.FrameInferenceProc := unsafeCast r
+@[implemented_by VCGen.FrameInferenceProcRef.toProcImpl]
+public opaque VCGen.FrameInferenceProcRef.toProc (r : VCGen.FrameInferenceProcRef) :
+    VCGen.FrameInferenceProc
+
 namespace VCGen
 
 public def Scope.registerJP (s : Scope) (fv : FVarId) (info : JumpSiteInfo) : Scope :=

src/Lean/Elab/Tactic/Do/Internal/VCGen/Driver.lean

@@ -147,18 +147,18 @@ Return the VCs and invariant goals.
 `stepLimit?`, when `some n`, seeds the fuel counter to `n`; when `none`, fuel is unlimited.
 -/
 public partial def run (goal : Grind.Goal) (ctx : Context) (scope : VCGen.Scope)
-    (stepLimit? : Option Nat := none) (frameDB? : Option (Deferred FrameDB) := none) :
+    (stepLimit? : Option Nat := none) (frameDB : Deferred FrameDB := default) :
     Grind.GrindM Result := do
   let initState : State :=
-    { fuel := match stepLimit? with | some n => .limited n | none => .unlimited, frameDB? }
+    { fuel := match stepLimit? with | some n => .limited n | none => .unlimited, frameDB }
   let ((), state) ← StateRefT'.run (ReaderT.run (work scope goal) ctx) initState
   _ ← state.invariants.mapIdxM fun idx mv => do
     mv.setTag (Name.mkSimple ("inv" ++ toString (idx + 1)))
   _ ← state.vcs.mapIdxM fun idx g => do
     g.mvarId.setTag (Name.mkSimple ("vc" ++ toString (idx + 1)) ++ (← g.mvarId.getTag).eraseMacroScopes)
   let vcs ← state.vcs.filterM (not <$> ·.mvarId.isAssigned)
-  let unmatchedFrames := match state.frameDB? with
-    | some (.elaborated db) => db.entries.filterMap fun e => if e.retired then none else some e.frameStx
+  let unmatchedFrames := match state.frameDB with
+    | .elaborated db => db.entries.filterMap fun e => if e.retired then none else some e.frameStx
     | _ => #[]
   return {
     invariants := state.invariants,

src/Lean/Elab/Tactic/Do/Internal/VCGen/Entails.lean

@@ -7,6 +7,7 @@ module
 
 prelude
 public import Lean.Elab.Tactic.Do.Internal.VCGen.Context
+public import Lean.Elab.Tactic.Do.Internal.VCGen.FrameProc
 public import Lean.Elab.Tactic.Do.Internal.VCGen.EPost
 public import Lean.Elab.Tactic.Do.Internal.VCGen.RuleCache
 public import Lean.Elab.Tactic.Do.Internal.VCGen.Util
@@ -76,11 +77,25 @@ public def splitLatticeOp? (goal : MVarId) (rhs : Expr) :
     VCGenM (Option (List MVarId)) :=
   rhs.withApp fun head args => do
     let some headName := head.constName? | return none
-    let some c := latticeSplits[headName]? | return none
-    let as := args.extract 2 (2 + c.numOperands)
-    let excessArgs := args.drop (2 + c.numOperands)
-    let resultType? := if c.needApplyArgs then none else args[0]?
-    let rule ← mkBackwardRuleForLatticeCached c as excessArgs resultType?
+    let ctx ← read
+    -- For a residual `PreservesSup.upperAdjoint f b`, dispatch on the head of the slice `f`: a custom
+    -- frame's magic wand (`f = conj F`, head `conj`) goes to its own `impSplit`. The meet `⇨` is its
+    -- own head `himp`, handled by the branch below.
+    let custom? :=
+      if headName == ``Lean.Order.PreservesSup.upperAdjoint then
+        -- `@PreservesSup.upperAdjoint α inst f b`: the slice `f` is at index 2.
+        (args[2]?.bind (·.getAppFn.constName?)).bind (ctx.customImpSplits[·]?)
+      else
+        ctx.customLatticeSplits[headName]?
+    let some c := custom? <|> latticeSplits[headName]? | return none
+    let rule ← match c.applyLemma with
+      | none => mkBackwardRuleForLatticeDirectCached c
+      | some _ =>
+        let params := args.extract 2 (2 + c.numParams)
+        let as := args.extract (2 + c.numParams) (2 + c.numParams + c.numOperands)
+        let excessArgs := args.drop (2 + c.numParams + c.numOperands)
+        let resultType? := if c.needApplyArgs then none else args[0]?
+        mkBackwardRuleForLatticeCached c params as excessArgs resultType?
     match ← rule.applyChecked goal with
     | .goals goals => return some goals
     | .failed => return none

src/Lean/Elab/Tactic/Do/Internal/VCGen/FrameProc.lean

@@ -0,0 +1,124 @@
+/-
+Copyright (c) 2026 Lean FRO LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Graf
+-/
+module
+
+prelude
+public import Lean.Elab.Tactic.Do.Internal.VCGen.Context
+
+/-!
+The `@[frameproc]` attribute: register a frame inference metaprogram for `vcgen`.
+
+A `FrameProc` bundles the inference procedure, the head constant of the frame operator it frames
+with, and the lattice split that decomposes that operator. The mechanism mirrors `@[simproc]`: the
+attribute records the declaration name, and the compiled `FrameProc` value is recovered from the
+environment via `evalConst`, so it works across modules.
+-/
+
+open Lean Meta Elab Sym
+open Lean.Elab.Tactic.Do.Internal
+
+public section
+
+namespace Lean.Elab.Tactic.Do.Internal.VCGen
+
+/-- A frame inference procedure registered with `@[frameproc]`, together with its frame operator. -/
+structure FrameProc where
+  /-- Head constant of the program type (the monad) whose `wp` this procedure frames. Keys the
+  procedure in the registry; `vcgen` consults it for a program with that head. -/
+  prog : Name
+  /-- The frame inference metaprogram. -/
+  proc : FrameInferenceProc
+  /-- Head constant of the frame operator framed with; keys `split`/`impSplit` for `splitLatticeOp?`. -/
+  conj : Name
+  /-- Builds the frame operator (head constant `conj`) for the goal's assertion type. -/
+  op : WPInfo → MetaM Expr
+  /-- The lattice split decomposing `conj F R` on the RHS of an entailment. -/
+  split : LatticeSplit
+  /-- The lattice split decomposing the residual `PreservesSup.upperAdjoint conj F R` (the frame's magic wand)
+  on the RHS of an entailment, so the wand never surfaces in a VC. -/
+  impSplit : LatticeSplit
+  /-- Optional equation `conj r b = <built-in connective>` (e.g. a `meet`). A fallback to *reduce*
+  `conj r b` over a nested-base lattice, where the direct `split` cannot peel the extra state
+  coordinate; reducing it to the built-in connective lets the meet/ofProp splits decompose it over all
+  coordinates. `none` for frames whose flat `split` always suffices. -/
+  conjReduce : Option Name := none
+  /-- Optional equation `upperAdjoint (conj r) b = <closed form>` (e.g. a cost shift `fun m => b (m+r)`).
+  A fallback to *reduce* the residual wand over a nested-base lattice, where the direct `impSplit`
+  cannot peel the extra state coordinate; reducing it exposes the body `b` (with its inner `wp`) so the
+  normal spec step runs. `none` for frames whose flat `impSplit` always suffices. -/
+  impReduce : Option Name := none
+
+unsafe def getFrameProcFromDeclImpl (declName : Name) : ImportM FrameProc := do
+  let ctx ← read
+  match ctx.env.find? declName with
+  | none => throw <| .userError s!"unknown constant '{declName}'"
+  | some info =>
+    unless info.type.isConstOf ``FrameProc do
+      throw <| .userError s!"`{declName}` is not a `FrameProc`"
+    IO.ofExcept <| ctx.env.evalConst FrameProc ctx.opts declName
+
+/-- Recover the compiled `FrameProc` value of a `@[frameproc]`-annotated declaration. -/
+@[implemented_by getFrameProcFromDeclImpl]
+opaque getFrameProcFromDecl (declName : Name) : ImportM FrameProc
+
+/-- The registered frame inference procedures, keyed by the head constant of the program type they
+frame, and their splits keyed by frame-operator head constant. -/
+structure FrameProcs where
+  procs : Std.HashMap Name FrameProc := {}
+  /-- Splits for the frame operators `conj F R`, keyed by `conj` head. -/
+  splits : Std.HashMap Name LatticeSplit := {}
+  /-- Splits for the residual wands `PreservesSup.upperAdjoint conj F R`, keyed by `conj` head. -/
+  impSplits : Std.HashMap Name LatticeSplit := {}
+  /-- Optional conj-reduction equations (`FrameProc.conjReduce`), keyed by `conj` head. -/
+  conjReduces : Std.HashMap Name Name := {}
+  /-- Optional wand-reduction equations (`FrameProc.impReduce`), keyed by `conj` head. -/
+  impReduces : Std.HashMap Name Name := {}
+
+instance : Inhabited FrameProcs := ⟨{}⟩
+
+def FrameProcs.insert (s : FrameProcs) (_declName : Name) (fp : FrameProc) : FrameProcs :=
+  { procs := s.procs.insert fp.prog fp
+    splits := s.splits.insert fp.conj fp.split
+    impSplits := s.impSplits.insert fp.conj fp.impSplit
+    conjReduces := match fp.conjReduce with
+      | some eq => s.conjReduces.insert fp.conj eq
+      | none => s.conjReduces
+    impReduces := match fp.impReduce with
+      | some eq => s.impReduces.insert fp.conj eq
+      | none => s.impReduces }
+
+abbrev FrameProcExtension := ScopedEnvExtension Name (Name × FrameProc) FrameProcs
+
+def toFrameProcEntry (declName : Name) : ImportM (Name × FrameProc) := do
+  return (declName, ← getFrameProcFromDecl declName)
+
+builtin_initialize frameProcExt : FrameProcExtension ←
+  registerScopedEnvExtension {
+    name         := by exact decl_name%
+    mkInitial    := return {}
+    ofOLeanEntry := fun _ declName => toFrameProcEntry declName
+    toOLeanEntry := fun (declName, _) => declName
+    addEntry     := fun s (declName, fp) => s.insert declName fp
+  }
+
+/-- The frame inference procedures in scope. -/
+def getFrameProcs {m : Type → Type} [Monad m] [MonadEnv m] : m FrameProcs :=
+  return frameProcExt.getState (← getEnv)
+
+def addFrameProcAttr (declName : Name) (kind : AttributeKind) : AttrM Unit := do
+  ensureAttrDeclIsMeta `frameproc declName kind
+  let fp ← getFrameProcFromDecl declName
+  frameProcExt.add (declName, fp) kind
+
+builtin_initialize registerBuiltinAttribute {
+  ref             := by exact decl_name%
+  name            := `frameproc
+  descr           := "register a frame inference procedure for `vcgen`"
+  applicationTime := .afterCompilation
+  add             := fun declName _stx kind => addFrameProcAttr declName kind
+}
+
+end Lean.Elab.Tactic.Do.Internal.VCGen

src/Lean/Elab/Tactic/Do/Internal/VCGen/Frontend.lean

@@ -116,7 +116,13 @@ public def mkContext (lemmas : Syntax) (goal : MVarId) (ignoreStarArg := false)
         catch _ => continue
   let backwardRules ← VCGen.mkBackwardRules
   let allSpecThms ← extendWithSimpSpecs specThms simpThms
-  let ctx : VCGen.Context := { backwardRules }
+  -- Registered `@[frameproc]` procedures contribute the lattice splits for their frame operators and
+  -- the residual wands of those operators.
+  let frameProcs ← VCGen.getFrameProcs
+  let ctx : VCGen.Context :=
+    { backwardRules, frameInferenceProc := VCGen.matchFrame?.toRef,
+      customLatticeSplits := frameProcs.splits, customImpSplits := frameProcs.impSplits,
+      customConjReduces := frameProcs.conjReduces, customImpReduces := frameProcs.impReduces }
   return (ctx, { specs := allSpecThms })
 
 end VCGen
@@ -251,7 +257,7 @@ private structure ParsedArgs where
   ctx : VCGen.Context
   scope : VCGen.Scope
   invariantAlts? : Option (Std.HashMap Nat Syntax)
-  frameDB? : Option (Deferred FrameDB)
+  frameDB : Deferred FrameDB
 
 /-- Build a `Sym.Pattern` from `e` by abstracting the metavariables `xs` into pattern variables.
 `checkTypeMask?` is `none` because `until` holes appear as function arguments, whose types the
@@ -339,7 +345,7 @@ private def parseArgs (stx : Syntax) (goal : MVarId) : TermElabM ParsedArgs := g
   -- distinguish "default true" from "user-set true"); not yet wired.
   let (ctx, scope) ← VCGen.mkContext stx[2] goal
   let untilPat? ← if stx[3].isNone then pure none else some <$> elabUntilPattern ⟨stx[3][1]⟩
-  let frameDB? ← if stx[4].isNone then pure none else some <$> elabFrameDB stx[4][1].getArgs
+  let frameDB ← elabFrameDB (if stx[4].isNone then #[] else stx[4][1].getArgs)
   let invariantAlts? ← parseInvariantMap stx[5]
   let hypSimpMethods ← elabSimplifyingAssumptions stx[6]
   let ctx := { ctx with
@@ -351,15 +357,15 @@ private def parseArgs (stx : Syntax) (goal : MVarId) : TermElabM ParsedArgs := g
     internalize := config.internalize,
     invariantAlts := invariantAlts?.getD {},
     untilPat? }
-  return { config, ctx, scope, invariantAlts?, frameDB? }
+  return { config, ctx, scope, invariantAlts?, frameDB }
 
 /-- `vcgen` step inside `sym => …` blocks. -/
 @[builtin_grind_tactic Lean.Parser.Tactic.Grind.vcgen]
 def evalSymVCGen : Lean.Elab.Tactic.Grind.GrindTactic := fun stx => do
   let goal ← Lean.Elab.Tactic.Grind.getMainGoal
   let args ← parseArgs stx goal.mvarId
   let result ← Lean.Elab.Tactic.Grind.liftGrindM do
-    let result ← VCGen.run goal args.ctx args.scope args.config.stepLimit (frameDB? := args.frameDB?)
+    let result ← VCGen.run goal args.ctx args.scope args.config.stepLimit (frameDB := args.frameDB)
     if let some alts := args.invariantAlts? then
       elabRemainingInvariants alts result.invariants result.inlineHandledInvariants
     return result