Recursively delete unused nodes after Remove (#36)

* Recursively delete unused nodes after Remove

Addresses #35

* Delete nodes without values or mid children

Extends deletion code to delete all nodes without values or middle
children, rather than deleting only nodes whose children are all nil.

trie/trie.go

@@ -18,6 +18,52 @@ type node[V any] struct {
 	valid            bool
 }
 
+func (n *node[V]) isUnused() bool {
+	return !n.valid && n.mid == nil
+}
+
+func (n *node[V]) delete() *node[V] {
+	if n == nil {
+		return nil
+	}
+	if n.right == nil {
+		return n.left
+	}
+	if n.left == nil {
+		return n.right
+	}
+
+	deleted := n
+	n = deleted.right.minChild()
+	n.right = deleted.right.deleteMinChild()
+	n.left = deleted.left
+	return n
+}
+
+func (n *node[V]) minChild() *node[V] {
+	if n == nil {
+		return nil
+	}
+
+	for n.left != nil {
+		n = n.left
+	}
+	return n
+}
+
+func (n *node[V]) deleteMinChild() *node[V] {
+	if n == nil {
+		return nil
+	}
+
+	if n.left == nil { // n is the min node
+		return n.right
+	}
+
+	n.left = n.left.deleteMinChild()
+	return n
+}
+
 // New returns an empty trie.
 func New[V any]() *Trie[V] {
 	return &Trie[V]{}
@@ -73,37 +119,62 @@ func (t *Trie[V]) Put(key string, val V) {
 	if !t.Contains(key) {
 		t.n++
 	}
-	t.root = t.put(t.root, key, val, 0, true)
+	t.root = t.put(t.root, key, val, 0)
+}
+
+func (t *Trie[V]) put(x *node[V], key string, val V, d int) *node[V] {
+	c := key[d]
+	if x == nil {
+		x = &node[V]{
+			c: c,
+		}
+	}
+	if c < x.c {
+		x.left = t.put(x.left, key, val, d)
+	} else if c > x.c {
+		x.right = t.put(x.right, key, val, d)
+	} else if d < len(key)-1 {
+		x.mid = t.put(x.mid, key, val, d+1)
+	} else {
+		x.val = val
+		x.valid = true
+	}
+	return x
 }
 
-// Remove removes the value associated with 'key'.
+// Remove removes the value associated with 'key', along with any nodes of the key that are no
+// longer used.
 func (t *Trie[V]) Remove(key string) {
-	if len(key) == 0 || !t.Contains(key) {
+	if len(key) == 0 {
 		return
 	}
-	var v V
+
+	t.root = t.remove(t.root, key, 0)
 	t.n--
-	// put a tombstone into the deleted value's node
-	t.root = t.put(t.root, key, v, 0, false)
 }
 
-func (t *Trie[V]) put(x *node[V], key string, val V, d int, valid bool) *node[V] {
-	c := key[d]
+func (t *Trie[V]) remove(x *node[V], key string, d int) *node[V] {
 	if x == nil {
-		x = &node[V]{
-			c: c,
-		}
+		return nil
 	}
+
+	c := key[d]
 	if c < x.c {
-		x.left = t.put(x.left, key, val, d, valid)
+		x.left = t.remove(x.left, key, d)
 	} else if c > x.c {
-		x.right = t.put(x.right, key, val, d, valid)
+		x.right = t.remove(x.right, key, d)
 	} else if d < len(key)-1 {
-		x.mid = t.put(x.mid, key, val, d+1, valid)
+		x.mid = t.remove(x.mid, key, d+1)
 	} else {
-		x.val = val
-		x.valid = valid
+		var v V
+		x.val = v
+		x.valid = false
 	}
+
+	if x.isUnused() {
+		return x.delete()
+	}
+
 	return x
 }