From 6166c410a4b3353377e640acbae9f56e877a9118 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Tue, 11 Jul 2023 15:23:49 +0200
Subject: Work on the progress tactic

---
 backends/lean/Base/Progress/Base.lean | 175 ++++++++++++++++++++++++++++++++++
 1 file changed, 175 insertions(+)
 create mode 100644 backends/lean/Base/Progress/Base.lean

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
new file mode 100644
index 00000000..3f44f46c
--- /dev/null
+++ b/backends/lean/Base/Progress/Base.lean
@@ -0,0 +1,175 @@
+import Lean
+import Base.Utils
+import Base.Primitives
+
+namespace Progress
+
+open Lean Elab Term Meta
+open Utils
+
+-- We can't define and use trace classes in the same file
+initialize registerTraceClass `Progress
+
+-- Return the first conjunct if the expression is a conjunction, or the
+-- expression itself otherwise. Also return the second conjunct if it is a
+-- conjunction.
+def getFirstConj (e : Expr) : MetaM (Expr × Option Expr) := do
+  e.withApp fun f args =>
+  if f.isConstOf ``And ∧ args.size = 2 then pure (args.get! 0, some (args.get! 1))
+  else pure (e, none)
+
+-- Destruct an equaliy and return the two sides
+def destEq (e : Expr) : MetaM (Expr × Expr) := do
+  e.withApp fun f args =>
+  if f.isConstOf ``Eq ∧ args.size = 3 then pure (args.get! 1, args.get! 2)
+  else throwError "Not an equality: {e}"
+
+-- Return the set of FVarIds in the expression
+partial def getFVarIds (e : Expr) (hs : HashSet FVarId := HashSet.empty) : MetaM (HashSet FVarId) := do
+  e.withApp fun body args => do
+  let hs := if body.isFVar then hs.insert body.fvarId! else hs
+  args.foldlM (fun hs arg => getFVarIds arg hs) hs
+
+/- # Progress tactic -/
+
+structure PSpecDesc where
+  -- The universally quantified variables
+  fvars : Array Expr
+  -- The existentially quantified variables
+  evars : Array Expr
+  -- The function
+  fName : Name
+  -- The function arguments
+  fLevels : List Level
+  args : Array Expr
+  -- The universally quantified variables which appear in the function arguments
+  argsFVars : Array FVarId
+  -- The returned value
+  ret : Expr
+  -- The postcondition (if there is)
+  post : Option Expr
+
+section Methods
+  variable [MonadLiftT MetaM m] [MonadControlT MetaM m] [Monad m] [MonadOptions m]
+  variable [MonadTrace m] [MonadLiftT IO m] [MonadRef m] [AddMessageContext m]
+  variable [MonadError m]
+  variable {a : Type}
+
+  /- Analyze a pspec theorem to decompose its arguments.
+
+     PSpec theorems should be of the following shape:
+     ```
+     ∀ x1 ... xn, H1 → ... Hn → ∃ y1 ... ym. f x1 ... xn = .ret ... ∧ Post1 ∧ ... ∧ Postk
+     ```
+
+     The continuation `k` receives the following inputs:
+     - universally quantified variables
+     - assumptions
+     - existentially quantified variables
+     - function name
+     - function arguments
+     - return
+     - postconditions
+
+     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) :
+    m a := do
+    trace[Progress] "Theorem: {th}"
+    -- Dive into the quantified variables and the assumptions
+    forallTelescope th fun fvars th => do
+    trace[Progress] "All argumens: {fvars}"
+  /-  -- Filter the argumens which are not propositions
+    let rec getFirstPropIdx (i : Nat) : MetaM Nat := do
+      if i ≥ fargs.size then pure i
+      else do
+        let x := fargs.get! i
+        if ← Meta.isProp (← inferType x) then pure i
+        else getFirstPropIdx (i + 1)
+    let i ← getFirstPropIdx 0
+    let fvars := fargs.extract 0 i
+    let hyps := fargs.extract i fargs.size
+    trace[Progress] "Quantified variables: {fvars}"
+    trace[Progress] "Assumptions: {hyps}"
+    -- Sanity check: all hypotheses are propositions (in particular, all the
+    -- quantified variables are at the beginning)
+    let hypsAreProp ← hyps.allM fun x => do Meta.isProp (← inferType x)
+    if ¬ hypsAreProp then
+      throwError "The theorem doesn't have the proper shape: all the quantified arguments should be at the beginning"
+      -/
+    -- Dive into the existentials
+    existsTelescope th fun evars th => do
+    trace[Progress] "Existentials: {evars}"
+    -- Take the first conjunct
+    let (th, post) ← getFirstConj th
+    -- Destruct the equality
+    let (th, ret) ← destEq th
+    -- Destruct the application to get the name
+    th.withApp fun f args => do
+    if ¬ f.isConst then throwError "Not a constant: {f}"
+    -- Compute the set of universally quantified variables which appear in the function arguments
+    let allArgsFVars ← args.foldlM (fun hs arg => getFVarIds arg hs) HashSet.empty
+    -- Sanity check
+    if sanityChecks then
+      let fvarsSet : HashSet FVarId := HashSet.ofArray (fvars.map (fun x => x.fvarId!))
+      let filtArgsFVars := allArgsFVars.toArray.filter (fun fvar => ¬ fvarsSet.contains fvar)
+      if ¬ filtArgsFVars.isEmpty then
+        let filtArgsFVars := filtArgsFVars.map (fun fvarId => Expr.fvar fvarId)
+        throwError "Some of the function inputs are not universally quantified: {filtArgsFVars}"
+    let argsFVars := fvars.map (fun x => x.fvarId!)
+    let argsFVars := argsFVars.filter (fun fvar => allArgsFVars.contains fvar)
+    -- Return
+    trace[Progress] "Function: {f.constName!}";
+    let thDesc := {
+      fvars := fvars
+      evars := evars
+      fName := f.constName!
+      fLevels := f.constLevels!
+      args := args
+      argsFVars
+      ret := ret
+      post := post
+    }
+    k thDesc
+end Methods
+
+
+def getPSpecFunName (th : Expr) : MetaM Name :=
+  withPSpec th (fun d => do pure d.fName) true
+
+structure PSpecAttr where
+  attr : AttributeImpl
+  ext  : MapDeclarationExtension Name
+  deriving Inhabited
+
+/- The persistent map from function to pspec theorems. -/
+initialize pspecAttr : PSpecAttr ← do
+  let ext ← mkMapDeclarationExtension `pspecMap
+  let attrImpl := {
+    name := `pspec
+    descr := "Marks theorems to use 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 'pspec', must be global"
+      -- 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 }
+
+def PSpecAttr.find? (s : PSpecAttr) (name : Name) : MetaM (Option Name) := do
+  return (s.ext.getState (← getEnv)).find? name
+  --return s.ext.find? (← getEnv) name
+
+
+end Progress
-- 
cgit v1.2.3


From a18d899a2c2b9bdd36f4a5a4b70472c12a835a96 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Wed, 12 Jul 2023 14:34:55 +0200
Subject: Finish a first version of the progress tactic

---
 backends/lean/Base/Progress/Base.lean | 24 +-----------------------
 1 file changed, 1 insertion(+), 23 deletions(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 3f44f46c..613f38f8 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -10,26 +10,6 @@ open Utils
 -- We can't define and use trace classes in the same file
 initialize registerTraceClass `Progress
 
--- Return the first conjunct if the expression is a conjunction, or the
--- expression itself otherwise. Also return the second conjunct if it is a
--- conjunction.
-def getFirstConj (e : Expr) : MetaM (Expr × Option Expr) := do
-  e.withApp fun f args =>
-  if f.isConstOf ``And ∧ args.size = 2 then pure (args.get! 0, some (args.get! 1))
-  else pure (e, none)
-
--- Destruct an equaliy and return the two sides
-def destEq (e : Expr) : MetaM (Expr × Expr) := do
-  e.withApp fun f args =>
-  if f.isConstOf ``Eq ∧ args.size = 3 then pure (args.get! 1, args.get! 2)
-  else throwError "Not an equality: {e}"
-
--- Return the set of FVarIds in the expression
-partial def getFVarIds (e : Expr) (hs : HashSet FVarId := HashSet.empty) : MetaM (HashSet FVarId) := do
-  e.withApp fun body args => do
-  let hs := if body.isFVar then hs.insert body.fvarId! else hs
-  args.foldlM (fun hs arg => getFVarIds arg hs) hs
-
 /- # Progress tactic -/
 
 structure PSpecDesc where
@@ -103,7 +83,7 @@ section Methods
     existsTelescope th fun evars th => do
     trace[Progress] "Existentials: {evars}"
     -- Take the first conjunct
-    let (th, post) ← getFirstConj th
+    let (th, post) ← optSplitConj th
     -- Destruct the equality
     let (th, ret) ← destEq th
     -- Destruct the application to get the name
@@ -169,7 +149,5 @@ initialize pspecAttr : PSpecAttr ← do
 
 def PSpecAttr.find? (s : PSpecAttr) (name : Name) : MetaM (Option Name) := do
   return (s.ext.getState (← getEnv)).find? name
-  --return s.ext.find? (← getEnv) name
-
 
 end Progress
-- 
cgit v1.2.3


From eb97bdb6761437e492bcf1a95b4fa43d2b69601b Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Wed, 12 Jul 2023 18:04:19 +0200
Subject: Improve progress to use assumptions and start working on a nice
 syntax

---
 backends/lean/Base/Progress/Base.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 613f38f8..a288d889 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -60,7 +60,7 @@ section Methods
     trace[Progress] "Theorem: {th}"
     -- Dive into the quantified variables and the assumptions
     forallTelescope th fun fvars th => do
-    trace[Progress] "All argumens: {fvars}"
+    trace[Progress] "All arguments: {fvars}"
   /-  -- Filter the argumens which are not propositions
     let rec getFirstPropIdx (i : Nat) : MetaM Nat := do
       if i ≥ fargs.size then pure i
-- 
cgit v1.2.3


From 2fa3cb8ee04dd7ff4184e3e1000fdc025abc50a4 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Mon, 17 Jul 2023 23:37:48 +0200
Subject: Start proving theorems for primitive definitions

---
 backends/lean/Base/Progress/Base.lean | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index a288d889..00b0a478 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -1,6 +1,7 @@
 import Lean
+import Std.Lean.HashSet
 import Base.Utils
-import Base.Primitives
+import Base.Primitives.Base
 
 namespace Progress
 
-- 
cgit v1.2.3


From e07177ee2de3fd1346ab6b1fc09aefbcb0e24459 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Tue, 18 Jul 2023 12:22:59 +0200
Subject: Improve progress

---
 backends/lean/Base/Progress/Base.lean | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 00b0a478..7eace667 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -58,10 +58,10 @@ section Methods
   def withPSpec [Inhabited (m a)] [Nonempty (m a)] (th : Expr) (k : PSpecDesc → m a)
     (sanityChecks : Bool := false) :
     m a := do
-    trace[Progress] "Theorem: {th}"
+    trace[Progress] "Proposition: {th}"
     -- Dive into the quantified variables and the assumptions
     forallTelescope th fun fvars th => do
-    trace[Progress] "All arguments: {fvars}"
+    trace[Progress] "Universally quantified arguments and assumptions: {fvars}"
   /-  -- Filter the argumens which are not propositions
     let rec getFirstPropIdx (i : Nat) : MetaM Nat := do
       if i ≥ fargs.size then pure i
@@ -83,12 +83,16 @@ section Methods
     -- Dive into the existentials
     existsTelescope th fun evars th => do
     trace[Progress] "Existentials: {evars}"
+    trace[Progress] "Proposition after stripping the quantifiers: {th}"
     -- Take the first conjunct
     let (th, post) ← optSplitConj th
+    trace[Progress] "After splitting the conjunction:\n- eq: {th}\n- post: {post}"
     -- Destruct the equality
     let (th, ret) ← destEq th
+    trace[Progress] "After splitting the equality:\n- lhs: {th}\n- rhs: {ret}"
     -- Destruct the application to get the name
-    th.withApp fun f args => do
+    th.consumeMData.withApp fun f args => do
+    trace[Progress] "After stripping the arguments:\n- f: {f}\n- args: {args}"
     if ¬ f.isConst then throwError "Not a constant: {f}"
     -- Compute the set of universally quantified variables which appear in the function arguments
     let allArgsFVars ← args.foldlM (fun hs arg => getFVarIds arg hs) HashSet.empty
@@ -114,6 +118,7 @@ section Methods
       post := post
     }
     k thDesc
