diff options
Diffstat (limited to 'backends/lean/Base/Diverge/ElabBase.lean')
-rw-r--r-- | backends/lean/Base/Diverge/ElabBase.lean | 63 |
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 |