From 3c092169efcbc36a9b435c68c590b36f69204f94 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Fri, 8 Dec 2023 12:38:55 +0100
Subject: Update the progress tactic to use discrimination trees

---
 backends/lean/Base/Primitives/Scalar.lean | 126 ++++++++---------
 backends/lean/Base/Progress/Base.lean     | 228 ++++++++----------------------
 backends/lean/Base/Progress/Progress.lean |  93 ++++++------
 3 files changed, 175 insertions(+), 272 deletions(-)

(limited to 'backends/lean/Base')

diff --git a/backends/lean/Base/Primitives/Scalar.lean b/backends/lean/Base/Primitives/Scalar.lean
index f74fecd4..db522df2 100644
--- a/backends/lean/Base/Primitives/Scalar.lean
+++ b/backends/lean/Base/Primitives/Scalar.lean
@@ -528,7 +528,7 @@ instance {ty} : HAnd (Scalar ty) (Scalar ty) (Scalar ty) where
   hAnd x y := Scalar.and x y
 
 -- Generic theorem - shouldn't be used much
-@[cpspec]
+@[pspec]
 theorem Scalar.add_spec {ty} {x y : Scalar ty}
   (hmin : Scalar.min ty ≤ x.val + y.val)
   (hmax : x.val + y.val ≤ Scalar.max ty) :
@@ -550,62 +550,62 @@ theorem Scalar.add_unsigned_spec {ty} (s: ¬ ty.isSigned) {x y : Scalar ty}
   apply add_spec <;> assumption
 
 /- Fine-grained theorems -/