+
 end Methods
 
 
-- 
cgit v1.2.3


From 36258c9ba583f19b5ddcb3b90e6521f9845b8944 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Wed, 19 Jul 2023 16:41:22 +0200
Subject: Start implementing support for some type classes for progress

---
 backends/lean/Base/Progress/Base.lean | 64 +++++++++++++++++++++++++++++++++--
 1 file changed, 62 insertions(+), 2 deletions(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 7eace667..0032c33d 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -121,19 +121,42 @@ section Methods
 
 end Methods
 
-
 def getPSpecFunName (th : Expr) : MetaM Name :=
   withPSpec th (fun d => do pure d.fName) true
 
+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
+
+-- "Regular" 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 ) x y`. As a consequence,
+   we store the theorem through the bindings: HAdd.add → Scalar → ...
+-/
+structure PSpecClassAttr where
+  attr : AttributeImpl
+  ext  : MapDeclarationExtension (NameMap Name)
+  deriving Inhabited
+
 /- The persistent map from function to pspec theorems. -/
 initialize pspecAttr : PSpecAttr ← do
   let ext ← mkMapDeclarationExtension `pspecMap
-  let attrImpl := {
+  let attrImpl : AttributeImpl := {
     name := `pspec
     descr := "Marks theorems to use with the `progress` tactic"
     add := fun thName stx attrKind => do
@@ -153,7 +176,44 @@ initialize pspecAttr : PSpecAttr ← do
   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) ← mkMapDeclarationExtension `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 }
+
+
 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
+
 end Progress
-- 
cgit v1.2.3


From abee28555eb9f95b1c548cc17b9fe746bc982b56 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Wed, 19 Jul 2023 18:50:19 +0200
Subject: Add some utilities for progress

---
 backends/lean/Base/Progress/Base.lean | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 0032c33d..785b9362 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -211,9 +211,28 @@ initialize pspecClassAttr : PSpecClassAttr ← do
 def PSpecAttr.find? (s : PSpecAttr) (name : Name) : MetaM (Option Name) := do
   return (s.ext.getState (← getEnv)).find? name
 
+def PSpecAttr.getState (s : PSpecAttr) : MetaM (NameMap Name) := do
+  pure (s.ext.getState (← getEnv))
+
 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 PSpecClassAttr.getState (s : PSpecClassAttr) : MetaM (NameMap (NameMap Name)) := 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
+
 end Progress
-- 
cgit v1.2.3


From 821b09b14794ebc2fe7b7047fc60fd56fb2cd107 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Wed, 19 Jul 2023 19:03:17 +0200
Subject: Fix a small issue with the persistent state of progress

---
 backends/lean/Base/Progress/Base.lean | 9 +++++++++
 1 file changed, 9 insertions(+)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 785b9362..72438d40 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -153,6 +153,15 @@ structure PSpecClassAttr where
   ext  : MapDeclarationExtension (NameMap Name)
   deriving Inhabited
 
+-- TODO: the original function doesn't define correctly the `addImportedFn`. Do a PR?
+def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) : IO (MapDeclarationExtension α) :=
+  registerSimplePersistentEnvExtension {
+    name          := name,
+    addImportedFn := fun a => a.foldl (fun s a => a.foldl (fun s (k, v) => s.insert k v) s) RBMap.empty,
+    addEntryFn    := fun s n => s.insert n.1 n.2 ,
+    toArrayFn     := fun es => es.toArray.qsort (fun a b => Name.quickLt a.1 b.1)
+  }
+
 /- The persistent map from function to pspec theorems. -/
 initialize pspecAttr : PSpecAttr ← do
   let ext ← mkMapDeclarationExtension `pspecMap
-- 
cgit v1.2.3


From 975b7c555cbffef2648a6469b777d1f1760d926d Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Thu, 20 Jul 2023 11:22:18 +0200
Subject: Improve progress further

---
 backends/lean/Base/Progress/Base.lean | 102 ++++++++++++++++++++++++++--------
 1 file changed, 80 insertions(+), 22 deletions(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 72438d40..3599d866 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -62,24 +62,6 @@ section Methods
     -- Dive into the quantified variables and the assumptions
     forallTelescope th fun fvars th => do
     trace[Progress] "Universally quantified arguments and assumptions: {fvars}"
-  /-  -- Filter the argumens which are not propositions
-    let rec getFirstPropIdx (i : Nat) : MetaM Nat := do
-      if i ≥ fargs.size then pure i
-      else do
-        let x := fargs.get! i
-        if ← Meta.isProp (← inferType x) then pure i
-        else getFirstPropIdx (i + 1)
-    let i ← getFirstPropIdx 0
-    let fvars := fargs.extract 0 i
-    let hyps := fargs.extract i fargs.size
-    trace[Progress] "Quantified variables: {fvars}"
-    trace[Progress] "Assumptions: {hyps}"
-    -- Sanity check: all hypotheses are propositions (in particular, all the
-    -- quantified variables are at the beginning)
-    let hypsAreProp ← hyps.allM fun x => do Meta.isProp (← inferType x)
-    if ¬ hypsAreProp then
-      throwError "The theorem doesn't have the proper shape: all the quantified arguments should be at the beginning"
-      -/
     -- Dive into the existentials
     existsTelescope th fun evars th => do
     trace[Progress] "Existentials: {evars}"
@@ -98,6 +80,8 @@ section Methods
     let allArgsFVars ← args.foldlM (fun hs arg => getFVarIds arg hs) HashSet.empty
     -- Sanity check
     if sanityChecks then
+      -- All the variables which appear in the inputs given to the function are
+      -- universally quantified (in particular, they are not *existentially* quantified)
       let fvarsSet : HashSet FVarId := HashSet.ofArray (fvars.map (fun x => x.fvarId!))
       let filtArgsFVars := allArgsFVars.toArray.filter (fun fvar => ¬ fvarsSet.contains fvar)
       if ¬ filtArgsFVars.isEmpty then
@@ -132,6 +116,12 @@ def getPSpecClassFunNames (th : Expr) : MetaM (Name × Name) :=
     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
+
 -- "Regular" pspec attribute
 structure PSpecAttr where
   attr : AttributeImpl
@@ -145,14 +135,26 @@ structure PSpecAttr where
    Example:
    ========
    We use type classes for addition. For instance, the addition between two
-   U32 is written (without syntactic sugar) as `HAdd.add (Scalar ) x y`. As a consequence,
+   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.
 -/
 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)
+  deriving Inhabited
+
 -- TODO: the original function doesn't define correctly the `addImportedFn`. Do a PR?
 def mkMapDeclarationExtension [Inhabited α] (name : Name := by exact decl_name%) : IO (MapDeclarationExtension α) :=
   registerSimplePersistentEnvExtension {
@@ -216,21 +218,69 @@ initialize pspecClassAttr : PSpecClassAttr ← do
   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) ← mkMapDeclarationExtension `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' (reduce 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)
