Refactor access list search (#637)

* Add pointer btrees for access lists * Debugging block 4 * Reset access lists ptrs on init_access_lists * Remove linked list use * Clean code * Address review * Minor * Reset access lists pointers using PROVER_INPUT * Rustfmt * Address reviews * Refactor cfg(test) --------- Co-authored-by: Robin Salen <[email protected]>
0xPolygonZero · Sep 26, 2024 · fc28337 · fc28337
1 parent c9caaa1
commit fc28337
Show file tree

Hide file tree

Showing 8 changed files with 253 additions and 215 deletions.
diff --git a/evm_arithmetization/src/cpu/kernel/asm/core/access_lists.asm b/evm_arithmetization/src/cpu/kernel/asm/core/access_lists.asm
@@ -45,6 +45,11 @@ global init_access_lists:
     // Store the segment scaled length
     %increment
     %mstore_global_metadata(@GLOBAL_METADATA_ACCESSED_STORAGE_KEYS_LEN)
+    
+    // Reset the access lists pointers in the `GenerationState`
+    PROVER_INPUT(access_lists::reset)
+    POP // reset pushed a 0
+
     JUMP
 
 %macro init_access_lists

diff --git a/evm_arithmetization/src/cpu/kernel/interpreter.rs b/evm_arithmetization/src/cpu/kernel/interpreter.rs
@@ -5,7 +5,7 @@
 //! the future execution and generate nondeterministically the corresponding
 //! jumpdest table, before the actual CPU carries on with contract execution.
 
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeSet, HashMap};
 
 use anyhow::anyhow;
 use ethereum_types::{BigEndianHash, U256};
