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