+      setEnv env
+      pure ()
+  }
+  registerBuiltinAttribute attrImpl
+  pure { attr := attrImpl, ext := ext }
+
 
 def PSpecAttr.find? (s : PSpecAttr) (name : Name) : MetaM (Option Name) := do
   return (s.ext.getState (← getEnv)).find? name
 
-def PSpecAttr.getState (s : PSpecAttr) : MetaM (NameMap Name) := do
-  pure (s.ext.getState (← getEnv))
-
 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
+  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!""
@@ -244,4 +294,12 @@ def showStoredPSpecClass : MetaM Unit := do
     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!""
+  IO.println s
+
 end Progress
-- 
cgit v1.2.3


From d87e35e1a53b2252cc2c8c554216115773fd9678 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Thu, 20 Jul 2023 11:38:55 +0200
Subject: Add fine-grained lemmas for the arithmetic operations

---
 backends/lean/Base/Progress/Base.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 3599d866..2fbd24dd 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -240,7 +240,7 @@ initialize pspecClassExprAttr : PSpecClassExprAttr ← do
       -- We store two bindings:
       -- - arg to theorem name
       -- - reduced arg to theorem name
-      let rarg ← MetaM.run' (reduce arg)
+      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 :=
-- 
cgit v1.2.3


From 03492250b45855fe9db5e0a28a96166607cd84a1 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Thu, 20 Jul 2023 14:14:34 +0200
Subject: Make some proofs in Hashmap/Properties.lean and improve progress

