Forward.State: make insertion lazy

Previously, the forward state was eagerly updated whenever a hypothesis
was added. Now, we use a two-phase protocol. Initially, hypotheses are
added to (and removed from) per-rule queues, which were already in place as
part of a different laziness scheme. They are properly processed only
once we need the forward matches, i.e. when selecting forward rules.

The processing is also staggered by phase, so before norm rule selection
only norm rules are processed. This yields a nice speedup when
hypotheses both appear and disappear within the same phase, e.g. due to
repeated `simp`s or `subst`s during normalisation. Previously, safe and
unsafe forward rules would have been updated for each new hypotheses.
Now, we only process those hypotheses that remain after normalisation.

TODO docs

Author: JLimperg

Aesop/Forward/State.lean

@@ -7,42 +7,37 @@ Authors: Jannis Limperg
 import Aesop.Forward.State
 import Aesop.RuleSet
 
-namespace Aesop
+namespace Aesop.ForwardState
 
 open Lean Lean.Meta
 
 /-- Apply a goal diff to the state, adding and removing hypotheses as indicated
 by the diff. -/
-def ForwardState.applyGoalDiff (rs : LocalRuleSet) (diff : GoalDiff)
-    (fs : ForwardState) : BaseM (ForwardState × Array ForwardRuleMatch) :=
+def applyGoalDiff (rs : LocalRuleSet) (diff : GoalDiff) (fs : ForwardState) :
+    BaseM ForwardState :=
   profilingForwardState do
   if ! aesop.dev.statefulForward.get (← getOptions) then
