Skip to content

Commit

Permalink
Simplify the data structure
Browse files Browse the repository at this point in the history
  • Loading branch information
PatrickCronin committed Dec 31, 2022
1 parent 1bb8216 commit daea4e0
Show file tree
Hide file tree
Showing 6 changed files with 395 additions and 565 deletions.
8 changes: 4 additions & 4 deletions pkg/routesum/routesum.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,21 @@ import (
"strings"

"github.com/PatrickCronin/routesum/pkg/routesum/bitslice"
"github.com/PatrickCronin/routesum/pkg/routesum/rstrie"
"github.com/PatrickCronin/routesum/pkg/routesum/rstree"
"github.com/pkg/errors"
"inet.af/netaddr"
)

// RouteSum has methods supporting route summarization of networks and hosts
type RouteSum struct {
ipv4, ipv6 *rstrie.RSTrie
ipv4, ipv6 *rstree.RSTree
}

// NewRouteSum returns an initialized RouteSum object
func NewRouteSum() *RouteSum {
rs := new(RouteSum)
rs.ipv4 = rstrie.NewRSTrie()
rs.ipv6 = rstrie.NewRSTrie()
rs.ipv4 = rstree.NewRSTree()
rs.ipv6 = rstree.NewRSTree()

return rs
}
Expand Down
12 changes: 6 additions & 6 deletions pkg/routesum/routesum_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -456,7 +456,7 @@ func TestSummarize(t *testing.T) { //nolint: funlen
}
}

