device: remove recursion from insertion and connect parent pointers

This makes the insertion algorithm a bit more efficient, while also now
taking on the additional task of connecting up parent pointers. This
will be handy in the following commit.

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

device/allowedips.go

@@ -14,9 +14,15 @@ import (
 	"unsafe"
 )
 
+type parentIndirection struct {
+	parentBit     **trieEntry
+	parentBitType uint8
+}
+
 type trieEntry struct {
 	peer        *Peer
 	child       [2]*trieEntry
+	parent      parentIndirection
 	cidr        uint8
 	bitAtByte   uint8
 	bitAtShift  uint8
@@ -114,43 +120,45 @@ func (node *trieEntry) maskSelf() {
 	}
 }
 
-func (node *trieEntry) insert(ip net.IP, cidr uint8, peer *Peer) *trieEntry {
+func (node *trieEntry) nodePlacement(ip net.IP, cidr uint8) (parent *trieEntry, exact bool) {
+	for node != nil && node.cidr <= cidr && commonBits(node.bits, ip) >= node.cidr {
+		parent = node
+		if parent.cidr == cidr {
+			exact = true
+			return
+		}
+		bit := node.choose(ip)
+		node = node.child[bit]
+	}
+	return
+}
 
-	// at leaf
+func (trie parentIndirection) insert(ip net.IP, cidr uint8, peer *Peer) {
-
+	if *trie.parentBit == nil {
-	if node == nil {
 		node := &trieEntry{
-			bits:       ip,
 			peer:       peer,
+			parent:     trie,
+			bits:       ip,
 			cidr:       cidr,
 			bitAtByte:  cidr / 8,
 			bitAtShift: 7 - (cidr % 8),
 		}
 		node.maskSelf()
 		node.addToPeerEntries()
-		return node
+		*trie.parentBit = node
+		return
 	}
-
+	node, exact := (*trie.parentBit).nodePlacement(ip, cidr)
-	// traverse deeper
+	if exact {
-
+		node.removeFromPeerEntries()
-	common := commonBits(node.bits, ip)
+		node.peer = peer
-	if node.cidr <= cidr && common >= node.cidr {
+		node.addToPeerEntries()
-		if node.cidr == cidr {
+		return
-			node.removeFromPeerEntries()
-			node.peer = peer
-			node.addToPeerEntries()
-			return node
-		}
-		bit := node.choose(ip)
-		node.child[bit] = node.child[bit].insert(ip, cidr, peer)
-		return node
 	}
 
-	// split node
-
 	newNode := &trieEntry{
-		bits:       ip,
 		peer:       peer,
+		bits:       ip,
 		cidr:       cidr,
 		bitAtByte:  cidr / 8,
 		bitAtShift: 7 - (cidr % 8),
@@ -158,34 +166,61 @@ func (node *trieEntry) insert(ip net.IP, cidr uint8, peer *Peer) *trieEntry {
 	newNode.maskSelf()
 	newNode.addToPeerEntries()
 
+	var down *trieEntry
+	if node == nil {
+		down = *trie.parentBit
+	} else {
+		bit := node.choose(ip)
+		down = node.child[bit]
+		if down == nil {
+			newNode.parent = parentIndirection{&node.child[bit], bit}
+			node.child[bit] = newNode
+			return
+		}
+	}
+	common := commonBits(down.bits, ip)
 	if common < cidr {
 		cidr = common
 	}
-
+	parent := node
-	// check for shorter prefix
 
 	if newNode.cidr == cidr {
-		bit := newNode.choose(node.bits)
+		bit := newNode.choose(down.bits)
-		newNode.child[bit] = node
+		down.parent = parentIndirection{&newNode.child[bit], bit}
-		return newNode
+		newNode.child[bit] = down
+		if parent == nil {
+			newNode.parent = trie
+			*trie.parentBit = newNode
+		} else {
+			bit := parent.choose(newNode.bits)
+			newNode.parent = parentIndirection{&parent.child[bit], bit}
+			parent.child[bit] = newNode
+		}
+		return
 	}
 
-	// create new parent for node & newNode
+	node = &trieEntry{
-
+		bits:       append([]byte{}, newNode.bits...),
-	parent := &trieEntry{
-		bits:       append([]byte{}, ip...),
-		peer:       nil,
 		cidr:       cidr,
 		bitAtByte:  cidr / 8,
 		bitAtShift: 7 - (cidr % 8),
 	}
-	parent.maskSelf()
+	node.maskSelf()
 
-	bit := parent.choose(ip)
+	bit := node.choose(down.bits)
-	parent.child[bit] = newNode
+	down.parent = parentIndirection{&node.child[bit], bit}
-	parent.child[bit^1] = node
+	node.child[bit] = down
-
+	bit = node.choose(newNode.bits)
-	return parent
+	newNode.parent = parentIndirection{&node.child[bit], bit}
+	node.child[bit] = newNode
+	if parent == nil {
+		node.parent = trie
+		*trie.parentBit = node
+	} else {
+		bit := parent.choose(node.bits)
+		node.parent = parentIndirection{&parent.child[bit], bit}
+		parent.child[bit] = node
+	}
 }
 
 func (node *trieEntry) lookup(ip net.IP) *Peer {
@@ -236,9 +271,9 @@ func (table *AllowedIPs) Insert(ip net.IP, cidr uint8, peer *Peer) {
 
 	switch len(ip) {
 	case net.IPv6len:
-		table.IPv6 = table.IPv6.insert(ip, cidr, peer)
+		parentIndirection{&table.IPv6, 2}.insert(ip, cidr, peer)
 	case net.IPv4len:
-		table.IPv4 = table.IPv4.insert(ip, cidr, peer)
+		parentIndirection{&table.IPv4, 2}.insert(ip, cidr, peer)
 	default:
 		panic(errors.New("inserting unknown address type"))
 	}

device/allowedips_rand_test.go

@@ -65,9 +65,9 @@ func (r SlowRouter) Lookup(addr []byte) *Peer {
 }
 
 func TestTrieRandomIPv4(t *testing.T) {
-	var trie *trieEntry
 	var slow SlowRouter
 	var peers []*Peer
+	var allowedIPs AllowedIPs
 
 	rand.Seed(1)
 
@@ -82,7 +82,7 @@ func TestTrieRandomIPv4(t *testing.T) {
 		rand.Read(addr[:])
 		cidr := uint8(rand.Uint32() % (AddressLength * 8))
 		index := rand.Int() % NumberOfPeers
-		trie = trie.insert(addr[:], cidr, peers[index])
+		allowedIPs.Insert(addr[:], cidr, peers[index])
 		slow = slow.Insert(addr[:], cidr, peers[index])
 	}
 
@@ -90,7 +90,7 @@ func TestTrieRandomIPv4(t *testing.T) {
 		var addr [AddressLength]byte
 		rand.Read(addr[:])
 		peer1 := slow.Lookup(addr[:])
-		peer2 := trie.lookup(addr[:])
+		peer2 := allowedIPs.LookupIPv4(addr[:])
 		if peer1 != peer2 {
 			t.Error("Trie did not match naive implementation, for:", addr)
 		}
@@ -98,9 +98,9 @@ func TestTrieRandomIPv4(t *testing.T) {
 }
 
 func TestTrieRandomIPv6(t *testing.T) {
-	var trie *trieEntry
 	var slow SlowRouter
 	var peers []*Peer
+	var allowedIPs AllowedIPs
 
 	rand.Seed(1)
 
@@ -115,7 +115,7 @@ func TestTrieRandomIPv6(t *testing.T) {
 		rand.Read(addr[:])
 		cidr := uint8(rand.Uint32() % (AddressLength * 8))
 		index := rand.Int() % NumberOfPeers
-		trie = trie.insert(addr[:], cidr, peers[index])
+		allowedIPs.Insert(addr[:], cidr, peers[index])
 		slow = slow.Insert(addr[:], cidr, peers[index])
 	}
 
@@ -123,7 +123,7 @@ func TestTrieRandomIPv6(t *testing.T) {
 		var addr [AddressLength]byte
 		rand.Read(addr[:])
 		peer1 := slow.Lookup(addr[:])
-		peer2 := trie.lookup(addr[:])
+		peer2 := allowedIPs.LookupIPv6(addr[:])
 		if peer1 != peer2 {
 			t.Error("Trie did not match naive implementation, for:", addr)
 		}

device/allowedips_test.go

@@ -42,6 +42,7 @@ func TestCommonBits(t *testing.T) {
 func benchmarkTrie(peerNumber int, addressNumber int, addressLength int, b *testing.B) {
 	var trie *trieEntry
 	var peers []*Peer
+	root := parentIndirection{&trie, 2}
 
 	rand.Seed(1)
 
@@ -56,7 +57,7 @@ func benchmarkTrie(peerNumber int, addressNumber int, addressLength int, b *test
 		rand.Read(addr[:])
 		cidr := uint8(rand.Uint32() % (AddressLength * 8))
 		index := rand.Int() % peerNumber
-		trie = trie.insert(addr[:], cidr, peers[index])
+		root.insert(addr[:], cidr, peers[index])
 	}
 
 	for n := 0; n < b.N; n++ {
@@ -94,21 +95,21 @@ func TestTrieIPv4(t *testing.T) {
 	g := &Peer{}
 	h := &Peer{}
 
-	var trie *trieEntry
+	var allowedIPs AllowedIPs
 
 	insert := func(peer *Peer, a, b, c, d byte, cidr uint8) {
-		trie = trie.insert([]byte{a, b, c, d}, cidr, peer)
+		allowedIPs.Insert([]byte{a, b, c, d}, cidr, peer)
 	}
 
 	assertEQ := func(peer *Peer, a, b, c, d byte) {
-		p := trie.lookup([]byte{a, b, c, d})
+		p := allowedIPs.LookupIPv4([]byte{a, b, c, d})
 		if p != peer {
 			t.Error("Assert EQ failed")
 		}
 	}
 
 	assertNEQ := func(peer *Peer, a, b, c, d byte) {
-		p := trie.lookup([]byte{a, b, c, d})
+		p := allowedIPs.LookupIPv4([]byte{a, b, c, d})
 		if p == peer {
 			t.Error("Assert NEQ failed")
 		}
@@ -150,7 +151,7 @@ func TestTrieIPv4(t *testing.T) {
 	assertEQ(a, 192, 0, 0, 0)
 	assertEQ(a, 255, 0, 0, 0)
 
-	trie = trie.removeByPeer(a)
+	allowedIPs.RemoveByPeer(a)
 
 	assertNEQ(a, 1, 0, 0, 0)
 	assertNEQ(a, 64, 0, 0, 0)
@@ -158,12 +159,12 @@ func TestTrieIPv4(t *testing.T) {
 	assertNEQ(a, 192, 0, 0, 0)
 	assertNEQ(a, 255, 0, 0, 0)
 
-	trie = nil
+	allowedIPs = AllowedIPs{}
 
 	insert(a, 192, 168, 0, 0, 16)
 	insert(a, 192, 168, 0, 0, 24)
 
-	trie = trie.removeByPeer(a)
+	allowedIPs.RemoveByPeer(a)
 
 	assertNEQ(a, 192, 168, 0, 1)
 }
@@ -181,7 +182,7 @@ func TestTrieIPv6(t *testing.T) {
 	g := &Peer{}
 	h := &Peer{}
 
-	var trie *trieEntry
+	var allowedIPs AllowedIPs
 
 	expand := func(a uint32) []byte {
 		var out [4]byte
@@ -198,7 +199,7 @@ func TestTrieIPv6(t *testing.T) {
 		addr = append(addr, expand(b)...)
 		addr = append(addr, expand(c)...)
 		addr = append(addr, expand(d)...)
-		trie = trie.insert(addr, cidr, peer)
+		allowedIPs.Insert(addr, cidr, peer)
 	}
 
 	assertEQ := func(peer *Peer, a, b, c, d uint32) {
@@ -207,7 +208,7 @@ func TestTrieIPv6(t *testing.T) {
 		addr = append(addr, expand(b)...)
 		addr = append(addr, expand(c)...)
 		addr = append(addr, expand(d)...)
-		p := trie.lookup(addr)
+		p := allowedIPs.LookupIPv6(addr)
 		if p != peer {
 			t.Error("Assert EQ failed")
 		}