1 files changed, 56 insertions, 23 deletions

diff --git a/src/trie.go b/src/trie.go
index 7fd7c5f..31a4d92 100644
--- a/src/trie.go
+++ b/src/trie.go
@@ -1,9 +1,11 @@
 package main
 
-import "fmt"
-
-/* Syncronization must be done seperatly
+/* Binary trie
+ *
+ * Syncronization done seperatly
+ * See: routing.go
  *
+ * Todo: Better commenting
  */
 
 type Trie struct {
@@ -13,7 +15,6 @@ type Trie struct {
 	peer  *Peer
 
 	// Index of "branching" bit
-	// bit_at_shift
 	bit_at_byte  uint
 	bit_at_shift uint
 }
@@ -92,7 +93,14 @@ func (node *Trie) RemovePeer(p *Peer) *Trie {
 	return node.child[0]
 }
 
+func (node *Trie) choose(key []byte) byte {
+	return (key[node.bit_at_byte] >> node.bit_at_shift) & 1
+}
+
 func (node *Trie) Insert(key []byte, cidr uint, peer *Peer) *Trie {
+
+	// At leaf
+
 	if node == nil {
 		return &Trie{
 			bits:         key,
@@ -107,22 +115,17 @@ func (node *Trie) Insert(key []byte, cidr uint, peer *Peer) *Trie {
 
 	common := commonBits(node.bits, key)
 	if node.cidr <= cidr && common >= node.cidr {
-		// Check if match the t.bits[:t.cidr] exactly
 		if node.cidr == cidr {
 			node.peer = peer
 			return node
 		}
-
-		// Go to child
-		bit := (key[node.bit_at_byte] >> node.bit_at_shift) & 1
+		bit := node.choose(key)
 		node.child[bit] = node.child[bit].Insert(key, cidr, peer)
 		return node
 	}
 
 	// Split node
 
-	fmt.Println("new", common)
-
 	newNode := &Trie{
 		bits:         key,
 		peer:         peer,
@@ -132,23 +135,53 @@ func (node *Trie) Insert(key []byte, cidr uint, peer *Peer) *Trie {
 	}
 
 	cidr = min(cidr, common)
-	node.cidr = cidr
-	node.bit_at_byte = cidr / 8
-	node.bit_at_shift = 7 - (cidr % 8)
 
-	// bval := node.bits[node.bit_at_byte] >> node.bit_at_shift // todo : remember index
-	// Work in progress
-	node.child[0] = newNode
-	node.child[1] = newNode
+	// Check for shorter prefix
 
-	return node
-}
+	if newNode.cidr == cidr {
+		bit := newNode.choose(node.bits)
+		newNode.child[bit] = node
+		return newNode
+	}
+
+	// Create new parent for node & newNode
 
-func (t *Trie) Lookup(key []byte) *Peer {
-	if t == nil {
-		return nil
+	parent := &Trie{
+		bits:         key,
+		peer:         nil,
+		cidr:         cidr,
+		bit_at_byte:  cidr / 8,
+		bit_at_shift: 7 - (cidr % 8),
 	}
 
-	return nil
+	bit := parent.choose(key)
+	parent.child[bit] = newNode
+	parent.child[bit^1] = node
+
+	return parent
+}
+
+func (node *Trie) Lookup(key []byte) *Peer {
+	var found *Peer
+	size := uint(len(key))
+	for node != nil && commonBits(node.bits, key) >= node.cidr {
+		if node.peer != nil {
+			found = node.peer
+		}
+		if node.bit_at_byte == size {
+			break
+		}
+		bit := node.choose(key)
+		node = node.child[bit]
+	}
+	return found
+}
 
+func (node *Trie) Count() uint {
+	if node == nil {
+		return 0
+	}
+	l := node.child[0].Count()
+	r := node.child[1].Count()
+	return l + r
 }