reth_cli_commands/test_vectors/
tables.rs

Help
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                            println!("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        (HeaderTerminalDifficulties, PER_TABLE, TABLE),
73        (HeaderNumbers, PER_TABLE, TABLE),
74        (Headers<Header>, PER_TABLE, TABLE),
75        (BlockBodyIndices, PER_TABLE, TABLE),
76        (BlockOmmers<Header>, 100, TABLE),
77        (TransactionHashNumbers, PER_TABLE, TABLE),
78        (Transactions<TransactionSigned>, 100, TABLE),
79        (PlainStorageState, PER_TABLE, DUPSORT),
80        (PlainAccountState, PER_TABLE, TABLE)
81    ]);
82
83    Ok(())
84}
85
86/// Generates test-vectors for normal tables. Keys are sorted and not repeated.
87fn generate_table_vector<T>(runner: &mut TestRunner, per_table: usize) -> Result<()>
88where
89    T: Table,
90    T::Key: for<'a> Arbitrary<'a> + serde::Serialize + Ord + std::hash::Hash + Clone,
91    T::Value: for<'a> Arbitrary<'a> + serde::Serialize + Clone,
92{
93    let mut rows = vec![];
94    let mut seen_keys = HashSet::new();
95    let strategy =
96        proptest::collection::vec(arb::<TableRow<T>>(), per_table - rows.len()).no_shrink().boxed();
97
98    while rows.len() < per_table {
99        // Generate all `per_table` rows: (Key, Value)
100        rows.extend(
101            &mut strategy
102                .new_tree(runner)
103                .map_err(|e| eyre::eyre!("{e}"))?
104                .current()
105                .into_iter()
106                .filter(|e| seen_keys.insert(e.0.clone())),
107        );
108    }
109    // Sort them by `Key`
110    rows.sort_by(|a, b| a.0.cmp(&b.0));
111
112    save_to_file::<T>(rows)
113}
114
115/// Generates test-vectors for DUPSORT tables. Each key has multiple (subkey, value). Keys and
116/// subkeys are sorted.
117fn generate_dupsort_vector<T>(runner: &mut TestRunner, per_table: usize) -> Result<()>
118where
119    T: Table + DupSort,
120    T::Key: for<'a> Arbitrary<'a> + serde::Serialize + Ord + std::hash::Hash + Clone,
121    T::Value: for<'a> Arbitrary<'a> + serde::Serialize + Ord + Clone,
122{
123    let mut rows = vec![];
124
125    // We want to control our repeated keys
126    let mut seen_keys = HashSet::new();
127
128    let start_values = proptest::collection::vec(arb::<T::Value>(), 100..300).no_shrink().boxed();
129
130    let start_keys = arb::<T::Key>().no_shrink().boxed();
131
132    while rows.len() < per_table {
133        let key: T::Key = start_keys.new_tree(runner).map_err(|e| eyre::eyre!("{e}"))?.current();
134
135        if !seen_keys.insert(key.clone()) {
136            continue
137        }
138
139        let mut values: Vec<T::Value> =
140            start_values.new_tree(runner).map_err(|e| eyre::eyre!("{e}"))?.current();
141
142        values.sort();
143
144        for value in values {
145            rows.push((key.clone(), value));
146        }
147    }
148
149    // Sort them by `Key`
150    rows.sort_by(|a, b| a.0.cmp(&b.0));
151
152    save_to_file::<T>(rows)
153}
154
155/// Save rows to file.
156fn save_to_file<T: Table>(rows: Vec<TableRow<T>>) -> eyre::Result<()>
157where
158    T::Key: serde::Serialize,
159    T::Value: serde::Serialize,
160{
161    serde_json::to_writer_pretty(
162        std::io::BufWriter::new(
163            std::fs::File::create(format!("{VECTORS_FOLDER}/{}.json", T::NAME)).unwrap(),
164        ),
165        &rows,
166    )
167    .map_err(|e| eyre::eyre!({ e }))
168}
reth_cli_commands/test_vectors/tables.rs

reth_cli_commands/test_vectors/
tables.rs
