forked from open-policy-agent/opa
-
Notifications
You must be signed in to change notification settings - Fork 0
/
functrie.go
79 lines (69 loc) · 1.57 KB
/
functrie.go
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
package planner
import (
"github.com/open-policy-agent/opa/ast"
"github.com/open-policy-agent/opa/internal/ir"
)
// functrie implements a simple trie structure for organizing planned functions.
// The functions are organized to facilitate access when planning references
// against the data document.
type functrie struct {
children map[ast.Value]*functrie
val *functrieValue
}
type functrieValue struct {
Fn *ir.Func
Rules []*ast.Rule
}
func (val *functrieValue) Arity() int {
return len(val.Rules[0].Head.Args)
}
func newFunctrie() *functrie {
return &functrie{
children: map[ast.Value]*functrie{},
}
}
func (t *functrie) Insert(key ast.Ref, val *functrieValue) {
node := t
for _, elem := range key {
child, ok := node.children[elem.Value]
if !ok {
child = newFunctrie()
node.children[elem.Value] = child
}
node = child
}
node.val = val
}
func (t *functrie) Lookup(key ast.Ref) *functrieValue {
node := t
for _, elem := range key {
var ok bool
if node == nil {
return nil
} else if node, ok = node.children[elem.Value]; !ok {
return nil
}
}
return node.val
}
func (t *functrie) LookupOrInsert(key ast.Ref, orElse *functrieValue) *functrieValue {
if val := t.Lookup(key); val != nil {
return val
}
t.Insert(key, orElse)
return orElse
}
func (t *functrie) FuncMap() map[string]*ir.Func {
result := map[string]*ir.Func{}
t.toMap(result)
return result
}
func (t *functrie) toMap(result map[string]*ir.Func) {
if t.val != nil {
result[t.val.Fn.Name] = t.val.Fn
return
}
for _, node := range t.children {
node.toMap(result)
}
}