1 files changed, 22 insertions, 18 deletions

diff --git a/src/trie.go b/src/trie.go
index 31a4d92..746c1b4 100644
--- a/src/trie.go
+++ b/src/trie.go
@@ -1,5 +1,9 @@
 package main
 
+import (
+	"net"
+)
+
 /* Binary trie
  *
  * Syncronization done seperatly
@@ -22,13 +26,13 @@ type Trie struct {
 /* Finds length of matching prefix
  * Maybe there is a faster way
  *
- * Assumption: len(s1) == len(s2)
+ * Assumption: len(ip1) == len(ip2)
  */
-func commonBits(s1 []byte, s2 []byte) uint {
+func commonBits(ip1 net.IP, ip2 net.IP) uint {
 	var i uint
-	size := uint(len(s1))
+	size := uint(len(ip1))
 	for i = 0; i < size; i += 1 {
-		v := s1[i] ^ s2[i]
+		v := ip1[i] ^ ip2[i]
 		if v != 0 {
 			v >>= 1
 			if v == 0 {
@@ -93,17 +97,17 @@ 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) choose(ip net.IP) byte {
+	return (ip[node.bit_at_byte] >> node.bit_at_shift) & 1
 }
 
-func (node *Trie) Insert(key []byte, cidr uint, peer *Peer) *Trie {
+func (node *Trie) Insert(ip net.IP, cidr uint, peer *Peer) *Trie {
 
 	// At leaf
 
 	if node == nil {
 		return &Trie{
-			bits:         key,
+			bits:         ip,
 			peer:         peer,
 			cidr:         cidr,
 			bit_at_byte:  cidr / 8,
@@ -113,21 +117,21 @@ func (node *Trie) Insert(key []byte, cidr uint, peer *Peer) *Trie {
 
 	// Traverse deeper
 
-	common := commonBits(node.bits, key)
+	common := commonBits(node.bits, ip)
 	if node.cidr <= cidr && common >= node.cidr {
 		if node.cidr == cidr {
 			node.peer = peer
 			return node
 		}
-		bit := node.choose(key)
-		node.child[bit] = node.child[bit].Insert(key, cidr, peer)
+		bit := node.choose(ip)
+		node.child[bit] = node.child[bit].Insert(ip, cidr, peer)
 		return node
 	}
 
 	// Split node
 
 	newNode := &Trie{
-		bits:         key,
+		bits:         ip,
 		peer:         peer,
 		cidr:         cidr,
 		bit_at_byte:  cidr / 8,
@@ -147,31 +151,31 @@ func (node *Trie) Insert(key []byte, cidr uint, peer *Peer) *Trie {
 	// Create new parent for node & newNode
 
 	parent := &Trie{
-		bits:         key,
+		bits:         ip,
 		peer:         nil,
 		cidr:         cidr,
 		bit_at_byte:  cidr / 8,
 		bit_at_shift: 7 - (cidr % 8),
 	}
 
-	bit := parent.choose(key)
+	bit := parent.choose(ip)
 	parent.child[bit] = newNode
 	parent.child[bit^1] = node
 
 	return parent
 }
 
-func (node *Trie) Lookup(key []byte) *Peer {
+func (node *Trie) Lookup(ip net.IP) *Peer {
 	var found *Peer
-	size := uint(len(key))
-	for node != nil && commonBits(node.bits, key) >= node.cidr {
+	size := uint(len(ip))
+	for node != nil && commonBits(node.bits, ip) >= node.cidr {
 		if node.peer != nil {
 			found = node.peer
 		}
 		if node.bit_at_byte == size {
 			break
 		}
-		bit := node.choose(key)
+		bit := node.choose(ip)
 		node = node.child[bit]
 	}
 	return found