---
 backends/lean/Base/Progress/Base.lean | 25 ++++++++++++++++++-------
 1 file changed, 18 insertions(+), 7 deletions(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index 2fbd24dd..b54bdf7a 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -19,6 +19,7 @@ structure PSpecDesc where
   -- The existentially quantified variables
   evars : Array Expr
   -- The function
+  fExpr : Expr
   fName : Name
   -- The function arguments
   fLevels : List Level
@@ -60,21 +61,30 @@ section Methods
     m a := do
     trace[Progress] "Proposition: {th}"
     -- Dive into the quantified variables and the assumptions
-    forallTelescope th fun fvars th => do
+    forallTelescope th.consumeMData fun fvars th => do
     trace[Progress] "Universally quantified arguments and assumptions: {fvars}"
     -- Dive into the existentials
-    existsTelescope th fun evars th => do
+    existsTelescope th.consumeMData fun evars th => do
     trace[Progress] "Existentials: {evars}"
     trace[Progress] "Proposition after stripping the quantifiers: {th}"
     -- Take the first conjunct
-    let (th, post) ← optSplitConj th
+    let (th, post) ← optSplitConj th.consumeMData
     trace[Progress] "After splitting the conjunction:\n- eq: {th}\n- post: {post}"
     -- Destruct the equality
-    let (th, ret) ← destEq th
+    let (mExpr, ret) ← destEq th.consumeMData
     trace[Progress] "After splitting the equality:\n- lhs: {th}\n- rhs: {ret}"
-    -- Destruct the application to get the name
-    th.consumeMData.withApp fun f args => do
-    trace[Progress] "After stripping the arguments:\n- f: {f}\n- args: {args}"
+    -- Destruct the monadic application to dive into the bind, if necessary (this
+    -- is for when we use `withPSpec` inside of the `progress` tactic), and
+    -- 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
+      if mf.isConst ∧ mf.constName = ``Bind.bind then do
+        -- Dive into the bind
+        let fExpr := margs.get! 4
+        fExpr.consumeMData.withApp fun f args => pure (fExpr, f, args)
+      else pure (mExpr, mf, margs)
+    trace[Progress] "After stripping the arguments of the function call:\n- f: {f}\n- args: {args}"
     if ¬ f.isConst then throwError "Not a constant: {f}"
     -- Compute the set of universally quantified variables which appear in the function arguments
     let allArgsFVars ← args.foldlM (fun hs arg => getFVarIds arg hs) HashSet.empty
@@ -94,6 +104,7 @@ section Methods
     let thDesc := {
       fvars := fvars
       evars := evars
+      fExpr
       fName := f.constName!
       fLevels := f.constLevels!
       args := args
-- 
cgit v1.2.3


From 1854c631a6a7a3f8d45ad18e05547f9d3782c3ee Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Tue, 25 Jul 2023 16:26:08 +0200
Subject: Make progress on the hashmap properties

---
 backends/lean/Base/Progress/Base.lean | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'backends/lean/Base/Progress/Base.lean')

diff --git a/backends/lean/Base/Progress/Base.lean b/backends/lean/Base/Progress/Base.lean
index b54bdf7a..6f820a84 100644
--- a/backends/lean/Base/Progress/Base.lean
+++ b/backends/lean/Base/Progress/Base.lean
@@ -81,8 +81,8 @@ section Methods
     let (fExpr, f, args) ← do
       if mf.isConst ∧ mf.constName = ``Bind.bind then do
         -- Dive into the bind
-        let fExpr := margs.get! 4
-        fExpr.consumeMData.withApp fun f args => pure (fExpr, f, args)
+        let fExpr := (margs.get! 4).consumeMData
+        fExpr.withApp fun f args => pure (fExpr, f, args)
       else pure (mExpr, mf, margs)
     trace[Progress] "After stripping the arguments of the function call:\n- f: {f}\n- args: {args}"
     if ¬ f.isConst then throwError "Not a constant: {f}"
-- 
cgit v1.2.3