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
78pub 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 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.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 pub coins: Vec<CoinConfig>,
120 pub messages: Vec<MessageConfig>,
122 #[serde(default)]
124 pub blobs: Vec<BlobConfig>,
125 pub contracts: Vec<ContractConfig>,
127 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() }
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() }
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() }
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() }
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() }
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() }
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 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 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 let snapshot = writer
753 .write_state_config(state.clone(), &ChainConfig::local_testnet())
754 .unwrap();
755
756 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 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 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 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 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 let snapshot = writer
834 .close(Some(block_config), &ChainConfig::local_testnet())
835 .unwrap();
836
837 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 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 let coins = reader.read::<Coins>().unwrap();
856
857 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 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 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 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}