blob: 08ef96f7099d9b16a31c1886edcfc7bc8a9e388f (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
|
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
|
