Improve `scalar_tac` and `scalar_decr_tac` (#256)

* Fix an issue in a proof of the hashmap

* Improve scalar_decr_tac

* Improve the error message of scalar_tac and add the missing Termination.lean

7 files changed, 155 insertions, 50 deletions

diff --git a/backends/lean/Base.lean b/backends/lean/Base.lean
index 2077d410..53baae1e 100644
--- a/backends/lean/Base.lean
+++ b/backends/lean/Base.lean
@@ -1,6 +1,7 @@
-import Base.Utils
-import Base.Primitives
-import Base.Diverge
 import Base.Arith
-import Base.Progress
+import Base.Diverge
 import Base.IList
+import Base.Primitives
+import Base.Progress
+import Base.Utils
+import Base.Termination
diff --git a/backends/lean/Base/Arith/Int.lean b/backends/lean/Base/Arith/Int.lean
index b1927cfd..958e31c9 100644
--- a/backends/lean/Base/Arith/Int.lean
+++ b/backends/lean/Base/Arith/Int.lean
@@ -106,19 +106,26 @@ def collectInstancesFromMainCtx (k : Expr → MetaM (Option Expr)) : Tactic.Tact
   Tactic.withMainContext do
   -- Get the local context
   let ctx ← Lean.MonadLCtx.getLCtx
-  -- Just a matter of precaution
-  let ctx ← instantiateLCtxMVars ctx
   -- Initialize the hashset
   let hs := HashSet.empty
   -- Explore the declarations
   let decls ← ctx.getDecls
   let hs ← decls.foldlM (fun hs d => do
     -- Collect instances over all subexpressions in the context.
-    -- Note that we explore the *type* of the local declarations: if we have
+    -- Note that if the local declaration is
+    -- Note that we explore the *type* of propositions: if we have
     -- for instance `h : A ∧ B` in the context, the expression itself is simply
     -- `h`; the information we are interested in is its type.
-    let ty ← Lean.Meta.inferType d.toExpr
-    collectInstances k hs ty
+    -- However, if the decl is not a proposition, we explore it directly.
+    -- For instance: `x : U32`
+    -- TODO: case disjunction on whether the local decl is a Prop or not. If prop,
+    -- we need to explore its type.
+    let d := d.toExpr
+    if d.isProp then
+      collectInstances k hs d
+    else
+      let ty ← Lean.Meta.inferType d
+      collectInstances k hs ty
   ) hs
   -- Also explore the goal
   collectInstances k hs (← Tactic.getMainTarget)
@@ -296,7 +303,7 @@ def intTac (tacName : String) (splitGoalConjs : Bool) (extraPreprocess :  Tactic
     try do Tactic.Omega.omegaTactic {}
     catch _ =>
       let g ← Tactic.getMainGoal
-      throwError "{tacName} failed to prove the goal:\n{g}"
+      throwError "{tacName} failed to prove the goal below.\n\nNote that {tacName} is equivalent to:\n  {tacName}_preprocess; omega\n\nGoal: \n{g}"
 
 elab "int_tac" args:(" split_goal"?): tactic =>
   let split := args.raw.getArgs.size > 0
diff --git a/backends/lean/Base/Arith/Scalar.lean b/backends/lean/Base/Arith/Scalar.lean
index 31110b95..a36aadf3 100644
--- a/backends/lean/Base/Arith/Scalar.lean
+++ b/backends/lean/Base/Arith/Scalar.lean
@@ -44,17 +44,6 @@ def scalarTac (splitGoalConjs : Bool) : Tactic.TacticM Unit := do
 elab "scalar_tac" : tactic =>
   scalarTac false
 
--- For termination proofs
-syntax "scalar_decr_tac" : tactic
-macro_rules
-  | `(tactic| scalar_decr_tac) =>
-    `(tactic|
-      simp_wf;
-      -- TODO: don't use a macro (namespace problems)
-      (first | apply Arith.to_int_to_nat_lt
-             | apply Arith.to_int_sub_to_nat_lt) <;>
-      simp_all <;> scalar_tac)
-
 instance (ty : ScalarTy) : HasIntProp (Scalar ty) where
   -- prop_ty is inferred
   prop := λ x => And.intro x.hmin x.hmax
diff --git a/backends/lean/Base/IList/IList.lean b/backends/lean/Base/IList/IList.lean
index ab71daed..c77f075f 100644
--- a/backends/lean/Base/IList/IList.lean
+++ b/backends/lean/Base/IList/IList.lean
@@ -4,8 +4,29 @@
 import Base.Arith
 import Base.Utils
 
+-- TODO: move?
+-- This lemma is generally useful. It often happens that (because we
+-- make a split on a condition for instance) we have `x ≠ y` in the context
+-- and need to simplify `y ≠ x` somewhere.
+@[simp]
+theorem neq_imp {α : Type u} {x y : α} (h : ¬ x = y) : ¬ y = x := by intro; simp_all
+
 namespace List
 
+-- Small helper
+-- We cover a set of cases which might imply inequality, to make sure that using
+-- this as the precondition of a `simp` lemma will allow the lemma to get correctly
+-- triggered.
+-- TODO: there should be something more systematic to do, with discharged procedures
+-- or simprocs I guess.
+@[simp]
+abbrev Int.not_eq (i j : Int) : Prop :=
+  i ≠ j ∨ j ≠ i ∨ i < j ∨ j < i
+
+theorem Int.not_eq_imp_not_eq {i j} : Int.not_eq i j → i ≠ j := by
+  intro h g
+  simp_all
+
 def len (ls : List α) : Int :=
   match ls with
   | [] => 0
@@ -32,7 +53,7 @@ def indexOpt (ls : List α) (i : Int) : Option α :=
 
 @[simp] theorem indexOpt_nil : indexOpt ([] : List α) i = none := by simp [indexOpt]
 @[simp] theorem indexOpt_zero_cons : indexOpt ((x :: tl) : List α) 0 = some x := by simp [indexOpt]
-@[simp] theorem indexOpt_nzero_cons (hne : i ≠ 0) : indexOpt ((x :: tl) : List α) i = indexOpt tl (i - 1) := by simp [*, indexOpt]
+@[simp] theorem indexOpt_nzero_cons (hne : Int.not_eq i 0) : indexOpt ((x :: tl) : List α) i = indexOpt tl (i - 1) := by simp [indexOpt]; intro; simp_all
 
 -- Remark: if i < 0, then the result is the default element
 def index [Inhabited α] (ls : List α) (i : Int) : α :=
@@ -42,10 +63,7 @@ def index [Inhabited α] (ls : List α) (i : Int) : α :=
     if i = 0 then x else index tl (i - 1)
 
 @[simp] theorem index_zero_cons [Inhabited α] : index ((x :: tl) : List α) 0 = x := by simp [index]
-@[simp] theorem index_nzero_cons [Inhabited α] (hne : i ≠ 0) : index ((x :: tl) : List α) i = index tl (i - 1) := by simp [*, index]
-@[simp] theorem index_zero_lt_cons [Inhabited α] (hne : 0 < i) : index ((x :: tl) : List α) i = index tl (i - 1) := by
-  have : i ≠ 0 := by scalar_tac
-  simp [*, index]
+@[simp] theorem index_nzero_cons [Inhabited α] (hne : Int.not_eq i 0) : index ((x :: tl) : List α) i = index tl (i - 1) := by simp [index]; intro; simp_all
 
 theorem indexOpt_bounds (ls : List α) (i : Int) :
   ls.indexOpt i = none ↔ i < 0 ∨ ls.len ≤ i :=
@@ -128,15 +146,15 @@ decreasing_by int_decr_tac
 
 @[simp] theorem update_nil : update ([] : List α) i y = [] := by simp [update]
 @[simp] theorem update_zero_cons : update ((x :: tl) : List α) 0 y = y :: tl := by simp [update]
-@[simp] theorem update_nzero_cons (hne : i ≠ 0) : update ((x :: tl) : List α) i y = x :: update tl (i - 1) y := by simp [*, update]
+@[simp] theorem update_nzero_cons (hne : Int.not_eq i 0) : update ((x :: tl) : List α) i y = x :: update tl (i - 1) y := by simp [update]; intro; simp_all
 
 @[simp] theorem idrop_nil : idrop i ([] : List α) = [] := by simp [idrop]
 @[simp] theorem idrop_zero : idrop 0 (ls : List α) = ls := by cases ls <;> simp [idrop]
-@[simp] theorem idrop_nzero_cons (hne : i ≠ 0) : idrop i ((x :: tl) : List α) = idrop (i - 1) tl := by simp [*, idrop]
+@[simp] theorem idrop_nzero_cons (hne : Int.not_eq i 0) : idrop i ((x :: tl) : List α) = idrop (i - 1) tl := by simp [idrop]; intro; simp_all
 
 @[simp] theorem itake_nil : itake i ([] : List α) = [] := by simp [itake]
 @[simp] theorem itake_zero : itake 0 (ls : List α) = [] := by cases ls <;> simp [itake]
-@[simp] theorem itake_nzero_cons (hne : i ≠ 0) : itake i ((x :: tl) : List α) = x :: itake (i - 1) tl := by simp [*, itake]
+@[simp] theorem itake_nzero_cons (hne : Int.not_eq i 0) : itake i ((x :: tl) : List α) = x :: itake (i - 1) tl := by simp [itake]; intro; simp_all
 
 @[simp] theorem slice_nil : slice i j ([] : List α) = [] := by simp [slice]
 @[simp] theorem slice_zero : slice 0 0 (ls : List α) = [] := by cases ls <;> simp [slice]
@@ -146,20 +164,21 @@ decreasing_by int_decr_tac
   rw [ireplicate]; simp [*]
 
 @[simp]
-theorem slice_nzero_cons (i j : Int) (x : α) (tl : List α) (hne : i ≠ 0) : slice i j ((x :: tl) : List α) = slice (i - 1) (j - 1) tl :=
+theorem slice_nzero_cons (i j : Int) (x : α) (tl : List α) (hne : Int.not_eq i 0) :
+  slice i j ((x :: tl) : List α) = slice (i - 1) (j - 1) tl := by
+  apply Int.not_eq_imp_not_eq at hne
   match tl with
-  | [] => by simp [slice]; simp [*]
+  | [] => simp [slice]; simp [*]
   | hd :: tl =>
-    if h: i - 1 = 0 then by
+    if h: i - 1 = 0 then
       have : i = 1 := by int_tac
       simp [*, slice]
     else
-      have hi := slice_nzero_cons (i - 1) (j - 1) hd tl h
-      by
-        conv => lhs; simp [slice, *]
-        conv at hi => lhs; simp [slice, *]
-        simp [slice]
-        simp [*]
+      have hi := slice_nzero_cons (i - 1) (j - 1) hd tl (by simp_all)
+      conv => lhs; simp [slice, *]
+      conv at hi => lhs; simp [slice, *]
+      simp [slice]
+      simp [*]
 
 @[simp]
 theorem ireplicate_replicate {α : Type u} (l : ℤ) (x : α) (h : 0 ≤ l) :
@@ -319,7 +338,8 @@ theorem itake_len_le (i : Int) (ls : List α) : (ls.itake i).len ≤ ls.len :=
       by simp [*]
 
 @[simp]
-theorem itake_len (i : Int) (ls : List α) (_ : 0 ≤ i) (_ : i ≤ ls.len) : (ls.itake i).len = i :=
+theorem itake_len (i : Int) (ls : List α) (_ : 0 ≤ i) (_ : i ≤ ls.len) :
+  (ls.itake i).len = i :=
   match ls with
   | [] => by simp_all; int_tac
   | hd :: tl =>
@@ -359,7 +379,8 @@ theorem index_itake [Inhabited α] (i : Int) (j : Int) (ls : List α)
   | [] => by simp at *
   | hd :: tl =>
     have : ¬ 0 = i := by int_tac -- TODO: this is slightly annoying
-    if h: j = 0 then by simp [*] at *
+    if h: j = 0 then by
+      simp_all
     else by
       simp [*]
       -- TODO: rewriting rule: ¬ i = 0 → 0 ≤ i → 0 < i ?
@@ -422,7 +443,7 @@ theorem index_itake_append_end [Inhabited α] (i j : Int) (l0 l1 : List α)
 @[simp]
 theorem index_update_ne
   {α : Type u} [Inhabited α] (l: List α) (i: ℤ) (j: ℤ) (x: α) :
-   j ≠ i → (l.update i x).index j = l.index j
+   Int.not_eq i j → (l.update i x).index j = l.index j
   :=
   λ _ => match l with
   | [] => by simp at *
diff --git a/backends/lean/Base/Primitives/Vec.lean b/backends/lean/Base/Primitives/Vec.lean
index e584777a..12789fa9 100644
--- a/backends/lean/Base/Primitives/Vec.lean
+++ b/backends/lean/Base/Primitives/Vec.lean
@@ -48,10 +48,7 @@ abbrev Vec.len (α : Type u) (v : Vec α) : Usize :=
 theorem Vec.len_val {α : Type u} (v : Vec α) : (Vec.len α v).val = v.length :=
   by rfl
 
--- This shouldn't be used
-def Vec.push_fwd (α : Type u) (_ : Vec α) (_ : α) : Unit := ()
-
--- This is actually the backward function
+@[irreducible]
 def Vec.push (α : Type u) (v : Vec α) (x : α) : Result (Vec α)
   :=
   let nlen := List.length v.val + 1
@@ -68,7 +65,7 @@ def Vec.push (α : Type u) (v : Vec α) (x : α) : Result (Vec α)
 theorem Vec.push_spec {α : Type u} (v : Vec α) (x : α) (h : v.val.len < Usize.max) :
   ∃ v1, v.push α x = ok v1 ∧
   v1.val = v.val ++ [x] := by
-  simp [push]
+  rw [push]; simp
   split <;> simp_all [List.len_eq_length]
   scalar_tac
 
diff --git a/backends/lean/Base/Termination.lean b/backends/lean/Base/Termination.lean
new file mode 100644
index 00000000..de8e678f
--- /dev/null
+++ b/backends/lean/Base/Termination.lean
@@ -0,0 +1,94 @@
+/- Some utilities to prove termination -/
+import Lean
+import Mathlib.Tactic.Core
+import Base.Utils
+import Base.Arith
+
+namespace Utils
+
+open Lean Lean.Elab Command Term Lean.Meta Tactic
+
+-- Inspired by the `clear` tactic
+def clearFvarIds (fvarIds : Array FVarId) : TacticM Unit := do
+  let fvarIds ← withMainContext <| sortFVarIds fvarIds
+  for fvarId in fvarIds.reverse do
+    withMainContext do
+      let mvarId ← (← getMainGoal).clear fvarId
+      replaceMainGoal [mvarId]
+
+/- Utility function for proofs of termination (i.e., inside `decreasing_by`).
+
+   Clean up the local context by removing all assumptions containing occurrences
+   of `invImage` (those are introduced automatically when doing proofs of
+   termination). We do so because we often need to simplify the context in the
+   proofs, and if we simplify those assumptions they tend to make the context
+   blow up.
+-/
+def removeInvImageAssumptions : TacticM Unit := do
+  withMainContext do
+  -- Get the local declarations
+  let ctx ← Lean.MonadLCtx.getLCtx
+  let decls ← ctx.getDecls
+  -- Retrieve the list of local declarations which contain `invertImage`
+  let containsInvertImage (decl : LocalDecl) : MetaM Bool := do
+    let expr := decl.toExpr
+    reduceVisit (
+      fun _ b e =>
+      pure (
+        b ||
+        match e with
+        | .const name _ => name == ``invImage
+        | _ => false)) false (← inferType expr)
+  let filtDecls ← liftM (decls.filterM containsInvertImage)
+  -- It can happen that other variables depend on the variables we want to clear:
+  -- filter them.
+  let allFVarsInTypes ← decls.foldlM (fun hs d => do
+    let hs ← getFVarIds (← inferType d.toExpr) hs
+    -- Explore the body if it is not opaque
+    match d.value? with
+    | none => pure hs
+    | some e => getFVarIds e hs
+    ) HashSet.empty
+  let fvarIds := filtDecls.map fun d => d.fvarId
+  let fvarIds := fvarIds.filter (fun id => ¬ allFVarsInTypes.contains id)
+  -- Clear them
+  clearFvarIds fvarIds.toArray
+
+elab "remove_invImage_assumptions" : tactic =>
+  removeInvImageAssumptions
+
+-- Auxiliary function
+def scalarDecrTac_apply_lemmas : TacticM Unit := do
+  withMainContext do
+  let goal ← getMainGoal
+  let rec applyFirst (names : List Name) : TacticM (List MVarId) :=
+    match names with
+    | [] => pure [goal]
+    | name :: names =>
+      -- Should use try ... catch or Lean.observing?
+      -- Generally speaking we should use Lean.observing? to restore the state,
+      -- but with tactics the try ... catch variant seems to work
+      try do
+        let th ← mkConstWithFreshMVarLevels name
+        goal.apply th
+      catch _ => do applyFirst names
+  let ngoals ← applyFirst [``Arith.to_int_to_nat_lt, ``Arith.to_int_sub_to_nat_lt]
+  setGoals ngoals
+
+elab "scalar_decr_tac_apply_lemmas" : tactic =>
+  scalarDecrTac_apply_lemmas
+
+-- For termination proofs
+syntax "scalar_decr_tac" : tactic
+macro_rules
+  | `(tactic| scalar_decr_tac) =>
+    `(tactic|
+      simp_wf <;>
+      -- Simplify the context - otherwise simp_all below will blow up
+      remove_invImage_assumptions <;>
+      -- Transform the goal a bit
+      scalar_decr_tac_apply_lemmas <;>
+      -- Finish
+      simp_all <;> scalar_tac)
+
+end Utils
diff --git a/backends/lean/Base/Utils.lean b/backends/lean/Base/Utils.lean
index b9de2fd1..2352979b 100644
--- a/backends/lean/Base/Utils.lean
+++ b/backends/lean/Base/Utils.lean
@@ -42,10 +42,6 @@ namespace List
 
 end List
 
--- TODO: move?
-@[simp]
-theorem neq_imp {α : Type u} {x y : α} (h : ¬ x = y) : ¬ y = x := by intro; simp_all
-
 namespace Lean
 
 namespace LocalContext