clean

crates/fpvm/src/memory/mod.rs

@@ -1,10 +1,5 @@
 //! Contains the memory implementation for the Howitzer FPVM.
 
-use crate::mips::mips_isa::DoubleWord;
-
-// TODO: Remove, dead
-pub const MEMORY_PROOF_SIZE: usize = (DoubleWord::BITS - 4) as usize;
-
 /// An [Address] is a 64 bit address in the MIPS emulator's memory.
 pub type Address = u64;
 

crates/fpvm/src/memory/page.rs

@@ -1,18 +1,14 @@
 //! This module contains the [Page] type as well as associated page parameterization constants.
 
+/// An index of a [Page] within a [TrieMemory] structure.
+pub(crate) type PageIndex = u64;
+
+/// A page of memory, representing [PAGE_SIZE] bytes of data.
+pub(crate) type Page = [u8; PAGE_SIZE];
+
 pub(crate) const PAGE_ADDRESS_SIZE: usize = 12;
 pub(crate) const PAGE_SIZE: usize = 1 << PAGE_ADDRESS_SIZE;
 pub(crate) const PAGE_ADDRESS_MASK: usize = PAGE_SIZE - 1;
-// pub(crate) const MAX_PAGE_COUNT: usize = 1 << PAGE_KEY_SIZE;
-// pub(crate) const PAGE_SIZE_WORDS: usize = PAGE_SIZE >> 5;
-// pub(crate) const PAGE_KEY_SIZE: usize = 64 - PAGE_ADDRESS_SIZE;
-// pub(crate) const PAGE_KEY_MASK: usize = MAX_PAGE_COUNT - 1;
 
 /// An empty page of memory, zeroed out.
 pub(crate) const EMPTY_PAGE: Page = [0u8; PAGE_SIZE];
-
-/// An index of a [Page] within a [TrieMemory] structure.
-pub(crate) type PageIndex = u64;
-
-/// A page of memory, representing [PAGE_SIZE] bytes of data.
-pub(crate) type Page = [u8; PAGE_SIZE];

crates/fpvm/src/memory/trie.rs

@@ -48,8 +48,8 @@ pub(crate) const EMPTY_ROOT_HASH: B256 =
 /// allowing for RLP encoding and decoding of the types for storage and retrieval. The
 /// implementation of these traits will not blind the child nodes of the [TrieNode].
 ///
-/// The [Self::root] function will merkle-encode the [TrieNode] recursively and return the hash of the
-/// encoded bytes.
+/// The [Self::root] function will merkle-encode the [TrieNode] recursively and return the hash of
+/// the encoded bytes.
 ///
 /// ## SAFETY
 /// As this implementation only supports uniform key sizes, the [TrieNode] data structure will fail
@@ -393,12 +393,10 @@ where
                 }
                 encoded_key_len + node.length()
             }
-            TrieNode::Branch { stack, .. } => {
-                stack.iter().fold(0, |mut acc, node| {
-                    acc += node.length();
-                    acc
-                })
-            }
+            TrieNode::Branch { stack, .. } => stack.iter().fold(0, |mut acc, node| {
+                acc += node.length();
+                acc
+            }),
         }
     }
 
@@ -428,7 +426,11 @@ where
                 // that are longer than 32 bytes in length.
                 Header { list: true, payload_length }.encode(out);
                 stack.iter_mut().for_each(|node| {
-                    node.root().encode(out);
+                    if matches!(node, TrieNode::Empty) {
+                        out.put_u8(EMPTY_STRING_CODE)
+                    } else {
+                        node.root().encode(out)
+                    };
                 });
             }
         }
@@ -464,23 +466,17 @@ where
                 }
                 encoded_key_len + value.length()
             }
