From b3ebef29fe3f13a9004b39fcb89afb33fbbfd248 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Thu, 7 Dec 2023 20:52:04 +0100
Subject: Start working on a version of Diverge.FixI more suited to
 higher-order functions

---
 backends/lean/Base/Diverge/ElabBase.lean | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'backends/lean/Base/Diverge/ElabBase.lean')

diff --git a/backends/lean/Base/Diverge/ElabBase.lean b/backends/lean/Base/Diverge/ElabBase.lean
index fedb1c74..1a0676e2 100644
--- a/backends/lean/Base/Diverge/ElabBase.lean
+++ b/backends/lean/Base/Diverge/ElabBase.lean
@@ -1,8 +1,11 @@
 import Lean
+import Base.Utils
+import Base.Primitives.Base
 
 namespace Diverge
 
 open Lean Elab Term Meta
+open Utils
 
 -- We can't define and use trace classes in the same file
 initialize registerTraceClass `Diverge.elab
@@ -12,4 +15,7 @@ initialize registerTraceClass `Diverge.def.genBody
 initialize registerTraceClass `Diverge.def.valid
 initialize registerTraceClass `Diverge.def.unfold
 
+-- For the attribute (for higher-order functions)
+initialize registerTraceClass `Diverge.attr
+
 end Diverge
-- 
cgit v1.2.3


From c23a37617188a1bbf913b5c700522abc33bf39c9 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Mon, 11 Dec 2023 17:00:01 +0100
Subject: Update Diverge/Elab.lean to use the more general FixII definitions

---
 backends/lean/Base/Diverge/ElabBase.lean | 2 ++
 1 file changed, 2 insertions(+)

(limited to 'backends/lean/Base/Diverge/ElabBase.lean')

diff --git a/backends/lean/Base/Diverge/ElabBase.lean b/backends/lean/Base/Diverge/ElabBase.lean
index 1a0676e2..b818d5af 100644
--- a/backends/lean/Base/Diverge/ElabBase.lean
+++ b/backends/lean/Base/Diverge/ElabBase.lean
@@ -11,7 +11,9 @@ open Utils
 initialize registerTraceClass `Diverge.elab
 initialize registerTraceClass `Diverge.def
 initialize registerTraceClass `Diverge.def.sigmas
+initialize registerTraceClass `Diverge.def.prods
 initialize registerTraceClass `Diverge.def.genBody
+initialize registerTraceClass `Diverge.def.genBody.visit
 initialize registerTraceClass `Diverge.def.valid
 initialize registerTraceClass `Diverge.def.unfold
 
-- 
cgit v1.2.3


From 78367ef21c147b26040e0f6062a907fceab1f390 Mon Sep 17 00:00:00 2001
From: Son Ho
Date: Mon, 11 Dec 2023 18:34:10 +0100
Subject: Start working on higher-order examples for Diverge

---
 backends/lean/Base/Diverge/ElabBase.lean | 63 +++++++++++++++++++++++++++++++-
 1 file changed, 62 insertions(+), 1 deletion(-)

(limited to 'backends/lean/Base/Diverge/ElabBase.lean')

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
-- 
cgit v1.2.3

