reth_cli_commands/test_vectors/
tables.rs

1use alloy_consensus::Header;
2use alloy_primitives::{hex, B256};
3use arbitrary::Arbitrary;
4use eyre::Result;
5use proptest::{
6    prelude::ProptestConfig,
7    strategy::{Strategy, ValueTree},
8    test_runner::{TestRng, TestRunner},
9};
10use proptest_arbitrary_interop::arb;
11use reth_db_api::{
12    table::{DupSort, Table, TableRow},
13    tables,
14};
15use reth_ethereum_primitives::TransactionSigned;
16use reth_fs_util as fs;
17use std::collections::HashSet;
18use tracing::error;
19
20const VECTORS_FOLDER: &str = "testdata/micro/db";
21const PER_TABLE: usize = 1000;
22
23/// Generates test vectors for specified `tables`. If list is empty, then generate for all tables.
24pub fn generate_vectors(mut tables: Vec<String>) -> Result<()> {
25    // Prepare random seed for test (same method as used by proptest)
26    let seed = B256::random();
27    println!("Seed for table test vectors: {:?}", hex::encode_prefixed(seed));
28
29    // Start the runner with the seed
30    let config = ProptestConfig::default();
31    let rng = TestRng::from_seed(config.rng_algorithm, &seed.0);
32    let mut runner = TestRunner::new_with_rng(config, rng);
33
34    fs::create_dir_all(VECTORS_FOLDER)?;
35
36    macro_rules! generate_vector {
37        ($table_type:ident$(<$($generic:ident),+>)?, $per_table:expr, TABLE) => {
38            generate_table_vector::<tables::$table_type$(<$($generic),+>)?>(&mut runner, $per_table)?;
39        };
40        ($table_type:ident$(<$($generic:ident),+>)?, $per_table:expr, DUPSORT) => {
41            generate_dupsort_vector::<tables::$table_type$(<$($generic),+>)?>(&mut runner, $per_table)?;
42        };
43    }
44
45    macro_rules! generate {
46        ([$(($table_type:ident$(<$($generic:ident),+>)?, $per_table:expr, $table_or_dup:tt)),*]) => {
47            let all_tables = vec![$(stringify!($table_type).to_string(),)*];
48
49            if tables.is_empty() {
50                tables = all_tables;
51            }
52
53            for table in tables {
54                match table.as_str() {
55                    $(
56                        stringify!($table_type) => {
57                            tracing::info!(target: "reth::cli", "Generating test vectors for {} <{}>.", stringify!($table_or_dup), tables::$table_type$(::<$($generic),+>)?::NAME);
58
59                            generate_vector!($table_type$(<$($generic),+>)?, $per_table, $table_or_dup);
60                        },
61                    )*
62                    _ => {
63                        error!(target: "reth::cli", "Unknown table: {}", table);
64                    }
65                }
66            }
67        }
68    }
69
70    generate!([
71        (CanonicalHeaders, PER_TABLE, TABLE),
72        (HeaderNumbers, PER_TABLE, TABLE),
73        (Headers<Header>, PER_TABLE, TABLE),
74        (BlockBodyIndices, PER_TABLE, TABLE),
75        (BlockOmmers<Header>, 100, TABLE),
76        (TransactionHashNumbers, PER_TABLE, TABLE),
77        (Transactions<TransactionSigned>, 100, TABLE),
78        (PlainStorageState, PER_TABLE, DUPSORT),
79        (PlainAccountState, PER_TABLE, TABLE)
80    ]);
81
82    Ok(())
83}
84
85/// Generates test-vectors for normal tables. Keys are sorted and not repeated.
86fn generate_table_vector<T>(runner: &mut TestRunner, per_table: usize) -> Result<()>
87where
88    T: Table,
89    T::Key: for<'a> Arbitrary<'a> + serde::Serialize + Ord + std::hash::Hash + Clone,
90    T::Value: for<'a> Arbitrary<'a> + serde::Serialize + Clone,
91{
92    let mut rows = vec![];
93    let mut seen_keys = HashSet::new();
94    let strategy =
95        proptest::collection::vec(arb::<TableRow<T>>(), per_table - rows.len()).no_shrink().boxed();
96
97    while rows.len() < per_table {
98        // Generate all `per_table` rows: (Key, Value)
99        rows.extend(
100            &mut strategy
101                .new_tree(runner)
102                .map_err(|e| eyre::eyre!("{e}"))?
103                .current()
104                .into_iter()
105                .filter(|e| seen_keys.insert(e.0.clone())),
106        );
107    }
108    // Sort them by `Key`
109    rows.sort_by(|a, b| a.0.cmp(&b.0));
110
111    save_to_file::<T>(rows)
112}
113
114/// Generates test-vectors for DUPSORT tables. Each key has multiple (subkey, value). Keys and
115/// subkeys are sorted.
116fn generate_dupsort_vector<T>(runner: &mut TestRunner, per_table: usize) -> Result<()>
117where
118    T: Table + DupSort,
119    T::Key: for<'a> Arbitrary<'a> + serde::Serialize + Ord + std::hash::Hash + Clone,
120    T::Value: for<'a> Arbitrary<'a> + serde::Serialize + Ord + Clone,
121{
122    let mut rows = vec![];
123
124    // We want to control our repeated keys
125    let mut seen_keys = HashSet::new();
126
127    let start_values = proptest::collection::vec(arb::<T::Value>(), 100..300).no_shrink().boxed();
128
129    let start_keys = arb::<T::Key>().no_shrink().boxed();
130
131    while rows.len() < per_table {
132        let key: T::Key = start_keys.new_tree(runner).map_err(|e| eyre::eyre!("{e}"))?.current();
133
134        if !seen_keys.insert(key.clone()) {
135            continue
136        }
137
138        let mut values: Vec<T::Value> =
139            start_values.new_tree(runner).map_err(|e| eyre::eyre!("{e}"))?.current();
140
141        values.sort();
142
143        for value in values {
144            rows.push((key.clone(), value));
145        }
146    }
147
148    // Sort them by `Key`
149    rows.sort_by(|a, b| a.0.cmp(&b.0));
150
151    save_to_file::<T>(rows)
152}
153
154/// Save rows to file.
155fn save_to_file<T: Table>(rows: Vec<TableRow<T>>) -> eyre::Result<()>
156where
157    T::Key: serde::Serialize,
158    T::Value: serde::Serialize,
159{
160    serde_json::to_writer_pretty(
161        std::io::BufWriter::new(
162            std::fs::File::create(format!("{VECTORS_FOLDER}/{}.json", T::NAME)).unwrap(),
163        ),
164        &rows,
165    )
166    .map_err(|e| eyre::eyre!({ e }))
167}
reth_cli_commands/test_vectors/tables.rs

reth_cli_commands/test_vectors/
tables.rs