func TestRSTrieMemUsage(t *testing.T) { //nolint: funlen
func TestMemUsage(t *testing.T) { //nolint: funlen
tests := []struct {
name string
entries []string
Expand All @@ -473,7 +473,7 @@ func TestRSTrieMemUsage(t *testing.T) { //nolint: funlen
entries: []string{
"192.0.2.1",
},
expectedNumInternalNodes: 0,
expectedNumInternalNodes: 32,
expectedNumLeafNodes: 1,
},
{
Expand All @@ -482,7 +482,7 @@ func TestRSTrieMemUsage(t *testing.T) { //nolint: funlen
"192.0.2.1",
"192.0.2.0",
},
expectedNumInternalNodes: 0,
expectedNumInternalNodes: 31,
expectedNumLeafNodes: 1,
},
{
Expand All @@ -491,15 +491,15 @@ func TestRSTrieMemUsage(t *testing.T) { //nolint: funlen
"192.0.2.1",
"192.0.2.2",
},
expectedNumInternalNodes: 1,
expectedNumInternalNodes: 33,
expectedNumLeafNodes: 2,
},
{
name: "one item, IPv6",
entries: []string{
"2001:db8::1",
},
expectedNumInternalNodes: 0,
expectedNumInternalNodes: 128,
expectedNumLeafNodes: 1,
},
{
Expand All @@ -508,7 +508,7 @@ func TestRSTrieMemUsage(t *testing.T) { //nolint: funlen
"192.0.2.0",
"2001:db8::1",
},
expectedNumInternalNodes: 0,
expectedNumInternalNodes: 160,
expectedNumLeafNodes: 2,
},
}
Expand Down
229 changes: 229 additions & 0 deletions pkg/routesum/rstree/rstree.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,229 @@
// Package rstree provides a datatype that supports building a space-efficient summary of networks
// and IPs.
package rstree

import (
"container/list"
"unsafe"

"github.com/PatrickCronin/routesum/pkg/routesum/bitslice"
)

type node struct {
children *[2]*node
}

func (n *node) isLeaf() bool {
return n.children == nil
}

// RSTree is a binary tree that supports the storage and retrieval of networks and IPs for the
// purpose of route summarization.
type RSTree struct {
root *node
}

// NewRSTree returns an initialized RSTree for use
func NewRSTree() *RSTree {
return &RSTree{
root: nil,
}
}

// InsertRoute inserts a new BitSlice into the tree. Each insert results in a space-optimized tree
// structure. If a route being inserted is already covered by an existing route, it's simply
// ignored. If a route being inserted covers one or more routes already stored, those routes are
// replaced.
func (t *RSTree) InsertRoute(routeBits bitslice.BitSlice) {
// If the tree has no root node, create one.
if t.root == nil {
t.root = &node{children: nil}

if len(routeBits) > 0 {
t.root.children = new([2]*node)
}
}

t.root.insertRoute(routeBits)
}

// insertRoute returns whether or not a change was made to the tree that might require upwards
// optimization.
func (n *node) insertRoute(remainingRouteBits bitslice.BitSlice) bool {
// Does the current node cover the requested route? If so, we're done.
if n.isLeaf() {
return false
}

// Does the requested route cover the current node? If so, update the current node.
remainingRouteBitsLen := len(remainingRouteBits)
if remainingRouteBitsLen == 0 {
n.children = nil
return true
}

// Otherwise the requested route diverges from the current node.
nextBit := remainingRouteBits[0]
if n.children[nextBit] == nil {
// As an optimization, if we would create a node only to realize it is redundant, just
// trim the rundundant child now and return.
if remainingRouteBitsLen == 1 &&
n.children[^nextBit&1] != nil &&
n.children[^nextBit&1].isLeaf() {
n.children = nil
return true
}

// Otherwise we add a node
n.children[nextBit] = new(node)
if remainingRouteBitsLen > 1 {
n.children[nextBit].children = new([2]*node)
}
}

// Traverse to the new node
if n.children[nextBit].insertRoute(remainingRouteBits[1:]) {
return n.maybeRemoveRedundantChildren()
}

return false
}

// A node's children are redundant if they, taken together, represent a complete subtree from the
// node's perspective. This situation can be represented more simply as the node having a nil
// children pointer.
func (n *node) maybeRemoveRedundantChildren() bool {
if n.isLeaf() {
return false
}

if n.children[0] == nil ||
!n.children[0].isLeaf() ||
n.children[1] == nil ||
!n.children[1].isLeaf() {
return false
}

n.children = nil
return true
}

type traversalStep struct {
n *node
precedingRouteBits bitslice.BitSlice
}

// Contents returns the BitSlices contained in the RSTree.
func (t *RSTree) Contents() []bitslice.BitSlice {
// If the tree is empty
if t.root == nil {
return []bitslice.BitSlice{}
}

// Otherwise
remainingSteps := list.New()
remainingSteps.PushFront(traversalStep{
n: t.root,
precedingRouteBits: bitslice.BitSlice{},
})

contents := []bitslice.BitSlice{}
for remainingSteps.Len() > 0 {
step := remainingSteps.Remove(remainingSteps.Front()).(traversalStep)

if step.n.isLeaf() {
contents = append(contents, step.precedingRouteBits)
} else {
lenPrecedingRouteBits := len(step.precedingRouteBits)

if step.n.children[1] != nil {
highChildBits := make([]byte, lenPrecedingRouteBits+1)
copy(highChildBits, step.precedingRouteBits)
highChildBits[lenPrecedingRouteBits] = 1
remainingSteps.PushFront(traversalStep{
n: step.n.children[1],
precedingRouteBits: highChildBits,
})
}

if step.n.children[0] != nil {
lowChildBits := make([]byte, lenPrecedingRouteBits+1)
copy(lowChildBits, step.precedingRouteBits)
lowChildBits[lenPrecedingRouteBits] = 0
remainingSteps.PushFront(traversalStep{
n: step.n.children[0],
precedingRouteBits: lowChildBits,
})
}
}
}

return contents
}

func (t *RSTree) visitAll(cb func(*node)) {
// If the trie is empty
if t.root == nil {
return
}

// Otherwise
remainingSteps := list.New()
remainingSteps.PushFront(traversalStep{
n: t.root,
precedingRouteBits: bitslice.BitSlice{},
})

for remainingSteps.Len() > 0 {
curNode := remainingSteps.Remove(remainingSteps.Front()).(traversalStep)

// Act on this node
cb(curNode.n)

// Traverse the remainder of the nodes
if !curNode.n.isLeaf() {
lenPrecedingRouteBits := len(curNode.precedingRouteBits)

if curNode.n.children[1] != nil {
highChildBits := make([]byte, lenPrecedingRouteBits+1)
copy(highChildBits, curNode.precedingRouteBits)
highChildBits[lenPrecedingRouteBits] = 1
remainingSteps.PushFront(traversalStep{
n: curNode.n.children[1],
precedingRouteBits: highChildBits,
})
}

if curNode.n.children[0] != nil {
lowChildBits := make([]byte, lenPrecedingRouteBits+1)
copy(lowChildBits, curNode.precedingRouteBits)
lowChildBits[lenPrecedingRouteBits] = 0
remainingSteps.PushFront(traversalStep{
n: curNode.n.children[0],
precedingRouteBits: lowChildBits,
})
}
}
}
}

// MemUsage returns information about an RSTrie's current size in memory.
func (t *RSTree) MemUsage() (uint, uint, uintptr, uintptr) {
var numInternalNodes, numLeafNodes uint
var internalNodesTotalSize, leafNodesTotalSize uintptr

tallyNode := func(n *node) {
baseNodeSize := unsafe.Sizeof(node{}) //nolint: exhaustruct, gosec
if n.isLeaf() {
numLeafNodes++
leafNodesTotalSize += baseNodeSize
return
}

numInternalNodes++
internalNodesTotalSize += baseNodeSize + unsafe.Sizeof([2]*node{}) //nolint: gosec
}
t.visitAll(tallyNode)

return numInternalNodes, numLeafNodes, internalNodesTotalSize, leafNodesTotalSize
}
Loading

0 comments on commit daea4e0

Please sign in to comment.