summaryrefslogtreecommitdiff
path: root/backends/lean/Base/Diverge/ElabBase.lean
blob: 0d33e9d25f9dbbf6ab8c75675e3c9d8ae7e3373a (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
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 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