iota_genesis_builder/stardust/
parse.rs

1// Copyright (c) 2024 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4//! Types and logic to parse a full Stardust snapshot.
5use std::io::{BufReader, Read};
6
7use anyhow::Result;
8use iota_sdk::types::block::{
9    output::Output, payload::milestone::MilestoneOption, protocol::ProtocolParameters,
10};
11use iota_types::stardust::error::StardustError;
12use packable::{
13    Packable,
14    unpacker::{IoUnpacker, Unpacker},
15};
16
17use super::types::{output_header::OutputHeader, snapshot::FullSnapshotHeader};
18
19/// Parse a Hornet genesis snapshot using a [`BufReader`] internally.
20pub struct HornetSnapshotParser<R: Read> {
21    reader: IoUnpacker<BufReader<R>>,
22    /// The full-snapshot header
23    pub header: FullSnapshotHeader,
24}
25
26impl<R: Read> HornetSnapshotParser<R> {
27    /// Creates a new [`HornetSnapshotParser`].
28    ///
29    /// `VERIFY = true` ensures that only global snapshots parse successfully.
30    pub fn new<const VERIFY: bool>(reader: R) -> Result<Self> {
31        let mut reader = IoUnpacker::new(std::io::BufReader::new(reader));
32        let header = FullSnapshotHeader::unpack::<_, VERIFY>(&mut reader, &())?;
33        Ok(Self { reader, header })
34    }
35
36    /// Provide an iterator over the Stardust UTXOs recorded in the snapshot.
37    pub fn outputs(&mut self) -> impl Iterator<Item = anyhow::Result<(OutputHeader, Output)>> + '_ {
38        let protocol_params = self.protocol_parameters().unwrap_or_default();
39
40        (0..self.header.output_count()).map(move |_| {
41            Ok((
42                OutputHeader::unpack::<_, true>(&mut self.reader, &())?,
43                Output::unpack::<_, true>(&mut self.reader, &protocol_params)?,
44            ))
45        })
46    }
47
48    /// Get the bytes of the solid entry points.
49    pub fn solid_entry_points_bytes(mut self) -> anyhow::Result<Vec<u8>> {
50        let mut remaining_bytes = vec![];
51        // Workaround as .read_to_end() is not available
52        let mut next_byte = vec![0u8; 1];
53        while self.reader.unpack_bytes(&mut next_byte).is_ok() {
54            remaining_bytes.push(next_byte[0]);
55        }
56
57        let sep_bytes = remaining_bytes
58            .get(remaining_bytes.len() - self.header.sep_count() as usize * 32..)
59            .expect("missing SEP bytes")
60            .to_vec();
61
62        Ok(sep_bytes)
63    }
64
65    /// Provide the target milestone timestamp extracted from the snapshot
66    /// header.
67    pub fn target_milestone_timestamp(&self) -> u32 {
68        self.header.target_milestone_timestamp()
69    }
70
71    /// Provide the network main token total supply through the snapshot
72    /// protocol parameters.
73    pub fn total_supply(&self) -> Result<u64> {
74        Ok(self.protocol_parameters()?.token_supply())
75    }
76
77    /// Get the protocol parameters.
78    pub fn protocol_parameters(&self) -> Result<ProtocolParameters> {
79        if let MilestoneOption::Parameters(params) = self.header.parameters_milestone_option() {
80            Ok(
81                <ProtocolParameters as packable::PackableExt>::unpack_unverified(
82                    params.binary_parameters(),
83                )
84                .expect("invalid protocol params"),
85            )
86        } else {
87            Err(StardustError::HornetSnapshotParametersNotFound.into())
88        }
89    }
90}