|
| 1 | +// Package rstree provides a datatype that supports building a space-efficient summary of networks |
| 2 | +// and IPs. |
| 3 | +package rstree |
| 4 | + |
| 5 | +import ( |
| 6 | + "github.com/PatrickCronin/routesum/pkg/routesum/bitslice" |
| 7 | +) |
| 8 | + |
| 9 | +type node struct { |
| 10 | + children *[2]*node |
| 11 | +} |
| 12 | + |
| 13 | +func (n *node) isLeaf() bool { |
| 14 | + return n.children == nil |
| 15 | +} |
| 16 | + |
| 17 | +func (n *node) childrenAreCompleteSubtree() bool { |
| 18 | + if n.children[0] == nil || n.children[1] == nil || |
| 19 | + !n.children[0].isLeaf() || !n.children[1].isLeaf() { |
| 20 | + return false |
| 21 | + } |
| 22 | + |
| 23 | + return true |
| 24 | +} |
| 25 | + |
| 26 | +// RSTree is a binary tree that supports the storage and retrieval of networks and IPs for the |
| 27 | +// purpose of route summarization. |
| 28 | +type RSTree struct { |
| 29 | + root *node |
| 30 | +} |
| 31 | + |
| 32 | +// NewRSTree returns an initialized RSTree for use |
| 33 | +func NewRSTree() *RSTree { |
| 34 | + return &RSTree{ |
| 35 | + root: nil, |
| 36 | + } |
| 37 | +} |
| 38 | + |
| 39 | +// InsertRoute inserts a new BitSlice into the tree. Each insert results in a space-optimized tree |
| 40 | +// structure. If a route being inserted is already covered by an existing route, it's simply |
| 41 | +// ignored. If a route being inserted covers one or more routes already stored, those routes are |
| 42 | +// replaced. |
| 43 | +func (t *RSTree) InsertRoute(routeBits bitslice.BitSlice) { |
| 44 | + // If the tree has no root node, create one. |
| 45 | + if t.root == nil { |
| 46 | + t.root = &node{children: nil} |
| 47 | + |
| 48 | + if len(routeBits) > 0 { |
| 49 | + t.root.children = new([2]*node) |
| 50 | + } |
| 51 | + } |
| 52 | + |
| 53 | + // Perform a non-recursive search of the tree's nodes for the best place to insert the route, |
| 54 | + // and do so. |
| 55 | + visited := []*node{} |
| 56 | + curNode := t.root |
| 57 | + for i := range routeBits { |
| 58 | + if curNode.isLeaf() { |
| 59 | + // Does the current node cover the requested route? If so, we're done. |
| 60 | + return |
| 61 | + } |
| 62 | + |
| 63 | + if curNode.children[routeBits[i]] == nil { |
| 64 | + curNode.children[routeBits[i]] = new(node) |
| 65 | + if i < len(routeBits)-1 { |
| 66 | + curNode.children[routeBits[i]].children = new([2]*node) |
| 67 | + } |
| 68 | + } |
| 69 | + |
| 70 | + visited = append(visited, curNode) |
| 71 | + curNode = curNode.children[routeBits[i]] |
| 72 | + } |
| 73 | + |
| 74 | + if !curNode.isLeaf() { |
| 75 | + // Does the requested route cover the current node? If so, update the current node. |
| 76 | + curNode.children = nil |
| 77 | + } |
| 78 | + |
| 79 | + simplifyVisitedSubtrees(visited) |
| 80 | +} |
| 81 | + |
| 82 | +// A completed subtree is a node in the tree whose children when taken together represent the |
| 83 | +// complete subtree below the node. For example, if the root node had a leaf node child for 0 and a |
| 84 | +// leaf-node child for 1, the node would be representing the "0" and 1" routes. But that's the same |
| 85 | +// as representing every possible route, so we'd simplify this to replace the root node with a leaf |
| 86 | +// node. |
| 87 | +// simplifyVisitedSubtrees takes a stack of visited nodes and simplifies completed subtrees as far |
| 88 | +// up the stack as possible. If at any point in the stack we find a node representing an incomplete |
| 89 | +// subtree, we stop. |
| 90 | +func simplifyVisitedSubtrees(visited []*node) { |
| 91 | + for i := len(visited) - 1; i >= 0; i-- { |
| 92 | + if !visited[i].childrenAreCompleteSubtree() { |
| 93 | + return |
| 94 | + } |
| 95 | + |
| 96 | + visited[i].children = nil |
| 97 | + } |
| 98 | +} |
| 99 | + |
| 100 | +type traversalStep struct { |
| 101 | + n *node |
| 102 | + precedingRouteBits bitslice.BitSlice |
| 103 | +} |
| 104 | + |
| 105 | +// Contents returns the BitSlices contained in the RSTree. |
| 106 | +func (t *RSTree) Contents() []bitslice.BitSlice { |
| 107 | + // If the tree is empty |
| 108 | + if t.root == nil { |
| 109 | + return []bitslice.BitSlice{} |
| 110 | + } |
| 111 | + |
| 112 | + // Otherwise |
| 113 | + queue := []traversalStep{ |
| 114 | + { |
| 115 | + n: t.root, |
| 116 | + precedingRouteBits: bitslice.BitSlice{}, |
| 117 | + }, |
| 118 | + } |
| 119 | + |
| 120 | + contents := []bitslice.BitSlice{} |
| 121 | + for len(queue) > 0 { |
| 122 | + step := queue[0] |
| 123 | + queue = queue[1:] |
| 124 | + |
| 125 | + if step.n.isLeaf() { |
| 126 | + contents = append(contents, step.precedingRouteBits) |
| 127 | + } else { |
| 128 | + lenPrecedingRouteBits := len(step.precedingRouteBits) |
| 129 | + toPrepend := make([]traversalStep, 0, 2) |
| 130 | + |
| 131 | + if step.n.children[0] != nil { |
| 132 | + lowChildBits := make([]byte, lenPrecedingRouteBits+1) |
| 133 | + copy(lowChildBits, step.precedingRouteBits) |
| 134 | + lowChildBits[lenPrecedingRouteBits] = 0 |
| 135 | + toPrepend = append(toPrepend, traversalStep{ |
| 136 | + n: step.n.children[0], |
| 137 | + precedingRouteBits: lowChildBits, |
| 138 | + }) |
| 139 | + } |
| 140 | + if step.n.children[1] != nil { |
| 141 | + highChildBits := make([]byte, lenPrecedingRouteBits+1) |
| 142 | + copy(highChildBits, step.precedingRouteBits) |
| 143 | + highChildBits[lenPrecedingRouteBits] = 1 |
| 144 | + toPrepend = append(toPrepend, traversalStep{ |
| 145 | + n: step.n.children[1], |
| 146 | + precedingRouteBits: highChildBits, |
| 147 | + }) |
| 148 | + } |
| 149 | + |
| 150 | + queue = append(toPrepend, queue...) |
| 151 | + } |
| 152 | + } |
| 153 | + |
| 154 | + return contents |
| 155 | +} |
0 commit comments