feat: vcgen frames cost monads over a stateful base

Extend `@[frameproc]` frame inference to a *nested-base* assertion
lattice (e.g. `ComplexiT (StateM σ)`, lattice `Tick → σ → Prop`). There
the frame operator is applied to more coordinates than the registered
direct split lemma handles (`pre ⊑ cCostConj c R n s`, not `… n`), so
`splitLatticeOp?` returns `none` and the raw goal — with an un-specced
inner `wp` — is dumped on `finish`.

Add two fallback steps in `solve`, run after the direct split and the wp
phase (so flat lattices are untouched): for a registered frame `conj` or
its residual wand, `mkReduceRule` builds a backward rule from a
registered reduce-equation (`FrameProc.conjReduce` / `impReduce`) whose
conclusion matches the goal head, and `applyChecked`s it. The single
premise is the reduced subgoal — the built-in connective for the conj
(the meet split then peels it over all coordinates), or the shifted body
for the wand (exposing the inner `wp` for the spec step).

`mkReduceRule` reuses `mkBackwardRuleForLattice`'s composition tail
(`liftEqByArgs` + `Eq.mp` + `mkBackwardRuleFromExpr`), minus the
`relLemma`. `tests/bench/vcgen/mwe_nested_frame.lean` documents the
2-level setup. Flat `TickM` is unchanged (the fallbacks never fire).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: volodeyka

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

@@ -208,6 +208,14 @@ public structure VCGen.Context where
   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

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

@@ -40,6 +40,16 @@ structure FrameProc where
   /-- 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
@@ -62,13 +72,23 @@ structure FrameProcs where
   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 }
+    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
 

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

@@ -121,7 +121,8 @@ public def mkContext (lemmas : Syntax) (goal : MVarId) (ignoreStarArg := false)
   let frameProcs ← VCGen.getFrameProcs
   let ctx : VCGen.Context :=
     { backwardRules, frameInferenceProc := VCGen.matchFrame?.toRef,
-      customLatticeSplits := frameProcs.splits, customImpSplits := frameProcs.impSplits }
+      customLatticeSplits := frameProcs.splits, customImpSplits := frameProcs.impSplits,
+      customConjReduces := frameProcs.conjReduces, customImpReduces := frameProcs.impReduces }
   return (ctx, { specs := allSpecThms })
 
 end VCGen

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

@@ -186,6 +186,44 @@ public def LatticeSplit.mkBackwardRuleForLattice
   let res ← abstractMVars prf
   mkBackwardRuleFromExpr res.expr res.paramNames.toList
 