-            TrieNode::Extension { prefix, node, .. } => {
+            TrieNode::Extension { prefix, .. } => {
                 let mut encoded_key_len = prefix.len() / 2 + 1;
                 if encoded_key_len != 1 {
                     encoded_key_len += length_of_length(encoded_key_len);
                 }
-                encoded_key_len + node.merkle_encoding_length()
-            }
-            TrieNode::Branch { stack, .. } => {
-                // In branch nodes, if an element is longer than an encoded 32 byte string, it is
-                // blinded. Assuming we have an open trie node, we must re-hash the
-                // elements that are longer than an encoded 32 byte string
-                // in length.
-                stack.iter().fold(0, |mut acc, node| {
-                    acc += node.merkle_encoding_length();
-                    acc
-                })
+                encoded_key_len + B256::ZERO.length()
             }
+            TrieNode::Branch { stack, .. } => stack.iter().fold(0, |mut acc, node| {
+                acc += if matches!(node, TrieNode::Empty) { 1 } else { B256::ZERO.length() };
+                acc
+            }),
         }
     }
 

crates/fpvm/src/mips/instrumented.rs

@@ -2,9 +2,10 @@
 
 use super::mips_isa::DoubleWord;
 use crate::{
-    memory::{Address, MEMORY_PROOF_SIZE},
+    memory::Address,
     state::{State, StepWitness},
 };
+use alloy_primitives::Bytes;
 use anyhow::Result;
 use kona_preimage::{HintRouter, PreimageFetcher};
 use std::io::{BufWriter, Write};
@@ -25,12 +26,14 @@ where
     pub(crate) std_out: BufWriter<O>,
     /// The MIPS thread context's stderr buffer.
     pub(crate) std_err: BufWriter<E>,
-    /// The last address we accessed in memory.
-    pub(crate) last_mem_access: Option<Address>,
+
     /// Whether or not the memory proof generation is enabled.
     pub(crate) mem_proof_enabled: bool,
+    /// The last address we accessed in memory.
+    pub(crate) last_mem_access: Option<Address>,
     /// The memory proof, if it is enabled.
-    pub(crate) mem_proof: [u8; MEMORY_PROOF_SIZE * 32],
+    pub(crate) mem_proof: Option<Vec<Bytes>>,
+
     /// The [PreimageOracle] used to fetch preimages.
     pub(crate) preimage_oracle: P,
     /// Cached pre-image data, including 8 byte length prefix
@@ -55,7 +58,7 @@ where
             std_err: BufWriter::new(std_err),
             last_mem_access: None,
             mem_proof_enabled: false,
-            mem_proof: [0u8; MEMORY_PROOF_SIZE * 32],
+            mem_proof: Default::default(),
             preimage_oracle: oracle,
             last_preimage: Vec::default(),
             last_preimage_key: [0u8; 32],
@@ -78,15 +81,16 @@ where
         // state transition.
         let mut witness = proof
             .then(|| {
-                // let instruction_proof = self.state.memory.merkle_proof(self.state.pc as
-                // Address)?;
+                // Generate a merkle proof for the page containing the current instruction.
+                let instruction_proof = self.state.memory.merkle_proof(self.state.pc as Address)?;
+
+                // Allocate the proof vector and push the instruction proof.
+                let mut proof = Vec::with_capacity(2);
+                proof.push(instruction_proof);
 
-                let mem_proof = vec![0; MEMORY_PROOF_SIZE * 32 * 2];
-                // mem_proof[0..MEMORY_PROOF_SIZE *
-                // 32].copy_from_slice(instruction_proof.as_slice());
                 Ok::<_, anyhow::Error>(StepWitness {
                     state: self.state.encode_witness()?,
-                    mem_proof,
+                    proof,
                     ..Default::default()
                 })
             })
