Start implementing support for some type classes for progress

2 files changed, 117 insertions, 25 deletions

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
diff --git a/backends/lean/Base/Progress/Progress.lean b/backends/lean/Base/Progress/Progress.lean
index 9e2461a2..64d1c14a 100644
--- a/backends/lean/Base/Progress/Progress.lean
+++ b/backends/lean/Base/Progress/Progress.lean
@@ -152,6 +152,15 @@ def progressWith (fnExpr : Expr) (th : TheoremOrLocal) (ids : Array Name)
     --
     pure .Ok
 
+-- Small utility: if `args` is not empty, return the name of the app in the first
+-- arg, if it is a const.
+def getFirstArgAppName (args : Array Expr) : MetaM (Option Name) := do
+  if args.size = 0 then pure none
+  else
+    (args.get! 0).withApp fun f _ => do
+    if f.isConst then pure (some f.constName)
+    else pure none
+
 -- The array of ids are identifiers to use when introducing fresh variables
 def progressAsmsOrLookupTheorem (withTh : Option TheoremOrLocal) (ids : Array Name) (asmTac : TacticM Unit) : TacticM Unit := do
   withMainContext do
@@ -187,34 +196,57 @@ def progressAsmsOrLookupTheorem (withTh : Option TheoremOrLocal) (ids : Array Na
     -- 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 thName ← do
-      match ← pspecAttr.find? fName with
-      | none => throwError "Could not find a pspec theorem for {fName}"
-      | some thName => pure thName
-    trace[Progress] "Lookuped up theorem: {thName}"
-    -- Apply the theorem
-    let res ← do
-      try
-        let res ← progressWith fnExpr (.Theorem thName) ids asmTac
-        pure (some res)
-      catch _ => none
+    let res ←
+        match ← pspecAttr.find? fName with
+        | some thName =>
+          trace[Progress] "Lookuped up theorem: {thName}"
+          -- Apply the theorem
+          let res ← do
+            try
+              let res ← progressWith fnExpr (.Theorem thName) ids asmTac
+              pure (some res)
+            catch _ => none
+        | none =>
+          trace[Progress] "Could not find a pspec theorem for {fName}"
+          throwError "TODO"
     match res with
     | some .Ok => return ()
     | some (.Error msg) => throwError msg
     | none =>
-    -- Try a recursive call - we try the assumptions of kind "auxDecl"
-    let ctx ← Lean.MonadLCtx.getLCtx
-    let decls ← ctx.getAllDecls
-    let decls := decls.filter (λ decl => match decl.kind with
-      | .default | .implDetail => false | .auxDecl => true)
-    for decl in decls.reverse do
-      trace[Progress] "Trying recursive assumption: {decl.userName} : {decl.type}"
-      let res ← do try progressWith fnExpr (.Local decl) ids asmTac catch _ => continue
+      -- It failed: try to lookup a *class* spec theorem
+      let res ← do
+        match ← getFirstArgAppName args with
+        | none => none
+        | some argName => do
+          match ← pspecClassAttr.find? fName argName with
+          | some thName =>
+            trace[Progress] "Lookuped up class theorem: {thName}"
+            -- Apply the theorem
+            let res ← do
+              try
+                let res ← progressWith fnExpr (.Theorem thName) ids asmTac
+                pure (some res)
+              catch _ => none
+          | none =>
+            trace[Progress] "Could not find a pspec theorem for {fName}"
+            pure none
       match res with
-      | .Ok => return ()
-      | .Error msg => throwError msg
-    -- Nothing worked: failed
-    throwError "Progress failed"
+      | some .Ok => return ()
+      | some (.Error msg) => throwError msg
+      | none =>
+      -- Try a recursive call - we try the assumptions of kind "auxDecl"
+      let ctx ← Lean.MonadLCtx.getLCtx
+      let decls ← ctx.getAllDecls
+      let decls := decls.filter (λ decl => match decl.kind with
+        | .default | .implDetail => false | .auxDecl => true)
+      for decl in decls.reverse do
+        trace[Progress] "Trying recursive assumption: {decl.userName} : {decl.type}"
+        let res ← do try progressWith fnExpr (.Local decl) ids asmTac catch _ => continue
+        match res with
+        | .Ok => return ()
+        | .Error msg => throwError msg
+      -- Nothing worked: failed
+      throwError "Progress failed"
 
 syntax progressArgs := ("with" ident)? ("as" " ⟨ " (ident)+ " ⟩")?