/ ledger / tests / helpers / mod.rs
mod.rs
  1  // Copyright (c) 2019-2025 Alpha-Delta Network Inc.
  2  // This file is part of the alphavm library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  use alphastd::StorageMode;
 17  use alphavm_console::{
 18      account::{Address, PrivateKey},
 19      network::MainnetV0,
 20      prelude::*,
 21  };
 22  use alphavm_ledger::{Block, Ledger};
 23  use alphavm_ledger_narwhal::{BatchCertificate, BatchHeader, Subdag};
 24  use alphavm_ledger_store::ConsensusStore;
 25  use alphavm_synthesizer::vm::VM;
 26  
 27  use indexmap::{IndexMap, IndexSet};
 28  use std::collections::{BTreeMap, HashMap};
 29  use time::OffsetDateTime;
 30  
 31  pub type CurrentNetwork = MainnetV0;
 32  
 33  #[cfg(not(feature = "rocks"))]
 34  pub type LedgerType<N> = alphavm_ledger_store::helpers::memory::ConsensusMemory<N>;
 35  #[cfg(feature = "rocks")]
 36  pub type LedgerType<N> = alphavm_ledger_store::helpers::rocksdb::ConsensusDB<N>;
 37  
 38  /// Helper to build chains with custom structures for testing.
 39  pub struct TestChainBuilder {
 40      /// The keys of all validators.
 41      private_keys: Vec<PrivateKey<CurrentNetwork>>,
 42      /// The underlying ledger.
 43      ledger: Ledger<CurrentNetwork, LedgerType<CurrentNetwork>>,
 44      /// The round containing the leader certificate for the most recent block we generated.
 45      last_block_round: u64,
 46      /// The batch certificates of the last round we generated.
 47      round_to_certificates: HashMap<u64, IndexMap<usize, BatchCertificate<CurrentNetwork>>>,
 48      /// The batch certificate of the last leader (if any).
 49      previous_leader_certificate: Option<BatchCertificate<CurrentNetwork>>,
 50      /// The last round for each committee member where they created a batch.
 51      /// Invariant: for any validator i, last_batch[i] <= last_committed_batch[i]
 52      last_batch_round: HashMap<usize, u64>,
 53      /// The last batch of a validator that was included in a block.
 54      last_committed_batch_round: HashMap<usize, u64>,
 55      /// The start of the test chain.
 56      genesis_block: Block<CurrentNetwork>,
 57  }
 58  
 59  /// Additional options you can pass to the builder when generating blocks.
 60  #[derive(Default)]
 61  pub struct BlockOptions {
 62      /// Do not include votes to the previous leader certificate
 63      pub skip_votes: bool,
 64      /// Do not generate certificates for the specific node indices (to simulate a partition).
 65      pub skip_nodes: Vec<usize>,
 66  }
 67  
 68  impl TestChainBuilder {
 69      pub fn new(committee_size: usize, rng: &mut TestRng) -> Self {
 70          // Sample the genesis private key.
 71          let private_key = PrivateKey::<CurrentNetwork>::new(rng).unwrap();
 72          // Initialize the store.
 73          let store = ConsensusStore::<_, LedgerType<_>>::open(StorageMode::new_test(None)).unwrap();
 74          // Create a genesis block with a seeded RNG to reproduce the same genesis private keys.
 75          let seed: u64 = rng.r#gen();
 76          let genesis_rng = &mut TestRng::from_seed(seed);
 77          let genesis_block = VM::from(store).unwrap().genesis_beacon(&private_key, genesis_rng).unwrap();
 78  
 79          // Extract the private keys from the genesis committee by using the same RNG to sample private keys.
 80          let genesis_rng = &mut TestRng::from_seed(seed);
 81          let private_keys = (0..committee_size).map(|_| PrivateKey::new(genesis_rng).unwrap()).collect();
 82  
 83          Self::from_genesis(private_keys, genesis_block)
 84      }
 85  
 86      /// Initialize the builder with the specified committee and gensis block
 87      pub fn from_genesis(private_keys: Vec<PrivateKey<CurrentNetwork>>, genesis_block: Block<CurrentNetwork>) -> Self {
 88          // Initialize the ledger with the genesis block.
 89          let ledger = Ledger::<CurrentNetwork, LedgerType<CurrentNetwork>>::load(
 90              genesis_block.clone(),
 91              StorageMode::new_test(None),
 92          )
 93          .unwrap();
 94  
 95          Self {
 96              private_keys,
 97              ledger,
 98  
 99              genesis_block,
100              last_batch_round: Default::default(),
101              last_committed_batch_round: Default::default(),
102              last_block_round: 0,
103              round_to_certificates: Default::default(),
104              previous_leader_certificate: Default::default(),
105          }
106      }
107  
108      /// Create multiple blocks, with fully-connected DAGs.
109      #[allow(dead_code)]
110      pub fn generate_blocks(&mut self, num_blocks: usize, rng: &mut TestRng) -> Vec<Block<CurrentNetwork>> {
111          self.generate_blocks_with_opts(num_blocks, &BlockOptions::default(), rng)
112      }
113  
114      /// Create multiple blocks, with additional parameters.
115      pub fn generate_blocks_with_opts(
116          &mut self,
117          num_blocks: usize,
118          options: &BlockOptions,
119          rng: &mut TestRng,
120      ) -> Vec<Block<CurrentNetwork>> {
121          assert!(num_blocks > 0, "Need to build at least one block");
122  
123          (0..num_blocks).map(|_| self.generate_block_with_opts(options, rng)).collect()
124      }
125  
126      /// Create a new block, with a fully-connected DAG.
127      ///
128      /// This will "fill in" any gaps left in earlier rounds from non participating nodes.
129      pub fn generate_block(&mut self, rng: &mut TestRng) -> Block<CurrentNetwork> {
130          self.generate_block_with_opts(&BlockOptions::default(), rng)
131      }
132  
133      /// Same as `generate_block` but with additional options/parameters.
134      pub fn generate_block_with_opts(&mut self, options: &BlockOptions, rng: &mut TestRng) -> Block<CurrentNetwork> {
135          assert!(
136              options.skip_nodes.len() * 3 < self.private_keys.len(),
137              "Cannot mark more than f nodes as unavailable/skipped"
138          );
139  
140          let next_block_round = self.last_block_round + 2;
141  
142          // SubDAGs can be at most GC rounds long.
143          let mut round = if next_block_round < BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS as u64 {
144              // Batches from genesis round cannot be included in any block that isn't genesis
145              1
146          } else {
147              next_block_round - BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS as u64
148          };
149  
150          // =======================================
151          // Create certificates for the new block.
152          // =======================================
153          loop {
154              let mut created_anchor = false;
155  
156              let previous_certificate_ids = if round == 1 {
157                  IndexSet::default()
158              } else {
159                  self.round_to_certificates
160                      .get(&(round - 1))
161                      .unwrap()
162                      .iter()
163                      .filter_map(|(_, cert)| {
164                          // If votes are skipped, remove previous leader cert from the set.
165                          let skip = if let Some(leader) = &self.previous_leader_certificate {
166                              options.skip_votes && leader.id() == cert.id()
167                          } else {
168                              false
169                          };
170  
171                          if skip { None } else { Some(cert.id()) }
172                      })
173                      .collect()
174              };
175  
176              let committee = self.ledger.get_committee_lookback_for_round(round).unwrap().unwrap_or_else(|| {
177                  panic!("No committee for round {round}");
178              });
179  
180              for (key1_idx, private_key_1) in self.private_keys.iter().enumerate() {
181                  if options.skip_nodes.contains(&key1_idx) {
182                      continue;
183                  }
184                  // Don't recreate batches that already exist.
185                  if self.last_batch_round.get(&key1_idx).unwrap_or(&0) >= &round {
186                      continue;
187                  }
188  
189                  let batch_header = BatchHeader::new(
190                      private_key_1,
191                      round,
192                      OffsetDateTime::now_utc().unix_timestamp(),
193                      committee.id(),
194                      Default::default(),
195                      previous_certificate_ids.clone(),
196                      rng,
197                  )
198                  .unwrap();
199  
200                  // Add signatures for the batch header.
201                  let signatures = self
202                      .private_keys
203                      .iter()
204                      .enumerate()
205                      .filter(|&(key2_idx, _)| key1_idx != key2_idx)
206                      .map(|(_, private_key_2)| private_key_2.sign(&[batch_header.batch_id()], rng).unwrap())
207                      .collect();
208  
209                  // Update the round at which this validator last created a batch.
210                  self.last_batch_round.insert(key1_idx, round);
211  
212                  // Insert certificate into the round_to_certificates mapping.
213                  self.round_to_certificates
214                      .entry(round)
215                      .or_default()
216                      .insert(key1_idx, BatchCertificate::from(batch_header, signatures).unwrap());
217  
218                  // Check if this batch was an anchor.
219                  if round % 2 == 0 {
220                      let leader = committee.get_leader(round).unwrap();
221                      if leader == Address::try_from(private_key_1).unwrap() {
222                          created_anchor = true;
223                      }
224                  }
225              }
226  
227              // Anchor was confirmed by more than a third of the validators.
228              if created_anchor && round % 2 == 0 && self.last_block_round < round {
229                  self.last_block_round = round;
230                  break;
231              }
232  
233              round += 1;
234          }
235  
236          // ==============================================================
237          // Build a subdag from the new certificates and create the block.
238          // ==============================================================
239          let commit_round = round;
240  
241          let leader_committee = self.ledger.get_committee_lookback_for_round(round).unwrap().unwrap();
242          let leader = leader_committee.get_leader(commit_round).unwrap();
243          let (leader_idx, leader_certificate) =
244              self.round_to_certificates.get(&commit_round).unwrap().iter().find(|(_, c)| c.author() == leader).unwrap();
245          let leader_idx = *leader_idx;
246          let leader_certificate = leader_certificate.clone();
247  
248          // Construct the subdag for the new block.
249          let mut subdag_map = BTreeMap::new();
250  
251          // Figure out what the earliest round for the subDAG could be.
252          let start_round = if commit_round < BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS as u64 {
253              1
254          } else {
255              commit_round - BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS as u64 + 2
256          };
257  
258          for round in start_round..commit_round {
259              let mut to_insert = IndexSet::new();
260              for idx in 0..self.private_keys.len() {
261                  // Some of the batches we in previous rounds might not be new,
262                  // and already included in a previous block.
263                  let cround = self.last_committed_batch_round.entry(idx).or_default();
264                  // Batch already included in another block
265                  if *cround >= round {
266                      continue;
267                  }
268  
269                  if let Some(cert) = self.round_to_certificates.entry(round).or_default().get(&idx) {
270                      to_insert.insert(cert.clone());
271                      *cround = round;
272                  }
273              }
274              if !to_insert.is_empty() {
275                  subdag_map.insert(round, to_insert);
276              }
277          }
278  
279          // Add the leader certificate.
280          // (special case, because it is the only cert included from the commit round)
281          subdag_map.insert(commit_round, [leader_certificate.clone()].into());
282          self.last_committed_batch_round.insert(leader_idx, commit_round);
283  
284          // Construct the block.
285          let subdag = Subdag::from(subdag_map).unwrap();
286          let block = self.ledger.prepare_advance_to_next_quorum_block(subdag, Default::default(), rng).unwrap();
287          self.ledger.check_next_block(&block, rng).unwrap();
288  
289          // Update th ledger state.
290          self.ledger.advance_to_next_block(&block).unwrap();
291          self.previous_leader_certificate = Some(leader_certificate.clone());
292  
293          block
294      }
295  
296      /// Return the genesis block associated with the test chain
297      pub fn genesis_block(&self) -> &Block<CurrentNetwork> {
298          &self.genesis_block
299      }
300  }