@@ -19,6 +19,7 @@ use crate::cpu::columns::CpuColumnsView;
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::generation::debug_inputs;
+use crate::generation::linked_list::LinkedListsPtrs;
 use crate::generation::mpt::{load_linked_lists_and_txn_and_receipt_mpts, TrieRootPtrs};
 use crate::generation::rlp::all_rlp_prover_inputs_reversed;
 use crate::generation::state::{
@@ -115,8 +116,8 @@ pub(crate) struct ExtraSegmentData {
     pub(crate) ger_prover_inputs: Vec<U256>,
     pub(crate) trie_root_ptrs: TrieRootPtrs,
     pub(crate) jumpdest_table: Option<HashMap<usize, Vec<usize>>>,
-    pub(crate) accounts: BTreeMap<U256, usize>,
-    pub(crate) storage: BTreeMap<(U256, U256), usize>,
+    pub(crate) access_lists_ptrs: LinkedListsPtrs,
+    pub(crate) state_ptrs: LinkedListsPtrs,
     pub(crate) next_txn_index: usize,
 }
 
@@ -235,8 +236,8 @@ impl<F: RichField> Interpreter<F> {
         // Initialize the MPT's pointers.
         let (trie_root_ptrs, state_leaves, storage_leaves, trie_data) =
             load_linked_lists_and_txn_and_receipt_mpts(
-                &mut self.generation_state.accounts_pointers,
-                &mut self.generation_state.storage_pointers,
+                &mut self.generation_state.state_ptrs.accounts,
+                &mut self.generation_state.state_ptrs.storage,
                 &inputs.tries,
             )
             .expect("Invalid MPT data for preinitialization");

diff --git a/evm_arithmetization/src/cpu/kernel/tests/account_code.rs b/evm_arithmetization/src/cpu/kernel/tests/account_code.rs
@@ -32,8 +32,8 @@ pub(crate) fn initialize_mpts<F: RichField>(
     // Load all MPTs.
     let (mut trie_root_ptrs, state_leaves, storage_leaves, trie_data) =
         load_linked_lists_and_txn_and_receipt_mpts(
-            &mut interpreter.generation_state.accounts_pointers,
-            &mut interpreter.generation_state.storage_pointers,
+            &mut interpreter.generation_state.state_ptrs.accounts,
+            &mut interpreter.generation_state.state_ptrs.storage,
             trie_inputs,
         )
         .expect("Invalid MPT data for preinitialization");

diff --git a/evm_arithmetization/src/cpu/kernel/tests/mpt/linked_list.rs b/evm_arithmetization/src/cpu/kernel/tests/mpt/linked_list.rs
@@ -12,7 +12,7 @@ use rand::{thread_rng, Rng};
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::interpreter::Interpreter;
-use crate::generation::linked_list::LinkedList;
+use crate::generation::linked_list::testing::LinkedList;
 use crate::generation::linked_list::ACCOUNTS_LINKED_LIST_NODE_SIZE;
 use crate::generation::linked_list::STORAGE_LINKED_LIST_NODE_SIZE;
 use crate::memory::segments::Segment;

diff --git a/evm_arithmetization/src/generation/linked_list.rs b/evm_arithmetization/src/generation/linked_list.rs
@@ -1,42 +1,27 @@
-use std::fmt;
-use std::marker::PhantomData;
+use std::collections::BTreeMap;
 
-use anyhow::Result;
 use ethereum_types::U256;
+use serde::{Deserialize, Serialize};
 
 use crate::memory::segments::Segment;
-use crate::util::u256_to_usize;
-use crate::witness::errors::ProgramError;
-use crate::witness::errors::ProverInputError::InvalidInput;
 
 pub const ACCOUNTS_LINKED_LIST_NODE_SIZE: usize = 4;
 pub const STORAGE_LINKED_LIST_NODE_SIZE: usize = 5;
 
-pub(crate) trait LinkedListType {}
-#[derive(Clone)]
-/// A linked list that starts from the first node after the special node and
-/// iterates forever.
-pub(crate) struct Cyclic;
-#[derive(Clone)]
-/// A linked list that starts from the special node and iterates until the last
-/// node.
-pub(crate) struct Bounded;
-impl LinkedListType for Cyclic {}
-impl LinkedListType for Bounded {}
-
-// A linked list implemented using a vector `access_list_mem`.
-// In this representation, the values of nodes are stored in the range
-// `access_list_mem[i..i + node_size - 1]`, and `access_list_mem[i + node_size -
-// 1]` holds the address of the next node, where i = node_size * j.
-#[derive(Clone)]
-pub(crate) struct LinkedList<'a, const N: usize, T = Cyclic>
-where
-    T: LinkedListType,
-{
-    mem: &'a [Option<U256>],
-    offset: usize,
-    pos: usize,
-    _marker: PhantomData<T>,
+pub const DUMMYHEAD: (U256, U256) = (U256::MAX, U256::zero());
+
+// Provides quick access to pointers that reference the memory location
+// of a storage or accounts linked list node, containing a specific key.
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub(crate) struct LinkedListsPtrs {
+    /// Each entry contains the pair (key, ptr) where key is the (hashed) key
+    /// of an account in the accounts linked list, and ptr is the respective
+    /// node address in memory.
+    pub(crate) accounts: BTreeMap<U256, usize>,
+    /// Each entry contains the pair ((account_key, slot_key), ptr) where
+    /// account_key is the (hashed) key of an account, slot_key is the slot
+    /// key, and ptr is the respective node address in memory.
+    pub(crate) storage: BTreeMap<(U256, U256), usize>,
 }
 
 pub(crate) fn empty_list_mem<const N: usize>(segment: Segment) -> [Option<U256>; N] {
@@ -51,76 +36,99 @@ pub(crate) fn empty_list_mem<const N: usize>(segment: Segment) -> [Option<U256>;
     })
 }
 
-impl<'a, const N: usize, T: LinkedListType> LinkedList<'a, N, T> {
-    pub fn from_mem_and_segment(
+#[cfg(test)]
+pub(crate) mod testing {
+    use std::fmt;
+    use std::marker::PhantomData;
+
+    use anyhow::Result;
+
+    use super::*;
+    use crate::util::u256_to_usize;
+    use crate::witness::errors::ProgramError;
+    use crate::witness::errors::ProverInputError::InvalidInput;
+
+    pub const ADDRESSES_ACCESS_LIST_LEN: usize = 2;
+    pub(crate) trait LinkedListType {}
+    #[derive(Clone)]
+    /// A linked list that starts from the first node after the special node and
+    /// iterates forever.
+    pub(crate) struct Cyclic;
+    #[derive(Clone)]
+    /// A linked list that starts from the special node and iterates until the
+    /// last node.
+    pub(crate) struct Bounded;
+    impl LinkedListType for Cyclic {}
+    impl LinkedListType for Bounded {}
+
+    // A linked list implemented using a vector `access_list_mem`.
+    // In this representation, the values of nodes are stored in the range
+    // `access_list_mem[i..i + node_size - 1]`, and `access_list_mem[i + node_size -
+    // 1]` holds the address of the next node, where i = node_size * j.
+    #[derive(Clone)]
+    pub(crate) struct LinkedList<'a, const N: usize, T = Cyclic>
+    where
+        T: LinkedListType,
+    {
         mem: &'a [Option<U256>],
-        segment: Segment,
-    ) -> Result<Self, ProgramError> {
-        Self::from_mem_len_and_segment(mem, segment)
+        offset: usize,
+        pos: usize,
+        _marker: PhantomData<T>,
     }
 
-    pub fn from_mem_len_and_segment(
-        mem: &'a [Option<U256>],
-        segment: Segment,
-    ) -> Result<Self, ProgramError> {
-        if mem.is_empty() {
-            return Err(ProgramError::ProverInputError(InvalidInput));
+    impl<'a, const N: usize, T: LinkedListType> LinkedList<'a, N, T> {
+        pub fn from_mem_and_segment(
+            mem: &'a [Option<U256>],
+            segment: Segment,
+        ) -> Result<Self, ProgramError> {
+            Self::from_mem_len_and_segment(mem, segment)
         }
-        Ok(Self {
-            mem,
-            offset: segment as usize,
-            pos: 0,
-            _marker: PhantomData,
-        })
-    }
-}
 
-impl<'a, const N: usize> fmt::Debug for LinkedList<'a, N> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        writeln!(f, "Linked List {{")?;
-        let cloned_list = self.clone();
-        for (i, node) in cloned_list.enumerate() {
-            if i > 0 && node[0] == U256::MAX {
-                break;
+        pub fn from_mem_len_and_segment(
+            mem: &'a [Option<U256>],
+            segment: Segment,
+        ) -> Result<Self, ProgramError> {
+            if mem.len() % N != 0 {
+                return Err(ProgramError::ProverInputError(InvalidInput));
             }
-            writeln!(f, "{:?} ->", node)?;
+            Ok(Self {
+                mem,
+                offset: segment as usize,
+                pos: 0,
+                _marker: PhantomData,
+            })
         }
-        write!(f, "}}")
     }
-}
 
-impl<'a, const N: usize> fmt::Debug for LinkedList<'a, N, Bounded> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        writeln!(f, "Linked List {{")?;
-        let cloned_list = self.clone();
-        for node in cloned_list {
-            writeln!(f, "{:?} ->", node)?;
+    impl<'a, const N: usize> fmt::Debug for LinkedList<'a, N> {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            writeln!(f, "Linked List {{")?;
+            let cloned_list = self.clone();
+            for (i, node) in cloned_list.enumerate() {
+                if i > 0 && node[0] == U256::MAX {
+                    break;
+                }
+                writeln!(f, "{:?} ->", node)?;
+            }
+            write!(f, "}}")
         }
-        write!(f, "}}")
     }
-}
-
-impl<'a, const N: usize> Iterator for LinkedList<'a, N> {
-    type Item = [U256; N];
 
-    fn next(&mut self) -> Option<Self::Item> {
-        let node = Some(std::array::from_fn(|i| {
-            self.mem[self.pos + i].unwrap_or_default()
-        }));
-        if let Ok(new_pos) = u256_to_usize(self.mem[self.pos + N - 1].unwrap_or_default()) {
-            self.pos = new_pos - self.offset;
-            node
-        } else {
-            None
+    impl<'a, const N: usize> fmt::Debug for LinkedList<'a, N, Bounded> {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            writeln!(f, "Linked List {{")?;
+            let cloned_list = self.clone();
+            for node in cloned_list {
+                writeln!(f, "{:?} ->", node)?;
+            }
+            write!(f, "}}")
         }
     }
-}
 
-impl<'a, const N: usize> Iterator for LinkedList<'a, N, Bounded> {
-    type Item = [U256; N];
+    impl<'a, const N: usize> Iterator for LinkedList<'a, N> {
+        type Item = [U256; N];
 
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.mem[self.pos] != Some(U256::MAX) {
+        fn next(&mut self) -> Option<Self::Item> {
             let node = Some(std::array::from_fn(|i| {
                 self.mem[self.pos + i].unwrap_or_default()
             }));
@@ -130,8 +138,26 @@ impl<'a, const N: usize> Iterator for LinkedList<'a, N, Bounded> {
             } else {
                 None
             }
-        } else {
-            None
+        }
+    }
+
+    impl<'a, const N: usize> Iterator for LinkedList<'a, N, Bounded> {
+        type Item = [U256; N];
+
+        fn next(&mut self) -> Option<Self::Item> {
+            if self.mem[self.pos] != Some(U256::MAX) {
+                let node = Some(std::array::from_fn(|i| {
+                    self.mem[self.pos + i].unwrap_or_default()
+                }));
+                if let Ok(new_pos) = u256_to_usize(self.mem[self.pos + N - 1].unwrap_or_default()) {
+                    self.pos = new_pos - self.offset;
+                    node
+                } else {
+                    None
+                }
+            } else {
+                None
+            }
         }
     }
 }