Merge pull request #2 from 0xPolygonZero/zbrown/adjust-deserializatio…

…n-logic

Handle bytestrings and LegacyReceiptRlp decoding

Cargo.toml

@@ -19,6 +19,7 @@ thiserror = "1.0.49"
 rlp = "0.5.2"
 rlp-derive = "0.1.0"
 serde = "1.0.166"
+serde_with = "3.4.0"
 
 [dev-dependencies]
 pretty_env_logger = "0.5.0"

src/compact/compact_prestate_processing.rs

@@ -1268,7 +1268,7 @@ fn process_compact_prestate_common(
     state: TrieCompact,
     create_and_extract_header_f: fn(Vec<u8>) -> CompactParsingResult<(Header, ParserState)>,
 ) -> CompactParsingResult<ProcessedCompactOutput> {
-    let (header, parser) = create_and_extract_header_f(state.bytes)?;
+    let (header, parser) = create_and_extract_header_f(state.0)?;
     let witness_out = parser.parse()?;
 
     let out = ProcessedCompactOutput {

src/compact/complex_test_payloads.rs

@@ -37,9 +37,7 @@ impl TestProtocolInputAndRoot {
         let protocol_bytes = hex::decode(self.byte_str).unwrap();
         let expected_hash = TrieRootHash::from_slice(&hex::decode(self.root_str).unwrap());
 
-        let out = match process_compact_prestate_f(TrieCompact {
-            bytes: protocol_bytes,
-        }) {
+        let out = match process_compact_prestate_f(TrieCompact(protocol_bytes)) {
             Ok(x) => x,
             Err(err) => panic!("{}", err),
         };

src/deserializers.rs

@@ -0,0 +1,64 @@
+//! Custom deserializers for Serde.
+use hex::FromHex;
+use plonky2_evm::generation::mpt::LegacyReceiptRlp;
+use rlp::DecoderError;
+use serde::{
+    de::{Error, Visitor},
+    Deserialize, Deserializer,
+};
+
+#[derive(Clone, Debug, Default, Deserialize)]
+pub(crate) struct ByteString(#[serde(with = "self")] pub(crate) Vec<u8>);
+
+impl From<ByteString> for Vec<u8> {
+    fn from(v: ByteString) -> Self {
+        v.0
+    }
+}
+
+impl TryFrom<ByteString> for LegacyReceiptRlp {
+    type Error = DecoderError;
+
+    fn try_from(value: ByteString) -> Result<Self, Self::Error> {
+        rlp::decode(&value.0)
+    }
+}
+
+fn remove_hex_prefix_if_present(data: &str) -> &str {
+    let prefix = &data[..2];
+
+    match matches!(prefix, "0x" | "0X") {
+        false => data,
+        true => &data[2..],
+    }
+}
+
+// Gross, but there is no Serde crate that can both parse a hex string with a
+// prefix and also deserialize from a `Vec<u8>`.
+fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Vec<u8>, D::Error> {
+    struct PrefixHexStrVisitor();
+
+    impl<'de> Visitor<'de> for PrefixHexStrVisitor {
+        type Value = Vec<u8>;
+
+        fn visit_str<E>(self, data: &str) -> Result<Self::Value, E>
+        where
+            E: Error,
+        {
+            FromHex::from_hex(remove_hex_prefix_if_present(data)).map_err(Error::custom)
+        }
+
+        fn visit_borrowed_str<E>(self, data: &'de str) -> Result<Self::Value, E>
+        where
+            E: Error,
+        {
+            FromHex::from_hex(remove_hex_prefix_if_present(data)).map_err(Error::custom)
+        }
+
+        fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+            write!(f, "a hex encoded string with a prefix")
+        }
+    }
+
+    deserializer.deserialize_string(PrefixHexStrVisitor())
+}

src/lib.rs

@@ -1,8 +1,10 @@
 #![feature(linked_list_cursors)]
 #![feature(trait_alias)]
+#![feature(iter_array_chunks)]
 
 mod compact;
 pub mod decoding;
+mod deserializers;
 pub mod processed_block_trace;
 pub mod proof_gen_types;
 pub mod trace_protocol;

src/processed_block_trace.rs

@@ -178,6 +178,7 @@ impl TxnInfo {
     ) -> ProcessedTxnInfo {
         let mut nodes_used_by_txn = NodesUsedByTxn::default();
         let mut contract_code_accessed = HashMap::new();
+        let block_bloom = self.block_bloom();
 
         for (addr, trace) in self.traces {
             let hashed_addr = hash(addr.as_bytes());
@@ -248,7 +249,7 @@ impl TxnInfo {
         let new_meta_state = TxnMetaState {
             txn_bytes: self.meta.byte_code,
             gas_used: self.meta.gas_used,
-            block_bloom: self.meta.bloom,
+            block_bloom,
         };
 
         ProcessedTxnInfo {
@@ -257,6 +258,25 @@ impl TxnInfo {
             meta: new_meta_state,
         }
     }
+
+    fn block_bloom(&self) -> Bloom {
+        let mut bloom = [U256::zero(); 8];
+
+        // Note that bloom can be empty.
+        for (i, v) in self
+            .meta
+            .new_receipt_trie_node_byte
+            .bloom
+            .clone()
+            .into_iter()
+            .array_chunks::<32>()
+            .enumerate()
+        {
+            bloom[i] = U256::from_big_endian(v.as_slice());
+        }
+
+        bloom
+    }
 }
 
 /// Note that "*_accesses" includes writes.

src/trace_protocol.rs

@@ -25,17 +25,20 @@ use std::collections::HashMap;
 
 use eth_trie_utils::partial_trie::HashedPartialTrie;
 use ethereum_types::{Address, U256};
-use serde::{Deserialize, Serialize};
+use plonky2_evm::generation::mpt::LegacyReceiptRlp;
+use serde::Deserialize;
+use serde_with::{serde_as, FromInto, TryFromInto};
 
 use crate::{
-    types::{Bloom, CodeHash, HashedAccountAddr, StorageAddr, StorageVal},
+    deserializers::ByteString,
+    types::{CodeHash, HashedAccountAddr, StorageAddr, StorageVal},
     utils::hash,
 };
 
 /// Core payload needed to generate a proof for a block. Note that the scheduler
 /// may need to request some additional data from the client along with this in
 /// order to generate a proof.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 pub struct BlockTrace {
     /// The trie pre-images (state & storage) in multiple possible formats.
     pub trie_pre_images: BlockTraceTriePreImages,
@@ -46,54 +49,53 @@ pub struct BlockTrace {
 }
 
 /// Minimal hashed out tries needed by all txns in the block.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum BlockTraceTriePreImages {
     Separate(SeparateTriePreImages),
     Combined(CombinedPreImages),
 }
 
 /// State/Storage trie pre-images that are separate.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 pub struct SeparateTriePreImages {
     pub state: SeparateTriePreImage,
     pub storage: SeparateStorageTriesPreImage,
 }
 
 /// A trie pre-image where state & storage are separate.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum SeparateTriePreImage {
     Uncompressed(TrieUncompressed),
     Direct(TrieDirect),
 }
 
 /// A trie pre-image where both state & storage are combined into one payload.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum CombinedPreImages {
     Compact(TrieCompact),
 }
 
 // TODO
 /// Bulkier format that is quicker to process.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 pub struct TrieUncompressed {}
 
 // TODO
+#[serde_as]
 /// Compact representation of a trie (will likely be very close to https://github.com/ledgerwatch/erigon/blob/devel/docs/programmers_guide/witness_formal_spec.md)
-#[derive(Debug, Serialize, Deserialize)]
-pub struct TrieCompact {
-    pub bytes: Vec<u8>,
-}
+#[derive(Debug, Deserialize)]
+pub struct TrieCompact(#[serde_as(as = "FromInto<ByteString>")] pub Vec<u8>);
 
 // TODO
 /// Trie format that is in exactly the same format of our internal trie format.
 /// This is the fastest format for us to processes.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 pub struct TrieDirect(pub HashedPartialTrie);
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum SeparateStorageTriesPreImage {
     /// A single hash map that contains all node hashes from all storage tries
@@ -107,7 +109,7 @@ pub enum SeparateStorageTriesPreImage {
 }
 
 /// Info specific to txns in the block.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 pub struct TxnInfo {
     /// Trace data for the txn. This is used by the protocol to:
     /// - Mutate it's own trie state between txns to arrive at the correct trie
@@ -120,33 +122,34 @@ pub struct TxnInfo {
     pub meta: TxnMeta,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[serde_as]
+#[derive(Debug, Deserialize)]
 pub struct TxnMeta {
     /// Txn byte code.
+    #[serde_as(as = "FromInto<ByteString>")]
     pub byte_code: Vec<u8>,
 
     /// Rlped bytes of the new txn value inserted into the txn trie by
     /// this txn. Note that the key is not included and this is only the rlped
     /// value of the node!
+    #[serde_as(as = "FromInto<ByteString>")]
     pub new_txn_trie_node_byte: Vec<u8>,
 
     /// Rlped bytes of the new receipt value inserted into the receipt trie by
     /// this txn. Note that the key is not included and this is only the rlped
     /// value of the node!
-    pub new_receipt_trie_node_byte: Vec<u8>,
+    #[serde_as(as = "TryFromInto<ByteString>")]
+    pub new_receipt_trie_node_byte: LegacyReceiptRlp,
 
     /// Gas used by this txn (Note: not cumulative gas used).
     pub gas_used: u64,
-
-    /// Bloom after txn execution.
-    pub bloom: Bloom,
 }
 
 /// A "trace" specific to an account for a txn.
 ///
 /// Specifically, since we can not execute the txn before proof generation, we
 /// rely on a separate EVM to run the txn and supply this data for us.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 pub struct TxnTrace {
     /// If the balance changed, then the new balance will appear here. Will be
     /// `None` if no change.
@@ -171,7 +174,7 @@ pub struct TxnTrace {
 }
 
 /// Contract code access type. Used by txn traces.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Deserialize)]
 #[serde(rename_all = "snake_case")]
 pub enum ContractCodeUsage {
     /// Contract was read.