iota_types/
mock_checkpoint_builder.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::mem;
6
7use fastcrypto::traits::Signer;
8
9use crate::{
10    base_types::{AuthorityName, VerifiedExecutionData},
11    committee::Committee,
12    crypto::{AuthoritySignInfo, AuthoritySignature, IotaAuthoritySignature},
13    effects::{TransactionEffects, TransactionEffectsAPI},
14    gas::GasCostSummary,
15    messages_checkpoint::{
16        CertifiedCheckpointSummary, CheckpointContents, CheckpointSummary,
17        CheckpointVersionSpecificData, EndOfEpochData, FullCheckpointContents, VerifiedCheckpoint,
18        VerifiedCheckpointContents,
19    },
20    transaction::VerifiedTransaction,
21};
22
23pub trait ValidatorKeypairProvider {
24    fn get_validator_key(&self, name: &AuthorityName) -> &dyn Signer<AuthoritySignature>;
25    fn get_committee(&self) -> &Committee;
26}
27
28/// A utility to build consecutive checkpoints by adding transactions to the
29/// checkpoint builder. It's mostly used by simulations, tests and benchmarks.
30#[derive(Debug)]
31pub struct MockCheckpointBuilder {
32    previous_checkpoint: VerifiedCheckpoint,
33    transactions: Vec<VerifiedExecutionData>,
34    epoch_rolling_gas_cost_summary: GasCostSummary,
35    epoch: u64,
36}
37
38impl MockCheckpointBuilder {
39    pub fn new(previous_checkpoint: VerifiedCheckpoint) -> Self {
40        let epoch_rolling_gas_cost_summary =
41            previous_checkpoint.epoch_rolling_gas_cost_summary.clone();
42        let epoch = previous_checkpoint.epoch;
43
44        Self {
45            previous_checkpoint,
46            transactions: Vec::new(),
47            epoch_rolling_gas_cost_summary,
48            epoch,
49        }
50    }
51
52    pub fn size(&self) -> usize {
53        self.transactions.len()
54    }
55
56    pub fn epoch_rolling_gas_cost_summary(&self) -> &GasCostSummary {
57        &self.epoch_rolling_gas_cost_summary
58    }
59
60    pub fn push_transaction(
61        &mut self,
62        transaction: VerifiedTransaction,
63        effects: TransactionEffects,
64    ) {
65        self.epoch_rolling_gas_cost_summary += effects.gas_cost_summary();
66
67        self.transactions
68            .push(VerifiedExecutionData::new(transaction, effects))
69    }
70
71    /// Builds a checkpoint using internally buffered transactions.
72    pub fn build(
73        &mut self,
74        validator_keys: &impl ValidatorKeypairProvider,
75        timestamp_ms: u64,
76    ) -> (
77        VerifiedCheckpoint,
78        CheckpointContents,
79        VerifiedCheckpointContents,
80    ) {
81        self.build_internal(validator_keys, timestamp_ms, None)
82    }
83
84    pub fn build_end_of_epoch(
85        &mut self,
86        validator_keys: &impl ValidatorKeypairProvider,
87        timestamp_ms: u64,
88        new_epoch: u64,
89        end_of_epoch_data: EndOfEpochData,
90    ) -> (
91        VerifiedCheckpoint,
92        CheckpointContents,
93        VerifiedCheckpointContents,
94    ) {
95        self.build_internal(
96            validator_keys,
97            timestamp_ms,
98            Some((new_epoch, end_of_epoch_data)),
99        )
100    }
101
102    fn build_internal(
103        &mut self,
104        validator_keys: &impl ValidatorKeypairProvider,
105        timestamp_ms: u64,
106        new_epoch_data: Option<(u64, EndOfEpochData)>,
107    ) -> (
108        VerifiedCheckpoint,
109        CheckpointContents,
110        VerifiedCheckpointContents,
111    ) {
112        let contents =
113            CheckpointContents::new_with_causally_ordered_execution_data(self.transactions.iter());
114        let full_contents = VerifiedCheckpointContents::new_unchecked(
115            FullCheckpointContents::new_with_causally_ordered_transactions(
116                mem::take(&mut self.transactions)
117                    .into_iter()
118                    .map(|e| e.into_inner()),
119            ),
120        );
121
122        let (epoch, epoch_rolling_gas_cost_summary, end_of_epoch_data) =
123            if let Some((next_epoch, end_of_epoch_data)) = new_epoch_data {
124                let epoch = std::mem::replace(&mut self.epoch, next_epoch);
125                assert_eq!(next_epoch, epoch + 1);
126                let epoch_rolling_gas_cost_summary =
127                    std::mem::take(&mut self.epoch_rolling_gas_cost_summary);
128
129                (
130                    epoch,
131                    epoch_rolling_gas_cost_summary,
132                    Some(end_of_epoch_data),
133                )
134            } else {
135                (
136                    self.epoch,
137                    self.epoch_rolling_gas_cost_summary.clone(),
138                    None,
139                )
140            };
141
142        let summary = CheckpointSummary {
143            epoch,
144            sequence_number: self
145                .previous_checkpoint
146                .sequence_number
147                .checked_add(1)
148                .unwrap(),
149            network_total_transactions: self.previous_checkpoint.network_total_transactions
150                + contents.size() as u64,
151            content_digest: *contents.digest(),
152            previous_digest: Some(*self.previous_checkpoint.digest()),
153            epoch_rolling_gas_cost_summary,
154            end_of_epoch_data,
155            timestamp_ms,
156            version_specific_data: bcs::to_bytes(&CheckpointVersionSpecificData::empty_for_tests())
157                .unwrap(),
158            checkpoint_commitments: Default::default(),
159        };
160
161        let checkpoint = Self::create_certified_checkpoint(validator_keys, summary);
162        self.previous_checkpoint = checkpoint.clone();
163        (checkpoint, contents, full_contents)
164    }
165
166    fn create_certified_checkpoint(
167        validator_keys: &impl ValidatorKeypairProvider,
168        checkpoint: CheckpointSummary,
169    ) -> VerifiedCheckpoint {
170        let signatures = validator_keys
171            .get_committee()
172            .voting_rights
173            .iter()
174            .map(|(name, _)| {
175                let intent_msg = shared_crypto::intent::IntentMessage::new(
176                    shared_crypto::intent::Intent::iota_app(
177                        shared_crypto::intent::IntentScope::CheckpointSummary,
178                    ),
179                    &checkpoint,
180                );
181                let key = validator_keys.get_validator_key(name);
182                let signature = AuthoritySignature::new_secure(&intent_msg, &checkpoint.epoch, key);
183                AuthoritySignInfo {
184                    epoch: checkpoint.epoch,
185                    authority: *name,
186                    signature,
187                }
188            })
189            .collect();
190
191        let checkpoint_cert =
192            CertifiedCheckpointSummary::new(checkpoint, signatures, validator_keys.get_committee())
193                .unwrap();
194        VerifiedCheckpoint::new_unchecked(checkpoint_cert)
195    }
196}
iota_types/mock_checkpoint_builder.rs

iota_types/
mock_checkpoint_builder.rs
