From c69e29f4114522f0343aee27ec5256aec374f2ef Mon Sep 17 00:00:00 2001
From: Jannis Limperg <jannis@limperg.de>
Date: Thu, 16 Jan 2025 00:22:50 +0100
Subject: [PATCH] Fix Stats collection for rules

I messed up the propagation of the Stats data structure in 7e00a02, so
stats for rules were effectively not collected.
---
 Aesop/BaseM.lean                 |  5 +++--
 Aesop/Saturate.lean              |  2 +-
 Aesop/Search/Expansion/Norm.lean |  1 -
 Aesop/Search/Main.lean           |  4 ++--
 Aesop/Search/SearchM.lean        | 11 ++++-------
 AesopTest/LegacyForward.lean     |  2 +-
 6 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/Aesop/BaseM.lean b/Aesop/BaseM.lean
index a2ce6e2..f9d5cbf 100644
--- a/Aesop/BaseM.lean
+++ b/Aesop/BaseM.lean
@@ -33,8 +33,9 @@ abbrev BaseM := StateRefT BaseM.State MetaM
 namespace BaseM
 
 /-- Run a `BaseM` action. -/
-protected def run (x : BaseM α) : MetaM α :=
-  StateRefT'.run' x ∅
+protected def run (x : BaseM α) (stats : Stats := ∅) : MetaM (α × Stats) := do
+  let (a, s) ← StateRefT'.run x { stats, rulePatternCache := ∅, rpinfCache := ∅ }
+  return (a, s.stats)
 
 instance : MonadHashMapCacheAdapter Expr RulePatternCache.Entry BaseM where
   getCache := return (← get).rulePatternCache.map
diff --git a/Aesop/Saturate.lean b/Aesop/Saturate.lean
index 739865e..7f6d24f 100644
--- a/Aesop/Saturate.lean
+++ b/Aesop/Saturate.lean
@@ -36,7 +36,7 @@ namespace SaturateM
 
 def run (options : Aesop.Options') (x : SaturateM α) :
     MetaM (α × Array Script.LazyStep) :=
-  ReaderT.run x { options } |>.run' {} |>.run.run
+  (·.fst) <$> (ReaderT.run x { options } |>.run' {} |>.run.run)
 
 end SaturateM
 
diff --git a/Aesop/Search/Expansion/Norm.lean b/Aesop/Search/Expansion/Norm.lean
index 4eb42a0..b9e154d 100644
--- a/Aesop/Search/Expansion/Norm.lean
+++ b/Aesop/Search/Expansion/Norm.lean
@@ -25,7 +25,6 @@ structure Context where
   options : Options'
   ruleSet : LocalRuleSet
   normSimpContext : NormSimpContext
-  statsRef : StatsRef
 
 structure State where
   forwardState : ForwardState
diff --git a/Aesop/Search/Main.lean b/Aesop/Search/Main.lean
index 32c4dc2..39f9a8a 100644
--- a/Aesop/Search/Main.lean
+++ b/Aesop/Search/Main.lean
@@ -278,8 +278,8 @@ def search (goal : MVarId) (ruleSet? : Option LocalRuleSet := none)
     show SearchM Q _ from
     try searchLoop
     finally freeTree
-  let (goals, _, _, stats) ←
-    go.run ruleSet options simpConfig simpConfigSyntax? goal stats |>.run
+  let ((goals, _, _), stats) ←
+    go.run ruleSet options simpConfig simpConfigSyntax? goal |>.run stats
   return (goals, stats)
 
 end Aesop
diff --git a/Aesop/Search/SearchM.lean b/Aesop/Search/SearchM.lean
index c3ccd69..05ad1b3 100644
--- a/Aesop/Search/SearchM.lean
+++ b/Aesop/Search/SearchM.lean
@@ -29,7 +29,6 @@ structure Context where
   ruleSet : LocalRuleSet
   normSimpContext : NormSimpContext
   options : Aesop.Options'
-  statsRef : StatsRef
   deriving Nonempty
 
 structure State (Q) [Aesop.Queue Q] where
@@ -74,15 +73,14 @@ instance : MonadLift TreeM (SearchM Q) where
     liftM $ ReaderT.run x ctx
 
 protected def run' (ctx : SearchM.Context) (σ : SearchM.State Q) (tree : Tree)
-    (x : SearchM Q α) : BaseM (α × SearchM.State Q × Tree × Stats) := do
+    (x : SearchM Q α) : BaseM (α × SearchM.State Q × Tree) := do
   let ((a, σ), t) ←
     x.run ctx |>.run σ |>.run { tree }
-  return (a, σ, t.tree, ← ctx.statsRef.get)
+  return (a, σ, t.tree)
 
 protected def run (ruleSet : LocalRuleSet) (options : Aesop.Options')
     (simpConfig : Simp.Config) (simpConfigStx? : Option Term)
-    (goal : MVarId) (stats : Stats) (x : SearchM Q α) :
-    BaseM (α × State Q × Tree × Stats) := do
+    (goal : MVarId) (x : SearchM Q α) : BaseM (α × State Q × Tree) := do
   let t ← mkInitialTree goal ruleSet
   let normSimpContext := {
     toContext := ← Simp.mkContext simpConfig (simpTheorems := ruleSet.simpTheoremsArray.map (·.snd))
@@ -92,8 +90,7 @@ protected def run (ruleSet : LocalRuleSet) (options : Aesop.Options')
     enabled := options.enableSimp
     useHyps := options.useSimpAll
   }
-  let statsRef ← IO.mkRef stats
-  let ctx := { ruleSet, options, normSimpContext, statsRef }
+  let ctx := { ruleSet, options, normSimpContext }
   let #[rootGoal] := (← t.root.get).goals
     | throwError "aesop: internal error: root mvar cluster does not contain exactly one goal."
   let state := {
diff --git a/AesopTest/LegacyForward.lean b/AesopTest/LegacyForward.lean
index f6f2a7a..368fb40 100644
--- a/AesopTest/LegacyForward.lean
+++ b/AesopTest/LegacyForward.lean
@@ -26,7 +26,7 @@ def forwardTac (goal : MVarId) (id : Ident) (immediate : Option (Array Syntax))
   let ldecl ← getLocalDeclFromUserName userName
   let immediate ← RuleBuilder.getImmediatePremises ldecl.type none
     (immediate.map (·.map (·.getId)))
-  let (goal, _) ←
+  let ((goal, _), _) ←
     RuleTac.applyForwardRule goal (mkFVar ldecl.fvarId) none immediate clear
       (maxDepth? := none) ∅ |>.run.run
   return [goal.mvarId]