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
78pub 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 pub block_height: BlockHeight,
94 pub da_block_height: DaBlockHeight,
96 pub consensus_parameters_version: ConsensusParametersVersion,
98 pub state_transition_version: StateTransitionBytecodeVersion,
100 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 pub coins: Vec<CoinConfig>,
124 pub messages: Vec<MessageConfig>,
126 #[serde(default)]
128 pub blobs: Vec<BlobConfig>,
129 pub contracts: Vec<ContractConfig>,
131 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() }
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() }
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() }
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() }
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() }
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() }
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 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 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 let snapshot = writer
757 .write_state_config(state.clone(), &ChainConfig::local_testnet())
758 .unwrap();
759
760 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 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 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 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 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 let snapshot = writer
838 .close(Some(block_config), &ChainConfig::local_testnet())
839 .unwrap();
840
841 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 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 let coins = reader.read::<Coins>().unwrap();
860
861 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 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 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 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}