+/--
+Creates a backward rule that *reduces* an operator-applied RHS via a registered equation
+`eqName : LHS = RHS`, where `LHS` is the operator applied to its arguments (e.g. `cCostConj c R`, or
+`PreservesSup.upperAdjoint (cCostConj c) R`). Unlike `mkBackwardRuleForLattice` it chains no `relLemma`:
+the single premise `pre ⊑ RHS s₁…sₙ` *is* the reduced subgoal. Produces
+
+```
+∀ … (pre), pre ⊑ RHS s₁…sₙ → pre ⊑ LHS s₁…sₙ
+```
+
+with the operator's own arguments and the excess coordinates `s₁…sₙ` (read off from `goalRhs` past the
+`LHS` arity) as fresh metavariables — so `applyChecked` unifies the rule's conclusion (same head as the
+goal) directly, with no defeq/normalization. Returns `none` if `eqName` is not an equation or `goalRhs`
+has fewer arguments than `LHS`. -/
+public def mkReduceRule (eqName : Name) (goalRhs : Expr) : MetaM (Option BackwardRule) := do
+  let eqConst ← mkConstWithFreshMVarLevels eqName
+  let (argMVars, _, eqBody) ← forallMetaTelescope (← Meta.inferType eqConst)
+  let some (_, lhs, _) := eqBody.eq? | return none
+  let numLhsArgs := lhs.getAppNumArgs
+  if goalRhs.getAppNumArgs < numLhsArgs then return none
+  -- Pin the operator's own arguments (carrier `Pred`, instance, resource …) to the goal by unifying
+  -- the equation's `LHS` with the goal RHS's operator prefix, so `LHS`'s carrier type is concrete and
+  -- the coordinates below apply.
+  unless ← isDefEq lhs (mkAppN goalRhs.getAppFn (goalRhs.getAppArgs.extract 0 numLhsArgs)) do
+    return none
+  -- The excess coordinates are the goal RHS arguments past the operator's own; freshen them so the
+  -- rule is a pattern (like `mkBackwardRuleForLattice`'s `ss`).
+  let ss ← (goalRhs.getAppArgs.extract numLhsArgs goalRhs.getAppNumArgs).mapM fun arg => do
+    mkFreshExprMVar (← Meta.inferType arg)
+  let eqDistributed ← liftEqByArgs (mkAppN eqConst argMVars) ss.toList  -- LHS s₁…sₙ = RHS s₁…sₙ
+  let some (_, goal, _) := (← Meta.inferType eqDistributed).eq? | return none
+  let pre ← mkFreshExprMVar (userName := `Pre) (← Meta.inferType goal)
+  let relEq ← mkCongrArg (← mkAppM ``PartialOrder.rel #[pre]) eqDistributed
+  -- eqMp : (pre ⊑ RHS s₁…sₙ) → (pre ⊑ LHS s₁…sₙ)
+  let eqMp ← mkAppM ``Eq.mp #[← mkEqSymm relEq]
+  let res ← abstractMVars eqMp
+  return some (← mkBackwardRuleFromExpr res.expr res.paramNames.toList)
+
 /-! ## Spec rules -/
 
 /-- Build the explicit pointwise implication premise used to weaken a concrete `post`.

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

@@ -475,6 +475,58 @@ private def foldUpperAdjointMeet? (goal : MVarId) (target rhs : Expr) : VCGenM (
   let newTarget ← mkAppNS target.getAppFn (relArgs.set! (relArgs.size - 1) newRhs)
   return some (← goal.replaceTargetDefEq newTarget)
 
+/-- `unfoldReducible`-normalize a goal's target, so a later `applyChecked` unifies past reducible state
+types (`StateM σ`, `Tick`, …) that sit behind reducible definitions — the normalization that
+`BackwardRule.applyChecked`'s `+debug` retry (`Util.lean`) diagnoses. When `betaRhs` is set, the RHS of
+the `pre ⊑ rhs` target is also beta-reduced first (the wand reduction leaves a `(fun m => …) s` redex
+whose inner `wp` head must be exposed). -/
+private def normalizeReducedGoal (goal : MVarId) (betaRhs : Bool := false) : VCGenM MVarId := do
+  let ty ← goal.getType
+  unless betaRhs do return ← goal.replaceTargetDefEq (← unfoldReducible ty)
+  let relArgs := ty.getAppArgs
+  let some rhs := relArgs[relArgs.size - 1]? | return goal
+  let newTarget ← mkAppNS ty.getAppFn (relArgs.set! (relArgs.size - 1) (← unfoldReducible rhs.headBeta))
+  goal.replaceTargetDefEq newTarget
+
+/-- Fallback for a registered frame `conj` that the direct `splitLatticeOp?` could not peel over a
+*nested-base* assertion lattice: `vcgen` introduced an extra inner state coordinate, so the goal is
+`conj c R n s` — one application deeper than the registered direct split lemma expects — and
+`splitLatticeOp?` returns `none`.
+
+Build a backward rule from the registered `conjReduce` equation (`conj c R = <built-in connective>`,
+e.g. a `meet`) via `mkReduceRule` and `applyChecked` it. The rule's conclusion has the same head as the
+goal, so it unifies directly (no defeq/normalization); its single premise `pre ⊑ <connective> n s` is the
+reduced subgoal. `unfoldReducible`-normalize it so the next iteration's `splitLatticeOp?` decomposes the
+exposed connective over *all* coordinates and the operand `wp` flows into the normal spec step.
+
+Keyed on `customConjReduces`, so only registered frame `conj` heads are touched, and run *after* the
+direct split and the wp phase — a flat lattice keeps its direct split, and only the nested case that
+would otherwise stall on `.noProgress` reaches here. Terminates: each firing exposes a non-`conj`
+connective head, so it never re-fires on its own output. -/
+private def reduceFrameConj? (goal : MVarId) (rhs : Expr) : VCGenM (Option MVarId) := do
+  let some headName := rhs.getAppFn.constName? | return none
+  let some eqName := (← read).customConjReduces[headName]? | return none
+  let some rule ← mkReduceRule eqName rhs | return none
+  let .goals [g] ← rule.applyChecked goal | return none
+  return some (← normalizeReducedGoal g)
+
+/-- Wand companion to `reduceFrameConj?`: reduce a residual frame wand `PreservesSup.upperAdjoint
+(conj c) R` that the direct `impSplit` couldn't peel over a nested-base lattice (the wand is applied to
+an extra inner state coordinate). Build a backward rule from the registered `impReduce` equation
+(`upperAdjoint (conj c) R = <closed form>`, e.g. a cost shift `fun m => R (m + c)`) via `mkReduceRule`
+and `applyChecked` it; the premise `pre ⊑ <closed form> n s` is the reduced subgoal. Beta-reduce and
+`unfoldReducible`-normalize it so the body `R` — and its inner `wp` — is exposed to the normal spec step
+instead of stranding in a VC. Run as a fallback, so a flat lattice keeps its direct `impSplit`. -/
+private def reduceFrameImp? (goal : MVarId) (rhs : Expr) : VCGenM (Option MVarId) := do
+  unless rhs.isAppOf ``Lean.Order.PreservesSup.upperAdjoint do return none
+  -- `@PreservesSup.upperAdjoint α inst (conj c) R …`: the slice `conj c` is at index 2; reduce keyed on
+  -- its head.
+  let some sliceHead := rhs.getAppArgs[2]?.bind (·.getAppFn.constName?) | return none
+  let some eqName := (← read).customImpReduces[sliceHead]? | return none
+  let some rule ← mkReduceRule eqName rhs | return none
+  let .goals [g] ← rule.applyChecked goal | return none
+  return some (← normalizeReducedGoal g (betaRhs := true))
+
 /--
 The main VC generation step. Operates on a plain `MVarId` with no knowledge of grind.
 Returns `.goals subgoals` when the goal was decomposed, or a classification result
@@ -565,6 +617,13 @@ public def solve (scope : VCGen.Scope) (goal : MVarId) : VCGenM SolveResult := g
       | .notFramed goal info => return ← applySpec scope goal info
     throwError "Failed to decompose weakest precondition for {info.prog}. This should not happen."
 
+  -- Phase 5 (fallback): a registered frame `conj`/wand that the direct split couldn't peel over a
+  -- nested-base lattice. `reduceFrameConj?` exposes the `conj`'s `meet`-first connective; `reduceFrameImp?`
+  -- reduces the residual wand via its `impReduce` equation. Either way the next iteration's
+  -- `splitLatticeOp?`/spec step decomposes the result over all coordinates and specs the inner `wp`.
+  if let some g ← reduceFrameConj? goal rhs then return .goals scope [g]
+  if let some g ← reduceFrameImp? goal rhs then return .goals scope [g]
+
   return .stop (.noProgress pre rhs)
 
 end VCGen

tests/bench/vcgen/mwe_nested_frame.lean

@@ -0,0 +1,122 @@
+/-
+Copyright (c) 2026 Lean FRO LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Vladimir Gladshtein, Sebastian Graf
+-/
+import Lean
+import Std.Internal
+import Std.Tactic.Do
+
+set_option mvcgen.warning false
+set_option grind.warning false
+
+/-!
+# MWE: `vcgen` cost-framing over a *2-level* (stateful-base) assertion lattice
+
+`tests/elab/vcgenTickFrames.lean` shows `vcgen` cost-framing working for `TickM := StateM Nat`,
+whose assertion lattice is the **flat** `Nat → Prop` (the single cost coordinate). The *same*
+construction over a **stateful base monad** — e.g. `ComplexiT (StateM σ)` in the `loom-complexity`
+project, assertion lattice `Nat → (σ → Prop)` (a cost coordinate **and** a base-state coordinate) —
+*used to* fail (`vcgen [foo] with finish` reported `finish failed`). It is **now fixed** by the
+`reduceFrameConj?` / `reduceFrameImp?` fallbacks in `VCGen/Solve.lean` (see "The fix" below). This
+file kept the diagnosis that pinned the root cause down.
+
+Concretely, with `foo : ComplexiT (StateM Unit) Unit := do withTick ()` and
+`fooSpec : ⦃fun acc _ => acc = 0⦄ foo ⦃fun _ acc _ => acc ≤ 1⦄`, after the registered frameproc
+fires the raw stuck goal is (`s✝¹` = cost, `s✝` = base `Unit` state):
+
+    (s✝¹ = 0) ⊑ costConj s✝¹ (wp (Gadget.skipFrame foo) (fun a => (costConj s✝¹) -⋆ Q) ⊥) s✝¹ s✝
+                            └──────────────────── costConj applied to TWO coords ───────────┘
+
+The registered *direct* split `le_costConj_point_apply : pre ⊑ costConj c R n` matches `costConj`
+applied to **one** coordinate. Over the 2-level lattice `vcgen` has introduced the extra base-state
+arg, so the goal is `costConj c R n s` (one application deeper); the split can't unify, returns
+`none`, and the raw `costConj` goal — with the un-specced inner `wp` still inside it — is dumped on
+`finish`, which cannot evaluate it.
+
+## What this file pins down (each `example` below verifies a claim)
+
+* **§A — the built-in stateful `meet` split is NOT the culprit.** `vcgen` decomposes
+  `pre ⊑ (guard ⊓ shift) coords` over a *2-level* lattice for every operand shape we tried (opaque,
+  λ, cost-shifted, `⌜·⌝`-guarded — i.e. the exact `costConj` shape). In `loom-complexity` it was
+  further verified to split even when an operand is `wp (Gadget.skipFrame prog) …` (reaching the
+  inner `wp`, i.e. getting *past* the meet). So once the meet is exposed, the split fires fine,
+  over both coordinates.
+
+* **§B — `costConj` is NOT defeq to its `meet` form.** The pi-lattice meet is a genuinely different
+  term (`meet_apply` is a *lemma*, not `rfl`). So a fix cannot reduce the *current* `costConj` with
+  `Meta.unfoldDefinition?` + `MVarId.replaceTargetDefEq` (the `rfl` it relies on does not hold); the
+  meet form must be exposed either by *defining* `costConj` meet-first, or by a propositional
+  `replaceTargetEq`.
+
+## The fix (implemented in `VCGen/Solve.lean`)
+
+Two fallback steps in `solve`, run *after* `splitLatticeOp?` so flat lattices keep their direct split:
+
+* `reduceFrameConj?` — for a registered frameproc `conj` the direct split couldn't peel, delta-unfold
+  the `conj` (defined **meet-first**, §B) to expose its `meet`, then **`unfoldReducible`-normalize**.
+  That last step was the real blocker: the meet split's `applyChecked` only unifies against the
+  reducible-normal form (state types `StateM σ`, `Tick` sit behind reducible defs). `applyChecked`'s
+  `+debug` retry diagnoses exactly this ("succeeded after `unfoldReducible`-normalization"). The next
+  `solve` iteration's meet split then peels both coordinates and specs the inner `wp`. Fires every
+  iteration, so it recurses at each framing level.
+
+* `reduceFrameImp?` — for ≥2 cost calls, the residual wand `upperAdjoint (conj c) R` wraps the *next*
+  call's un-specced `wp` and the direct `impSplit` can't peel the extra coord. Rewrite it via a
+  registered `FrameProc.impReduce` equation (here `costConj_imp : upperAdjoint (costConj c) R =
+  fun m => R (m+c)`), then beta + `unfoldReducible`, exposing the inner `wp` for the spec step.
+
+The leftover VCs (budget guard `c ≤ n`, `WP.Frames`) close under `with finish`.
+-/
+
+open Lean Order Std Internal.Do
+
+/-- The 2-level cost-frame operator: a guard on the cost coordinate `n`, met — on the base-state
+lattice `σ → Prop` — with the body run on the remaining cost `n - r`. The analog of the flat
+`costConj` of `vcgenTickFrames.lean`, but over `Nat → (σ → Prop)`. -/
+def costConj {σ : Type} (r : Nat) (b : Nat → (σ → Prop)) : Nat → (σ → Prop) :=
+  fun n => ⌜r ≤ n⌝ ⊓ b (n - r)
+
+/-! ## §A — the built-in stateful `meet` split decomposes 2-level meets
+
+Each goal is a bare lattice entailment whose RHS is a `meet` of two `Nat → (Unit → Prop)` functions;
+`vcgen` introduces the cost arg *and* the base-state arg and fires the meet split across **both**
+coordinates, leaving one obligation per operand, which `finish` discharges. -/
+
+/-- Two `⊤` operands over the 2-level lattice `Nat → Unit → Prop`: the split introduces the cost arg
+`n` *and* the base-state arg `s`, then closes `pre ⊑ ⊤` on each side. -/
+example (pre : Nat → Unit → Prop) :
+    pre ⊑ (fun _ => (⊤ : Unit → Prop)) ⊓ (fun _ => (⊤ : Unit → Prop)) := by
+  vcgen with finish [top_apply]
+
+-- The exact `costConj` shape — `⌜·⌝`-guard ⊓ cost-shift — over the 2-level lattice. The split peels
+-- **both** coordinates (`s✝¹` the cost, `s✝` the base state), leaving the guard obligation (`vc1`)
+-- and the shifted-body obligation (`vc2`), each trivially true. (In `vcgenTickFrames.lean` the
+-- analogous split has only the single cost coordinate; here it must — and does — recurse into base.)
+/--
+error: unsolved goals
+case vc1
+b : Nat → PUnit → Prop
+c s✝¹ : Nat
+s✝ : PUnit
+⊢ (c ≤ s✝¹ ∧ b (s✝¹ - c) s✝) ⊑ ⌜c ≤ s✝¹⌝ s✝
+
+case vc2
+b : Nat → PUnit → Prop
+c s✝¹ : Nat
+s✝ : PUnit
+⊢ (c ≤ s✝¹ ∧ b (s✝¹ - c) s✝) ⊑ b (s✝¹ - c) s✝
+-/
+#guard_msgs in
+example (b : Nat → Unit → Prop) (c : Nat) :
+    (fun n s => c ≤ n ∧ b (n - c) s) ⊑ (fun n => ⌜c ≤ n⌝) ⊓ (fun n => b (n - c)) := by
+  vcgen
+
+/-! ## §B — `costConj` is not definitionally its `meet` form -/
+
+/-- The `rfl` that an `unfoldDefinition?` + `replaceTargetDefEq` reduction would rely on **fails**;
+the equation holds only *propositionally* (`meet_apply` is a lemma, not `rfl`). -/
+example {σ : Type} (c : Nat) (b : Nat → σ → Prop) :
+    costConj c b = (fun n => ⌜c ≤ n⌝) ⊓ (fun n => b (n - c)) := by
+  fail_if_success rfl
+  funext n; simp only [costConj, meet_apply]