add support for compact identifier types

crates/cli/commands/src/test_vectors/compact.rs

@@ -1,4 +1,4 @@
-use alloy_primitives::{hex, private::getrandom::getrandom};
+use alloy_primitives::{hex, private::getrandom::getrandom, TxKind};
 use arbitrary::Arbitrary;
 use eyre::Result;
 use proptest::{
@@ -17,7 +17,7 @@ use reth_db::{
     ClientVersion,
 };
 use reth_fs_util as fs;
-use reth_primitives::{Account, Log, LogData, StorageEntry, TransactionSignedNoHash};
+use reth_primitives::{Account, Log, LogData, StorageEntry, Transaction, TransactionSignedNoHash, TxType};
 use reth_prune_types::{PruneCheckpoint, PruneMode};
 use reth_stages_types::{
     AccountHashingCheckpoint, CheckpointBlockRange, EntitiesCheckpoint, ExecutionCheckpoint,
@@ -26,79 +26,100 @@ use reth_stages_types::{
 };
 use reth_trie::{hash_builder::HashBuilderValue, TrieMask};
 use reth_trie_common::{hash_builder::HashBuilderState, StoredNibbles, StoredNibblesSubKey};
-use std::{fs::File, io::BufReader};
+use std::{collections::HashSet, fs::File, io::BufReader, sync::LazyLock};
 
 const VECTORS_FOLDER: &str = "testdata/micro/compact";
 const VECTOR_SIZE: usize = 100;
 
 macro_rules! compact_types {
-    ($($ty:ty),*) => {
+    (regular: [$($regular_ty:ident),*], identifier: [$($id_ty:ident),*]) => {
         const GENERATE_VECTORS: &[fn(&mut TestRunner) -> Result<()>] = &[
             $(
-                generate_vector::<$ty> as fn(&mut TestRunner) -> Result<()>,
+                generate_vector::<$regular_ty> as fn(&mut TestRunner) -> Result<()>,
+            )*
+            $(
+                generate_vector::<$id_ty> as fn(&mut TestRunner) -> Result<()>,
             )*
         ];
 
         const READ_VECTORS: &[fn() -> Result<()>] = &[
             $(
-                read_vector::<$ty> as fn() -> Result<()>,
+                read_vector::<$regular_ty> as fn() -> Result<()>,
+            )*
+            $(
+                read_vector::<$id_ty> as fn() -> Result<()>,
             )*
         ];
+
+        static IDENTIFIER_TYPE: LazyLock<HashSet<String>> = LazyLock::new(|| {
+            let mut map = HashSet::new();
+            $(
+                map.insert(type_name::<$id_ty>());
+            )*
+            map
+        });
     };
 }
 
 // The type that **actually** implements `Compact` should go here. If it's an alloy type, import the
 // auxiliary type from reth_codecs::alloy instead.
 compact_types!(
-    // reth-primitives
-    Account,
-    // reth_codecs::alloy
-    Authorization,
-    GenesisAccount,
-    Header,
-    HeaderExt,
-    Withdrawal,
-    TxEip2930,
-    TxEip4844,
-    TxEip7702,
-    TxLegacy,
-    HashBuilderValue,
-    LogData,
-    Log,
-    // BranchNodeCompact, // todo requires arbitrary
-    TrieMask,
-    // TxDeposit, TODO(joshie): optimism
-    // reth_prune_types
-    PruneCheckpoint,
-    PruneMode,
-    // reth_stages_types
-    AccountHashingCheckpoint,
-    StorageHashingCheckpoint,
-    ExecutionCheckpoint,
-    HeadersCheckpoint,
-    IndexHistoryCheckpoint,
-    EntitiesCheckpoint,
-    CheckpointBlockRange,
-    StageCheckpoint,
-    StageUnitCheckpoint,
-    // reth_db_api
-    StoredBlockOmmers,
-    StoredBlockBodyIndices,
-    StoredBlockWithdrawals,
-    // Manual implementations
-    // Transaction & TxType & TxKind & Signature todo requires special handling, since from_compact
-    // receives a txidentifier,
-    TransactionSignedNoHash,
-    // Bytecode, // todo revm arbitrary
-    StorageEntry,
-    // MerkleCheckpoint, // todo storedsubnode -> branchnodecompact arbitrary
-    AccountBeforeTx,
-    ClientVersion,
-    StoredNibbles,
-    StoredNibblesSubKey,
-    // StorageTrieEntry, // todo branchnodecompact arbitrary
-    // StoredSubNode, // todo branchnodecompact arbitrary
-    HashBuilderState
+    regular: [
+        // reth-primitives
+        Account,
+        // reth_codecs::alloy
+        Authorization,
+        GenesisAccount,
+        Header,
+        HeaderExt,
+        Withdrawal,
+        TxEip2930,
+        TxEip4844,
+        TxEip7702,
+        TxLegacy,
+        HashBuilderValue,
+        LogData,
+        Log,
+        // BranchNodeCompact, // todo requires arbitrary
+        TrieMask,
+        // TxDeposit, TODO(joshie): optimism
+        // reth_prune_types
+        PruneCheckpoint,
+        PruneMode,
+        // reth_stages_types
+        AccountHashingCheckpoint,
+        StorageHashingCheckpoint,
+        ExecutionCheckpoint,
+        HeadersCheckpoint,
+        IndexHistoryCheckpoint,
+        EntitiesCheckpoint,
+        CheckpointBlockRange,
+        StageCheckpoint,
+        StageUnitCheckpoint,
+        // reth_db_api
+        StoredBlockOmmers,
+        StoredBlockBodyIndices,
+        StoredBlockWithdrawals,
+        // Manual implementations
+        TransactionSignedNoHash,
+        // Bytecode, // todo revm arbitrary
+        StorageEntry,
+        // MerkleCheckpoint, // todo storedsubnode -> branchnodecompact arbitrary
+        AccountBeforeTx,
+        ClientVersion,
+        StoredNibbles,
+        StoredNibblesSubKey,
+        // StorageTrieEntry, // todo branchnodecompact arbitrary
+        // StoredSubNode, // todo branchnodecompact arbitrary
+        HashBuilderState
+    ],
+    // These types require an extra identifier which is usually stored elsewhere (eg. parent type).
+    identifier: [
+        // Signature todo we for v we only store parity(true || false), while v can take more values
+        Transaction,
+        TxType,
+        TxKind
+    ]
 );
 
 pub(crate) fn generate_vectors() -> Result<()> {
@@ -141,6 +162,7 @@ where
         + std::fmt::Debug
         + 'static,
 {
+    let type_name = type_name::<T>();
     let mut bytes = std::iter::repeat(0u8).take(256).collect::<Vec<u8>>();
     let mut compact_buffer = vec![];
 
@@ -150,18 +172,23 @@ where
         compact_buffer.clear();
 
         let obj = T::arbitrary(&mut arbitrary::Unstructured::new(&bytes))?;
-        obj.to_compact(&mut compact_buffer);
+        let res = obj.to_compact(&mut compact_buffer);
+
+        if IDENTIFIER_TYPE.contains(&type_name) {
+            compact_buffer.push(res as u8);
+        }
+
         values.push((obj, hex::encode(&compact_buffer)));
     }
 
     serde_json::to_writer(
         std::io::BufWriter::new(
-            std::fs::File::create(format!("{VECTORS_FOLDER}/{}.json", type_name::<T>())).unwrap(),
+            std::fs::File::create(format!("{VECTORS_FOLDER}/{}.json", &type_name)).unwrap(),
         ),
         &values,
     )?;
 
-    println!("{} ✅", type_name::<T>());
+    println!("{} ✅", &type_name);
 
     Ok(())
 }
@@ -172,24 +199,33 @@ fn read_vector<T>() -> Result<()>
 where
     T: serde::de::DeserializeOwned + reth_codecs::Compact + PartialEq + Clone + std::fmt::Debug,
 {
+    let type_name = type_name::<T>();
+
     // Read the file where the vectors are stored
-    let file_path = format!("{VECTORS_FOLDER}/{}.json", type_name::<T>());
+    let file_path = format!("{VECTORS_FOLDER}/{}.json", &type_name);
     let file = File::open(&file_path)?;
     let reader = BufReader::new(file);
 
     let stored_values: Vec<(T, String)> = serde_json::from_reader(reader)?;
 
     for (original, hex_str) in stored_values {
-        let compact_bytes = hex::decode(hex_str)?;
-        let (reconstructed, _) = T::from_compact(&compact_bytes, compact_bytes.len());
+        let mut compact_bytes = hex::decode(hex_str)?;
+        let mut identifier = None;
+
+        if IDENTIFIER_TYPE.contains(&type_name) {
+            identifier = compact_bytes.pop().map(|b| b as usize);
+        }
+
+        let len_or_identifier = identifier.unwrap_or(compact_bytes.len());
+        let (reconstructed, _) = T::from_compact(&compact_bytes, len_or_identifier);
 
         if original != reconstructed {
-            println!("{} ❌", type_name::<T>());
+            println!("{} ❌", &type_name);
             panic!("mismatch found on {original:?} and {reconstructed:?}");
         }
     }
 
-    println!("{} ✅", type_name::<T>());
+    println!("{} ✅", &type_name);
 
     Ok(())
 }