Skip to main content

chia_sdk_test/full_node_simulator/
chain.rs

1use std::collections::VecDeque;
2
3use chia_consensus::validation_error::ErrorCode;
4use chia_protocol::{BlockRecord, Bytes32, SpendBundle};
5use indexmap::{IndexMap, IndexSet};
6use rand::Rng;
7
8use crate::{FullNodeSimulatorEvent, SimulatorError};
9
10use super::{
11    BLOCK_REWARD_AMOUNT, FullNodeSimulator, SimBlock, SimCoinRecord, ValidatedSpend,
12    state::{BlockDelta, CoinChange, HintChange, SpendChange},
13};
14
15impl FullNodeSimulator {
16    pub(super) fn create_block_from_mempool(&mut self) -> BlockRecord {
17        let previous_header_hash = self.header_hash();
18        let height = self.state.height + 1;
19        let timestamp = self.state.next_timestamp;
20        let header_hash = self.random_hash();
21
22        let mut included_tx_ids = Vec::new();
23        let mut included = Vec::new();
24        let mut included_spends_by_coin = IndexMap::<Bytes32, ValidatedSpend>::new();
25        for (tx_id, item) in self.mempool.clone() {
26            let Ok(validated) = self.validate_bundle(item.spend_bundle.clone()) else {
27                continue;
28            };
29            let has_conflict = validated.removals.iter().any(|coin_id| {
30                let Some(existing_spend) = included_spends_by_coin.get(coin_id) else {
31                    return false;
32                };
33                let Some(new_spend) = validated.spends.get(coin_id) else {
34                    return true;
35                };
36                !Self::spends_are_dedup_compatible(existing_spend, new_spend)
37            });
38            if has_conflict {
39                continue;
40            }
41            for coin_id in &validated.removals {
42                let Some(spend) = validated.spends.get(coin_id) else {
43                    continue;
44                };
45                included_spends_by_coin
46                    .entry(*coin_id)
47                    .or_insert_with(|| spend.clone());
48            }
49            included_tx_ids.push(tx_id);
50            included.push(validated);
51        }
52
53        let mut additions = Vec::new();
54        let mut removals = Vec::new();
55        let mut spends = Vec::new();
56        let mut transactions = Vec::new();
57        let mut fees = 0_u64;
58        let mut applied_removals = IndexSet::new();
59        let mut applied_additions = IndexSet::new();
60        let mut applied_spends = IndexSet::new();
61        let mut changed_coins = IndexSet::new();
62        let mut changed_spends = IndexSet::new();
63        let mut changed_hints = IndexSet::new();
64        let mut staged_coins = self.state.coins.clone();
65        let mut staged_coin_spends = self.state.coin_spends.clone();
66        let mut staged_coin_hints = self.state.coin_hints.clone();
67        let reward_coin = Self::reward_coin(
68            header_hash,
69            height,
70            0,
71            self.farming_puzzle_hash,
72            BLOCK_REWARD_AMOUNT,
73        );
74        let reward_coin_id = reward_coin.coin_id();
75        staged_coins.insert(
76            reward_coin_id,
77            SimCoinRecord {
78                coin: reward_coin,
79                coinbase: true,
80                confirmed_block_index: height,
81                spent_block_index: None,
82                timestamp,
83            },
84        );
85        changed_coins.insert(reward_coin_id);
86        additions.push(reward_coin_id);
87
88        for item in included {
89            fees = fees.saturating_add(item.fee);
90            transactions.push(item.spend_bundle);
91
92            for (coin, hint) in item.additions {
93                let coin_id = coin.coin_id();
94                if !applied_additions.insert(coin_id) {
95                    continue;
96                }
97                staged_coins.insert(
98                    coin_id,
99                    SimCoinRecord {
100                        coin,
101                        coinbase: false,
102                        confirmed_block_index: height,
103                        spent_block_index: None,
104                        timestamp,
105                    },
106                );
107                changed_coins.insert(coin_id);
108                if let Some(hint) = hint {
109                    staged_coin_hints.insert(coin_id, hint);
110                    changed_hints.insert(coin_id);
111                }
112                additions.push(coin_id);
113            }
114
115            for coin_id in item.removals {
116                if !applied_removals.insert(coin_id) {
117                    continue;
118                }
119                if let Some(record) = staged_coins.get_mut(&coin_id) {
120                    record.spent_block_index = Some(height);
121                    changed_coins.insert(coin_id);
122                    removals.push(coin_id);
123                }
124            }
125
126            for (coin_id, spend) in item.spends {
127                if applied_spends.insert(coin_id) {
128                    staged_coin_spends.insert(coin_id, spend.coin_spend.clone());
129                    changed_spends.insert(coin_id);
130                    spends.push(spend.coin_spend);
131                }
132            }
133        }
134
135        let delta = BlockDelta {
136            coins: changed_coins
137                .into_iter()
138                .map(|coin_id| CoinChange {
139                    coin_id,
140                    before: self.state.coins.get(&coin_id).copied(),
141                    after: staged_coins.get(&coin_id).copied(),
142                })
143                .collect(),
144            spends: changed_spends
145                .into_iter()
146                .map(|coin_id| SpendChange {
147                    coin_id,
148                    before: self.state.coin_spends.get(&coin_id).cloned(),
149                    after: staged_coin_spends.get(&coin_id).cloned(),
150                })
151                .collect(),
152            hints: changed_hints
153                .into_iter()
154                .map(|coin_id| HintChange {
155                    coin_id,
156                    before: self.state.coin_hints.get(&coin_id).copied(),
157                    after: staged_coin_hints.get(&coin_id).copied(),
158                })
159                .collect(),
160        };
161
162        let record = Self::make_block_record(
163            header_hash,
164            previous_header_hash,
165            height,
166            timestamp,
167            self.header_hash_of(height.saturating_sub(1))
168                .unwrap_or_default(),
169            fees,
170            height.saturating_sub(1),
171            self.farming_puzzle_hash,
172            vec![reward_coin],
173        );
174        let block = SimBlock {
175            record: record.clone(),
176            additions: additions.clone(),
177            removals: removals.clone(),
178            spends,
179            transactions,
180            delta,
181        };
182
183        self.state
184            .apply_block(block)
185            .expect("locally built block delta must apply atomically");
186        for tx_id in included_tx_ids {
187            self.mempool.swap_remove(&tx_id);
188        }
189
190        self.events.push(FullNodeSimulatorEvent::Block {
191            height,
192            header_hash,
193            previous_header_hash,
194            additions: self.records_for_ids(&additions),
195            removals: self.records_for_ids(&removals),
196        });
197
198        record
199    }
200
201    pub(super) fn revert_canonical_blocks(&mut self, blocks: u32) -> Vec<SimBlock> {
202        let mut reverted = Vec::new();
203        for _ in 0..blocks {
204            let Some(block) = self
205                .state
206                .revert_tip()
207                .expect("canonical tip delta must revert exactly")
208            else {
209                break;
210            };
211            reverted.push(block);
212        }
213        reverted.reverse();
214        reverted
215    }
216
217    pub(super) fn requeue_transactions(
218        &mut self,
219        transactions: impl IntoIterator<Item = SpendBundle>,
220    ) {
221        let mut pending = transactions.into_iter().collect::<VecDeque<_>>();
222        while !pending.is_empty() {
223            let mut deferred = VecDeque::new();
224            let mut made_progress = false;
225            while let Some(spend_bundle) = pending.pop_front() {
226                match self.normalize_and_insert(spend_bundle.clone()) {
227                    Ok(()) => made_progress = true,
228                    Err(SimulatorError::Validation(ErrorCode::UnknownUnspent)) => {
229                        deferred.push_back(spend_bundle);
230                    }
231                    Err(_) => {}
232                }
233            }
234            if !made_progress {
235                break;
236            }
237            pending = deferred;
238        }
239    }
240
241    pub(super) fn prune_mempool(&mut self) {
242        let spend_bundles = self
243            .mempool
244            .values()
245            .map(|item| item.spend_bundle.clone())
246            .collect::<Vec<_>>();
247        self.mempool.clear();
248        self.requeue_transactions(spend_bundles);
249    }
250
251    fn random_hash(&mut self) -> Bytes32 {
252        let mut bytes = [0; 32];
253        self.rng.fill(&mut bytes);
254        bytes.into()
255    }
256
257    pub fn drain_events(&mut self) -> Vec<FullNodeSimulatorEvent> {
258        std::mem::take(&mut self.events)
259    }
260
261    pub fn set_farming_ph(&mut self, puzzle_hash: Bytes32) {
262        self.farming_puzzle_hash = puzzle_hash;
263    }
264
265    pub fn farm_block(&mut self, blocks: u32) -> Vec<BlockRecord> {
266        let count = blocks.max(1);
267        let mut records = Vec::new();
268        for _ in 0..count {
269            records.push(self.create_block_from_mempool());
270        }
271        records
272    }
273
274    pub fn revert_blocks(&mut self, blocks: u32) -> Vec<Bytes32> {
275        let reverted = self.revert_canonical_blocks(blocks);
276        self.requeue_transactions(reverted.iter().flat_map(|block| block.transactions.clone()));
277        reverted
278            .iter()
279            .map(|block| block.record.header_hash)
280            .collect()
281    }
282
283    pub fn reorg_blocks(
284        &mut self,
285        num_of_blocks_to_rev: u32,
286        num_of_new_blocks: u32,
287    ) -> Vec<BlockRecord> {
288        let old_peak_hash = self.header_hash();
289        let reverted = self.revert_canonical_blocks(num_of_blocks_to_rev);
290        let fork_height = self.state.height;
291        let reverted_header_hashes = reverted
292            .iter()
293            .map(|block| block.record.header_hash)
294            .collect::<Vec<_>>();
295        let replacement_block_count =
296            num_of_new_blocks.saturating_add(u32::try_from(reverted.len()).unwrap());
297
298        for block in reverted {
299            self.orphaned_blocks.insert(block.record.header_hash, block);
300        }
301
302        let mut records = Vec::new();
303        let mut new_header_hashes = Vec::new();
304        for _ in 0..replacement_block_count {
305            let record = self.create_block_from_mempool();
306            new_header_hashes.push(record.header_hash);
307            records.push(record);
308        }
309
310        self.prune_mempool();
311        self.events.push(FullNodeSimulatorEvent::Reorg {
312            fork_height,
313            old_peak_hash,
314            new_peak_hash: self.header_hash(),
315            reverted_header_hashes,
316            new_header_hashes,
317        });
318
319        records
320    }
321}