prune correctly stem

Signed-off-by: Karim Taam <karim.t2am@gmail.com>

Author: matkt

src/main/java/org/hyperledger/besu/ethereum/stateless/bintrie/SimpleBinTrie.java

@@ -20,6 +20,7 @@
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.LeafNode;
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.Node;
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.NullNode;
+import org.hyperledger.besu.ethereum.stateless.bintrie.pruning.StemPrunableNodeRegistry;
 import org.hyperledger.besu.ethereum.stateless.bintrie.visitor.CommitVisitor;
 import org.hyperledger.besu.ethereum.stateless.bintrie.visitor.FlattenVisitor;
 import org.hyperledger.besu.ethereum.stateless.bintrie.visitor.GetVisitor;
@@ -30,6 +31,7 @@
 
 import java.util.Optional;
 
+import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
 
 /**
@@ -39,11 +41,13 @@
  * @param <V> The type of values in the Bin Trie.
  */
 public class SimpleBinTrie<K extends BitSequence<K>, V> implements BinTrie<K, V> {
+
+  private final StemPrunableNodeRegistry<K> stemPrunableNodeRegistry;
   protected Node<K, V> root;
 
   /** Creates a new Bin Trie with a null node as the root. */
   public SimpleBinTrie() {
-    this.root = NullNode.node();
+    this(NullNode.node());
   }
 
   /**
@@ -52,7 +56,7 @@ public SimpleBinTrie() {
    * @param root The root node of the Bin Trie.
    */
   public SimpleBinTrie(final Optional<Node<K, V>> root) {
-    this.root = root.orElse(NullNode.node());
+    this(root.orElse(NullNode.node()));
   }
 
   /**
@@ -62,6 +66,7 @@ public SimpleBinTrie(final Optional<Node<K, V>> root) {
    */
   public SimpleBinTrie(final Node<K, V> root) {
     this.root = root;
+    this.stemPrunableNodeRegistry = new StemPrunableNodeRegistry<>();
   }
 
   /**
@@ -99,7 +104,7 @@ public Optional<V> get(final K key) {
   public Optional<V> put(final K key, final V value) {
     checkNotNull(key);
     checkNotNull(value);
-    final PutVisitor<K, V> kvPutVisitor = new PutVisitor<>(key, value);
+    final PutVisitor<K, V> kvPutVisitor = new PutVisitor<>(key, value, stemPrunableNodeRegistry);
     this.root = root.accept(kvPutVisitor);
     return kvPutVisitor.getOldValue();
   }
@@ -118,7 +123,7 @@ public void remove(final K key) {
   /** Restructure tree to get minimal representation. */
   @Override
   public void flatten() {
-    this.root = root.accept(new FlattenVisitor<K, V>());
+    this.root = root.accept(new FlattenVisitor<K, V>(stemPrunableNodeRegistry));
   }
 
   /**
@@ -128,7 +133,7 @@ public void flatten() {
    */
   @Override
   public Bytes32 getRootHash() {
-    root = root.accept(new FlattenVisitor<K, V>());
+    root = root.accept(new FlattenVisitor<K, V>(stemPrunableNodeRegistry));
     root = root.accept(new HashVisitor<K, V>());
     assert root.commitment.isPresent() : "HashVisitor should produce a rootHash";
     return root.commitment.get();
@@ -152,6 +157,14 @@ public String toString() {
   @Override
   public void commit(final NodeUpdater nodeUpdater) {
     getRootHash();
+    stemPrunableNodeRegistry
+        .getPrunableStems()
+        .forEach(
+            stem -> {
+              nodeUpdater.store(Bytes.wrap(stem.encode()), null, null);
+            });
+
+    stemPrunableNodeRegistry.clear();
     root = root.accept(new CommitVisitor<K, V>(nodeUpdater));
   }
 

src/main/java/org/hyperledger/besu/ethereum/stateless/bintrie/pruning/StemPrunableNodeRegistry.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright Hyperledger Besu Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package org.hyperledger.besu.ethereum.stateless.bintrie.pruning;
+
+import org.hyperledger.besu.ethereum.stateless.bintrie.BitSequence;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Registry responsible for tracking {@code StemNode} stems that are eligible for pruning.
+ *
+ * <p>When a {@code StemNode} becomes obsolete (e.g., replaced by a {@code NullNode}), its stem can
+ * be marked as prunable and later removed during a cleanup pass. This class centralizes that
+ * tracking mechanism.
+ */
+public class StemPrunableNodeRegistry<K extends BitSequence<K>> {
+
+  private final List<K> prunableStems = new ArrayList<>();
+
+  /** Creates a new registry for tracking prunable stem nodes. */
+  public StemPrunableNodeRegistry() {}
+
+  /**
+   * Marks a stem as prunable, indicating that the corresponding {@code StemNode} can be safely
+   * removed if no longer needed.
+   *
+   * @param key the stem key to mark as prunable
+   */
+  public void markPrunableStem(final K key) {
+    prunableStems.add(key);
+  }
+
+  /**
+   * Removes a stem from the prunable set, usually because it has been revived or re-added to the
+   * trie.
+   *
+   * @param key the stem key to remove from pruning consideration
+   */
+  public void removePrunableStem(final K key) {
+    prunableStems.remove(key);
+  }
+
+  /**
+   * Returns the list of all stems currently marked as prunable.
+   *
+   * @return list of stems eligible for pruning
+   */
+  public List<K> getPrunableStems() {
+    return prunableStems;
+  }
+
+  /** Clears all tracked prunable stems. Typically used after a batch prune operation. */
+  public void clear() {
+    prunableStems.clear();
+  }
+}

src/main/java/org/hyperledger/besu/ethereum/stateless/bintrie/visitor/FlattenVisitor.java

@@ -20,6 +20,7 @@
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.Node;
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.NullNode;
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.StemNode;
+import org.hyperledger.besu.ethereum.stateless.bintrie.pruning.StemPrunableNodeRegistry;
 
 /**
  * Class representing a visitor for flattening a node in a Trie tree.
@@ -32,6 +33,13 @@
  * @param <V> The type of node values.
  */
 public class FlattenVisitor<K extends BitSequence<K>, V> implements NodeVisitor<K, V> {
+
+  private final StemPrunableNodeRegistry<K> stemPrunableNodeRegistry;
+
+  public FlattenVisitor(final StemPrunableNodeRegistry<K> stemPrunableNodeRegistry) {
+    this.stemPrunableNodeRegistry = stemPrunableNodeRegistry;
+  }
+
   @Override
   public Node<K, V> visit(InternalNode<K, V> internalNode) {
 
@@ -71,6 +79,7 @@ public Node<K, V> visit(StemNode<K, V> stemNode) {
       return stemNode;
     }
     if (stemNode.allLeavesAreNull()) {
+      stemPrunableNodeRegistry.markPrunableStem(stemNode.stem);
       return NullNode.node();
     }
     return stemNode;

src/main/java/org/hyperledger/besu/ethereum/stateless/bintrie/visitor/PutVisitor.java

@@ -23,6 +23,7 @@
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.NullNode;
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.StemNode;
 import org.hyperledger.besu.ethereum.stateless.bintrie.node.ValueNode;
+import org.hyperledger.besu.ethereum.stateless.bintrie.pruning.StemPrunableNodeRegistry;
 
 import java.util.Optional;
 
@@ -35,6 +36,9 @@
  * @param <V> The type of values to insert or update.
  */
 public class PutVisitor<K extends BitSequence<K>, V> implements NodeVisitor<K, V> {
+
+  private final StemPrunableNodeRegistry<K> stemPrunableNodeRegistry;
+
   public final K path;
   private Optional<V> oldValue;
   public final V value;
@@ -44,12 +48,16 @@ public class PutVisitor<K extends BitSequence<K>, V> implements NodeVisitor<K, V
    * Constructs a new PutVisitor with the provided value to insert or update.
    *
    * @param value The value to be inserted or updated in the Verkle Trie.
+   * @param stemPrunableNodeRegistry Tracker of stem to delete
    */
-  public PutVisitor(final K path, final V value) {
+  public PutVisitor(
+      final K path, final V value, final StemPrunableNodeRegistry<K> stemPrunableNodeRegistry) {
     assert path.length() <= Node.KEY_SIZE;
     this.path = path;
     this.value = value;
     this.oldValue = Optional.empty();
+    this.stemPrunableNodeRegistry = stemPrunableNodeRegistry;
+
     // System.out.println(String.format("PutVisit path=%s, value=%s",
     // Bytes.wrap(path.toBytes()),
     // value));
@@ -134,8 +142,9 @@ public Node<K, V> visit(final StemNode<K, V> stemNode) {
   public Node<K, V> visit(final NullNode<K, V> nullNode) {
     // System.out.println(String.format("Put visit Null: depth=%s, loc=%s, stem=%s", depth+1,
     // path.slice(0, depth+1).toHexString(), path.slice(0, Node.STEM_SIZE).toHexString()));
-    return new StemNode<K, V>(Optional.of(path.slice(0, depth + 1)), path.slice(0, Node.STEM_SIZE))
-        .accept(this);
+    final K stem = path.slice(0, Node.STEM_SIZE);
+    stemPrunableNodeRegistry.removePrunableStem(stem);
+    return new StemNode<K, V>(Optional.of(path.slice(0, depth + 1)), stem).accept(this);
   }
 
   /**

src/test/java/org/hyperledger/besu/ethereum/stateless/bintrie/StemPrunableNodeRegistryTest.java

@@ -0,0 +1,114 @@
+/*
+ * Copyright Hyperledger Besu Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package org.hyperledger.besu.ethereum.stateless.bintrie;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.hyperledger.besu.ethereum.stateless.bintrie.factory.StoredNodeFactory;
+
+import org.apache.tuweni.bytes.Bytes;
+import org.apache.tuweni.bytes.Bytes32;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+@DisplayName("StemPrunableNodeRegistryTest – stem pruning behaviour")
+class StemPrunableNodeRegistryTest {
+
+  private static final BytesBitSequenceFactory KEY_FACTORY = new BytesBitSequenceFactory();
+  private static final BytesBitSequence KEY =
+      KEY_FACTORY.fromHexString(
+          "0x00112233445566778899aabbccddeeff00112233445566778899aabbccddeeff");
+  private static final Bytes32 VALUE =
+      Bytes32.fromHexString("0x1000000000000000000000000000000000000000000000000000000000000000");
+  private static final Bytes STEM =
+      Bytes.fromHexString(
+          "0x00112233445566778899aabbccddeeff00112233445566778899aabbccddee00"); // first 31 bytes
+
+  private org.hyperledger.besu.ethereum.stateless.bintrie.NodeUpdaterMock nodeUpdater;
+  private StoredBinTrie<BytesBitSequence, Bytes32> trie;
+
+  @BeforeEach
+  void setUp() {
+    nodeUpdater = new NodeUpdaterMock();
+    NodeLoaderMock nodeLoader = new NodeLoaderMock(nodeUpdater.storage);
+    StoredNodeFactory<BytesBitSequence, Bytes32> nodeFactory =
+        new StoredNodeFactory<>(nodeLoader, KEY_FACTORY, x -> (Bytes32) x);
+    trie = new StoredBinTrie<>(nodeFactory);
+  }
+
+  @Test
+  @DisplayName("prunes stem after key is removed and committed")
+  void shouldPruneStemAfterRemoveAndCommit() {
+    // put + commit
+    trie.put(KEY, VALUE);
+    trie.commit(nodeUpdater);
+    assertThat(stemExists()).isTrue();
+    assertThat(trie.getRootHash()).isNotEqualTo(Bytes32.ZERO);
+
+    // remove + second commit
+    trie.remove(KEY);
+    trie.commit(nodeUpdater);
+
+    assertThat(stemExists()).isFalse();
+    assertThat(trie.getRootHash()).isEqualTo(Bytes32.ZERO);
+  }
+
+  @Test
+  @DisplayName("keeps stem when key is removed then re-added before commit")
+  void shouldKeepStemWhenKeyRemovedAndReaddedBeforeSameCommit() {
+    // put + commit
+    trie.put(KEY, VALUE);
+    trie.commit(nodeUpdater);
+    assertThat(stemExists()).isTrue();
+
+    // remove + put (same transaction) + commit
+    trie.remove(KEY);
+    trie.put(KEY, VALUE);
+    trie.commit(nodeUpdater);
+
+    assertThat(stemExists()).isTrue();
+    assertThat(trie.getRootHash()).isNotEqualTo(Bytes32.ZERO);
+  }
+
+  @Test
+  @DisplayName("keeps stem when key is re-added in a later transaction")
+  void shouldKeepStemWhenKeyReaddedAfterSeparateCommit() {
+    // put + commit
+    trie.put(KEY, VALUE);
+    trie.commit(nodeUpdater);
+    assertThat(stemExists()).isTrue();
+
+    // remove + commit  ➜ stem prunable
+    trie.remove(KEY);
+    trie.commit(nodeUpdater);
+    assertThat(stemExists()).isFalse();
+
+    // put again + commit ➜ stem revived
+    trie.put(KEY, VALUE);
+    trie.commit(nodeUpdater);
+
+    assertThat(stemExists()).isTrue();
+    assertThat(trie.getRootHash()).isNotEqualTo(Bytes32.ZERO);
+  }
+
+  /**
+   * @return {@code true} if the current in-memory DB still contains the stem.
+   */
+  private boolean stemExists() {
+    return nodeUpdater.storage.keySet().stream().anyMatch(k -> k.equals(STEM));
+  }
+}