changes with witness tracking

core/state/database.go

@@ -130,11 +130,7 @@ type Trie interface {
 
 	// Witness returns a set containing all trie nodes that have been accessed.
 	// The returned map could be nil if the witness is empty.
-	Witness() map[string]struct{}
-
-	// WitnessPaths returns a set of paths for all trie nodes. For future reference,
-	// witness can be deprecated and used as a replacement to witness.
-	WitnessPaths() map[string]struct{}
+	Witness() map[string][]byte
 
 	// NodeIterator returns an iterator that returns nodes of the trie. Iteration
 	// starts at the key after the given start key. And error will be returned

core/state/statedb.go

@@ -841,7 +841,7 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 				// If witness building is enabled and the state object has a trie,
 				// gather the witnesses for its specific storage trie
 				if s.witness != nil && obj.trie != nil {
-					s.witness.AddState(obj.trie.Witness(), obj.trie.WitnessPaths())
+					s.witness.AddState(obj.trie.Witness())
 				}
 			}
 			return nil
@@ -858,9 +858,9 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 				continue
 			}
 			if trie := obj.getPrefetchedTrie(); trie != nil {
-				s.witness.AddState(trie.Witness(), trie.WitnessPaths())
+				s.witness.AddState(trie.Witness())
 			} else if obj.trie != nil {
-				s.witness.AddState(obj.trie.Witness(), obj.trie.WitnessPaths())
+				s.witness.AddState(obj.trie.Witness())
 			}
 		}
 		// Pull in only-read and non-destructed trie witnesses
@@ -874,9 +874,9 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 				continue
 			}
 			if trie := obj.getPrefetchedTrie(); trie != nil {
-				s.witness.AddState(trie.Witness(), trie.WitnessPaths())
+				s.witness.AddState(trie.Witness())
 			} else if obj.trie != nil {
-				s.witness.AddState(obj.trie.Witness(), obj.trie.WitnessPaths())
+				s.witness.AddState(obj.trie.Witness())
 			}
 		}
 	}
@@ -942,7 +942,7 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 
 	// If witness building is enabled, gather the account trie witness
 	if s.witness != nil {
-		s.witness.AddState(s.trie.Witness(), nil)
+		s.witness.AddState(s.trie.Witness())
 	}
 	return hash
 }

core/stateless/witness.go

@@ -91,16 +91,15 @@ func (w *Witness) AddCode(code []byte) {
 }
 
 // AddState inserts a batch of MPT trie nodes into the witness.
-func (w *Witness) AddState(nodes map[string]struct{}, paths map[string]struct{}) {
-	if len(nodes) == 0 {
+func (w *Witness) AddState(nodemap map[string][]byte) {
+	if len(nodemap) == 0 {
 		return
 	}
 	w.lock.Lock()
 	defer w.lock.Unlock()
-
-	maps.Copy(w.State, nodes)
-	if paths != nil {
-		maps.Copy(w.Paths, paths)
+	for path, value := range nodemap {
+		w.State[string(value)] = struct{}{}
+		w.Paths[string(path)] = struct{}{} // Also track the path for the node
 	}
 }
 

trie/secure_trie.go

@@ -273,7 +273,7 @@ func (t *StateTrie) GetKey(shaKey []byte) []byte {
 }
 
 // Witness returns a set containing all trie nodes that have been accessed.
-func (t *StateTrie) Witness() map[string]struct{} {
+func (t *StateTrie) Witness() map[string][]byte {
 	return t.trie.Witness()
 }
 

trie/tracer.go

@@ -158,6 +158,10 @@ func (t *prevalueTracer) keys() []string {
 	return slices.Collect(maps.Keys(t.data))
 }
 
+func (t *prevalueTracer) getMap() map[string][]byte {
+	return maps.Clone(t.data)
+}
+
 // reset resets the cached content in the prevalueTracer.
 func (t *prevalueTracer) reset() {
 	t.lock.Lock()

trie/trie.go

@@ -133,6 +133,12 @@ func (t *Trie) NodeIterator(start []byte) (NodeIterator, error) {
 	return newNodeIterator(t, start), nil
 }
 
+// Owner returns the owner of the trie, allowing us to distinguis between
+// storage and account tries.
+func (t *Trie) Owner() common.Hash {
+	return t.owner
+}
+
 // MustGet is a wrapper of Get and will omit any encountered error but just
 // print out an error message.
 func (t *Trie) MustGet(key []byte) []byte {
@@ -517,9 +523,6 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 			// always creates a new slice) instead of append to
 			// avoid modifying n.Key since it might be shared with
 			// other nodes.
-			if t.owner == (common.Hash{}) {
-				stateDepthAggregator.record(int64(len(prefix) + len(key)))
-			}
 			return true, &shortNode{concat(n.Key, child.Key...), child.Val, t.newFlag()}, nil
 		default:
 			return true, &shortNode{n.Key, child, t.newFlag()}, nil
@@ -577,10 +580,7 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 					// Replace the entire full node with the short node.
 					// Mark the original short node as deleted since the
 					// value is embedded into the parent now.
-					t.tracer.onDelete(append(prefix, byte(pos)))
-					if t.owner == (common.Hash{}) {
-						stateDepthAggregator.record(int64(len(prefix) + 1))
-					}
+					t.opTracer.onDelete(append(prefix, byte(pos)))
 
 					k := append([]byte{byte(pos)}, cnode.Key...)
 					return true, &shortNode{k, cnode.Val, t.newFlag()}, nil
@@ -618,6 +618,13 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 	}
 }
 
+func concat(s1 []byte, s2 ...byte) []byte {
+	r := make([]byte, len(s1)+len(s2))
+	copy(r, s1)
+	copy(r[len(s1):], s2)
+	return r
+}
+
 // copyNode deep-copies the supplied node along with its children recursively.
 func copyNode(n node) node {
 	switch n := (n).(type) {
@@ -757,31 +764,12 @@ func (t *Trie) hashRoot() []byte {
 }
 
 // Witness returns a set containing all trie nodes that have been accessed.
-func (t *Trie) Witness() map[string]struct{} {
-	values := t.prevalueTracer.values()
-	if len(values) == 0 {
+func (t *Trie) Witness() map[string][]byte {
+	dataMap := t.prevalueTracer.getMap()
+	if len(dataMap) == 0 {
 		return nil
 	}
-	witnessStates := make(map[string]struct{}, len(values))
-	for _, val := range values {
-		witnessStates[string(val)] = struct{}{}
-	}
-
-	return witnessStates
-}
-
-func (t *Trie) WitnessPaths() map[string]struct{} {
-	// Return the paths of all nodes that have been accessed.
-	// The paths are the keys of the prevalue tracer.
-	keys := t.prevalueTracer.keys()
-	if len(keys) == 0 {
-		return nil
-	}
-	witnessPaths := make(map[string]struct{}, len(keys))
-	for _, key := range keys {
-		witnessPaths[string(key)] = struct{}{}
-	}
-	return witnessPaths
+	return dataMap
 }
 
 // Reset drops the referenced root node and cleans all internal state.

trie/verkle.go

@@ -452,11 +452,6 @@ func (t *VerkleTrie) nodeResolver(path []byte) ([]byte, error) {
 }
 
 // Witness returns a set containing all trie nodes that have been accessed.
-func (t *VerkleTrie) Witness() map[string]struct{} {
-	panic("not implemented")
-}
-
-// Witness returns a set containing all trie nodes that have been accessed.
-func (t *VerkleTrie) WitnessPaths() map[string]struct{} {
+func (t *VerkleTrie) Witness() map[string][]byte {
 	panic("not implemented")
 }