perf: internalize matched spec patterns during spec lookup (#14142)

This PR speeds up `mvcgen'` spec lookup by internalizing each matched
spec pattern into the `SymM` share table on first lookup, so its
instance arguments become pointer-equal to the program's and need not be
re-internalized on every later lookup.

`findSpecs` selects which spec applies by matching the program against
the candidate patterns held in the discrimination tree. Those patterns
are built at `@[spec]` registration and not internalized into the
`mvcgen'`-ambient `SymM` share table, so their instance arguments, the
`MonadStateOf`/`MonadReaderOf`/`MonadExceptOf` instances of
`get`/`set`/`read`/`throw` and the like, are structurally equal to the
program's but pointer-distinct, and the matcher re-internalizes them on
each match.

The first time a spec is tried, `findSpecs` internalizes its pattern and
stores it back in the database. The internalized pattern is structurally
unchanged, so it keys the discrimination tree at the same place and
replaces the old entry in place. Later lookups detect the internalized
pattern in O(1) with `Sym.isShared` (a membership query on the share
table) and skip the work. The updated database is threaded through the
`Scope`, which `solve` already carries.

`Sym.isShared` and `Pattern.shareCommon` are defined in the
spec-database module so the spec database and the backward-rule caches
share one definition.

Benchmark wall-clock for the `mvcgen'` step (`tests/bench/mvcgen/sym`,
single-threaded, min of 5):

| case | before | after |
| --- | --- | --- |
| ReaderState(1000) | 197 ms | 171 ms |
| AddSubCancelDeep(700) | 178 ms | 144 ms |
| GetThrowSet(600) | 126 ms | 119 ms |
| MatchIota(150) | 36 ms | 33 ms |
| AddSubCancel(1000) | 179 ms | 176 ms |

Author: sgraf812

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

@@ -294,7 +294,13 @@ and apply its cached backward rule. -/
 private def applySpec (scope : VCGen.Scope) (goal : MVarId) (target : Expr) (info : WPInfo) :
     VCGenM SolveResult := do
   trace[Elab.Tactic.Do.vcgen] "Applying a spec for {info.prog}. Excess args: {info.excessArgs}"
-  match ← SpecTheorems.findSpecs scope.specs info.prog with
+  -- Hand `findSpecs` the sole reference to the database so its in-place pattern internalization
+  -- does not copy the discrimination tree, then thread the updated database back into the scope.
+  let specs := scope.specs
+  let scope := { scope with specs := default }
+  let (result, specs) ← SpecTheorems.findSpecs specs info.prog
+  let scope := { scope with specs }
+  match result with
   | .error thms => stopOrErrorOnMissingSpec info.prog info.m thms
   | .ok thm =>
   trace[Elab.Tactic.Do.vcgen] "Spec for {info.prog}: {thm.proof}"

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

@@ -26,6 +26,24 @@ namespace Lean.Elab.Tactic.Do.Internal
 
 open SpecAttr
 
+/-- Returns `true` if `e` is already internalized into the current `SymM` share table, in which case
+`shareCommon e` returns `e` unchanged. -/
+public def _root_.Lean.Meta.Sym.isShared (e : Expr) : SymM Bool :=
+  return (← get).share.set.contains { expr := e }
+
+/--
+Internalizes the pattern's expressions into the current `SymM` share table.
+
+A pattern is built in `MetaM`, outside the `SymM` thread whose table the targets are internalized
+into, so its closed subterms (the instance telescope of a `wp` application, for example) are
+structurally equal to but distinct from the targets'. Internalizing the pattern once makes those
+subterms pointer-equal to every target internalized afterwards, so matching them is `O(1)`.
+-/
+public def _root_.Lean.Meta.Sym.Pattern.shareCommon (p : Pattern) : SymM Pattern := do
+  return { p with
+    pattern := ← Sym.shareCommon p.pattern
+    varTypes := ← p.varTypes.mapM Sym.shareCommon }
+
 /--
 Instantiates a spec theorem's proof.
 
@@ -110,19 +128,30 @@ Look up `SpecTheorem`s in the `@[spec]` database.
 Takes all specs that match the given program `e` and sorts by descending priority.
 -/
 public def SpecAttr.SpecTheorems.findSpecs (database : SpecTheorems) (e : Expr) :
-    SymM (Except (Array SpecTheorem) SpecTheorem) := do
+    SymM (Except (Array SpecTheorem) SpecTheorem × SpecTheorems) := do
   let e ← instantiateMVars e
   let e ← shareCommon e
   let candidates := Sym.getMatch database.specs e
   let candidates := candidates.filter fun spec => !database.erased.contains spec.proof
   if h : candidates.size = 1 then
     have : 0 < candidates.size := h ▸ Nat.zero_lt_one
-    return .ok candidates[0]
+    return (.ok candidates[0], database)
   -- It appears that insertion sort is *much* faster than qsort here.
   let candidates := candidates.insertionSort (·.priority > ·.priority)
+  let mut database := database
   for spec in candidates do
+    -- Match against the internalized pattern so its instance arguments are pointer-equal to the
+    -- program's. A spec is internalized the first time it is tried and stored back in the database,
+    -- so later lookups find it already in the share table and skip the work.
+    let mut spec := spec
+    unless ← isShared spec.pattern.pattern do
+      spec := { spec with pattern := ← spec.pattern.shareCommon }
+      -- Take sole ownership of the discrimination tree before inserting so the update is in place.
+      let specs := database.specs
+      database := { database with specs := default }
+      database := { database with specs := Sym.insertPattern specs spec.pattern spec }
     let some _res ← spec.pattern.match? e | continue
-    return .ok spec
-  return .error candidates
+    return (.ok spec, database)
+  return (.error candidates, database)
 
 end Lean.Elab.Tactic.Do.Internal

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

@@ -26,22 +26,6 @@ decomposes.
 
 namespace Lean.Elab.Tactic.Do.Internal
 
-/--
-Internalize a pattern's expressions into the current `SymM` share table.
-
-A pattern is built in `MetaM`, outside the `SymM` thread whose table holds the targets it matches
-against, so its closed subterms (the instance telescope of a `wp` application, for example) are
-equal to but distinct from the targets'. Internalizing the pattern once makes those subterms
-pointer-equal to every target internalized afterwards, so matching them no longer re-internalizes
-the same constant terms on every application.
-
-Designed for dot notation: `pattern.shareCommon`. Requires `open Lean.Elab.Tactic.Do.Internal`.
--/
-public def _root_.Lean.Meta.Sym.Pattern.shareCommon (p : Pattern) : SymM Pattern := do
-  return { p with
-    pattern := ← Sym.shareCommon p.pattern
-    varTypes := ← p.varTypes.mapM Sym.shareCommon }
-
 /-- Internalize a backward rule's pattern into the current `SymM` share table. See
 `Pattern.shareCommon`. Designed for dot notation: `rule.shareCommon`. -/
 public def _root_.Lean.Meta.Sym.BackwardRule.shareCommon (rule : BackwardRule) : SymM BackwardRule :=