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