add DeleteAtStem command for state rollbacks

doc.go

@@ -30,6 +30,8 @@ import "errors"
 var (
 	errInsertIntoHash         = errors.New("trying to insert into hashed node")
 	errDeleteHash             = errors.New("trying to delete from a hashed subtree")
+	errDeleteMissing          = errors.New("trying to delete a missing group")
+	errDeleteUnknown          = errors.New("trying to delete an out-of-view node")
 	errReadFromInvalid        = errors.New("trying to read from an invalid child")
 	errSerializeHashedNode    = errors.New("trying to serialize a hashed internal node")
 	errInsertIntoOtherStem    = errors.New("insert splits a stem where it should not happen")

tree.go

@@ -631,6 +631,75 @@ func (n *InternalNode) Delete(key []byte, resolver NodeResolverFn) (bool, error)
 	}
 }
 
+func (n *InternalNode) DeleteAtStem(key []byte, resolver NodeResolverFn) (bool, error) {
+	nChild := offset2key(key, n.depth)
+	switch child := n.children[nChild].(type) {
+	case Empty:
+		return false, nil
+	case HashedNode:
+		if resolver == nil {
+			return false, errDeleteHash
+		}
+		payload, err := resolver(key[:n.depth+1])
+		if err != nil {
+			return false, err
+		}
+		// deserialize the payload and set it as the child
+		c, err := ParseNode(payload, n.depth+1)
+		if err != nil {
+			return false, err
+		}
+		n.children[nChild] = c
+		return n.DeleteAtStem(key, resolver)
+	case *LeafNode:
+		if !bytes.Equal(child.stem, key[:31]) {
+			return false, errDeleteMissing
+		}
+
+		n.cowChild(nChild)
+		n.children[nChild] = Empty{}
+
+		// Check if all children are gone, if so
+		// signal that this node should be deleted
+		// as well.
+		for _, c := range n.children {
+			if _, ok := c.(Empty); !ok {
+				return false, nil
+			}
+		}
+
+		return true, nil
+	case *InternalNode:
+		n.cowChild(nChild)
+		del, err := child.DeleteAtStem(key, resolver)
+		if err != nil {
+			return false, err
+		}
+
+		// delete the entire child if instructed to by
+		// the recursive algorigthm.
+		if del {
+			n.children[nChild] = Empty{}
+
+			// Check if all children are gone, if so
+			// signal that this node should be deleted
+			// as well.
+			for _, c := range n.children {
+				if _, ok := c.(Empty); !ok {
+					return false, nil
+				}
+			}
+
+			return true, nil
+		}
+
+		return false, nil
+	default:
+		// only unknown nodes are left
+		return false, errDeleteUnknown
+	}
+}
+
 // Flush hashes the children of an internal node and replaces them
 // with HashedNode. It also sends the current node on the flush channel.
 func (n *InternalNode) Flush(flush NodeFlushFn) {