1  // Copyright (c) 2025 ADnet Contributors
  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 crate::{
 17      Block,
 18      Ledger,
 19      Transaction,
 20      Transmission,
 21      TransmissionID,
 22      narwhal::{BatchCertificate, BatchHeader, Subdag},
 23      puzzle::Solution,
 24      store::{ConsensusStore, helpers::memory::ConsensusMemory},
 25  };
 26  use alphavm_console::{
 27      account::{Address, PrivateKey},
 28      network::MainnetV0,
 29      prelude::*,
 30  };
 31  use alphavm_synthesizer::vm::VM;
 32  
 33  use alpha_std::StorageMode;
 34  
 35  use anyhow::{Context, Result};
 36  use indexmap::{IndexMap, IndexSet};
 37  use rand::SeedableRng;
 38  use rand_chacha::ChaChaRng;
 39  use std::collections::{BTreeMap, HashMap};
 40  use time::OffsetDateTime;
 41  
 42  pub type CurrentNetwork = MainnetV0;
 43  
 44  #[cfg(not(feature = "rocks"))]
 45  pub type LedgerType<N> = alphavm_ledger_store::helpers::memory::ConsensusMemory<N>;
 46  #[cfg(feature = "rocks")]
 47  pub type LedgerType<N> = alphavm_ledger_store::helpers::rocksdb::ConsensusDB<N>;
 48  
 49  /// Helper to build chains with custom structures for testing.
 50  pub struct TestChainBuilder<N: Network> {
 51      /// The keys of all validators.
 52      private_keys: Vec<PrivateKey<N>>,
 53      /// The underlying ledger.
 54      ledger: Ledger<N, ConsensusMemory<N>>,
 55      /// The round containing the leader certificate for the most recent block we generated.
 56      last_block_round: u64,
 57      /// The batch certificates of the last round we generated.
 58      round_to_certificates: HashMap<u64, IndexMap<usize, BatchCertificate<N>>>,
 59      /// The batch certificate of the last leader (if any).
 60      previous_leader_certificate: Option<BatchCertificate<N>>,
 61      /// The last round for each committee member where they created a batch.
 62      /// Invariant: for any validator i, last_batch[i] <= last_committed_batch[i]
 63      last_batch_round: HashMap<usize, u64>,
 64      /// The last batch of a validator that was included in a block.
 65      last_committed_batch_round: HashMap<usize, u64>,
 66      /// The start of the test chain.
 67      genesis_block: Block<N>,
 68  }
 69  
 70  /// Additional options you can pass to the builder when generating a set of blocks.
 71  #[derive(Clone)]
 72  pub struct GenerateBlocksOptions<N: Network> {
 73      /// Do not include votes to the previous leader certificate
 74      pub skip_votes: bool,
 75      /// Do not generate certificates for the specific node indices (to simulate a partition).
 76      pub skip_nodes: Vec<usize>,
 77      /// A flag indicating that a number of initial "placeholder blocks" should be baked
 78      /// wthout transactions in order to skip to the latest version of consensus.
 79      pub skip_to_current_version: bool,
 80      /// The number of validators.
 81      pub num_validators: usize,
 82      /// Preloaded transactions to populate the blocks with.
 83      pub transactions: Vec<Transaction<N>>,
 84  }
 85  
 86  impl<N: Network> Default for GenerateBlocksOptions<N> {
 87      fn default() -> Self {
 88          Self {
 89              skip_votes: false,
 90              skip_nodes: Default::default(),
 91              skip_to_current_version: false,
 92              num_validators: 0,
 93              transactions: Default::default(),
 94          }
 95      }
 96  }
 97  
 98  /// Additional options you can pass to the builder when generating a single block.
 99  /// Note: As of now, all certificates for this block will have the given timestamp and contain listed transmissions.
