Start working on higher-order examples for Diverge

1 files changed, 62 insertions, 1 deletions

diff --git a/backends/lean/Base/Diverge/ElabBase.lean b/backends/lean/Base/Diverge/ElabBase.lean
index b818d5af..0d33e9d2 100644
--- a/backends/lean/Base/Diverge/ElabBase.lean
+++ b/backends/lean/Base/Diverge/ElabBase.lean
@@ -1,13 +1,15 @@
 import Lean
 import Base.Utils
 import Base.Primitives.Base
+import Base.Extensions
 
 namespace Diverge
 
 open Lean Elab Term Meta
-open Utils
+open Utils Extensions
 
 -- We can't define and use trace classes in the same file
+initialize registerTraceClass `Diverge
 initialize registerTraceClass `Diverge.elab
 initialize registerTraceClass `Diverge.def
 initialize registerTraceClass `Diverge.def.sigmas
@@ -20,4 +22,63 @@ initialize registerTraceClass `Diverge.def.unfold
 -- For the attribute (for higher-order functions)
 initialize registerTraceClass `Diverge.attr
 
+-- Attribute
+
+-- divspec attribute
+structure DivSpecAttr where
+  attr : AttributeImpl
+  ext  : DiscrTreeExtension Name true
+  deriving Inhabited
+
+/- The persistent map from expressions to divspec theorems. -/
+initialize divspecAttr : DivSpecAttr ← do
+  let ext ← mkDiscrTreeExtention `divspecMap true
+  let attrImpl : AttributeImpl := {
+    name := `divspec
+    descr := "Marks theorems to use with the `divergent` encoding"
+    add := fun thName stx attrKind => do
+      Attribute.Builtin.ensureNoArgs stx
+      -- TODO: use the attribute kind
+      unless attrKind == AttributeKind.global do
+        throwError "invalid attribute divspec, must be global"
+      -- Lookup the theorem
+      let env ← getEnv
+      let thDecl := env.constants.find! thName
+      let fKey : Array (DiscrTree.Key true) ← MetaM.run' (do
+        /- The theorem should have the shape:
+           `∀ ..., is_valid_p k (λ k => ...)`
+
+           Dive into the ∀:
+         -/
+        let (_, _, fExpr) ← forallMetaTelescope thDecl.type.consumeMData
+        /- Dive into the argument of `is_valid_p`: -/
+        fExpr.consumeMData.withApp fun _ args => do
+        if args.size ≠ 7 then throwError "Invalid number of arguments to is_valid_p"
+        let fExpr := args.get! 6
+        /- Dive into the lambda: -/
+        let (_, _, fExpr) ← lambdaMetaTelescope fExpr.consumeMData
+        trace[Diverge] "Registering divspec theorem for {fExpr}"
+        -- Convert the function expression to a discrimination tree key
+        DiscrTree.mkPath fExpr)
+      let env := ext.addEntry env (fKey, thName)
+      setEnv env
+      trace[Diverge] "Saved the environment"
+      pure ()
+  }
+  registerBuiltinAttribute attrImpl
+  pure { attr := attrImpl, ext := ext }
+
+def DivSpecAttr.find? (s : DivSpecAttr) (e : Expr) : MetaM (Array Name) := do
+  (s.ext.getState (← getEnv)).getMatch e
+
+def DivSpecAttr.getState (s : DivSpecAttr) : MetaM (DiscrTree Name true) := do
+  pure (s.ext.getState (← getEnv))
+
+def showStoredDivSpec : MetaM Unit := do
+  let st ← divspecAttr.getState
+  -- 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 Diverge