-    return (fs, #[])
+    return fs
   let fs ← diff.oldGoal.withContext do
     diff.removedFVars.foldM (init := fs) λ fs h => eraseHyp h fs
   diff.newGoal.withContext do
-    let (fs, ruleMatches) ←
-      diff.addedFVars.foldM (init := (fs, ∅)) λ (fs, ruleMatches) h =>
-        addHyp h fs ruleMatches
+    let fs ← diff.addedFVars.foldM (init := fs) λ fs h => addHyp h fs
     if diff.targetChanged then
-      updateTarget fs ruleMatches
+      updateTarget fs
     else
-      return (fs, ruleMatches)
+      return fs
 where
   eraseHyp (h : FVarId) (fs : ForwardState) : BaseM ForwardState :=
     withConstAesopTraceNode .forward (return m!"erase hyp {Expr.fvar h} ({h.name})") do
       return fs.eraseHyp h
 
-  addHyp (h : FVarId) (fs : ForwardState)
-      (ruleMatches : Array ForwardRuleMatch) :
-      BaseM (ForwardState × Array ForwardRuleMatch) := do
+  addHyp (h : FVarId) (fs : ForwardState) : BaseM ForwardState := do
     let rules ← rs.applicableForwardRules (← h.getType)
     let patInsts ← rs.forwardRulePatternSubstsInLocalDecl (← h.getDecl)
-    fs.addHypWithPatSubstsCore ruleMatches diff.newGoal h rules patInsts
+    return fs.enqueueHypWithPatSubsts h rules patInsts
 
-  updateTarget (fs : ForwardState) (ruleMatches : Array ForwardRuleMatch) :
-      BaseM (ForwardState × Array ForwardRuleMatch) := do
+  updateTarget (fs : ForwardState) : BaseM ForwardState := do
     let patInsts ← rs.forwardRulePatternSubstsInExpr (← diff.newGoal.getType)
-    fs.updateTargetPatSubstsCore ruleMatches diff.newGoal patInsts
+    return fs.enqueueTargetPatSubsts patInsts
 
-end Aesop
+end Aesop.ForwardState

Aesop/Forward/State/ApplyGoalDiff.lean

@@ -18,20 +18,18 @@ def mkInitialForwardState (goal : MVarId) (rs : LocalRuleSet) :
     if ! aesop.dev.statefulForward.get (← getOptions) then
       return (∅, #[])
     let mut fs : ForwardState := ∅
-    let mut ruleMatches := rs.constForwardRuleMatches
-    aesop_trace[forward] do
-      for m in ruleMatches do
-        aesop_trace![forward] "match for constant rule {m.rule.name}"
     for ldecl in ← show MetaM _ from getLCtx do
       if ldecl.isImplementationDetail then
         continue
       let rules ← rs.applicableForwardRules ldecl.type
       let patInsts ← rs.forwardRulePatternSubstsInLocalDecl ldecl
-      let (fs', ruleMatches') ←
-        fs.addHypWithPatSubstsCore ruleMatches goal ldecl.fvarId rules patInsts
-      fs := fs'
-      ruleMatches := ruleMatches'
+      fs := fs.enqueueHypWithPatSubsts ldecl.fvarId rules patInsts
     let patInsts ← rs.forwardRulePatternSubstsInExpr (← goal.getType)
-    fs.addPatSubstsCore ruleMatches goal patInsts
+    fs := fs.enqueueTargetPatSubsts patInsts
+    let ruleMatches := rs.constForwardRuleMatches
+    aesop_trace[forward] do
+      for m in ruleMatches do
+        aesop_trace![forward] "match for constant rule {m.rule.name}"
+    return (fs, ruleMatches)
 
 end Aesop.LocalRuleSet

Aesop/Forward/State/Initial.lean

@@ -0,0 +1,26 @@
+/-
+Copyright (c) 2025 Jannis Limperg. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Jannis Limperg
+-/
+
+import Aesop.Forward.State
+import Aesop.Tree.TreeM
+import Aesop.Tree.RunMetaM
+
+open Lean
+
+namespace Aesop.GoalRef
+
+def updateForwardState (phase : PhaseName) (gref : GoalRef) : TreeM Unit := do
+  let g ← gref.get
+  if phase == .norm then
+    throwError "aesop: internal error: {decl_name%}: at goal {g.id}: norm phase not supported"
+  if ! g.isNormal then
+    throwError "aesop: internal error: {decl_name%}: attempt to update forward state of non-normal goal {g.id} to phase {phase}"
+  let (fs, ms) ← g.runMetaMInPostNormState' fun goal =>
+    g.forwardState.update goal phase
+  let ms := g.forwardRuleMatches.update ms ∅ ∅
+  gref.set <| g.setForwardState fs |>.setForwardRuleMatches ms
+
+end Aesop.GoalRef

Aesop/Forward/State/UpdateGoal.lean

@@ -48,22 +48,18 @@ initialize
 partial def saturateCore (rs : LocalRuleSet) (goal : MVarId) : SaturateM MVarId :=
   withExceptionPrefix "saturate: internal error: " do
   goal.checkNotAssigned `saturate
-  -- We use the forward state only to track the hypotheses present in the goal.
-  let (fs, _) ← rs.mkInitialForwardState goal
-  go goal fs
+  go goal
 where
-  go (goal : MVarId) (fs : ForwardState) : SaturateM MVarId :=
+  go (goal : MVarId) : SaturateM MVarId :=
     withIncRecDepth do
     checkSystem "saturate"
     trace[saturate] "goal {goal.name}:{indentD goal}"
     let mvars := UnorderedArraySet.ofHashSet $ ← goal.getMVarDependencies
     let preState ← show MetaM _ from saveState
     if let some diff ← tryNormRules goal mvars preState then
-      let (fs, _) ← fs.applyGoalDiff rs diff
-      go diff.newGoal fs
+      go diff.newGoal
     else if let some diff ← trySafeRules goal mvars preState then
-      let (fs, _) ← fs.applyGoalDiff rs diff
-      go diff.newGoal fs
+      go diff.newGoal
     else
       clearForwardImplDetailHyps goal
 
@@ -131,8 +127,13 @@ partial def saturateCore (rs : LocalRuleSet) (goal : MVarId) : SaturateM MVarId
   withExceptionPrefix "saturate: internal error: " do
   goal.withContext do
     goal.checkNotAssigned `saturate
+    let mut queue := ∅
     let (fs, ruleMatches) ← rs.mkInitialForwardState goal
-    let queue := ruleMatches.foldl (init := ∅) λ queue m => queue.insert m
+    for m in ruleMatches do
+      queue := queue.insert m
+    let (fs, ruleMatches) ← fs.update goal (phase? := none)
+    for m in ruleMatches do
+      queue := queue.insert m
     go ∅ fs queue ∅ goal
 where
   go (hypDepths : Std.HashMap FVarId Nat) (fs : ForwardState) (queue : Queue)
@@ -166,10 +167,11 @@ where
           else
             let rules ← profilingRuleSelection do
               rs.applicableForwardRules type
-            let (fs, ruleMatches) ← profilingForwardState do
+            let fs ← profilingForwardState do
               let patInsts ←
                 rs.forwardRulePatternSubstsInLocalDecl (← hyp.getDecl)
-              fs.addHypWithPatSubsts goal hyp rules patInsts
+              return fs.enqueueHypWithPatSubsts hyp rules patInsts
+            let (fs, ruleMatches) ← fs.update goal (phase? := none)
             let queue :=
               ruleMatches.foldl (init := queue) λ queue m => queue.insert m
             go hypDepths fs queue erasedHyps goal

Aesop/Saturate.lean

@@ -6,6 +6,7 @@ Authors: Jannis Limperg
 
 import Aesop.Search.Expansion.Norm
 import Aesop.Tree.AddRapp
+import Aesop.Forward.State.UpdateGoal
 
 open Lean
 open Lean.Meta
@@ -173,11 +174,12 @@ def SafeRulesResult.toEmoji : SafeRulesResult → String
   | skipped => ruleSkippedEmoji
 
 def runFirstSafeRule (gref : GoalRef) : SearchM Q SafeRulesResult := do
-  let g ← gref.get
-  if g.unsafeRulesSelected then
+  if (← gref.get).unsafeRulesSelected then
     return .skipped
     -- If the unsafe rules have been selected, we have already tried all the
     -- safe rules.
+  gref.updateForwardState .safe
+  let g ← gref.get
   let rules ← selectSafeRules g
   let mut postponedRules := {}
   for r in rules do
@@ -197,6 +199,7 @@ def applyPostponedSafeRule (r : PostponedSafeRule) (parentRef : GoalRef) :
 
 partial def runFirstUnsafeRule (postponedSafeRules : Array PostponedSafeRule)
     (parentRef : GoalRef) : SearchM Q RuleResult := do
+  parentRef.updateForwardState .unsafe
   let queue ← selectUnsafeRules postponedSafeRules parentRef
   let (remainingQueue, result) ← loop queue
   parentRef.modify λ g => g.setUnsafeQueue remainingQueue

Aesop/Search/Expansion.lean

@@ -41,7 +41,7 @@ def getForwardState : NormM ForwardState :=
 def getResetForwardState : NormM ForwardState := do
   modifyGetThe NormM.State λ s => (s.forwardState, { s with forwardState := ∅ })
 
-def updateForwardState (fs : ForwardState) (newMatches : Array ForwardRuleMatch)
+def modifyForwardState (fs : ForwardState) (newMatches : Array ForwardRuleMatch)
     (erasedHyps : Std.HashSet FVarId) : NormM Unit :=
   modifyThe NormM.State λ s => { s with
     forwardState := fs
@@ -50,13 +50,18 @@ def updateForwardState (fs : ForwardState) (newMatches : Array ForwardRuleMatch)
         (consumedForwardRuleMatches := #[]) -- We erase the consumed matches separately.
   }
 
+def updateForwardState (goal : MVarId) : NormM Unit := do
+  let fs ← getResetForwardState
+  let (fs, ms) ← fs.update goal (phase? := some .norm)
+  modifyForwardState fs ms ∅
+
 def eraseForwardRuleMatch (m : ForwardRuleMatch) : NormM Unit := do
   modifyThe NormM.State λ s => { s with forwardRuleMatches := s.forwardRuleMatches.erase m }
 
 def applyDiffToForwardState (diff : GoalDiff) : NormM Unit := do
   let fs ← getResetForwardState
-  let (fs, ms) ← fs.applyGoalDiff (← read).ruleSet diff
-  updateForwardState fs ms diff.removedFVars
+  let fs ← fs.applyGoalDiff (← read).ruleSet diff
+  modifyForwardState fs #[] diff.removedFVars
 
 inductive NormRuleResult
   | succeeded (goal : MVarId) (steps? : Option (Array Script.LazyStep))
@@ -97,7 +102,7 @@ returns the matches as well. -/
 def runNormRuleTac (rule : NormRule) (input : RuleTacInput) (fs : ForwardState)
     (rs : LocalRuleSet) :
     NormM $
-      Option (NormRuleResult × ForwardState × Array ForwardRuleMatch × Std.HashSet FVarId) ×
+      Option (NormRuleResult × ForwardState × Std.HashSet FVarId) ×
       Array ForwardRuleMatch := do
   let preMetaState ← show MetaM _ from saveState
   let result? ← runRuleTac rule.tac.run rule.name preMetaState input
@@ -111,18 +116,18 @@ def runNormRuleTac (rule : NormRule) (input : RuleTacInput) (fs : ForwardState)
       | err m!"rule did not produce exactly one rule application."
     show MetaM _ from restoreState rapp.postState
     if rapp.goals.isEmpty then
-      return (some (.proved rapp.scriptSteps?, fs, #[], ∅), forwardRuleMatches)
+      return (some (.proved rapp.scriptSteps?, fs, ∅), forwardRuleMatches)
     let (#[{ diff }]) := rapp.goals
       | err m!"rule produced more than one subgoal."
-    let (fs, ms) ← fs.applyGoalDiff rs diff
+    let fs ← fs.applyGoalDiff rs diff
     let g := diff.newGoal
     if ← Check.rules.isEnabled then
       let mvars := .ofArray input.mvars.toArray
       let actualMVars ← rapp.postState.runMetaM' g.getMVarDependencies
       if ! actualMVars == mvars then
          err "the goal produced by the rule depends on different metavariables than the original goal."
     let result := .succeeded g rapp.scriptSteps?
-    return (some (result, fs, ms, diff.removedFVars), forwardRuleMatches)
+    return (some (result, fs, diff.removedFVars), forwardRuleMatches)
   where
     err {α} (msg : MessageData) : MetaM α := throwError
       "aesop: error while running norm rule {rule.name}: {msg}\nThe rule was run on this goal:{indentD $ MessageData.ofGoal input.goal}"
@@ -142,9 +147,9 @@ def runNormRule (goal : MVarId) (mvars : UnorderedArraySet MVarId)
         runNormRuleTac rule.rule ruleInput fs (← read).ruleSet
       for m in consumedForwardRuleMatches do
         eraseForwardRuleMatch m
-      let (some (result, fs, ms, removedFVars)) := result?
+      let (some (result, fs, removedFVars)) := result?
         | return none
-      updateForwardState fs ms removedFVars
+      modifyForwardState fs #[] removedFVars
       return result
 
 def runFirstNormRule (goal : MVarId) (mvars : UnorderedArraySet MVarId)
@@ -323,6 +328,7 @@ def runNormSteps (goal : MVarId) (steps : Array NormStep)
   let mut anySuccess := false
   while iteration < maxIterations do
     if step.val == 0 then
+      updateForwardState goal
       let rules ←
         selectNormRules ctx.ruleSet (← getThe NormM.State).forwardRuleMatches
           goal

Aesop/Search/Expansion/Norm.lean

@@ -93,9 +93,9 @@ unsafe def copyGoals (assignedMVars : UnorderedArraySet MVarId)
       runInMetaState parentMetaState do
         let start ← start.get
         let diff ← diffGoals start.currentGoal g.preNormGoal
-        let (forwardState, ms) ← start.forwardState.applyGoalDiff rs diff
+        let forwardState ← start.forwardState.applyGoalDiff rs diff
         let forwardRuleMatches :=
-          start.forwardRuleMatches.update ms diff.removedFVars
+          start.forwardRuleMatches.update #[] diff.removedFVars
             (consumedForwardRuleMatches := #[]) -- TODO unsure whether this is correct
         let mvars ← .ofHashSet <$> g.preNormGoal.getMVarDependencies
         pure (forwardState, forwardRuleMatches, mvars)
@@ -129,9 +129,9 @@ def makeInitialGoal (goal : Subgoal) (mvars : UnorderedArraySet MVarId)
     (successProbability : Percent) (origin : GoalOrigin) : TreeM Goal := do
   let rs := (← read).ruleSet
   let (forwardState, forwardRuleMatches) ← runInMetaState parentMetaState do
-    let (fs, newMatches) ← parentForwardState.applyGoalDiff rs goal.diff
+    let fs ← parentForwardState.applyGoalDiff rs goal.diff
     let ms :=
-      parentForwardMatches.update newMatches goal.diff.removedFVars
+      parentForwardMatches.update #[] goal.diff.removedFVars
         consumedForwardRuleMatches
     pure (fs, ms)
   return Goal.mk {