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,78 @@ func (n *InternalNode) Delete(key []byte, resolver NodeResolverFn) (bool, error)
 	}
 }
 
+// DeleteAtStem delete a full stem. Unlike Delete, it will error out if the stem that is to
+// be deleted does not exist in the tree, because it's meant to be used by rollback code,
+// that should only delete things that exist.
+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, errDeleteMissing
+	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) {

tree_test.go

@@ -357,6 +357,53 @@ func TestDelLeaf(t *testing.T) { // skipcq: GO-R1005
 	}
 }
 
+func TestDeleteAtStem(t *testing.T) {
+	t.Parallel()
+
+	key1, _ := hex.DecodeString("0105000000000000000000000000000000000000000000000000000000000000")
+	key1p, _ := hex.DecodeString("0105000000000000000000000000000000000000000000000000000000000001")  // same Cn group as key1
+	key1pp, _ := hex.DecodeString("0105000000000000000000000000000000000000000000000000000000000081") // Other Cn group as key1
+	key2, _ := hex.DecodeString("0107000000000000000000000000000000000000000000000000000000000000")
+	tree := New()
+	if err := tree.Insert(key1, fourtyKeyTest, nil); err != nil {
+		t.Fatalf("inserting into the original failed: %v", err)
+	}
+	if err := tree.Insert(key1p, fourtyKeyTest, nil); err != nil {
+		t.Fatalf("inserting into the original failed: %v", err)
+	}
+	if err := tree.Insert(key1pp, fourtyKeyTest, nil); err != nil {
+		t.Fatalf("inserting into the original failed: %v", err)
+	}
+	if err := tree.Insert(key2, fourtyKeyTest, nil); err != nil {
+		t.Fatalf("inserting into the original failed: %v", err)
+	}
+	var init Point
+	init.Set(tree.Commit())
+
+	if _, err := tree.(*InternalNode).DeleteAtStem(key1[:31], nil); err != err {
+		t.Error(err)
+	}
+
+	res, err := tree.Get(key1, nil)
+	if err != nil {
+		t.Error(err)
+	}
+	if len(res) > 0 {
+		t.Error("leaf hasnt been deleted")
+	}
+	res, err = tree.Get(key1pp, nil)
+	if err != nil {
+		t.Error(err)
+	}
+	if len(res) > 0 {
+		t.Error("leaf hasnt been deleted")
+	}
+
+	if _, err := tree.(*InternalNode).DeleteAtStem(zeroKeyTest[:31], nil); err != errDeleteMissing {
+		t.Fatal(err)
+	}
+}
+
 func TestDeleteNonExistent(t *testing.T) {
 	t.Parallel()