import Lean import Base.Utils import Base.Primitives.Base import Base.Extensions namespace Diverge open Lean Elab Term Meta 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 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 -- For the attribute (for higher-order functions) initialize registerTraceClass `Diverge.attr -- Attribute -- divspec attribute structure DivSpecAttr where attr : AttributeImpl ext : DiscrTreeExtension Name deriving Inhabited /- The persistent map from expressions to divspec theorems. -/ initialize divspecAttr : DivSpecAttr ← do let ext ← mkDiscrTreeExtention `divspecMap 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) ← 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 -- We use the default configuration let config : WhnfCoreConfig := {} DiscrTree.mkPath fExpr config) 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 -- We use the default configuration let config : WhnfCoreConfig := {} (s.ext.getState (← getEnv)).getMatch e config def DivSpecAttr.getState (s : DivSpecAttr) : MetaM (DiscrTree Name) := 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