move trie key utils methods from besu to besu-verkle-trie (#80)

Signed-off-by: Karim Taam <[email protected]>
hyperledger · Jan 23, 2025 · a5cac6c · a5cac6c
1 parent 7f0789e
commit a5cac6c
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 18 deletions.
diff --git a/src/main/java/org/hyperledger/besu/ethereum/trie/verkle/adapter/TrieKeyFactory.java b/src/main/java/org/hyperledger/besu/ethereum/trie/verkle/adapter/TrieKeyFactory.java
@@ -74,7 +74,7 @@ public Bytes getHeaderStem(final Bytes address) {
    * @return The generated storage stem.
    */
   public Bytes getStorageStem(final Bytes address, final Bytes32 storageKey) {
-    final UInt256 trieIndex = TrieKeyUtils.getStorageKeyTrieIndex(storageKey);
+    final Bytes32 trieIndex = TrieKeyUtils.getStorageKeyTrieIndex(storageKey);
     return hasher.computeStem(address, trieIndex);
   }
 
@@ -86,7 +86,7 @@ public Bytes getStorageStem(final Bytes address, final Bytes32 storageKey) {
    * @return The generated code chunk stem.
    */
   public Bytes getCodeChunkStem(final Bytes address, final UInt256 chunkId) {
-    final UInt256 trieIndex = TrieKeyUtils.getCodeChunkKeyTrieIndex(chunkId.toBytes());
+    final Bytes32 trieIndex = TrieKeyUtils.getCodeChunkKeyTrieIndex(chunkId.toBytes());
     return hasher.computeStem(address, trieIndex);
   }
 
@@ -101,9 +101,9 @@ public Map<Bytes32, Bytes> manyStems(
     if (!headerKeys.isEmpty()) {
       trieIndex.add(UInt256.ZERO);
     }
-    for (Bytes32 storageKey : storageKeys) {
-      trieIndex.add(TrieKeyUtils.getStorageKeyTrieIndex(storageKey));
-    }
+
+    trieIndex.addAll(TrieKeyUtils.getStorageKeyTrieIndexes(storageKeys));
+
     for (Bytes32 codeChunkId : codeChunkIds) {
       trieIndex.add(TrieKeyUtils.getCodeChunkKeyTrieIndex(codeChunkId));
     }

diff --git a/src/main/java/org/hyperledger/besu/ethereum/trie/verkle/adapter/TrieKeyUtils.java b/src/main/java/org/hyperledger/besu/ethereum/trie/verkle/adapter/TrieKeyUtils.java
@@ -22,8 +22,12 @@
 import static org.hyperledger.besu.ethereum.trie.verkle.util.Parameters.VERKLE_NODE_WIDTH;
 import static org.hyperledger.besu.ethereum.trie.verkle.util.Parameters.VERKLE_NODE_WIDTH_LOG2;
 
+import org.hyperledger.besu.ethereum.trie.verkle.util.Parameters;
+
 import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
@@ -36,19 +40,41 @@ public class TrieKeyUtils {
 
   private TrieKeyUtils() {}
 
-  public static UInt256 getStorageKeyTrieIndex(final Bytes32 storageKey) {
+  public static Bytes32 getAccountKeyTrieIndex() {
+    return Parameters.BASIC_DATA_LEAF_KEY;
+  }
+
+  public static Bytes32 getStorageKeyTrieIndex(final Bytes32 storageKey) {
     final UInt256 uintStorageKey = UInt256.fromBytes(storageKey);
     if (uintStorageKey.compareTo(HEADER_STORAGE_SIZE) < 0) {
       return uintStorageKey.add(HEADER_STORAGE_OFFSET).divide(VERKLE_NODE_WIDTH);
     } else {
       // We divide by VerkleNodeWidthLog2 to make space and prevent any potential overflow
       // Then, we increment, a step that is safeguarded against overflow.
-      return uintStorageKey
-          .shiftRight(VERKLE_NODE_WIDTH_LOG2.intValue())
-          .add(MAIN_STORAGE_OFFSET_SHIFT_LEFT_VERKLE_NODE_WIDTH);
+      return Bytes32.wrap(
+          uintStorageKey
+              .shiftRight(VERKLE_NODE_WIDTH_LOG2.intValue())
+              .add(MAIN_STORAGE_OFFSET_SHIFT_LEFT_VERKLE_NODE_WIDTH));
     }
   }
 
+  public static List<Bytes32> getStorageKeyTrieIndexes(final List<Bytes32> storageSlotKeys) {
+    return storageSlotKeys.stream()
+        .map(TrieKeyUtils::getStorageKeyTrieIndex)
+        .map(Bytes32::wrap)
+        .toList();
+  }
+
+  public static Bytes32 getCodeChunkKeyTrieIndex(final Bytes32 chunkId) {
+    return Bytes32.wrap(CODE_OFFSET.add(UInt256.fromBytes(chunkId)).divide(VERKLE_NODE_WIDTH));
+  }
+
+  public static List<Bytes32> getCodeChunkKeyTrieIndexes(final Bytes code) {
+    return IntStream.range(0, TrieKeyUtils.getNbChunk(code))
+        .mapToObj(UInt256::valueOf)
+        .collect(Collectors.toUnmodifiableList());
+  }
+
   public static Bytes getStorageKeySuffix(final Bytes32 storageKey) {
     final UInt256 uintStorageKey = UInt256.fromBytes(storageKey);
     if (uintStorageKey.compareTo(HEADER_STORAGE_SIZE) < 0) {
@@ -58,10 +84,6 @@ public static Bytes getStorageKeySuffix(final Bytes32 storageKey) {
     }
   }
 
-  public static UInt256 getCodeChunkKeyTrieIndex(final Bytes32 chunkId) {
-    return CODE_OFFSET.add(UInt256.fromBytes(chunkId)).divide(VERKLE_NODE_WIDTH);
-  }
-
   public static Bytes getCodeChunkKeySuffix(final Bytes32 chunkId) {
     return getLastByte(CODE_OFFSET.add(UInt256.fromBytes(chunkId)).mod(VERKLE_NODE_WIDTH));
   }

diff --git a/src/test/java/org/hyperledger/besu/ethereum/trie/verkle/hasher/StemHasherTest.java b/src/test/java/org/hyperledger/besu/ethereum/trie/verkle/hasher/StemHasherTest.java
@@ -114,8 +114,8 @@ void testComputeStem_CacheMiss() {
   @Test
   void testManyStems_AllKeysInCache() {
     final Bytes address = Bytes.random(32);
-    final Bytes32 index1 = Bytes32.random();
-    final Bytes32 index2 = Bytes32.random();
+    final UInt256 index1 = UInt256.fromBytes(Bytes32.random());
+    final UInt256 index2 = UInt256.fromBytes(Bytes32.random());
     final Bytes value1 = Bytes.random(31);
     final Bytes value2 = Bytes.random(31);
 
@@ -134,9 +134,9 @@ void testManyStems_AllKeysInCache() {
   @Test
   void testManyStems_MissingKeys() {
     final Bytes address = Bytes.random(32);
-    final Bytes32 index1 = Bytes32.leftPad(Bytes.of(1));
-    final Bytes32 index2 = Bytes32.leftPad(Bytes.of(2));
-    final Bytes32 index3 = Bytes32.leftPad(Bytes.of(3));
+    final UInt256 index1 = UInt256.fromBytes(Bytes32.leftPad(Bytes.of(1)));
+    final UInt256 index2 = UInt256.fromBytes(Bytes32.leftPad(Bytes.of(2)));
+    final UInt256 index3 = UInt256.fromBytes(Bytes32.leftPad(Bytes.of(3)));
     final Bytes value1 = Bytes.random(31);
 
     when(mockStemCache.getIfPresent(index1)).thenReturn(value1);