feat: change AxEffects.toSimpTheorems to generate more specific dis…

…crtree keys (#148)

### Description:

The method we were using to construct simp-lemmas sets a `noIndexAtArgs`
option, meaning that a lemma like
`r (.GPR 0#5) s10` would get indexed *only* on `r`, and the arguments
were ignored. This caused a lot of attempts at unification.

Ideally, we'd want `mkSimpTheorems` to expose `noIndexAtArgs` as an
option, so we could disable it.
For now, we workaround the issue with a postprocessing step that
regenerates the discrtree keys with this behaviour turned off. Sadly,
the failed unifications turn out to not be the main expense: with this
change we can simulate SHA512 for 190 steps, which is a relatively small
improvement over before.

Then, I set the priority of the field-specfic lemmas higher than the
default priority used for the nonEffectProof, to prevent simp from
trying to discharge the expensive side-conditions when we already know a
cheaper lemma that is applicable. This did give an impressive speedup.

All together, we are now able to simulate 250 steps of SHA512 with the
default heartbeat limit (when tracing is turned off! Tracing does
consume some heartbeats, so with tracing the limit is lower)

### Testing:

`make all` succeeded locally

### License:

By submitting this pull request, I confirm that my contribution is
made under the terms of the Apache 2.0 license.

---------

Co-authored-by: Shilpi Goel <[email protected]>

Arm/MinTheory.lean

@@ -8,35 +8,59 @@ attribute [minimal_theory] ite_false
 attribute [minimal_theory] dite_true
 attribute [minimal_theory] dite_false
 attribute [minimal_theory] ite_self
-attribute [minimal_theory] and_true
-attribute [minimal_theory] true_and
-attribute [minimal_theory] and_false
-attribute [minimal_theory] false_and
+
+/-
+Notice how both `and_true : ?p ∧ True = ?p` and `and_self : ?p ∧ ?p = ?p` may
+be attempted by `simp` on a goal of the shape `_ ∧ True`, as they are both
+a match where the discrimination tree is concerned.
+
+However, assuming the left conjunct is not def-eq to `True`, an attempt of
+`and_self` will fail to unify, which causes a fallback to `and_true`.
+For the latter simp lemma, on the other hand, we know that if the discrimination
+tree gives a match, then the lemma should be applicable.
+This is because the variable `?p` is used only once in the pattern (i.e.,
+the pattern is linear), and the first unification of an unassigned metavariable
+is always successful.
+
+Another possibly expensive lemma is
+  `and_iff_left_of_imp : {a b : Prop} (ha : a → b) : a ∧ b ↔ a`
+The pattern of this lemma (`?a ∧ ?b`) is perfectly linear, but to apply such
+a lemma, `simp` first has to discharge the `(ha : a → b)` side-condition,
+which might fail and might be expensive.
+
+Thus, the rationale we follow is that only linear, side-condition free lemmas
+get the `high` priority, and everything else gets the default prioriry.
+This ensures that `and_true` gets tried before `and_self` or
+`and_iff_left_of_imp` -/
+attribute [minimal_theory high] and_true
+attribute [minimal_theory high] true_and
+attribute [minimal_theory high] and_false
+attribute [minimal_theory high] false_and
 attribute [minimal_theory] and_self
 attribute [minimal_theory] and_not_self
 attribute [minimal_theory] not_and_self
 attribute [minimal_theory] and_imp
-attribute [minimal_theory] not_and
+attribute [minimal_theory high] not_and
 attribute [minimal_theory] or_self
-attribute [minimal_theory] or_true
-attribute [minimal_theory] true_or
-attribute [minimal_theory] or_false
-attribute [minimal_theory] false_or
-attribute [minimal_theory] if_true_left
-attribute [minimal_theory] if_true_right
-attribute [minimal_theory] if_false_left
-attribute [minimal_theory] if_false_right
+attribute [minimal_theory high] or_true
+attribute [minimal_theory high] true_or
+attribute [minimal_theory high] or_false
+attribute [minimal_theory high] false_or
+attribute [minimal_theory high] if_true_left
+attribute [minimal_theory high] if_true_right
+attribute [minimal_theory high] if_false_left
+attribute [minimal_theory high] if_false_right
 attribute [minimal_theory] iff_self
-attribute [minimal_theory] iff_true
-attribute [minimal_theory] true_iff
-attribute [minimal_theory] iff_false
-attribute [minimal_theory] false_iff
-attribute [minimal_theory] eq_iff_iff
-attribute [minimal_theory] false_implies
-attribute [minimal_theory] implies_true
-attribute [minimal_theory] true_implies
-attribute [minimal_theory] not_false_eq_true
-attribute [minimal_theory] not_true_eq_false
+attribute [minimal_theory high] iff_true
+attribute [minimal_theory high] true_iff
+attribute [minimal_theory high] iff_false
+attribute [minimal_theory high] false_iff
+attribute [minimal_theory high] eq_iff_iff
+attribute [minimal_theory high] false_implies
+attribute [minimal_theory high] implies_true
+attribute [minimal_theory high] true_implies
+attribute [minimal_theory high] not_false_eq_true
+attribute [minimal_theory high] not_true_eq_false
 attribute [minimal_theory] not_iff_self
 attribute [minimal_theory] and_self_left
 attribute [minimal_theory] and_self_right
@@ -52,56 +76,61 @@ attribute [minimal_theory] or_iff_right_of_imp
 attribute [minimal_theory] or_iff_left_of_imp
 attribute [minimal_theory] or_iff_left_iff_imp
 attribute [minimal_theory] or_iff_right_iff_imp
-attribute [minimal_theory] Bool.or_false
-attribute [minimal_theory] Bool.or_true
-attribute [minimal_theory] Bool.false_or
-attribute [minimal_theory] Bool.false_eq_true
-attribute [minimal_theory] Bool.true_or
+
+attribute [minimal_theory high] Bool.or_false
+attribute [minimal_theory high] Bool.or_true
+attribute [minimal_theory high] Bool.false_or
+attribute [minimal_theory high] Bool.false_eq_true
+attribute [minimal_theory high] Bool.true_or
 attribute [minimal_theory] Bool.or_self
-attribute [minimal_theory] Bool.or_eq_true
-attribute [minimal_theory] Bool.and_false
-attribute [minimal_theory] Bool.and_true
-attribute [minimal_theory] Bool.false_and
-attribute [minimal_theory] Bool.true_and
+attribute [minimal_theory high] Bool.or_eq_true
+attribute [minimal_theory high] Bool.and_false
+attribute [minimal_theory high] Bool.and_true
+attribute [minimal_theory high] Bool.false_and
+attribute [minimal_theory high] Bool.true_and
 attribute [minimal_theory] Bool.and_self
-attribute [minimal_theory] Bool.and_eq_true
-attribute [minimal_theory] Bool.not_not
-attribute [minimal_theory] Bool.not_true
-attribute [minimal_theory] Bool.not_false
-attribute [minimal_theory] beq_true
-attribute [minimal_theory] beq_false
-attribute [minimal_theory] Bool.not_eq_true'
-attribute [minimal_theory] Bool.not_eq_false'
-attribute [minimal_theory] Bool.beq_to_eq
-attribute [minimal_theory] Bool.not_beq_to_not_eq
-attribute [minimal_theory] Bool.not_eq_true
-attribute [minimal_theory] Bool.not_eq_false
-attribute [minimal_theory] decide_eq_true_eq
-attribute [minimal_theory] decide_not
-attribute [minimal_theory] not_decide_eq_true
+attribute [minimal_theory high] Bool.and_eq_true
+attribute [minimal_theory high] Bool.not_not
+attribute [minimal_theory high] Bool.not_true
+attribute [minimal_theory high] Bool.not_false
+attribute [minimal_theory high] beq_true
+attribute [minimal_theory high] beq_false
+attribute [minimal_theory high] Bool.not_eq_true'
+attribute [minimal_theory high] Bool.not_eq_false'
+attribute [minimal_theory high] Bool.beq_to_eq
+attribute [minimal_theory high] Bool.not_beq_to_not_eq
+attribute [minimal_theory high] Bool.not_eq_true
+attribute [minimal_theory high] Bool.not_eq_false
+attribute [minimal_theory high] decide_eq_true_eq
+attribute [minimal_theory high] decide_not
+attribute [minimal_theory high] not_decide_eq_true
+
+-- NOTE: `heq_eq_eq` might look linear, but if we consider implicit variables,
+-- the pattern is `@HEq ?α ?a ?α ?b`; `?α` is used non-linearly
 attribute [minimal_theory] heq_eq_eq
-attribute [minimal_theory] cond_true
-attribute [minimal_theory] cond_false
+
+attribute [minimal_theory high] cond_true
+attribute [minimal_theory high] cond_false
 attribute [minimal_theory] beq_self_eq_true
 attribute [minimal_theory] beq_self_eq_true'
 attribute [minimal_theory] bne_self_eq_false
 attribute [minimal_theory] bne_self_eq_false'
-attribute [minimal_theory] decide_False
-attribute [minimal_theory] decide_True
-attribute [minimal_theory] decide_eq_false_iff_not
-attribute [minimal_theory] decide_eq_true_iff
-attribute [minimal_theory] bne_iff_ne
-attribute [minimal_theory] Bool.false_eq
-attribute [minimal_theory] Bool.and_eq_false_imp
+attribute [minimal_theory high] decide_False
+attribute [minimal_theory high] decide_True
+attribute [minimal_theory high] decide_eq_false_iff_not
+attribute [minimal_theory high] decide_eq_true_iff
+attribute [minimal_theory high] bne_iff_ne
+attribute [minimal_theory high] Bool.false_eq
+attribute [minimal_theory high] Bool.and_eq_false_imp
 
-attribute [minimal_theory] Decidable.not_not
+attribute [minimal_theory high] Decidable.not_not
 
-attribute [minimal_theory] Nat.le_zero_eq
-attribute [minimal_theory] Nat.zero_add
-attribute [minimal_theory] Nat.zero_eq
-attribute [minimal_theory] Nat.succ.injEq
-attribute [minimal_theory] Nat.succ_ne_zero
-attribute [minimal_theory] Nat.sub_zero
+attribute [minimal_theory high] Nat.le_zero_eq
+attribute [minimal_theory high] Nat.zero_add
+attribute [minimal_theory high] Nat.zero_eq
+attribute [minimal_theory high] Nat.succ.injEq
+attribute [minimal_theory high] Nat.succ_ne_zero
+attribute [minimal_theory high] Nat.sub_zero
 
 attribute [minimal_theory] Nat.le_refl
 
@@ -111,8 +140,8 @@ theorem option_get_bang_of_some [Inhabited α] (v : α) :
 attribute [minimal_theory] Option.isNone_some
 
 attribute [minimal_theory] Fin.isValue
-attribute [minimal_theory] Fin.zero_eta
-attribute [minimal_theory] Fin.mk.injEq
+attribute [minimal_theory high] Fin.zero_eta
+attribute [minimal_theory high] Fin.mk.injEq
 
 -- attribute [minimal_theory] ↓reduceIte
 attribute [minimal_theory] reduceCtorEq

Tactics/Reflect/AxEffects.lean

@@ -645,6 +645,7 @@ where
         let ⟨fvar, goal⟩ ← goal.note h v t?
         return ⟨fvar, goal⟩
 
+open Meta.DiscrTree in
 /-- Return a `SimpTheorems` with the proofs contained in the given `AxEffects`
 
 Note this adds *only* the (non-)effect proofs, to get a simp configuration with
@@ -661,19 +662,53 @@ def toSimpTheorems (eff : AxEffects) : MetaM SimpTheorems := do
         none
     let baseName := baseName?.getD (Name.mkSimple "AxEffects")
 
-    let add (thms : SimpTheorems) (e : Expr) (name : String) := do
-      trace[Tactic.sym] "adding {e} with name {name}"
-      let origin : Origin :=
-        if e.isFVar then
-          .fvar e.fvarId!
-        else
-          .other <| Name.str baseName name
-      thms.add origin #[] e
+    let add (thms : SimpTheorems) (e : Expr) (name : String)
+        (prio : Nat := 1000) :=
+      let msg := m!"adding {e} with name {name}"
+      withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+        let origin : Origin :=
+          if e.isFVar then
+            .fvar e.fvarId!
+          else
+            .other <| Name.str baseName name
+        let newThms ← mkSimpTheorems origin #[] e (prio := prio)
+        let newThms ← newThms.mapM (fun thm => do
+          /- `mkSimpTheorems` sets `noIndexAtArgs := true`, meaning that all
+          our (non-effects) theorems will have `[r, *, *]` as key.
+          causing many unneeded attempts at unification.
+          The following code fixes up the keys, so that for example
+            `h_x0_s10 : r (.GPR 0#5) s10`
+          will get a key that includes the state and statefield constant
+
+          FIXME: we should suggest a PR upstream that adds `noIndexAtArgs` as
+          an optional argument to `mkSimpTheorems`, so that we can control
+          this properly -/
+          let type ← inferType thm.proof
+          let ⟨_, _, type⟩ ← forallMetaTelescope type
+          match type.eq? with
+            | none =>
+                trace[Tactic.sym] "{Lean.crossEmoji} {type} is not an equality\
+                  , giving up on fixing the discrtree keys.
+
+                  Currently, the keys are:\n{thm.keys}"
+                pure thm
+            | some (_eqType, lhs, _rhs) =>
+                let keys ← mkPath lhs simpDtConfig (noIndexAtArgs := false)
+                pure { thm with keys }
+        )
+
+        pure <| newThms.foldl addSimpTheoremEntry thms
 
     let mut thms : SimpTheorems := {}
 
     for ⟨field, {proof, ..}⟩ in eff.fields do
-      thms ← add thms proof s!"field_{field}"
+      /- We give the field-specific lemmas a high priority, since their
+      applicability is determined entirely by discrtree matching.
+      This is important for performance, because it avoids unneccesary
+      (expensive!) attempts to discharge the `field ≠ otherField`
+      side-conditions of the non-effect proof -/
+      thms ← add thms proof s!"field_{field}" (prio := 1500)
+
 
     thms ← add thms eff.nonEffectProof "nonEffectProof"
     thms ← add thms eff.memoryEffectProof "memoryEffectProof"