100  #[derive(Clone)]
101  pub struct GenerateBlockOptions<N: Network> {
102      /// Do not include votes to the previous leader certificate
103      pub skip_votes: bool,
104      /// Do not generate certificates for the specific node indices (to simulate a partition).
105      pub skip_nodes: Vec<usize>,
106      /// The timestamp for this block.
107      pub timestamp: i64,
108      /// The transmissions to be included in the block.
109      pub solutions: Vec<Solution<N>>,
110      pub transactions: Vec<Transaction<N>>,
111  }
112  
113  impl<N: Network> Default for GenerateBlockOptions<N> {
114      fn default() -> Self {
115          Self {
116              skip_votes: false,
117              skip_nodes: Default::default(),
118              transactions: Default::default(),
119              solutions: Default::default(),
120              timestamp: OffsetDateTime::now_utc().unix_timestamp(),
121          }
122      }
123  }
124  
125  impl<N: Network> TestChainBuilder<N> {
126      /// Generate a new committee and genesis block.
127      pub fn initialize_components(committee_size: usize, rng: &mut TestRng) -> Result<(Vec<PrivateKey<N>>, Block<N>)> {
128          // Sample the genesis private key.
129          let private_key = PrivateKey::<N>::new(rng)?;
130  
131          // Initialize the store.
132          let store = ConsensusStore::<_, ConsensusMemory<_>>::open(StorageMode::new_test(None))
133              .with_context(|| "Failed to initialize consensus store")?;
134          // Create a genesis block with a seeded RNG to reproduce the same genesis private keys.
135          let seed: u64 = rng.r#gen();
136          trace!("Using seed {seed} and key {} for genesis RNG", private_key);
137          let genesis_rng = &mut TestRng::from_seed(seed);
138          let genesis_block = VM::from(store).unwrap().genesis_beacon(&private_key, genesis_rng)?;
139  
140          // Extract the private keys from the genesis committee by using the same RNG to sample private keys.
141          let genesis_rng = &mut TestRng::from_seed(seed);
142          let private_keys = (0..committee_size).map(|_| PrivateKey::new(genesis_rng).unwrap()).collect();
143  
144          trace!(
145              "Generated genesis block ({}) and private keys for all {committee_size} committee members",
146              genesis_block.hash()
147          );
148  
149          Ok((private_keys, genesis_block))
150      }
151  
152      /// Initialize the builder with the default quorum size.
153      pub fn new(rng: &mut TestRng) -> Result<Self> {
154          Self::new_with_quorum_size(4, rng)
155      }
156  
157      /// Initialize the builder with the specified quorum size.
158      pub fn new_with_quorum_size(num_validators: usize, rng: &mut TestRng) -> Result<Self> {
159          let (private_keys, genesis) = Self::initialize_components(num_validators, rng)?;
160          Self::from_components(private_keys, genesis)
161      }
162  
163      /// Initialize the builder with the specified genesis block..
164      /// Note: this function mirrors the way the private keys are sampled in snarkOS `fn parse_genesis`.
165      pub fn new_with_quorum_size_and_genesis_block(num_validators: usize, genesis_path: String) -> Result<Self> {
166          // Attempts to load the genesis block file.
167          let buffer = std::fs::read(genesis_path)?;
168          // Return the genesis block.
169          let genesis = Block::from_bytes_le(&buffer)?;
170          /// The development mode RNG seed.
171          pub const DEVELOPMENT_MODE_RNG_SEED: u64 = 1234567890u64;
172          // Initialize the (fixed) RNG.
173          let mut rng = ChaChaRng::seed_from_u64(DEVELOPMENT_MODE_RNG_SEED);
174          // Initialize the development private keys.
175          let private_keys = (0..num_validators).map(|_| PrivateKey::new(&mut rng).unwrap()).collect();
176          // Initialize the builder with the specified committee and genesis block.
177          Self::from_components(private_keys, genesis)
178      }
179  
180      /// Initialize the builder with the specified committee and genesis block
181      pub fn from_components(private_keys: Vec<PrivateKey<N>>, genesis_block: Block<N>) -> Result<Self> {
182          // Initialize the ledger with the genesis block.
183          let ledger = Ledger::<N, ConsensusMemory<N>>::load(genesis_block.clone(), StorageMode::new_test(None))
184              .with_context(|| "Failed to set up ledger for test chain")?;
185  
186          ensure!(ledger.genesis_block == genesis_block);
187  
188          Self::from_genesis(private_keys, genesis_block)
189      }
190  
191      /// Initialize the builder with the specified committee and gensis block
192      pub fn from_genesis(private_keys: Vec<PrivateKey<N>>, genesis_block: Block<N>) -> Result<Self> {
193          // Initialize the ledger with the genesis block.
194          let ledger = Ledger::<N, ConsensusMemory<N>>::load(genesis_block.clone(), StorageMode::new_test(None))
195              .with_context(|| "Failed to set up ledger for test chain")?;
196  
197          Ok(Self {
198              private_keys,
199              ledger,
200              genesis_block,
201              last_batch_round: Default::default(),
202              last_committed_batch_round: Default::default(),
203              last_block_round: 0,
204              round_to_certificates: Default::default(),
205              previous_leader_certificate: Default::default(),
206          })
207      }
208  
209      /// Create multiple blocks, with fully-connected DAGs.
210      pub fn generate_blocks(&mut self, num_blocks: usize, rng: &mut TestRng) -> Result<Vec<Block<N>>> {
211          let num_validators = self.private_keys.len();
212  
213          self.generate_blocks_with_opts(num_blocks, GenerateBlocksOptions { num_validators, ..Default::default() }, rng)
214      }
215  
216      /// Create multiple blocks, with additional parameters.
217      ///
218      /// # Panics
219      /// This function panics if called from an async context.
220      pub fn generate_blocks_with_opts(
221          &mut self,
222          num_blocks: usize,
223          mut options: GenerateBlocksOptions<N>,
224          rng: &mut TestRng,
225      ) -> Result<Vec<Block<N>>> {
226          assert!(num_blocks > 0, "Need to build at least one block");
227  
228          let mut result = vec![];
229  
230          // If configured, skip enough blocks to reach the current consensus version.
231          if options.skip_to_current_version {
232              let (version, target_height) = TEST_CONSENSUS_VERSION_HEIGHTS.last().unwrap();
233              let mut current_height = self.ledger.latest_height();
234  
235              let diff = target_height.saturating_sub(current_height);
236  
237              if diff > 0 {
238                  println!("Skipping {diff} blocks to reach {version}");
239  
240                  while current_height < *target_height && result.len() < num_blocks {
241                      let options = GenerateBlockOptions {
242                          skip_votes: options.skip_votes,
243                          skip_nodes: options.skip_nodes.clone(),
244                          ..Default::default()
245                      };
246  
247                      let block = self.generate_block_with_opts(options, rng)?;
248                      current_height = block.height();
249                      result.push(block);
250                  }
251  
252                  println!("Advanced to the current consensus version at height {target_height}");
253              } else {
254                  debug!("Already at the current consensus version. No blocks to skip.");
255              }
256          }
257  
258          while result.len() < num_blocks {
259              let num_txs = (BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH * options.num_validators)
260                  .min(options.transactions.len());
261  
262              let options = GenerateBlockOptions {
263                  skip_votes: options.skip_votes,
264                  skip_nodes: options.skip_nodes.clone(),
265                  transactions: options.transactions.drain(..num_txs).collect(),
266                  ..Default::default()
267              };
268  
269              let block = self.generate_block_with_opts(options, rng)?;
270              result.push(block);
271          }
272  
273          Ok(result)
274      }
275  
276      /// Create a new block, with a fully-connected DAG.
277      ///
278      /// This will "fill in" any gaps left in earlier rounds from non participating nodes.
279      pub fn generate_block(&mut self, rng: &mut TestRng) -> Result<Block<N>> {
280          self.generate_block_with_opts(GenerateBlockOptions::default(), rng)
281      }
282  
283      /// Same as `generate_block` but with additional options/parameters.
284      pub fn generate_block_with_opts(
285          &mut self,
286          options: GenerateBlockOptions<N>,
287          rng: &mut TestRng,
288      ) -> Result<Block<N>> {
289          assert!(
290              options.skip_nodes.len() * 3 < self.private_keys.len(),
291              "Cannot mark more than f nodes as unavailable/skipped"
292          );
293  
294          let next_block_round = self.last_block_round + 2;
295          let mut cert_count = 0;
296  
297          // SubDAGs can be at most GC rounds long.
298          // Batches from genesis round cannot be included in any block that isn't genesis
299          let mut round = next_block_round.checked_sub(BatchHeader::<N>::MAX_GC_ROUNDS as u64).unwrap_or(1).max(1);
300  
301          let mut transmissions = IndexMap::default();
302  
303          for txn in options.transactions {
304              let txn_id = txn.id();
305              let transmission = Transmission::from(txn);
306              let transmission_id = TransmissionID::Transaction(txn_id, transmission.to_checksum().unwrap().unwrap());
307  
308              transmissions.insert(transmission_id, transmission);
309          }
310  
311          for solution in options.solutions {
312              let transmission = Transmission::from(solution);
313              let transmission_id = TransmissionID::Solution(solution.id(), transmission.to_checksum().unwrap().unwrap());
314  
315              transmissions.insert(transmission_id, transmission);
316          }
317  
318          // =======================================
319          // Create certificates for the new block.
320          // =======================================
321          loop {
322              let mut created_anchor = false;
323  
324              let previous_certificate_ids = if round == 1 {
325                  IndexSet::default()
326              } else {
327                  self.round_to_certificates
328                      .get(&(round - 1))
329                      .unwrap()
330                      .iter()
331                      .filter_map(|(_, cert)| {
332                          // If votes are skipped, remove previous leader cert from the set.
333                          let skip = if let Some(leader) = &self.previous_leader_certificate {
334                              options.skip_votes && leader.id() == cert.id()
335                          } else {
336                              false
337                          };
338  
339                          if skip { None } else { Some(cert.id()) }
340                      })
341                      .collect()
342              };
343  
344              let committee = self.ledger.get_committee_lookback_for_round(round).unwrap().unwrap_or_else(|| {
345                  panic!("No committee for round {round}");
346              });
347  
348              for (key1_idx, private_key_1) in self.private_keys.iter().enumerate() {
349                  if options.skip_nodes.contains(&key1_idx) {
350                      continue;
351                  }
352                  // Don't recreate batches that already exist.
353                  if self.last_batch_round.get(&key1_idx).unwrap_or(&0) >= &round {
354                      continue;
355                  }
356  
357                  let transmission_ids: IndexSet<_> = transmissions
358                      .keys()
359                      .skip(key1_idx * BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH)
360                      .take(BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH)
361                      .copied()
362                      .collect();
363  
364                  let batch_header = BatchHeader::new(
365                      private_key_1,
366                      round,
367                      options.timestamp,
368                      committee.id(),
369                      transmission_ids.clone(),
370                      previous_certificate_ids.clone(),
371                      rng,
372                  )
373                  .unwrap();
374  
375                  // Add signatures for the batch header.
376                  let signatures = self
377                      .private_keys
378                      .iter()
379                      .enumerate()
380                      .filter(|&(key2_idx, _)| key1_idx != key2_idx)
381                      .map(|(_, private_key_2)| private_key_2.sign(&[batch_header.batch_id()], rng).unwrap())
382                      .collect();
383  
384                  // Update the round at which this validator last created a batch.
385                  self.last_batch_round.insert(key1_idx, round);
386  
387                  // Insert certificate into the round_to_certificates mapping.
388                  self.round_to_certificates
389                      .entry(round)
390                      .or_default()
391                      .insert(key1_idx, BatchCertificate::from(batch_header, signatures).unwrap());
392  
393                  cert_count += 1;
394  
395                  // Check if this batch was an anchor.
396                  if round % 2 == 0 {
397                      let leader = committee.get_leader(round).unwrap();
398                      if leader == Address::try_from(private_key_1).unwrap() {
399                          created_anchor = true;
400                      }
401                  }
402              }
403  
404              // Anchor was confirmed by more than a third of the validators.
405              if created_anchor && round % 2 == 0 && self.last_block_round < round {
406                  self.last_block_round = round;
407                  break;
408              }
409  
410              round += 1;
411          }
412  
413          // ==============================================================
414          // Build a subdag from the new certificates and create the block.
415          // ==============================================================
416          let commit_round = round;
417  
418          let leader_committee = self.ledger.get_committee_lookback_for_round(round).unwrap().unwrap();
419          let leader = leader_committee.get_leader(commit_round).unwrap();
420          let (leader_idx, leader_certificate) =
421              self.round_to_certificates.get(&commit_round).unwrap().iter().find(|(_, c)| c.author() == leader).unwrap();
422          let leader_idx = *leader_idx;
423          let leader_certificate = leader_certificate.clone();
424  
425          // Construct the subdag for the new block.
426          let mut subdag_map = BTreeMap::new();
427  
428          // Figure out what the earliest round for the subDAG could be.
429          let start_round = if commit_round < BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS as u64 {
430              1
431          } else {
432              commit_round - BatchHeader::<CurrentNetwork>::MAX_GC_ROUNDS as u64 + 2
433          };
434  
435          for round in start_round..commit_round {
436              let mut to_insert = IndexSet::new();
437              for idx in 0..self.private_keys.len() {
438                  // Some of the batches we in previous rounds might not be new,
439                  // and already included in a previous block.
440                  let cround = self.last_committed_batch_round.entry(idx).or_default();
441                  // Batch already included in another block
442                  if *cround >= round {
443                      continue;
444                  }
445  
446                  if let Some(cert) = self.round_to_certificates.entry(round).or_default().get(&idx) {
447                      to_insert.insert(cert.clone());
448                      *cround = round;
449                  }
450              }
451              if !to_insert.is_empty() {
452                  subdag_map.insert(round, to_insert);
453              }
454          }
455  
456          // Add the leader certificate.
457          // (special case, because it is the only cert included from the commit round)
458          subdag_map.insert(commit_round, [leader_certificate.clone()].into());
459          self.last_committed_batch_round.insert(leader_idx, commit_round);
460  
461          trace!("Generated {cert_count} certificates for the next block");
462  
463          // Construct the block.
464          let subdag = Subdag::from(subdag_map).unwrap();
465  
466          let block = self.ledger.prepare_advance_to_next_quorum_block(subdag, transmissions, rng)?;
467  
468          // Skip to increase performance.
469          //self.ledger.check_next_block(&block, rng).with_context(|| "Failed to (internally) check generated block")?;
470  
471          trace!("Generated new block {} at height {}", block.hash(), block.height());
472  
473          // Update the ledger state.
474          self.ledger
475              .advance_to_next_block(&block)
476              .with_context(|| "Failed to (internally) advance to generated block")?;
477          self.previous_leader_certificate = Some(leader_certificate.clone());
478  
479          trace!("Updated internal ledger to height {}", block.height());
480          Ok(block)
481      }
482  
483      /// Return the genesis block associated with the test chain
484      pub fn genesis_block(&self) -> &Block<N> {
485          &self.genesis_block
486      }
487  
488      /// Returns the private keys of the genesis committee of this test chain
489      pub fn private_keys(&self) -> &[PrivateKey<N>] {
490          &self.private_keys
491      }
492  
493      /// Returns the private keys of the genesis committee of this test chain
494      pub fn validator_key(&self, index: usize) -> &PrivateKey<N> {
495          &self.private_keys[index]
496      }
497  
498      /// Returns the address of the specified validator.
499      pub fn validator_address(&self, index: usize) -> Address<N> {
500          Address::try_from(*self.validator_key(index)).unwrap()
501      }
502  
503      /// Create a test ledger with this builder's genesis block.
504      pub fn instantiate_ledger(&self) -> Ledger<N, LedgerType<N>> {
505          Ledger::load(self.genesis_block().clone(), StorageMode::new_test(None)).unwrap()
506      }
507  }