-
Notifications
You must be signed in to change notification settings - Fork 0
/
tree.go
88 lines (75 loc) · 1.6 KB
/
tree.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
80
81
82
83
84
85
86
87
88
package main
import (
"net"
"sync"
)
// Tree is a binary tree that stores network prefixes
type Tree struct {
mu *sync.RWMutex // only used at the root node
children [2]*Tree
value interface{}
}
func popcount(i uint8) int {
i -= ((i >> 1) & 0x55)
i = (i & 0x33) + ((i >> 2) & 0x33)
return int((i + (i >> 4)) & 0x0F)
}
// New returns an initialized Tree
func New() *Tree {
return &Tree{mu: new(sync.RWMutex)}
}
// Insert adds a network and the associated value to the tree. The passed
// net.IPNet is assumed to be valid.
func (t *Tree) Insert(n *net.IPNet, value interface{}) {
t.mu.Lock()
defer t.mu.Unlock()
maskLen := 0
for _, b := range n.Mask {
maskLen += popcount(b)
}
node := t
for i := 0; i < maskLen; i++ {
bit := (n.IP[i/8] >> uint8(7-i%8)) & 0x01
c := node.children[bit]
if c == nil {
c = &Tree{}
node.children[bit] = c
}
node = c
}
node.value = value
return
}
// Lookup looks up the value associated with the most specific network that
// contains the passed net.IP. If no applicable network is found nil is
// returned.
func (t *Tree) Lookup(ip net.IP) interface{} {
if ip == nil {
return nil
}
if x := ip.To4(); x != nil {
ip = x
}
t.mu.RLock()
defer t.mu.RUnlock()
node := t
var longestPrefix *Tree
for i := 0; i < len(ip)*8; i++ {
bit := (ip[i/8] >> uint8(7-i%8)) & 0x01
child := node.children[bit]
if child == nil {
break
}
node = child
if node.value != nil {
longestPrefix = node
}
}
if longestPrefix == nil {
return nil
}
return longestPrefix.value
}
// path compression would probably be nice...
// func (t *Tree) Compact() {
// }