Trie related segments and metadata

Summary of the design:
- Tries are stored as immutable, copy-on-write trees
- All trie data is stored in the `TrieData` segment. Since it's immutable, data is never modified/deleted, new versions are just appended at the end.
- In order to support reverts, each context stores a pointer to the initial state trie version, plus the initial version of each storage trie.

One variation which may be worth considering is storing the whole state trie as one big trie (with sub-tries for storage). Reverts would then be simpler - we'd replace a single pointer. I thought that approach might make hashing the trie a bit more complex, as the node associated with an account would be a special type of node, rather than just another leaf. Either approach seems reasonable though.

evm/src/cpu/kernel/context_metadata.rs

@@ -20,10 +20,13 @@ pub(crate) enum ContextMetadata {
     /// Whether this context was created by `STATICCALL`, in which case state changes are
     /// prohibited.
     Static = 8,
+    /// Pointer to the initial version of the state trie, at the creation of this context. Used when
+    /// we need to revert a context. See also `StorageTrieCheckpointPointers`.
+    StateTrieCheckpointPointer = 9,
 }
 
 impl ContextMetadata {
-    pub(crate) const COUNT: usize = 9;
+    pub(crate) const COUNT: usize = 10;
 
     pub(crate) fn all() -> [Self; Self::COUNT] {
         [
@@ -36,6 +39,7 @@ impl ContextMetadata {
             Self::Caller,
             Self::CallValue,
             Self::Static,
+            Self::StateTrieCheckpointPointer,
         ]
     }
 
@@ -51,6 +55,7 @@ impl ContextMetadata {
             ContextMetadata::Caller => "CTX_METADATA_CALLER",
             ContextMetadata::CallValue => "CTX_METADATA_CALL_VALUE",
             ContextMetadata::Static => "CTX_METADATA_STATIC",
+            ContextMetadata::StateTrieCheckpointPointer => "CTX_METADATA_STATE_TRIE_CHECKPOINT_PTR",
         }
     }
 }

evm/src/cpu/kernel/global_metadata.rs

@@ -9,13 +9,34 @@ pub(crate) enum GlobalMetadata {
     Origin = 1,
     /// The size of active memory, in bytes.
     MemorySize = 2,
+    /// The size of the `TrieData` segment, in bytes. In other words, the next address available for
+    /// appending additional trie data.
+    TrieDataSize = 3,
+    /// A pointer to the root of the state trie within the `TrieData` buffer.
+    StateTrieRoot = 4,
+    /// A pointer to the root of the transaction trie within the `TrieData` buffer.
+    TransactionTrieRoot = 5,
+    /// A pointer to the root of the receipt trie within the `TrieData` buffer.
+    ReceiptTrieRoot = 6,
+    /// The number of storage tries involved in this transaction. I.e. the number of values in
+    /// `StorageTrieAddresses`, `StorageTriePointers` and `StorageTrieCheckpointPointers`.
+    NumStorageTries = 7,
 }
 
 impl GlobalMetadata {
-    pub(crate) const COUNT: usize = 3;
+    pub(crate) const COUNT: usize = 8;
 
     pub(crate) fn all() -> [Self; Self::COUNT] {
-        [Self::LargestContext, Self::Origin, Self::MemorySize]
+        [
+            Self::LargestContext,
+            Self::Origin,
+            Self::MemorySize,
+            Self::TrieDataSize,
+            Self::StateTrieRoot,
+            Self::TransactionTrieRoot,
+            Self::ReceiptTrieRoot,
+            Self::NumStorageTries,
+        ]
     }
 
     /// The variable name that gets passed into kernel assembly code.
@@ -24,6 +45,11 @@ impl GlobalMetadata {
             GlobalMetadata::LargestContext => "GLOBAL_METADATA_LARGEST_CONTEXT",
             GlobalMetadata::Origin => "GLOBAL_METADATA_ORIGIN",
             GlobalMetadata::MemorySize => "GLOBAL_METADATA_MEMORY_SIZE",
+            GlobalMetadata::TrieDataSize => "GLOBAL_METADATA_TRIE_DATA_SIZE",
+            GlobalMetadata::StateTrieRoot => "GLOBAL_METADATA_STATE_TRIE_ROOT",
+            GlobalMetadata::TransactionTrieRoot => "GLOBAL_METADATA_TXN_TRIE_ROOT",
+            GlobalMetadata::ReceiptTrieRoot => "GLOBAL_METADATA_RECEIPT_TRIE_ROOT",
+            GlobalMetadata::NumStorageTries => "GLOBAL_METADATA_NUM_STORAGE_TRIES",
         }
     }
 }

evm/src/memory/segments.rs

@@ -22,12 +22,24 @@ pub(crate) enum Segment {
     TxnFields = 8,
     /// Contains the data field of a transaction.
     TxnData = 9,
-    /// Raw RLP data.
+    /// A buffer used to hold raw RLP data.
     RlpRaw = 10,
+    /// Contains all trie data. Tries are stored as immutable, copy-on-write trees, so this is an
+    /// append-only buffer. It is owned by the kernel, so it only lives on context 0.
+    TrieData = 11,
+    /// The account address associated with the `i`th storage trie. Only lives on context 0.
+    StorageTrieAddresses = 12,
+    /// A pointer to the `i`th storage trie within the `TrieData` buffer. Only lives on context 0.
+    StorageTriePointers = 13,
+    /// Like `StorageTriePointers`, except that these pointers correspond to the version of each
+    /// trie at the creation of a given context. This lets us easily revert a context by replacing
+    /// `StorageTriePointers` with `StorageTrieCheckpointPointers`.
+    /// See also `StateTrieCheckpointPointer`.
+    StorageTrieCheckpointPointers = 14,
 }
 
 impl Segment {
-    pub(crate) const COUNT: usize = 11;
+    pub(crate) const COUNT: usize = 15;
 
     pub(crate) fn all() -> [Self; Self::COUNT] {
         [
@@ -42,6 +54,10 @@ impl Segment {
             Self::TxnFields,
             Self::TxnData,
             Self::RlpRaw,
+            Self::TrieData,
+            Self::StorageTrieAddresses,
+            Self::StorageTriePointers,
+            Self::StorageTrieCheckpointPointers,
         ]
     }
 
@@ -59,6 +75,10 @@ impl Segment {
             Segment::TxnFields => "SEGMENT_NORMALIZED_TXN",
             Segment::TxnData => "SEGMENT_TXN_DATA",
             Segment::RlpRaw => "SEGMENT_RLP_RAW",
+            Segment::TrieData => "SEGMENT_TRIE_DATA",
+            Segment::StorageTrieAddresses => "SEGMENT_STORAGE_TRIE_ADDRS",
+            Segment::StorageTriePointers => "SEGMENT_STORAGE_TRIE_PTRS",
+            Segment::StorageTrieCheckpointPointers => "SEGMENT_STORAGE_TRIE_CHECKPOINT_PTRS",
         }
     }
 
@@ -76,6 +96,10 @@ impl Segment {
             Segment::TxnFields => 256,
             Segment::TxnData => 256,
             Segment::RlpRaw => 8,
+            Segment::TrieData => 256,
+            Segment::StorageTrieAddresses => 160,
+            Segment::StorageTriePointers => 32,
+            Segment::StorageTrieCheckpointPointers => 32,
         }
     }
 }