fuel_core_chain_config/config/
state.rs

1use super::{
2    blob::BlobConfig,
3    coin::CoinConfig,
4    contract::ContractConfig,
5    message::MessageConfig,
6    table_entry::TableEntry,
7};
8use crate::{
9    ContractBalanceConfig,
10    ContractStateConfig,
11};
12use fuel_core_storage::{
13    ContractsAssetKey,
14    ContractsStateKey,
15    Mappable,
16    structured_storage::TableWithBlueprint,
17    tables::{
18        Coins,
19        ContractsAssets,
20        ContractsLatestUtxo,
21        ContractsRawCode,
22        ContractsState,
23        FuelBlocks,
24        Messages,
25        ProcessedTransactions,
26        SealedBlockConsensus,
27        Transactions,
28    },
29};
30use fuel_core_types::{
31    blockchain::primitives::DaBlockHeight,
32    entities::contract::ContractUtxoInfo,
33    fuel_types::{
34        BlockHeight,
35        Bytes32,
36    },
37    fuel_vm::BlobData,
38};
39use itertools::Itertools;
40use serde::{
41    Deserialize,
42    Serialize,
43};
44
45#[cfg(feature = "std")]
46use crate::SnapshotMetadata;
47
48#[cfg(feature = "test-helpers")]
49use crate::coin_config_helpers::CoinConfigGenerator;
50#[cfg(feature = "test-helpers")]
51use bech32::{
52    ToBase32,
53    Variant::Bech32m,
54};
55#[cfg(feature = "test-helpers")]
56use core::str::FromStr;
57use fuel_core_storage::tables::merkle::{
58    FuelBlockMerkleData,
59    FuelBlockMerkleMetadata,
60};
61use fuel_core_types::blockchain::header::{
62    BlockHeader,
63    ConsensusParametersVersion,
64    StateTransitionBytecodeVersion,
65};
66#[cfg(feature = "test-helpers")]
67use fuel_core_types::{
68    fuel_types::Address,
69    fuel_vm::SecretKey,
70};
71
72#[cfg(feature = "parquet")]
73mod parquet;
74mod reader;
75#[cfg(feature = "std")]
76mod writer;
77
78// Fuel Network human-readable part for bech32 encoding
79pub const FUEL_BECH32_HRP: &str = "fuel";
80pub const TESTNET_INITIAL_BALANCE: u64 = 10_000_000_000;
81
82pub const TESTNET_WALLET_SECRETS: [&str; 5] = [
83    "0xde97d8624a438121b86a1956544bd72ed68cd69f2c99555b08b1e8c51ffd511c",
84    "0x37fa81c84ccd547c30c176b118d5cb892bdb113e8e80141f266519422ef9eefd",
85    "0x862512a2363db2b3a375c0d4bbbd27172180d89f23f2e259bac850ab02619301",
86    "0x976e5c3fa620092c718d852ca703b6da9e3075b9f2ecb8ed42d9f746bf26aafb",
87    "0x7f8a325504e7315eda997db7861c9447f5c3eff26333b20180475d94443a10c6",
88];
89
90#[derive(Default, Copy, Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
91pub struct LastBlockConfig {
92    /// The block height of the last block.
93    pub block_height: BlockHeight,
94    /// The da height used in the last block.
95    pub da_block_height: DaBlockHeight,
96    /// The version of consensus parameters used to produce last block.
97    pub consensus_parameters_version: ConsensusParametersVersion,
98    /// The version of state transition function used to produce last block.
99    pub state_transition_version: StateTransitionBytecodeVersion,
100    /// The Merkle root of all blocks before regenesis.
101    pub blocks_root: Bytes32,
102}
103
104impl LastBlockConfig {
105    pub fn from_header(header: &BlockHeader, blocks_root: Bytes32) -> Self {
106        Self {
107            block_height: *header.height(),
108            da_block_height: header.da_height(),
109            consensus_parameters_version: header.consensus_parameters_version(),
110            state_transition_version: header.state_transition_bytecode_version(),
111            blocks_root,
112        }
113    }
114}
115
116#[derive(Default, Clone, Debug, Deserialize, Serialize, Eq, PartialEq)]
117pub struct StateConfig {
118    /// Spendable coins
119    pub coins: Vec<CoinConfig>,
120    /// Messages from Layer 1
121    pub messages: Vec<MessageConfig>,
122    /// Blobs
123    #[serde(default)]
124    pub blobs: Vec<BlobConfig>,
125    /// Contracts
126    pub contracts: Vec<ContractConfig>,
127    /// Last block config.
128    pub last_block: Option<LastBlockConfig>,
129}
130
131#[derive(Debug, Clone, Default)]
132pub struct StateConfigBuilder {
133    coins: Vec<TableEntry<Coins>>,
134    messages: Vec<TableEntry<Messages>>,
135    blobs: Vec<TableEntry<BlobData>>,
136    contract_state: Vec<TableEntry<ContractsState>>,
137    contract_balance: Vec<TableEntry<ContractsAssets>>,
138    contract_code: Vec<TableEntry<ContractsRawCode>>,
139    contract_utxo: Vec<TableEntry<ContractsLatestUtxo>>,
140}
141
142impl StateConfigBuilder {
143    pub fn merge(&mut self, builder: Self) -> &mut Self {
144        self.coins.extend(builder.coins);
145        self.messages.extend(builder.messages);
146        self.blobs.extend(builder.blobs);
147        self.contract_state.extend(builder.contract_state);
148        self.contract_balance.extend(builder.contract_balance);
149        self.contract_code.extend(builder.contract_code);
150        self.contract_utxo.extend(builder.contract_utxo);
151
152        self
153    }
154
155    #[cfg(feature = "std")]
156    pub fn build(
157        self,
158        latest_block_config: Option<LastBlockConfig>,
159    ) -> anyhow::Result<StateConfig> {
160        use std::collections::HashMap;
161
162        let coins = self.coins.into_iter().map(|coin| coin.into()).collect();
163        let messages = self
164            .messages
165            .into_iter()
166            .map(|message| message.into())
167            .collect();
168        let blobs = self.blobs.into_iter().map(|blob| blob.into()).collect();
169        let contract_ids = self
170            .contract_code
171            .iter()
172            .map(|entry| entry.key)
173            .collect::<Vec<_>>();
174        let mut state: HashMap<_, _> = self
175            .contract_state
176            .into_iter()
177            .map(|state| {
178                (
179                    *state.key.contract_id(),
180                    ContractStateConfig {
181                        key: *state.key.state_key(),
182                        value: state.value.into(),
183                    },
184                )
185            })
186            .into_group_map();
187
188        let mut balance: HashMap<_, _> = self
189            .contract_balance
190            .into_iter()
191            .map(|balance| {
192                (
193                    *balance.key.contract_id(),
194                    ContractBalanceConfig {
195                        asset_id: *balance.key.asset_id(),
196                        amount: balance.value,
197                    },
198                )
199            })
200            .into_group_map();
201
202        let mut contract_code: HashMap<_, Vec<u8>> = self
203            .contract_code
204            .into_iter()
205            .map(|entry| (entry.key, entry.value.into()))
206            .collect();
207
208        let mut contract_utxos: HashMap<_, _> = self
209            .contract_utxo
210            .into_iter()
211            .map(|entry| match entry.value {
212                ContractUtxoInfo::V1(utxo) => {
213                    (entry.key, (utxo.utxo_id, utxo.tx_pointer))
214                }
215                _ => unreachable!(),
216            })
217            .collect();
218
219        let contracts = contract_ids
220            .into_iter()
221            .map(|id| -> anyhow::Result<_> {
222                let code = contract_code
223                    .remove(&id)
224                    .ok_or_else(|| anyhow::anyhow!("Missing code for contract: {id}"))?;
225                let (utxo_id, tx_pointer) = contract_utxos
226                    .remove(&id)
227                    .ok_or_else(|| anyhow::anyhow!("Missing utxo for contract: {id}"))?;
228                let states = state.remove(&id).unwrap_or_default();
229                let balances = balance.remove(&id).unwrap_or_default();
230
231                Ok(ContractConfig {
232                    contract_id: id,
233                    code,
234                    tx_id: *utxo_id.tx_id(),
235                    output_index: utxo_id.output_index(),
236                    tx_pointer_block_height: tx_pointer.block_height(),
237                    tx_pointer_tx_idx: tx_pointer.tx_index(),
238                    states,
239                    balances,
240                })
241            })
242            .try_collect()?;
243
244        Ok(StateConfig {
245            coins,
246            messages,
247            blobs,
248            contracts,
249            last_block: latest_block_config,
250        })
251    }
252}
253
254pub trait AddTable<T>
255where
256    T: Mappable,
257{
258    fn add(&mut self, _entries: Vec<TableEntry<T>>);
259}
260
261impl AddTable<Coins> for StateConfigBuilder {
262    fn add(&mut self, entries: Vec<TableEntry<Coins>>) {
263        self.coins.extend(entries);
264    }
265}
266
267impl AsTable<Coins> for StateConfig {
268    fn as_table(&self) -> Vec<TableEntry<Coins>> {
269        self.coins
270            .clone()
271            .into_iter()
272            .map(|coin| coin.into())
273            .collect()
274    }
275}
276
277#[cfg(feature = "test-helpers")]
278impl crate::Randomize for StateConfig {
279    fn randomize(mut rng: impl rand::Rng) -> Self {
280        let amount = 2;
281        fn rand_collection<T: crate::Randomize>(
282            mut rng: impl rand::Rng,
283            amount: usize,
284        ) -> Vec<T> {
285            std::iter::repeat_with(|| crate::Randomize::randomize(&mut rng))
286                .take(amount)
287                .collect()
288        }
289
290        Self {
291            coins: rand_collection(&mut rng, amount),
292            messages: rand_collection(&mut rng, amount),
293            blobs: rand_collection(&mut rng, amount),
294            contracts: rand_collection(&mut rng, amount),
295            last_block: Some(LastBlockConfig {
296                block_height: rng.r#gen(),
297                da_block_height: rng.r#gen(),
298                consensus_parameters_version: rng.r#gen(),
299                state_transition_version: rng.r#gen(),
300                blocks_root: rng.r#gen(),
301            }),
302        }
303    }
304}
305
306pub trait AsTable<T>
307where
308    T: TableWithBlueprint,
309{
310    fn as_table(&self) -> Vec<TableEntry<T>>;
311}
312
313impl AsTable<Messages> for StateConfig {
314    fn as_table(&self) -> Vec<TableEntry<Messages>> {
315        self.messages
316            .clone()
317            .into_iter()
318            .map(|message| message.into())
319            .collect()
320    }
321}
322
323impl AddTable<Messages> for StateConfigBuilder {
324    fn add(&mut self, entries: Vec<TableEntry<Messages>>) {
325        self.messages.extend(entries);
326    }
327}
328
329impl AsTable<BlobData> for StateConfig {
330    fn as_table(&self) -> Vec<TableEntry<BlobData>> {
331        self.blobs
332            .clone()
333            .into_iter()
334            .map(|blob| blob.into())
335            .collect()
336    }
337}
338
339impl AddTable<BlobData> for StateConfigBuilder {
340    fn add(&mut self, entries: Vec<TableEntry<BlobData>>) {
341        self.blobs.extend(entries);
342    }
343}
344
345impl AsTable<ContractsState> for StateConfig {
346    fn as_table(&self) -> Vec<TableEntry<ContractsState>> {
347        self.contracts
348            .iter()
349            .flat_map(|contract| {
350                contract.states.iter().map(
351                    |ContractStateConfig {
352                         key: state_key,
353                         value: state_value,
354                     }| TableEntry {
355                        key: ContractsStateKey::new(&contract.contract_id, state_key),
356                        value: state_value.clone().into(),
357                    },
358                )
359            })
360            .collect()
361    }
362}
363
364impl AddTable<ContractsState> for StateConfigBuilder {
365    fn add(&mut self, entries: Vec<TableEntry<ContractsState>>) {
366        self.contract_state.extend(entries);
367    }
368}
369
370impl AsTable<ContractsAssets> for StateConfig {
371    fn as_table(&self) -> Vec<TableEntry<ContractsAssets>> {
372        self.contracts
373            .iter()
374            .flat_map(|contract| {
375                contract.balances.iter().map(
376                    |ContractBalanceConfig { asset_id, amount }| TableEntry {
377                        key: ContractsAssetKey::new(&contract.contract_id, asset_id),
378                        value: *amount,
379                    },
380                )
381            })
382            .collect()
383    }
384}
385
386impl AddTable<ContractsAssets> for StateConfigBuilder {
387    fn add(&mut self, entries: Vec<TableEntry<ContractsAssets>>) {
388        self.contract_balance.extend(entries);
389    }
390}
391
392impl AsTable<ContractsRawCode> for StateConfig {
393    fn as_table(&self) -> Vec<TableEntry<ContractsRawCode>> {
394        self.contracts
395            .iter()
396            .map(|config| TableEntry {
397                key: config.contract_id,
398                value: config.code.as_slice().into(),
399            })
400            .collect()
401    }
402}
403
404impl AddTable<ContractsRawCode> for StateConfigBuilder {
405    fn add(&mut self, entries: Vec<TableEntry<ContractsRawCode>>) {
406        self.contract_code.extend(entries);
407    }
408}
409
410impl AsTable<ContractsLatestUtxo> for StateConfig {
411    fn as_table(&self) -> Vec<TableEntry<ContractsLatestUtxo>> {
412        self.contracts
413            .iter()
414            .map(|config| TableEntry {
415                key: config.contract_id,
416                value: ContractUtxoInfo::V1(
417                    fuel_core_types::entities::contract::ContractUtxoInfoV1 {
418                        utxo_id: config.utxo_id(),
419                        tx_pointer: config.tx_pointer(),
420                    },
421                ),
422            })
423            .collect()
424    }
425}
426
427impl AddTable<ContractsLatestUtxo> for StateConfigBuilder {
428    fn add(&mut self, entries: Vec<TableEntry<ContractsLatestUtxo>>) {
429        self.contract_utxo.extend(entries);
430    }
431}
432
433impl AsTable<Transactions> for StateConfig {
434    fn as_table(&self) -> Vec<TableEntry<Transactions>> {
435        Vec::new() // Do not include these for now
436    }
437}
438
439impl AddTable<Transactions> for StateConfigBuilder {
440    fn add(&mut self, _entries: Vec<TableEntry<Transactions>>) {}
441}
442
443impl AsTable<FuelBlocks> for StateConfig {
444    fn as_table(&self) -> Vec<TableEntry<FuelBlocks>> {
445        Vec::new() // Do not include these for now
446    }
447}
448
449impl AddTable<FuelBlocks> for StateConfigBuilder {
450    fn add(&mut self, _entries: Vec<TableEntry<FuelBlocks>>) {}
451}
452
453impl AsTable<SealedBlockConsensus> for StateConfig {
454    fn as_table(&self) -> Vec<TableEntry<SealedBlockConsensus>> {
455        Vec::new() // Do not include these for now
456    }
457}
458
459impl AddTable<SealedBlockConsensus> for StateConfigBuilder {
460    fn add(&mut self, _entries: Vec<TableEntry<SealedBlockConsensus>>) {}
461}
462
463impl AsTable<FuelBlockMerkleData> for StateConfig {
464    fn as_table(&self) -> Vec<TableEntry<FuelBlockMerkleData>> {
465        Vec::new() // Do not include these for now
466    }
467}
468
469impl AddTable<FuelBlockMerkleData> for StateConfigBuilder {
470    fn add(&mut self, _entries: Vec<TableEntry<FuelBlockMerkleData>>) {}
471}
472
473impl AsTable<FuelBlockMerkleMetadata> for StateConfig {
474    fn as_table(&self) -> Vec<TableEntry<FuelBlockMerkleMetadata>> {
475        Vec::new() // Do not include these for now
476    }
477}
478
479impl AddTable<FuelBlockMerkleMetadata> for StateConfigBuilder {
480    fn add(&mut self, _entries: Vec<TableEntry<FuelBlockMerkleMetadata>>) {}
481}
482
483impl AddTable<ProcessedTransactions> for StateConfigBuilder {
484    fn add(&mut self, _: Vec<TableEntry<ProcessedTransactions>>) {}
485}
486
487impl AsTable<ProcessedTransactions> for StateConfig {
488    fn as_table(&self) -> Vec<TableEntry<ProcessedTransactions>> {
489        Vec::new() // Do not include these for now
490    }
491}
492
493impl StateConfig {
494    pub fn sorted(mut self) -> Self {
495        self.coins = self
496            .coins
497            .into_iter()
498            .sorted_by_key(|c| c.utxo_id())
499            .collect();
500
501        self.messages = self
502            .messages
503            .into_iter()
504            .sorted_by_key(|m| m.nonce)
505            .collect();
506
507        self.blobs = self
508            .blobs
509            .into_iter()
510            .sorted_by_key(|b| b.blob_id)
511            .collect();
512
513        self.contracts = self
514            .contracts
515            .into_iter()
516            .sorted_by_key(|c| c.contract_id)
517            .collect();
518
519        self
520    }
521
522    pub fn extend(&mut self, other: Self) {
523        self.coins.extend(other.coins);
524        self.messages.extend(other.messages);
525        self.contracts.extend(other.contracts);
526    }
527
528    #[cfg(feature = "std")]
529    pub fn from_snapshot_metadata(
530        snapshot_metadata: SnapshotMetadata,
531    ) -> anyhow::Result<Self> {
532        let reader = crate::SnapshotReader::open(snapshot_metadata)?;
533        Self::from_reader(&reader)
534    }
535
536    #[cfg(feature = "std")]
537    pub fn from_reader(reader: &SnapshotReader) -> anyhow::Result<Self> {
538        let mut builder = StateConfigBuilder::default();
539
540        let coins = reader
541            .read::<Coins>()?
542            .into_iter()
543            .flatten_ok()
544            .try_collect()?;
545
546        builder.add(coins);
547
548        let messages = reader
549            .read::<Messages>()?
550            .into_iter()
551            .flatten_ok()
552            .try_collect()?;
553
554        builder.add(messages);
555
556        let blobs = reader
557            .read::<BlobData>()?
558            .into_iter()
559            .flatten_ok()
560            .try_collect()?;
561
562        builder.add(blobs);
563
564        let contract_state = reader
565            .read::<ContractsState>()?
566            .into_iter()
567            .flatten_ok()
568            .try_collect()?;
569
570        builder.add(contract_state);
571
572        let contract_balance = reader
573            .read::<ContractsAssets>()?
574            .into_iter()
575            .flatten_ok()
576            .try_collect()?;
577
578        builder.add(contract_balance);
579
580        let contract_code = reader
581            .read::<ContractsRawCode>()?
582            .into_iter()
583            .flatten_ok()
584            .try_collect()?;
585
586        builder.add(contract_code);
587
588        let contract_utxo = reader
589            .read::<ContractsLatestUtxo>()?
590            .into_iter()
591            .flatten_ok()
592            .try_collect()?;
593
594        builder.add(contract_utxo);
595
596        builder.build(reader.last_block_config().cloned())
597    }
598
599    #[cfg(feature = "test-helpers")]
600    pub fn local_testnet() -> Self {
601        // endow some preset accounts with an initial balance
602        tracing::info!("Initial Accounts");
603
604        let mut coin_generator = CoinConfigGenerator::new();
605        let coins = TESTNET_WALLET_SECRETS
606            .into_iter()
607            .map(|secret| {
608                let secret = SecretKey::from_str(secret).expect("Expected valid secret");
609                let address = Address::from(*secret.public_key().hash());
610                let bech32_data = Bytes32::new(*address).to_base32();
611                let bech32_encoding =
612                    bech32::encode(FUEL_BECH32_HRP, bech32_data, Bech32m).unwrap();
613                tracing::info!(
614                    "PrivateKey({:#x}), Address({:#x} [bech32: {}]), Balance({})",
615                    secret,
616                    address,
617                    bech32_encoding,
618                    TESTNET_INITIAL_BALANCE
619                );
620                coin_generator.generate_with(secret, TESTNET_INITIAL_BALANCE)
621            })
622            .collect_vec();
623
624        Self {
625            coins,
626            ..StateConfig::default()
627        }
628    }
629
630    #[cfg(feature = "test-helpers")]
631    pub fn random_testnet() -> Self {
632        tracing::info!("Initial Accounts");
633        let mut rng = rand::thread_rng();
634        let mut coin_generator = CoinConfigGenerator::new();
635        let coins = (0..5)
636            .map(|_| {
637                let secret = SecretKey::random(&mut rng);
638                let address = Address::from(*secret.public_key().hash());
639                let bech32_data = Bytes32::new(*address).to_base32();
640                let bech32_encoding =
641                    bech32::encode(FUEL_BECH32_HRP, bech32_data, Bech32m).unwrap();
642                tracing::info!(
643                    "PrivateKey({:#x}), Address({:#x} [bech32: {}]), Balance({})",
644                    secret,
645                    address,
646                    bech32_encoding,
647                    TESTNET_INITIAL_BALANCE
648                );
649                coin_generator.generate_with(secret, TESTNET_INITIAL_BALANCE)
650            })
651            .collect_vec();
652
653        Self {
654            coins,
655            ..StateConfig::default()
656        }
657    }
658}
659
660pub use reader::{
661    GroupIter,
662    Groups,
663    SnapshotReader,
664};
665#[cfg(feature = "parquet")]
666pub use writer::ZstdCompressionLevel;
667#[cfg(feature = "std")]
668pub use writer::{
669    SnapshotFragment,
670    SnapshotWriter,
671};
672pub const MAX_GROUP_SIZE: usize = usize::MAX;
673
674#[cfg(test)]
675mod tests {
676    use std::path::Path;
677
678    use crate::{
679        ChainConfig,
680        Randomize,
681    };
682
683    use rand::{
684        SeedableRng,
685        rngs::StdRng,
686    };
687
688    use super::*;
689
690    #[test]
691    fn parquet_roundtrip() {
692        let writer = given_parquet_writer;
693
694        let reader = |metadata: SnapshotMetadata, _: usize| {
695            SnapshotReader::open(metadata).unwrap()
696        };
697
698        macro_rules! test_tables {
699                ($($table:ty),*) => {
700                    $(assert_roundtrip::<$table>(writer, reader);)*
701                };
702            }
703
704        test_tables!(
705            Coins,
706            BlobData,
707            ContractsAssets,
708            ContractsLatestUtxo,
709            ContractsRawCode,
710            ContractsState,
711            Messages
712        );
713    }
714
715    fn given_parquet_writer(path: &Path) -> SnapshotWriter {
716        SnapshotWriter::parquet(path, writer::ZstdCompressionLevel::Level1).unwrap()
717    }
718
719    fn given_json_writer(path: &Path) -> SnapshotWriter {
720        SnapshotWriter::json(path)
721    }
722
723    #[test]
724    fn json_roundtrip_non_contract_related_tables() {
725        let writer = |temp_dir: &Path| SnapshotWriter::json(temp_dir);
726        let reader = |metadata: SnapshotMetadata, group_size: usize| {
727            SnapshotReader::open_w_config(metadata, group_size).unwrap()
728        };
729
730        assert_roundtrip::<Coins>(writer, reader);
731        assert_roundtrip::<Messages>(writer, reader);
732        assert_roundtrip::<BlobData>(writer, reader);
733    }
734
735    #[test]
736    fn json_roundtrip_contract_related_tables() {
737        // given
738        let mut rng = StdRng::seed_from_u64(0);
739        let contracts = std::iter::repeat_with(|| ContractConfig::randomize(&mut rng))
740            .take(4)
741            .collect_vec();
742
743        let tmp_dir = tempfile::tempdir().unwrap();
744        let writer = SnapshotWriter::json(tmp_dir.path());
745
746        let state = StateConfig {
747            contracts,
748            ..Default::default()
749        };
750
751        // when
752        let snapshot = writer
753            .write_state_config(state.clone(), &ChainConfig::local_testnet())
754            .unwrap();
755
756        // then
757        let reader = SnapshotReader::open(snapshot).unwrap();
758        let read_state = StateConfig::from_reader(&reader).unwrap();
759
760        pretty_assertions::assert_eq!(state, read_state);
761    }
762
763    #[test_case::test_case(given_parquet_writer)]
764    #[test_case::test_case(given_json_writer)]
765    fn writes_in_fragments_correctly(writer: impl Fn(&Path) -> SnapshotWriter + Copy) {
766        // given
767        let temp_dir = tempfile::tempdir().unwrap();
768        let create_writer = || writer(temp_dir.path());
769
770        let mut rng = StdRng::seed_from_u64(0);
771        let state_config = StateConfig::randomize(&mut rng);
772
773        let chain_config = ChainConfig::local_testnet();
774
775        macro_rules! write_in_fragments {
776            ($($fragment_ty: ty,)*) => {[
777                $({
778                    let mut writer = create_writer();
779                    writer
780                        .write(AsTable::<$fragment_ty>::as_table(&state_config))
781                        .unwrap();
782                    writer.partial_close().unwrap()
783                }),*
784            ]}
785        }
786
787        let fragments = write_in_fragments!(
788            Coins,
789            Messages,
790            BlobData,
791            ContractsState,
792            ContractsAssets,
793            ContractsRawCode,
794            ContractsLatestUtxo,
795        );
796
797        // when
798        let snapshot = fragments
799            .into_iter()
800            .reduce(|fragment, next_fragment| fragment.merge(next_fragment).unwrap())
801            .unwrap()
802            .finalize(state_config.last_block, &chain_config)
803            .unwrap();
804
805        // then
806        let reader = SnapshotReader::open(snapshot).unwrap();
807
808        let read_state_config = StateConfig::from_reader(&reader).unwrap();
809        assert_eq!(read_state_config, state_config);
810        assert_eq!(reader.chain_config(), &chain_config);
811    }
812
813    #[test_case::test_case(given_parquet_writer)]
814    #[test_case::test_case(given_json_writer)]
815    fn roundtrip_block_heights(writer: impl FnOnce(&Path) -> SnapshotWriter) {
816        // given
817        let temp_dir = tempfile::tempdir().unwrap();
818        let block_height = 13u32.into();
819        let da_block_height = 14u64.into();
820        let consensus_parameters_version = 321u32;
821        let state_transition_version = 123u32;
822        let blocks_root = Bytes32::from([123; 32]);
823        let block_config = LastBlockConfig {
824            block_height,
825            da_block_height,
826            consensus_parameters_version,
827            state_transition_version,
828            blocks_root,
829        };
830        let writer = writer(temp_dir.path());
831
832        // when
833        let snapshot = writer
834            .close(Some(block_config), &ChainConfig::local_testnet())
835            .unwrap();
836
837        // then
838        let reader = SnapshotReader::open(snapshot).unwrap();
839
840        let block_config_decoded = reader.last_block_config().cloned();
841        pretty_assertions::assert_eq!(Some(block_config), block_config_decoded);
842    }
843
844    #[test_case::test_case(given_parquet_writer)]
845    #[test_case::test_case(given_json_writer)]
846    fn missing_tables_tolerated(writer: impl FnOnce(&Path) -> SnapshotWriter) {
847        // given
848        let temp_dir = tempfile::tempdir().unwrap();
849        let writer = writer(temp_dir.path());
850        let snapshot = writer.close(None, &ChainConfig::local_testnet()).unwrap();
851
852        let reader = SnapshotReader::open(snapshot).unwrap();
853
854        // when
855        let coins = reader.read::<Coins>().unwrap();
856
857        // then
858        assert_eq!(coins.into_iter().count(), 0);
859    }
860
861    fn assert_roundtrip<T>(
862        writer: impl FnOnce(&Path) -> SnapshotWriter,
863        reader: impl FnOnce(SnapshotMetadata, usize) -> SnapshotReader,
864    ) where
865        T: TableWithBlueprint,
866        T::OwnedKey: Randomize
867            + serde::Serialize
868            + serde::de::DeserializeOwned
869            + core::fmt::Debug
870            + PartialEq,
871        T::OwnedValue: Randomize
872            + serde::Serialize
873            + serde::de::DeserializeOwned
874            + core::fmt::Debug
875            + PartialEq,
876        StateConfig: AsTable<T>,
877        TableEntry<T>: Randomize,
878        StateConfigBuilder: AddTable<T>,
879    {
880        // given
881        let skip_n_groups = 3;
882        let temp_dir = tempfile::tempdir().unwrap();
883
884        let num_groups = 4;
885        let group_size = 1;
886        let mut group_generator =
887            GroupGenerator::new(StdRng::seed_from_u64(0), group_size, num_groups);
888        let mut snapshot_writer = writer(temp_dir.path());
889
890        // when
891        let expected_groups = group_generator
892            .write_groups::<T>(&mut snapshot_writer)
893            .into_iter()
894            .collect_vec();
895        let snapshot = snapshot_writer
896            .close(None, &ChainConfig::local_testnet())
897            .unwrap();
898
899        let actual_groups = reader(snapshot, group_size)
900            .read()
901            .unwrap()
902            .into_iter()
903            .collect_vec();
904
905        // then
906        assert_groups_identical(&expected_groups, actual_groups, skip_n_groups);
907    }
908
909    struct GroupGenerator<R> {
910        rand: R,
911        group_size: usize,
912        num_groups: usize,
913    }
914
915    impl<R: ::rand::RngCore> GroupGenerator<R> {
916        fn new(rand: R, group_size: usize, num_groups: usize) -> Self {
917            Self {
918                rand,
919                group_size,
920                num_groups,
921            }
922        }
923
924        fn write_groups<T>(
925            &mut self,
926            encoder: &mut SnapshotWriter,
927        ) -> Vec<Vec<TableEntry<T>>>
928        where
929            T: TableWithBlueprint,
930            T::OwnedKey: serde::Serialize,
931            T::OwnedValue: serde::Serialize,
932            TableEntry<T>: Randomize,
933            StateConfigBuilder: AddTable<T>,
934        {
935            let groups = self.generate_groups();
936            for group in &groups {
937                encoder.write(group.clone()).unwrap();
938            }
939            groups
940        }
941
942        fn generate_groups<T>(&mut self) -> Vec<Vec<T>>
943        where
944            T: Randomize,
945        {
946            ::std::iter::repeat_with(|| T::randomize(&mut self.rand))
947                .chunks(self.group_size)
948                .into_iter()
949                .map(|chunk| chunk.collect_vec())
950                .take(self.num_groups)
951                .collect()
952        }
953    }
954
955    fn assert_groups_identical<T>(
956        original: &[Vec<T>],
957        read: impl IntoIterator<Item = Result<Vec<T>, anyhow::Error>>,
958        skip: usize,
959    ) where
960        Vec<T>: PartialEq,
961        T: PartialEq + std::fmt::Debug,
962    {
963        pretty_assertions::assert_eq!(
964            original[skip..],
965            read.into_iter()
966                .skip(skip)
967                .collect::<Result<Vec<_>, _>>()
968                .unwrap()
969        );
970    }
971}