Add options and persistent extensions to add more rule sets for scala…

…r_tac

backends/lean/Base/Arith.lean

@@ -1,2 +1,3 @@
 import Base.Arith.Int
 import Base.Arith.Scalar
+import Base.Arith.Lemmas

backends/lean/Base/Arith/Init.lean

@@ -1,11 +1,40 @@
+import Base.Extensions
 import Aesop
+open Lean
 
 /-!
-# ScalarTac rule set
+# Scalar tac rules sets
 
-This module defines the `ScalarTac` Aesop rule set which is used by the
+This module defines several Aesop rule sets and options which are used by the
 `scalar_tac` tactic. Aesop rule sets only become visible once the file in which
 they're declared is imported, so we must put this declaration into its own file.
 -/
 
-declare_aesop_rule_sets [Aeneas.ScalarTac]
+namespace Arith
+
+declare_aesop_rule_sets [Aeneas.ScalarTac, Aeneas.ScalarTacNonLin]
+
+#check Lean.Option.register
+register_option scalarTac.nonLin : Bool := {
+  defValue := false
+  group := ""
+  descr := "Activate the use of a set of lemmas to reason about non-linear arithmetic by `scalar_tac`"
+}
+
+-- The sets of rules that `scalar_tac` should use
+open Extensions in
+initialize scalarTacRuleSets : ListDeclarationExtension Name ← do
+  mkListDeclarationExtension `scalarTacRuleSetsList
+
+def scalarTacRuleSets.get : MetaM (List Name) := do
+  pure (scalarTacRuleSets.getState (← getEnv))
+
+-- Note that the changes are not persistent
+def scalarTacRuleSets.set (names : List Name) : MetaM Unit := do
+  let _ := scalarTacRuleSets.setState (← getEnv) names
+
+-- Note that the changes are not persistent
+def scalarTacRuleSets.add (name : Name) : MetaM Unit := do
+  let _ := scalarTacRuleSets.modifyState (← getEnv) (fun ls => name :: ls)
+
+end Arith

backends/lean/Base/Arith/Int.lean

@@ -25,6 +25,10 @@ macro "int_tac" pat:term : attr =>
 macro "scalar_tac" pat:term : attr =>
   `(attr|aesop safe forward (rule_sets := [$(Lean.mkIdent `Aeneas.ScalarTac):ident]) (pattern := $pat))
 
+/-- The `nonlin_scalar_tac` attribute used to tag forward theorems for the `int_tac` and `scalar_tac` tactics. -/
+macro "nonlin_scalar_tac" pat:term : attr =>
+  `(attr|aesop safe forward (rule_sets := [$(Lean.mkIdent `Aeneas.ScalarTacNonLin):ident]) (pattern := $pat))
+
 /- Check if a proposition is a linear integer proposition.
    We notably use this to check the goals: this is useful to filter goals that
    are unlikely to be solvable with arithmetic tactics. -/
@@ -70,7 +74,14 @@ def intTacSaturateForward : Tactic.TacticM Unit := do
   -- Use a forward max depth of 0 to prevent recursively applying forward rules on the assumptions
   -- introduced by the forward rules themselves.
   let options ← options.toOptions' (some 0)
-  evalAesopSaturate options #[`Aeneas.ScalarTac]
+  -- We always use the rule set `Aeneas.ScalarTac`, but also need to add other rule sets locally
+  -- activated by the user. The `Aeneas.ScalarTacNonLin` rule set has a special treatment as
+  -- it is activated through an option.
+  let ruleSets :=
+    let ruleSets := `Aeneas.ScalarTac :: (← scalarTacRuleSets.get)
+    if scalarTac.nonLin.get (← getOptions) then `Aeneas.ScalarTacNonLin :: ruleSets
+    else ruleSets
+  evalAesopSaturate options ruleSets.toArray
 
 /- Boosting a bit the `omega` tac.
  -/

backends/lean/Base/Arith/Lemmas.lean

@@ -0,0 +1,31 @@
+import Base.Arith.Int
+import Base.Arith.Scalar
+
+@[nonlin_scalar_tac n % m]
+theorem Int.emod_of_pos_disj (n m : Int) : m ≤ 0 ∨ (0 ≤ n % m ∧ n % m < m) := by
+  if h: 0 < m then
+    right; constructor
+    . apply Int.emod_nonneg; omega
+    . apply Int.emod_lt_of_pos; omega
+  else left; omega
+
+theorem Int.pos_mul_pos_is_pos (n m : Int) (hm : 0 ≤ m) (hn : 0 ≤ n): 0 ≤ m * n := by
+  have h : (0 : Int) = 0 * 0 := by simp
+  rw [h]
+  apply mul_le_mul <;> norm_cast
+
+@[nonlin_scalar_tac m * n]
+theorem Int.pos_mul_pos_is_pos_disj (n m : Int) : m < 0 ∨ n < 0 ∨ 0 ≤ m * n := by
+  cases h: (m < 0 : Bool) <;> simp_all
+  cases h: (n < 0 : Bool) <;> simp_all
+  right; right; apply pos_mul_pos_is_pos <;> tauto
+
+-- Some tests
+section
+
+  -- Activate the rule set for non linear arithmetic
+  set_option scalarTac.nonLin true
+
+  example (x y : Int) (h : 0 ≤ x ∧ 0 ≤ y) : 0 ≤ x * y := by scalar_tac
+
+end

backends/lean/Base/Extensions.lean

@@ -11,6 +11,20 @@ namespace Extensions
 open Lean Elab Term Meta
 open Utils
 
+def ListDeclarationExtension (α : Type) := SimplePersistentEnvExtension α (List α)
+
+instance : Inhabited (ListDeclarationExtension α) :=
+  inferInstanceAs (Inhabited (SimplePersistentEnvExtension ..))
+
+def mkListDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) :
+  IO (ListDeclarationExtension α) :=
+  registerSimplePersistentEnvExtension {
+    name          := name,
+    addImportedFn := fun entries => entries.foldl (fun s l => l.data ++ s) [],
+    addEntryFn    := fun l x => x :: l,
+    toArrayFn     := fun l => l.toArray
+  }
+
 -- This is not used anymore but we keep it here.
 -- TODO: the original function doesn't define correctly the `addImportedFn`. Do a PR?
 def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) :

tests/lean/Hashmap/Properties.lean

@@ -141,17 +141,11 @@ attribute [-simp] Bool.exists_bool
 
 attribute [local simp] List.lookup
 
-/- Adding some theorems for `scalar_tac`-/
-theorem Int.emod_of_pos (n m : Int) : m ≤ 0 ∨ (0 ≤ n % m ∧ n % m < m) := by
-  if h: 0 < m then
-    right; constructor
-    . apply Int.emod_nonneg; scalar_tac
-    . apply Int.emod_lt_of_pos; scalar_tac
-  else left; scalar_tac
+/- Adding some theorems for `scalar_tac` -/
+-- We first activate the rule set for non linear arithmetic - this is needed because of the modulo operations
+set_option scalarTac.nonLin true
 
-attribute [local scalar_tac n % m] Int.emod_of_pos
-
--- TODO: make this rule local
+-- Custom, local rule
 @[local scalar_tac h]
 theorem inv_imp_eqs_ineqs {hm : HashMap α} (h : hm.inv) : 0 < hm.slots.length ∧ hm.num_entries = hm.al_v.len := by
   simp_all [inv]