add node iterator, delete leaf nodes

core/trie2/trie.go

@@ -96,9 +96,6 @@ func (t *Trie) Update(key, value *felt.Felt) error {
 // Retrieves the value associated with the given key.
 // Returns felt.Zero if the key doesn't exist.
 // May update the trie's internal structure if nodes need to be resolved.
-// TODO(weiihann):
-// The State should keep track of the modified key and values, so that we can avoid traversing the trie
-// No action needed for the Trie, remove this once State provides the functionality
 func (t *Trie) Get(key *felt.Felt) (felt.Felt, error) {
 	if t.committed {
 		return felt.Zero, ErrCommitted
@@ -198,6 +195,10 @@ func (t *Trie) Commit() (felt.Felt, error) {
 	return rootHash, nil
 }
 
+func (t *Trie) NodeIterator() (db.Iterator, error) {
+	return t.db.NewIterator(t.owner)
+}
+
 func (t *Trie) Copy() *Trie {
 	return &Trie{
 		height:         t.height,
@@ -373,9 +374,10 @@ func (t *Trie) delete(n Node, prefix, key *Path) (bool, Node, error) {
 		if match.Len() < n.Path.Len() {
 			return false, n, nil
 		}
-		// If the whole key matches, remove the entire edge node
+		// If the whole key matches, remove the entire edge node and its child
 		if match.Len() == key.Len() {
-			t.nodeTracer.onDelete(prefix)
+			t.nodeTracer.onDelete(prefix)                        // delete edge node
+			t.nodeTracer.onDelete(new(Path).Append(prefix, key)) // delete value node
 			return true, nil, nil
 		}
 
@@ -438,6 +440,7 @@ func (t *Trie) delete(n Node, prefix, key *Path) (bool, Node, error) {
 		// containing the other child as the child
 		return true, &EdgeNode{Path: bitPrefix, Child: n.Children[other], flags: newFlag()}, nil
 	case *ValueNode:
+		t.nodeTracer.onDelete(key)
 		return true, nil, nil
 	case nil:
 		return false, nil, nil

core/trie2/triedb/database.go

@@ -27,6 +27,7 @@ type TrieDB interface {
 	Get(buf *bytes.Buffer, owner felt.Felt, path trieutils.BitArray) (int, error)
 	Put(owner felt.Felt, path trieutils.BitArray, blob []byte) error
 	Delete(owner felt.Felt, path trieutils.BitArray) error
+	NewIterator(owner felt.Felt) (db.Iterator, error)
 }
 
 type Database struct {
@@ -91,6 +92,30 @@ func (d *Database) Delete(owner felt.Felt, path trieutils.BitArray) error {
 	return d.txn.Delete(buffer.Bytes())
 }
 
+func (d *Database) NewIterator(owner felt.Felt) (db.Iterator, error) {
+	buffer := dbBufferPool.Get().(*bytes.Buffer)
+	buffer.Reset()
+	defer func() {
+		buffer.Reset()
+		dbBufferPool.Put(buffer)
+	}()
+
+	_, err := buffer.Write(d.prefix.Key())
+	if err != nil {
+		return nil, err
+	}
+
+	if owner != (felt.Felt{}) {
+		oBytes := owner.Bytes()
+		_, err := buffer.Write(oBytes[:])
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return d.txn.NewIterator(buffer.Bytes(), true)
+}
+
 func (d *Database) dbKey(buf *bytes.Buffer, owner felt.Felt, path trieutils.BitArray) error {
 	_, err := buf.Write(d.prefix.Key())
 	if err != nil {
@@ -126,3 +151,7 @@ func (EmptyDatabase) Put(owner felt.Felt, path trieutils.BitArray, blob []byte)
 func (EmptyDatabase) Delete(owner felt.Felt, path trieutils.BitArray) error {
 	return ErrCallEmptyDatabase
 }
+
+func (EmptyDatabase) NewIterator(owner felt.Felt) (db.Iterator, error) {
+	return nil, ErrCallEmptyDatabase
+}