@@ -99,7 +103,10 @@ where
         // preimage read within the state transition.
         if proof {
             witness = witness.map(|mut wit| {
-                wit.mem_proof[MEMORY_PROOF_SIZE * 32..].copy_from_slice(self.mem_proof.as_slice());
+                if let Some(mem_proof) = self.mem_proof.take() {
+                    wit.proof.push(mem_proof);
+                }
+
                 if self.last_preimage_offset.is_some() {
                     wit.preimage_key = Some(self.last_preimage_key);
                     wit.preimage_value = Some(self.last_preimage.clone());

crates/fpvm/src/mips/mips_vm.rs

@@ -664,7 +664,7 @@ where
                 "Unexpected last memory access with existing access buffered."
             );
             self.last_mem_access = Some(effective_address);
-            // self.mem_proof = self.state.memory.merkle_proof(effective_address)?;
+            self.mem_proof = Some(self.state.memory.merkle_proof(effective_address)?);
         }
         Ok(())
     }

crates/fpvm/src/state/step_witness.rs

@@ -1,19 +1,19 @@
 //! This module contains the various witness types.
 
 use super::{StateWitness, STATE_WITNESS_SIZE};
-use crate::memory::MEMORY_PROOF_SIZE;
 use alloy_primitives::{Bytes, B256, U256};
 use alloy_sol_types::{sol, SolCall};
 use kona_preimage::PreimageKeyType;
 
 /// A [StepWitness] is produced after each instruction step of the MIPS emulator. It contains
 /// the encoded [StateWitness], the proof of memory access, and the preimage key, value, and
 /// offset.
+#[derive(Debug, Clone, Eq, PartialEq)]
 pub struct StepWitness {
     /// The encoded state witness
     pub state: StateWitness,
     /// The proof of memory access
-    pub mem_proof: Vec<u8>,
+    pub proof: Vec<Vec<Bytes>>,
     /// The preimage key
     pub preimage_key: Option<[u8; 32]>,
     /// The preimage value
@@ -26,7 +26,7 @@ impl Default for StepWitness {
     fn default() -> Self {
         Self {
             state: [0u8; STATE_WITNESS_SIZE],
-            mem_proof: Vec::with_capacity(MEMORY_PROOF_SIZE * 32 * 2),
+            proof: Vec::with_capacity(2),
             preimage_key: Default::default(),
             preimage_value: Default::default(),
             preimage_offset: Default::default(),
@@ -104,8 +104,8 @@ impl StepWitness {
     pub fn encode_step_input(&self) -> Bytes {
         let call = stepCall {
             _0: self.state.to_vec().into(),
-            _1: self.mem_proof.to_vec().into(),
-            _2: B256::ZERO, // constant local context
+            _1: Default::default(), // TODO: Load proof
+            _2: B256::ZERO,         // constant local context
         };
 
         call.abi_encode().into()

crates/kernel/src/kernel.rs

@@ -157,7 +157,7 @@ where
                                 post: poststate_hash,
                                 state_data: step_witness.state,
                                 step_input: step_witness.encode_step_input().to_vec(),
-                                proof_data: step_witness.mem_proof,
+                                proof_data: step_witness.proof,
                                 oracle_input: preimage_input.map(|k| k.to_vec()),
                                 oracle_key: step_witness.preimage_key.map(|k| k.to_vec()),
                                 oracle_value: step_witness.preimage_value,

crates/kernel/src/types.rs

@@ -1,5 +1,6 @@
 //! This module contains the types for the `howitzer` interface.
 
+use alloy_primitives::Bytes;
 use howitzer_fpvm::state::StateWitness;
 use serde::{Deserialize, Serialize};
 
@@ -12,7 +13,7 @@ pub struct Proof {
     pub post: [u8; 32],
     #[serde(with = "howitzer_fpvm::utils::ser::state_witness_hex")]
     pub state_data: StateWitness,
-    pub proof_data: Vec<u8>,
+    pub proof_data: Vec<Vec<Bytes>>,
     pub step_input: Vec<u8>,
     pub oracle_key: Option<Vec<u8>>,
     pub oracle_value: Option<Vec<u8>>,