-@[cepspec] theorem Usize.add_spec {x y : Usize} (hmax : x.val + y.val ≤ Usize.max) :
+@[pspec] theorem Usize.add_spec {x y : Usize} (hmax : x.val + y.val ≤ Usize.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
   apply Scalar.add_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U8.add_spec {x y : U8} (hmax : x.val + y.val ≤ U8.max) :
+@[pspec] theorem U8.add_spec {x y : U8} (hmax : x.val + y.val ≤ U8.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
   apply Scalar.add_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U16.add_spec {x y : U16} (hmax : x.val + y.val ≤ U16.max) :
+@[pspec] theorem U16.add_spec {x y : U16} (hmax : x.val + y.val ≤ U16.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
   apply Scalar.add_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U32.add_spec {x y : U32} (hmax : x.val + y.val ≤ U32.max) :
+@[pspec] theorem U32.add_spec {x y : U32} (hmax : x.val + y.val ≤ U32.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
   apply Scalar.add_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U64.add_spec {x y : U64} (hmax : x.val + y.val ≤ U64.max) :
+@[pspec] theorem U64.add_spec {x y : U64} (hmax : x.val + y.val ≤ U64.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
   apply Scalar.add_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U128.add_spec {x y : U128} (hmax : x.val + y.val ≤ U128.max) :
+@[pspec] theorem U128.add_spec {x y : U128} (hmax : x.val + y.val ≤ U128.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
   apply Scalar.add_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem Isize.add_spec {x y : Isize}
+@[pspec] theorem Isize.add_spec {x y : Isize}
   (hmin : Isize.min ≤ x.val + y.val) (hmax : x.val + y.val ≤ Isize.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val :=
   Scalar.add_spec hmin hmax
 
-@[cepspec] theorem I8.add_spec {x y : I8}
+@[pspec] theorem I8.add_spec {x y : I8}
   (hmin : I8.min ≤ x.val + y.val) (hmax : x.val + y.val ≤ I8.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val :=
   Scalar.add_spec hmin hmax
 
-@[cepspec] theorem I16.add_spec {x y : I16}
+@[pspec] theorem I16.add_spec {x y : I16}
   (hmin : I16.min ≤ x.val + y.val) (hmax : x.val + y.val ≤ I16.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val :=
   Scalar.add_spec hmin hmax
 
-@[cepspec] theorem I32.add_spec {x y : I32}
+@[pspec] theorem I32.add_spec {x y : I32}
   (hmin : I32.min ≤ x.val + y.val) (hmax : x.val + y.val ≤ I32.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val :=
   Scalar.add_spec hmin hmax
 
-@[cepspec] theorem I64.add_spec {x y : I64}
+@[pspec] theorem I64.add_spec {x y : I64}
   (hmin : I64.min ≤ x.val + y.val) (hmax : x.val + y.val ≤ I64.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val :=
   Scalar.add_spec hmin hmax
 
-@[cepspec] theorem I128.add_spec {x y : I128}
+@[pspec] theorem I128.add_spec {x y : I128}
   (hmin : I128.min ≤ x.val + y.val) (hmax : x.val + y.val ≤ I128.max) :
   ∃ z, x + y = ret z ∧ z.val = x.val + y.val :=
   Scalar.add_spec hmin hmax
 
 -- Generic theorem - shouldn't be used much
-@[cpspec]
+@[pspec]
 theorem Scalar.sub_spec {ty} {x y : Scalar ty}
   (hmin : Scalar.min ty ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ Scalar.max ty) :
@@ -629,56 +629,56 @@ theorem Scalar.sub_unsigned_spec {ty} (s: ¬ ty.isSigned) {x y : Scalar ty}
   apply sub_spec <;> assumption
 
 /- Fine-grained theorems -/
-@[cepspec] theorem Usize.sub_spec {x y : Usize} (hmin : Usize.min ≤ x.val - y.val) :
+@[pspec] theorem Usize.sub_spec {x y : Usize} (hmin : Usize.min ≤ x.val - y.val) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val := by
   apply Scalar.sub_unsigned_spec <;> simp only [Scalar.min, *]
 
-@[cepspec] theorem U8.sub_spec {x y : U8} (hmin : U8.min ≤ x.val - y.val) :
+@[pspec] theorem U8.sub_spec {x y : U8} (hmin : U8.min ≤ x.val - y.val) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val := by
   apply Scalar.sub_unsigned_spec <;> simp only [Scalar.min, *]
 
-@[cepspec] theorem U16.sub_spec {x y : U16} (hmin : U16.min ≤ x.val - y.val) :
+@[pspec] theorem U16.sub_spec {x y : U16} (hmin : U16.min ≤ x.val - y.val) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val := by
   apply Scalar.sub_unsigned_spec <;> simp only [Scalar.min, *]
 
-@[cepspec] theorem U32.sub_spec {x y : U32} (hmin : U32.min ≤ x.val - y.val) :
+@[pspec] theorem U32.sub_spec {x y : U32} (hmin : U32.min ≤ x.val - y.val) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val := by
   apply Scalar.sub_unsigned_spec <;> simp only [Scalar.min, *]
 
-@[cepspec] theorem U64.sub_spec {x y : U64} (hmin : U64.min ≤ x.val - y.val) :
+@[pspec] theorem U64.sub_spec {x y : U64} (hmin : U64.min ≤ x.val - y.val) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val := by
   apply Scalar.sub_unsigned_spec <;> simp only [Scalar.min, *]
 
-@[cepspec] theorem U128.sub_spec {x y : U128} (hmin : U128.min ≤ x.val - y.val) :
+@[pspec] theorem U128.sub_spec {x y : U128} (hmin : U128.min ≤ x.val - y.val) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val := by
   apply Scalar.sub_unsigned_spec <;> simp only [Scalar.min, *]
 
-@[cepspec] theorem Isize.sub_spec {x y : Isize} (hmin : Isize.min ≤ x.val - y.val)
+@[pspec] theorem Isize.sub_spec {x y : Isize} (hmin : Isize.min ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ Isize.max) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val :=
   Scalar.sub_spec hmin hmax
 
-@[cepspec] theorem I8.sub_spec {x y : I8} (hmin : I8.min ≤ x.val - y.val)
+@[pspec] theorem I8.sub_spec {x y : I8} (hmin : I8.min ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ I8.max) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val :=
   Scalar.sub_spec hmin hmax
 
-@[cepspec] theorem I16.sub_spec {x y : I16} (hmin : I16.min ≤ x.val - y.val)
+@[pspec] theorem I16.sub_spec {x y : I16} (hmin : I16.min ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ I16.max) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val :=
   Scalar.sub_spec hmin hmax
 
-@[cepspec] theorem I32.sub_spec {x y : I32} (hmin : I32.min ≤ x.val - y.val)
+@[pspec] theorem I32.sub_spec {x y : I32} (hmin : I32.min ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ I32.max) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val :=
   Scalar.sub_spec hmin hmax
 
-@[cepspec] theorem I64.sub_spec {x y : I64} (hmin : I64.min ≤ x.val - y.val)
+@[pspec] theorem I64.sub_spec {x y : I64} (hmin : I64.min ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ I64.max) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val :=
   Scalar.sub_spec hmin hmax
 
-@[cepspec] theorem I128.sub_spec {x y : I128} (hmin : I128.min ≤ x.val - y.val)
+@[pspec] theorem I128.sub_spec {x y : I128} (hmin : I128.min ≤ x.val - y.val)
   (hmax : x.val - y.val ≤ I128.max) :
   ∃ z, x - y = ret z ∧ z.val = x.val - y.val :=
   Scalar.sub_spec hmin hmax
@@ -705,62 +705,62 @@ theorem Scalar.mul_unsigned_spec {ty} (s: ¬ ty.isSigned) {x y : Scalar ty}
   apply mul_spec <;> assumption
 
 /- Fine-grained theorems -/
-@[cepspec] theorem Usize.mul_spec {x y : Usize} (hmax : x.val * y.val ≤ Usize.max) :
+@[pspec] theorem Usize.mul_spec {x y : Usize} (hmax : x.val * y.val ≤ Usize.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val := by
   apply Scalar.mul_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U8.mul_spec {x y : U8} (hmax : x.val * y.val ≤ U8.max) :
+@[pspec] theorem U8.mul_spec {x y : U8} (hmax : x.val * y.val ≤ U8.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val := by
   apply Scalar.mul_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U16.mul_spec {x y : U16} (hmax : x.val * y.val ≤ U16.max) :
+@[pspec] theorem U16.mul_spec {x y : U16} (hmax : x.val * y.val ≤ U16.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val := by
   apply Scalar.mul_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U32.mul_spec {x y : U32} (hmax : x.val * y.val ≤ U32.max) :
+@[pspec] theorem U32.mul_spec {x y : U32} (hmax : x.val * y.val ≤ U32.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val := by
   apply Scalar.mul_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U64.mul_spec {x y : U64} (hmax : x.val * y.val ≤ U64.max) :
+@[pspec] theorem U64.mul_spec {x y : U64} (hmax : x.val * y.val ≤ U64.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val := by
   apply Scalar.mul_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem U128.mul_spec {x y : U128} (hmax : x.val * y.val ≤ U128.max) :
+@[pspec] theorem U128.mul_spec {x y : U128} (hmax : x.val * y.val ≤ U128.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val := by
   apply Scalar.mul_unsigned_spec <;> simp only [Scalar.max, *]
 
-@[cepspec] theorem Isize.mul_spec {x y : Isize} (hmin : Isize.min ≤ x.val * y.val)
+@[pspec] theorem Isize.mul_spec {x y : Isize} (hmin : Isize.min ≤ x.val * y.val)
   (hmax : x.val * y.val ≤ Isize.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val :=
   Scalar.mul_spec hmin hmax
 
-@[cepspec] theorem I8.mul_spec {x y : I8} (hmin : I8.min ≤ x.val * y.val)
+@[pspec] theorem I8.mul_spec {x y : I8} (hmin : I8.min ≤ x.val * y.val)
   (hmax : x.val * y.val ≤ I8.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val :=
   Scalar.mul_spec hmin hmax
 
-@[cepspec] theorem I16.mul_spec {x y : I16} (hmin : I16.min ≤ x.val * y.val)
+@[pspec] theorem I16.mul_spec {x y : I16} (hmin : I16.min ≤ x.val * y.val)
   (hmax : x.val * y.val ≤ I16.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val :=
   Scalar.mul_spec hmin hmax
 
-@[cepspec] theorem I32.mul_spec {x y : I32} (hmin : I32.min ≤ x.val * y.val)
+@[pspec] theorem I32.mul_spec {x y : I32} (hmin : I32.min ≤ x.val * y.val)
   (hmax : x.val * y.val ≤ I32.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val :=
   Scalar.mul_spec hmin hmax
 
-@[cepspec] theorem I64.mul_spec {x y : I64} (hmin : I64.min ≤ x.val * y.val)
+@[pspec] theorem I64.mul_spec {x y : I64} (hmin : I64.min ≤ x.val * y.val)
   (hmax : x.val * y.val ≤ I64.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val :=
   Scalar.mul_spec hmin hmax
 
-@[cepspec] theorem I128.mul_spec {x y : I128} (hmin : I128.min ≤ x.val * y.val)
+@[pspec] theorem I128.mul_spec {x y : I128} (hmin : I128.min ≤ x.val * y.val)
   (hmax : x.val * y.val ≤ I128.max) :
   ∃ z, x * y = ret z ∧ z.val = x.val * y.val :=
   Scalar.mul_spec hmin hmax
 
 -- Generic theorem - shouldn't be used much
-@[cpspec]
+@[pspec]
 theorem Scalar.div_spec {ty} {x y : Scalar ty}
   (hnz : y.val ≠ 0)
   (hmin : Scalar.min ty ≤ scalar_div x.val y.val)
@@ -788,66 +788,66 @@ theorem Scalar.div_unsigned_spec {ty} (s: ¬ ty.isSigned) (x : Scalar ty) {y : S
   apply hs
 
 /- Fine-grained theorems -/
-@[cepspec] theorem Usize.div_spec (x : Usize) {y : Usize} (hnz : y.val ≠ 0) :
+@[pspec] theorem Usize.div_spec (x : Usize) {y : Usize} (hnz : y.val ≠ 0) :
   ∃ z, x / y = ret z ∧ z.val = x.val / y.val := by
   apply Scalar.div_unsigned_spec <;> simp [*]
 
-@[cepspec] theorem U8.div_spec (x : U8) {y : U8} (hnz : y.val ≠ 0) :
+@[pspec] theorem U8.div_spec (x : U8) {y : U8} (hnz : y.val ≠ 0) :
   ∃ z, x / y = ret z ∧ z.val = x.val / y.val := by
   apply Scalar.div_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U16.div_spec (x : U16) {y : U16} (hnz : y.val ≠ 0) :
+@[pspec] theorem U16.div_spec (x : U16) {y : U16} (hnz : y.val ≠ 0) :
   ∃ z, x / y = ret z ∧ z.val = x.val / y.val := by
   apply Scalar.div_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U32.div_spec (x : U32) {y : U32} (hnz : y.val ≠ 0) :
+@[pspec] theorem U32.div_spec (x : U32) {y : U32} (hnz : y.val ≠ 0) :
   ∃ z, x / y = ret z ∧ z.val = x.val / y.val := by
   apply Scalar.div_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U64.div_spec (x : U64) {y : U64} (hnz : y.val ≠ 0) :
+@[pspec] theorem U64.div_spec (x : U64) {y : U64} (hnz : y.val ≠ 0) :
   ∃ z, x / y = ret z ∧ z.val = x.val / y.val := by
   apply Scalar.div_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U128.div_spec (x : U128) {y : U128} (hnz : y.val ≠ 0) :
+@[pspec] theorem U128.div_spec (x : U128) {y : U128} (hnz : y.val ≠ 0) :
   ∃ z, x / y = ret z ∧ z.val = x.val / y.val := by
   apply Scalar.div_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem Isize.div_spec (x : Isize) {y : Isize}
+@[pspec] theorem Isize.div_spec (x : Isize) {y : Isize}
   (hnz : y.val ≠ 0)
   (hmin : Isize.min ≤ scalar_div x.val y.val)
   (hmax : scalar_div x.val y.val ≤ Isize.max):
   ∃ z, x / y = ret z ∧ z.val = scalar_div x.val y.val :=
   Scalar.div_spec hnz hmin hmax
 
-@[cepspec] theorem I8.div_spec (x : I8) {y : I8}
+@[pspec] theorem I8.div_spec (x : I8) {y : I8}
   (hnz : y.val ≠ 0)
   (hmin : I8.min ≤ scalar_div x.val y.val)
   (hmax : scalar_div x.val y.val ≤ I8.max):
   ∃ z, x / y = ret z ∧ z.val = scalar_div x.val y.val :=
   Scalar.div_spec hnz hmin hmax
 
-@[cepspec] theorem I16.div_spec (x : I16) {y : I16}
+@[pspec] theorem I16.div_spec (x : I16) {y : I16}
   (hnz : y.val ≠ 0)
   (hmin : I16.min ≤ scalar_div x.val y.val)
   (hmax : scalar_div x.val y.val ≤ I16.max):
   ∃ z, x / y = ret z ∧ z.val = scalar_div x.val y.val :=
   Scalar.div_spec hnz hmin hmax
 
-@[cepspec] theorem I32.div_spec (x : I32) {y : I32}
+@[pspec] theorem I32.div_spec (x : I32) {y : I32}
   (hnz : y.val ≠ 0)
   (hmin : I32.min ≤ scalar_div x.val y.val)
   (hmax : scalar_div x.val y.val ≤ I32.max):
   ∃ z, x / y = ret z ∧ z.val = scalar_div x.val y.val :=
   Scalar.div_spec hnz hmin hmax
 
-@[cepspec] theorem I64.div_spec (x : I64) {y : I64}
+@[pspec] theorem I64.div_spec (x : I64) {y : I64}
   (hnz : y.val ≠ 0)
   (hmin : I64.min ≤ scalar_div x.val y.val)
   (hmax : scalar_div x.val y.val ≤ I64.max):
   ∃ z, x / y = ret z ∧ z.val = scalar_div x.val y.val :=
   Scalar.div_spec hnz hmin hmax
 
-@[cepspec] theorem I128.div_spec (x : I128) {y : I128}
+@[pspec] theorem I128.div_spec (x : I128) {y : I128}
   (hnz : y.val ≠ 0)
   (hmin : I128.min ≤ scalar_div x.val y.val)
   (hmax : scalar_div x.val y.val ≤ I128.max):
@@ -855,7 +855,7 @@ theorem Scalar.div_unsigned_spec {ty} (s: ¬ ty.isSigned) (x : Scalar ty) {y : S
   Scalar.div_spec hnz hmin hmax
 
 -- Generic theorem - shouldn't be used much
-@[cpspec]
+@[pspec]
 theorem Scalar.rem_spec {ty} {x y : Scalar ty}
   (hnz : y.val ≠ 0)
   (hmin : Scalar.min ty ≤ scalar_rem x.val y.val)
@@ -883,59 +883,59 @@ theorem Scalar.rem_unsigned_spec {ty} (s: ¬ ty.isSigned) (x : Scalar ty) {y : S
   simp [*] at hs
   simp [*]
 
-@[cepspec] theorem Usize.rem_spec (x : Usize) {y : Usize} (hnz : y.val ≠ 0) :
+@[pspec] theorem Usize.rem_spec (x : Usize) {y : Usize} (hnz : y.val ≠ 0) :
   ∃ z, x % y = ret z ∧ z.val = x.val % y.val := by
   apply Scalar.rem_unsigned_spec <;> simp [*]
 
-@[cepspec] theorem U8.rem_spec (x : U8) {y : U8} (hnz : y.val ≠ 0) :
+@[pspec] theorem U8.rem_spec (x : U8) {y : U8} (hnz : y.val ≠ 0) :
   ∃ z, x % y = ret z ∧ z.val = x.val % y.val := by
   apply Scalar.rem_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U16.rem_spec (x : U16) {y : U16} (hnz : y.val ≠ 0) :
+@[pspec] theorem U16.rem_spec (x : U16) {y : U16} (hnz : y.val ≠ 0) :
   ∃ z, x % y = ret z ∧ z.val = x.val % y.val := by
   apply Scalar.rem_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U32.rem_spec (x : U32) {y : U32} (hnz : y.val ≠ 0) :
+@[pspec] theorem U32.rem_spec (x : U32) {y : U32} (hnz : y.val ≠ 0) :
   ∃ z, x % y = ret z ∧ z.val = x.val % y.val := by
   apply Scalar.rem_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U64.rem_spec (x : U64) {y : U64} (hnz : y.val ≠ 0) :
+@[pspec] theorem U64.rem_spec (x : U64) {y : U64} (hnz : y.val ≠ 0) :
   ∃ z, x % y = ret z ∧ z.val = x.val % y.val := by
   apply Scalar.rem_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem U128.rem_spec (x : U128) {y : U128} (hnz : y.val ≠ 0) :
+@[pspec] theorem U128.rem_spec (x : U128) {y : U128} (hnz : y.val ≠ 0) :
   ∃ z, x % y = ret z ∧ z.val = x.val % y.val := by
   apply Scalar.rem_unsigned_spec <;> simp [Scalar.max, *]
 
-@[cepspec] theorem I8.rem_spec (x : I8) {y : I8}
+@[pspec] theorem I8.rem_spec (x : I8) {y : I8}
   (hnz : y.val ≠ 0)
   (hmin : I8.min ≤ scalar_rem x.val y.val)
   (hmax : scalar_rem x.val y.val ≤ I8.max):
   ∃ z, x % y = ret z ∧ z.val = scalar_rem x.val y.val :=
   Scalar.rem_spec hnz hmin hmax
 
-@[cepspec] theorem I16.rem_spec (x : I16) {y : I16}
+@[pspec] theorem I16.rem_spec (x : I16) {y : I16}
   (hnz : y.val ≠ 0)
   (hmin : I16.min ≤ scalar_rem x.val y.val)
   (hmax : scalar_rem x.val y.val ≤ I16.max):
   ∃ z, x % y = ret z ∧ z.val = scalar_rem x.val y.val :=
   Scalar.rem_spec hnz hmin hmax
 
-@[cepspec] theorem I32.rem_spec (x : I32) {y : I32}
+@[pspec] theorem I32.rem_spec (x : I32) {y : I32}
   (hnz : y.val ≠ 0)
   (hmin : I32.min ≤ scalar_rem x.val y.val)
   (hmax : scalar_rem x.val y.val ≤ I32.max):
   ∃ z, x % y = ret z ∧ z.val = scalar_rem x.val y.val :=
   Scalar.rem_spec hnz hmin hmax
 
-@[cepspec] theorem I64.rem_spec (x : I64) {y : I64}
+@[pspec] theorem I64.rem_spec (x : I64) {y : I64}
   (hnz : y.val ≠ 0)
   (hmin : I64.min ≤ scalar_rem x.val y.val)
   (hmax : scalar_rem x.val y.val ≤ I64.max):
   ∃ z, x % y = ret z ∧ z.val = scalar_rem x.val y.val :=
   Scalar.rem_spec hnz hmin hmax
 
-@[cepspec] theorem I128.rem_spec (x : I128) {y : I128}
+@[pspec] theorem I128.rem_spec (x : I128) {y : I128}
   (hnz : y.val ≠ 0)
   (hmin : I128.min ≤ scalar_rem x.val y.val)
   (hmax : scalar_rem x.val y.val ≤ I128.max):
diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 573f0cc5..d50c357c 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -4,6 +4,9 @@ import Base.Utils
 import Base.Primitives.Base
 import Base.Lookup.Base
 
+import Lean.Meta.DiscrTree
+import Lean.Meta.Tactic.Simp
+
 namespace Progress
 
 open Lean Elab Term Meta
@@ -14,13 +17,35 @@ initialize registerTraceClass `Progress
 
 /- # Progress tactic -/
 
+/- Discrimination trees map expressions to values. When storing an expression
+   in a discrimination tree, the expression is first converted to an array
+   of `DiscrTree.Key`, which are the keys actually used by the discrimination
+   trees. The conversion operation is monadic, however, and extensions require
+   all the operations to be pure. For this reason, in the state extension, we
+   store the keys from *after* the transformation (i.e., the `DiscrTreeKey`
+   below).
+ -/
+abbrev DiscrTreeKey (simpleReduce : Bool) := Array (DiscrTree.Key simpleReduce)
+
+abbrev DiscrTreeExtension (α : Type) (simpleReduce : Bool) :=
+  SimplePersistentEnvExtension (DiscrTreeKey simpleReduce × α) (DiscrTree α simpleReduce)
+
+def mkDiscrTreeExtention [Inhabited α] [BEq α] (name : Name := by exact decl_name%) (simpleReduce : Bool) :
+  IO (DiscrTreeExtension α simpleReduce) :=
+  registerSimplePersistentEnvExtension {
+    name          := name,
+    addImportedFn := fun a => a.foldl (fun s a => a.foldl (fun s (k, v) => s.insertCore k v) s) DiscrTree.empty,
+    addEntryFn    := fun s n => s.insertCore n.1 n.2 ,
+  }
+
 structure PSpecDesc where
   -- The universally quantified variables
   fvars : Array Expr
   -- The existentially quantified variables
   evars : Array Expr
+  -- The function applied to its arguments
+  fArgsExpr : Expr
   -- The function
-  fExpr : Expr
   fName : Name
   -- The function arguments
   fLevels : List Level
@@ -38,7 +63,7 @@ section Methods
   variable [MonadError m]
   variable {a : Type}
 
-  /- Analyze a pspec theorem to decompose its arguments.
+  /- Analyze a goal or a pspec theorem to decompose its arguments.
 
      PSpec theorems should be of the following shape:
      ```
@@ -57,12 +82,20 @@ section Methods
      TODO: generalize for when we do inductive proofs
   -/
   partial
-  def withPSpec [Inhabited (m a)] [Nonempty (m a)] (th : Expr) (k : PSpecDesc → m a)
-    (sanityChecks : Bool := false) :
+  def withPSpec [Inhabited (m a)] [Nonempty (m a)] (sanityChecks : Bool := false)
+    (isGoal : Bool) (th : Expr) (k : PSpecDesc → m a) :
     m a := do
     trace[Progress] "Proposition: {th}"
     -- Dive into the quantified variables and the assumptions
-    forallTelescope th.consumeMData fun fvars th => do
+    -- Note that if we analyze a pspec theorem to register it in a database (i.e.
+    -- a discrimination tree), we need to introduce *meta-variables* for the
+    -- quantified variables.
+    let telescope (k : Array Expr → Expr → m a) : m a :=
+      if isGoal then forallTelescope th.consumeMData (fun fvars th => k fvars th)
+      else do
+        let (fvars, _, th) ← forallMetaTelescope th.consumeMData
+        k fvars th
+    telescope fun fvars th => do
     trace[Progress] "Universally quantified arguments and assumptions: {fvars}"
     -- Dive into the existentials
     existsTelescope th.consumeMData fun evars th => do
@@ -79,7 +112,7 @@ section Methods
     -- destruct the application to get the function name
     mExpr.consumeMData.withApp fun mf margs => do
     trace[Progress] "After stripping the arguments of the monad expression:\n- mf: {mf}\n- margs: {margs}"
-    let (fExpr, f, args) ← do
+    let (fArgsExpr, f, args) ← do
       if mf.isConst ∧ mf.constName = ``Bind.bind then do
         -- Dive into the bind
         let fExpr := (margs.get! 4).consumeMData
@@ -101,11 +134,11 @@ section Methods
     let argsFVars := fvars.map (fun x => x.fvarId!)
     let argsFVars := argsFVars.filter (fun fvar => allArgsFVars.contains fvar)
     -- Return
-    trace[Progress] "Function: {f.constName!}";
+    trace[Progress] "Function with arguments: {fArgsExpr}";
     let thDesc := {
       fvars := fvars
       evars := evars
-      fExpr
+      fArgsExpr
       fName := f.constName!
       fLevels := f.constLevels!
       args := args
@@ -117,60 +150,21 @@ section Methods
 
 end Methods
 
-def getPSpecFunName (th : Expr) : MetaM Name :=
-  withPSpec th (fun d => do pure d.fName) true
+/-def getPSpecFunArgsExpr (th : Expr) : MetaM Expr :=
+  withPSpec true th (fun d => do pure d.fArgsExpr)
 
-def getPSpecClassFunNames (th : Expr) : MetaM (Name × Name) :=
-  withPSpec th (fun d => do
-    let arg0 := d.args.get! 0
-    arg0.withApp fun f _ => do
-    if ¬ f.isConst then throwError "Not a constant: {f}"
-    pure (d.fName, f.constName)
-    ) true
-
-def getPSpecClassFunNameArg (th : Expr) : MetaM (Name × Expr) :=
-  withPSpec th (fun d => do
-    let arg0 := d.args.get! 0
-    pure (d.fName, arg0)
-    ) true
+def getPSpecFunName (th : Expr) : MetaM Name :=
+  withPSpec true th (fun d => do pure d.fName)-/
 
--- "Regular" pspec attribute
+-- pspec attribute
 structure PSpecAttr where
   attr : AttributeImpl
-  ext  : MapDeclarationExtension Name
-  deriving Inhabited
-
-/- pspec attribute for type classes: we use the name of the type class to
-   lookup another map. We use the *first* argument of the type class to lookup
-   into this second map.
-
-   Example:
-   ========
-   We use type classes for addition. For instance, the addition between two
-   U32 is written (without syntactic sugar) as `HAdd.add (Scalar ty) x y`. As a consequence,
-   we store the theorem through the bindings: HAdd.add → Scalar → ...
-
-   SH: TODO: this (and `PSpecClassExprAttr`) is a bit ad-hoc. For now it works for the
-   specs of the scalar operations, which is what I really need, but I'm not sure it
-   applies well to other situations. A better way would probably to use type classes, but
-   I couldn't get them to work on those cases. It is worth retrying.
-   UPDATE: use discrimination trees (`DiscrTree`) from core Lean
--/
-structure PSpecClassAttr where
-  attr : AttributeImpl
-  ext  : MapDeclarationExtension (NameMap Name)
-  deriving Inhabited
-
-/- Same as `PSpecClassAttr` but we use the full first argument (it works when it
-   is a constant). -/
-structure PSpecClassExprAttr where
-  attr : AttributeImpl
-  ext  : MapDeclarationExtension (HashMap Expr Name)
+  ext  : DiscrTreeExtension Name true
   deriving Inhabited
 
 /- The persistent map from function to pspec theorems. -/
 initialize pspecAttr : PSpecAttr ← do
-  let ext ← Lookup.mkMapDeclarationExtension `pspecMap
+  let ext ← mkDiscrTreeExtention `pspecMap true
   let attrImpl : AttributeImpl := {
     name := `pspec
     descr := "Marks theorems to use with the `progress` tactic"
@@ -182,130 +176,30 @@ initialize pspecAttr : PSpecAttr ← do
       -- Lookup the theorem
       let env ← getEnv
       let thDecl := env.constants.find! thName
-      let fName ← MetaM.run' (getPSpecFunName thDecl.type)
-      trace[Progress] "Registering spec theorem for {fName}"
-      let env := ext.addEntry env (fName, thName)
-      setEnv env
-      pure ()
-  }
-  registerBuiltinAttribute attrImpl
-  pure { attr := attrImpl, ext := ext }
-
-/- The persistent map from type classes to pspec theorems -/
-initialize pspecClassAttr : PSpecClassAttr ← do
-  let ext : MapDeclarationExtension (NameMap Name) ←
-    Lookup.mkMapMapDeclarationExtension Name.quickCmp `pspecClassMap
-  let attrImpl : AttributeImpl  := {
-    name := `cpspec
-    descr := "Marks theorems to use for type classes with the `progress` tactic"
-    add := fun thName stx attrKind => do
-      Attribute.Builtin.ensureNoArgs stx
-      -- TODO: use the attribute kind
-      unless attrKind == AttributeKind.global do
-        throwError "invalid attribute 'cpspec', must be global"
-      -- Lookup the theorem
-      let env ← getEnv
-      let thDecl := env.constants.find! thName
-      let (fName, argName) ← MetaM.run' (getPSpecClassFunNames thDecl.type)
-      trace[Progress] "Registering class spec theorem for ({fName}, {argName})"
-      -- Update the entry if there is one, add an entry if there is none
-      let env :=
-        match (ext.getState (← getEnv)).find? fName with
-        | none =>
-          let m := RBMap.ofList [(argName, thName)]
-          ext.addEntry env (fName, m)
-        | some m =>
-          let m := m.insert argName thName
-          ext.addEntry env (fName, m)
-      setEnv env
-      pure ()
-  }
-  registerBuiltinAttribute attrImpl
-  pure { attr := attrImpl, ext := ext }
-
-/- The 2nd persistent map from type classes to pspec theorems -/
-initialize pspecClassExprAttr : PSpecClassExprAttr ← do
-  let ext : MapDeclarationExtension (HashMap Expr Name) ←
-    Lookup.mkMapHashMapDeclarationExtension `pspecClassExprMap
-  let attrImpl : AttributeImpl  := {
-    name := `cepspec
-    descr := "Marks theorems to use for type classes with the `progress` tactic"
-    add := fun thName stx attrKind => do
-      Attribute.Builtin.ensureNoArgs stx
-      -- TODO: use the attribute kind
-      unless attrKind == AttributeKind.global do
-        throwError "invalid attribute 'cpspec', must be global"
-      -- Lookup the theorem
-      let env ← getEnv
-      let thDecl := env.constants.find! thName
-      let (fName, arg) ← MetaM.run' (getPSpecClassFunNameArg thDecl.type)
-      -- Sanity check: no variables appear in the argument
-      MetaM.run' do
-        let fvars ← getFVarIds arg
-        if ¬ fvars.isEmpty then throwError "The first argument ({arg}) contains variables"
-      -- We store two bindings:
-      -- - arg to theorem name
-      -- - reduced arg to theorem name
-      let rarg ← MetaM.run' (reduceAll arg)
-      trace[Progress] "Registering class spec theorem for ({fName}, {arg}) and ({fName}, {rarg})"
-      -- Update the entry if there is one, add an entry if there is none
-      let env :=
-        match (ext.getState (← getEnv)).find? fName with
-        | none =>
-          let m := HashMap.ofList [(arg, thName), (rarg, thName)]
-          ext.addEntry env (fName, m)
-        | some m =>
-          let m := m.insert arg thName
-          let m := m.insert rarg thName
-          ext.addEntry env (fName, m)
+      let isGoal := false
+      let fKey ← MetaM.run' (withPSpec true isGoal thDecl.type fun d => do
+        let fExpr := d.fArgsExpr
+        trace[Progress] "Registering spec theorem for {fExpr}"
+        -- Convert the function expression to a discrimination tree key
+        DiscrTree.mkPath fExpr)
+      let env := ext.addEntry env (fKey, thName)
       setEnv env
       pure ()
   }
   registerBuiltinAttribute attrImpl
   pure { attr := attrImpl, ext := ext }
 
+def PSpecAttr.find? (s : PSpecAttr) (e : Expr) : MetaM (Array Name) := do
+  (s.ext.getState (← getEnv)).getMatch e
 
-def PSpecAttr.find? (s : PSpecAttr) (name : Name) : MetaM (Option Name) := do
-  return (s.ext.getState (← getEnv)).find? name
-
-def PSpecClassAttr.find? (s : PSpecClassAttr) (className argName : Name) : MetaM (Option Name) := do
-  match (s.ext.getState (← getEnv)).find? className with
-  | none => return none
-  | some map => return map.find? argName
-
-def PSpecClassExprAttr.find? (s : PSpecClassExprAttr) (className : Name) (arg : Expr) : MetaM (Option Name) := do
-  match (s.ext.getState (← getEnv)).find? className with
-  | none => return none
-  | some map => return map.find? arg
-
-def PSpecAttr.getState (s : PSpecAttr) : MetaM (NameMap Name) := do
-  pure (s.ext.getState (← getEnv))
-
-def PSpecClassAttr.getState (s : PSpecClassAttr) : MetaM (NameMap (NameMap Name)) := do
-  pure (s.ext.getState (← getEnv))
-
-def PSpecClassExprAttr.getState (s : PSpecClassExprAttr) : MetaM (NameMap (HashMap Expr Name)) := do
+def PSpecAttr.getState (s : PSpecAttr) : MetaM (DiscrTree Name true) := do
   pure (s.ext.getState (← getEnv))
 
 def showStoredPSpec : MetaM Unit := do
   let st ← pspecAttr.getState
-  let s := st.toList.foldl (fun s (f, th) => f!"{s}\n{f} → {th}") f!""
-  IO.println s
-
-def showStoredPSpecClass : MetaM Unit := do
-  let st ← pspecClassAttr.getState
-  let s := st.toList.foldl (fun s (f, m) =>
-    let ms := m.toList.foldl (fun s (f, th) =>
-      f!"{s}\n  {f} → {th}") f!""
-    f!"{s}\n{f} → [{ms}]") f!""
-  IO.println s
-
-def showStoredPSpecExprClass : MetaM Unit := do
-  let st ← pspecClassExprAttr.getState
-  let s := st.toList.foldl (fun s (f, m) =>
-    let ms := m.toList.foldl (fun s (f, th) =>
-      f!"{s}\n  {f} → {th}") f!""
-    f!"{s}\n{f} → [{ms}]") f!""
+  -- TODO: how can we iterate over (at least) the values stored in the tree?
+  --let s := st.toList.foldl (fun s (f, th) => f!"{s}\n{f} → {th}") f!""
+  let s := f!"{st}"
   IO.println s
 
 end Progress
diff --git a/backends/lean/Base/Progress/Progress.lean b/backends/lean/Base/Progress/Progress.lean
index ba63f09d..93b7d7d5 100644
--- a/backends/lean/Base/Progress/Progress.lean
+++ b/backends/lean/Base/Progress/Progress.lean
@@ -204,11 +204,11 @@ def getFirstArg (args : Array Expr) : Option Expr := do
   if args.size = 0 then none
   else some (args.get! 0)
 
-/- Helper: try to lookup a theorem and apply it, or continue with another tactic
-   if it fails -/
+/- Helper: try to lookup a theorem and apply it.
+   Return true if it succeeded. -/
 def tryLookupApply (keep : Option Name) (ids : Array (Option Name)) (splitPost : Bool)
   (asmTac : TacticM Unit) (fExpr : Expr)
-  (kind : String) (th : Option TheoremOrLocal) (x : TacticM Unit) : TacticM Unit := do
+  (kind : String) (th : Option TheoremOrLocal) : TacticM Bool := do
   let res ← do
     match th with
     | none =>
@@ -223,9 +223,9 @@ def tryLookupApply (keep : Option Name) (ids : Array (Option Name)) (splitPost :
           pure (some res)
         catch _ => none
   match res with
-  | some .Ok => return ()
+  | some .Ok => return true
   | some (.Error msg) => throwError msg
-  | none => x
+  | none => return false
 
 -- The array of ids are identifiers to use when introducing fresh variables
 def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrLocal)
@@ -236,11 +236,19 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
   let goalTy ← mgoal.getType
   trace[Progress] "goal: {goalTy}"
   -- Dive into the goal to lookup the theorem
-  let (fExpr, fName, args) ← do
-    withPSpec goalTy fun desc =>
-    -- TODO: check that no quantified variables in the arguments
-    pure (desc.fExpr, desc.fName, desc.args)
-  trace[Progress] "Function: {fName}"
+  -- Remark: if we don't isolate the call to `withPSpec` to immediately "close"
+  -- the terms immediately, we may end up with the error:
+  -- "(kernel) declaration has free variables"
+  -- I'm not sure I understand why.
+  -- TODO: we should also check that no quantified variable appears in fExpr.
+  -- If such variables appear, we should just fail because the goal doesn't
+  -- have the proper shape.
+  let fExpr ← do
+    let isGoal := true
+    withPSpec false isGoal goalTy fun desc => do
+    let fExpr := desc.fArgsExpr
+    trace[Progress] "Expression to match: {fExpr}"
+    pure fExpr
   -- If the user provided a theorem/assumption: use it.
   -- Otherwise, lookup one.
   match withTh with
@@ -258,36 +266,24 @@ def progressAsmsOrLookupTheorem (keep : Option Name) (withTh : Option TheoremOrL
       match res with
       | .Ok => return ()
       | .Error msg => throwError msg
-    -- It failed: try to lookup a theorem
-    -- TODO: use a list of theorems, and try them one by one?
-    trace[Progress] "No assumption succeeded: trying to lookup a theorem"
-    let pspec ← do
-      let thName ← pspecAttr.find? fName
-      pure (thName.map fun th => .Theorem th)
-    tryLookupApply keep ids splitPost asmTac fExpr "pspec theorem" pspec do
-    -- It failed: try to lookup a *class* expr spec theorem (those are more
-    -- specific than class spec theorems)
-    trace[Progress] "Failed using a pspec theorem: trying to lookup a pspec class expr theorem"
-    let pspecClassExpr ← do
-      match getFirstArg args with
-      | none => pure none
-      | some arg => do
-        trace[Progress] "Using: f:{fName}, arg: {arg}"
-        let thName ← pspecClassExprAttr.find? fName arg
-        pure (thName.map fun th => .Theorem th)
-    tryLookupApply keep ids splitPost asmTac fExpr "pspec class expr theorem" pspecClassExpr do
-    -- It failed: try to lookup a *class* spec theorem
-    trace[Progress] "Failed using a pspec class expr theorem: trying to lookup a pspec class theorem"
-    let pspecClass ← do
-      match ← getFirstArgAppName args with
-      | none => pure none
-      | some argName => do
-        trace[Progress] "Using: f: {fName}, arg: {argName}"
-        let thName ← pspecClassAttr.find? fName argName
-        pure (thName.map fun th => .Theorem th)
-    tryLookupApply keep ids splitPost asmTac fExpr "pspec class theorem" pspecClass do
-    trace[Progress] "Failed using a pspec class theorem: trying to use a recursive assumption"
-    -- Try a recursive call - we try the assumptions of kind "auxDecl"
+    -- It failed: lookup the pspec theorems which match the expression
+    trace[Progress] "No assumption succeeded: trying to lookup a pspec theorem"
+    let pspecs : Array TheoremOrLocal ← do
+      let thNames ← pspecAttr.find? fExpr
+      -- TODO: because of reduction, there may be several valid theorems (for
+      -- instance for the scalars). We need to sort them from most specific to
+      -- least specific. For now, we assume the most specific theorems are at
+      -- the end.
+      let thNames := thNames.reverse
+      trace[Progress] "Looked up pspec theorems: {thNames}"
+      pure (thNames.map fun th => TheoremOrLocal.Theorem th)
+    -- Try the theorems one by one
+    for pspec in pspecs do
+      if ← tryLookupApply keep ids splitPost asmTac fExpr "pspec theorem" pspec then return ()
+      else pure ()
+    -- It failed: try to use the recursive assumptions
+    trace[Progress] "Failed using a pspec theorem: trying to use a recursive assumption"
+    -- We try to apply the assumptions of kind "auxDecl"
     let ctx ← Lean.MonadLCtx.getLCtx
     let decls ← ctx.getAllDecls
     let decls := decls.filter (λ decl => match decl.kind with
@@ -381,8 +377,6 @@ namespace Test
 
   -- The following commands display the databases of theorems
   -- #eval showStoredPSpec
-  -- #eval showStoredPSpecClass
-  -- #eval showStoredPSpecExprClass
   open alloc.vec
 
   example {ty} {x y : Scalar ty}
@@ -392,6 +386,8 @@ namespace Test
     progress keep _ as ⟨ z, h1 .. ⟩
     simp [*, h1]
 
+  set_option trace.Progress false
+
   example {ty} {x y : Scalar ty}
     (hmin : Scalar.min ty ≤ x.val + y.val)
     (hmax : x.val + y.val ≤ Scalar.max ty) :
@@ -399,6 +395,19 @@ namespace Test
     progress keep h with Scalar.add_spec as ⟨ z ⟩
     simp [*, h]
 
+  example {x y : U32}
+    (hmax : x.val + y.val ≤ U32.max) :
+    ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
+    -- This spec theorem is suboptimal, but it is good to check that it works
+    progress with Scalar.add_spec as ⟨ z, h1 .. ⟩
+    simp [*, h1]
+ 
+  example {x y : U32}
+    (hmax : x.val + y.val ≤ U32.max) :
+    ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
+    progress with U32.add_spec as ⟨ z, h1 .. ⟩
+    simp [*, h1]
+
   example {x y : U32}
     (hmax : x.val + y.val ≤ U32.max) :
     ∃ z, x + y = ret z ∧ z.val = x.val + y.val := by
-- 
cgit v1.2.3