Add bit trie that supports variable-length keys

trie/trie.go

@@ -67,7 +67,7 @@ func (trie *Trie) Size() int {
 }
 
 func (trie *Trie) SizeAtDepth(depth int) int {
-	if trie.Branch[0] == nil && trie.Branch[1] == nil {
+	if trie.IsLeaf() {
 		if trie.IsEmpty() {
 			return 0
 		} else {

trie2/add2.go

@@ -0,0 +1,46 @@
+package trie2
+
+func (trie *Trie) Add(q Key) (insertedDepth int, insertedOK bool) {
+	return trie.AddAtDepth(0, q)
+}
+
+func (trie *Trie) AddAtDepth(depth int, q Key) (insertedDepth int, insertedOK bool) {
+	if q.Len() == depth {
+		if trie.Key == nil {
+			trie.Key = q
+			return depth, true
+		} else {
+			if q.Equal(trie.Key) {
+				return depth, false
+			} else {
+				// trie.Key.Len() > depth
+				trie.Key, q = q, trie.Key
+			}
+		}
+	}
+	// q.Len() > depth
+	switch {
+	case trie.IsEmptyLeaf():
+		trie.Key = q
+		return depth, true
+	case trie.IsNonEmptyLeaf():
+		if q.Equal(trie.Key) {
+			return depth, false
+		} else {
+			if trie.Key.Len() == depth {
+				// both branches are nil
+				trie.Branch[0], trie.Branch[1] = &Trie{}, &Trie{}
+				return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
+			} else { // trie.Key.Len() > depth
+				p := trie.Key
+				trie.Key = nil
+				// both branches are nil
+				trie.Branch[0], trie.Branch[1] = &Trie{}, &Trie{}
+				trie.Branch[p.BitAt(depth)].Key = p
+				return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
+			}
+		}
+	default:
+		return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
+	}
+}

trie2/find2.go

@@ -0,0 +1,41 @@
+package trie2
+
+// FindSubKeys finds all keys in the trie that are prefixes of the query key.
+func (trie *Trie) FindSubKeys(q Key) (reachedDepth int, found []Key) {
+	return trie.FindSubKeysAtDepth(0, q)
+}
+
+func (trie *Trie) FindSubKeysAtDepth(depth int, q Key) (reachedDepth int, found []Key) {
+	if trie.IsLeaf() {
+		if trie.IsEmpty() {
+			return depth, nil
+		} else {
+			return depth, []Key{trie.Key}
+		}
+	} else {
+		reachedDepth, found = trie.Branch[q.BitAt(depth)].FindSubKeysAtDepth(depth+1, q)
+		if trie.IsEmpty() {
+			return reachedDepth, found
+		} else {
+			return reachedDepth, append(found, trie.Key)
+		}
+	}
+}
+
+// FindSuperKeys finds all keys in the trie that are prefixed by the query key.
+func (trie *Trie) FindSuperKeys(q Key) (reachedDepth int, found []Key) {
+	return trie.FindSuperKeysAtDepth(0, q)
+}
+
+func (trie *Trie) FindSuperKeysAtDepth(depth int, q Key) (reachedDepth int, found []Key) {
+	d, p := trie.Walk(depth, q)
+	return d, p.List()
+}
+
+func (trie *Trie) Walk(depth int, q Key) (reachedDepth int, arrivedAt *Trie) {
+	if trie.IsLeaf() {
+		return depth, trie
+	} else {
+		return trie.Branch[q.BitAt(depth)].Walk(depth+1, q)
+	}
+}

trie2/list2.go

@@ -0,0 +1,19 @@
+package trie2
+
+// List returns a list of all keys in the trie.
+func (trie *Trie) List() []Key {
+	switch {
+	case trie.IsEmptyLeaf():
+		return nil
+	case trie.IsNonEmptyLeaf():
+		return []Key{trie.Key}
+	case trie.IsEmpty():
+		return append(trie.Branch[0].List(), trie.Branch[1].List()...)
+	case !trie.IsEmpty():
+		return append(
+			[]Key{trie.Key},
+			append(trie.Branch[0].List(), trie.Branch[1].List()...)...,
+		)
+	}
+	panic("unreachable")
+}

trie2/trie2.go

@@ -0,0 +1,50 @@
+package trie2
+
+type Key interface {
+	Equal(Key) bool
+	BitAt(int) byte
+	Len() int
+}
+
+type Trie struct {
+	Branch [2]*Trie
+	Key    Key
+}
+
+func (trie *Trie) IsEmpty() bool {
+	return trie.Key == nil
+}
+
+func (trie *Trie) IsLeaf() bool {
+	return trie.Branch[0] == nil && trie.Branch[1] == nil
+}
+
+func (trie *Trie) IsEmptyLeaf() bool {
+	return trie.IsEmpty() && trie.IsLeaf()
+}
+
+func (trie *Trie) IsNonEmptyLeaf() bool {
+	return !trie.IsEmpty() && trie.IsLeaf()
+}
+
+func (trie *Trie) Size() int {
+	return trie.SizeAtDepth(0)
+}
+
+func (trie *Trie) SizeAtDepth(depth int) int {
+	if trie.IsLeaf() {
+		if trie.IsEmpty() {
+			return 0
+		} else {
+			return 1
+		}
+	} else {
+		var here int
+		if trie.IsEmpty() {
+			here = 0
+		} else {
+			here = 1
+		}
+		return here + trie.Branch[0].SizeAtDepth(depth+1) + trie.Branch[1].SizeAtDepth(depth+1)
+	}
+}