misc(ledger/testchain-gen): minor cleanups

Author: kaimast

ledger/src/test_helpers/chain_builder.rs

@@ -36,10 +36,7 @@ use anyhow::{Context, Result};
 use indexmap::{IndexMap, IndexSet};
 use rand::SeedableRng;
 use rand_chacha::ChaChaRng;
-use std::{
-    cmp,
-    collections::{BTreeMap, HashMap},
-};
+use std::collections::{BTreeMap, HashMap};
 use time::OffsetDateTime;
 
 pub type CurrentNetwork = MainnetV0;
@@ -222,7 +219,7 @@ impl<N: Network> TestChainBuilder<N> {
     /// This function panics if called from an async context.
     pub fn generate_blocks_with_opts(
         &mut self,
-        mut num_blocks: usize,
+        num_blocks: usize,
         mut options: GenerateBlocksOptions<N>,
         rng: &mut TestRng,
     ) -> Result<Vec<Block<N>>> {
@@ -232,34 +229,35 @@ impl<N: Network> TestChainBuilder<N> {
 
         // If configured, skip enough blocks to reach the current consensus version.
         if options.skip_to_current_version {
-            let (version, height) = TEST_CONSENSUS_VERSION_HEIGHTS.last().unwrap();
+            let (version, target_height) = TEST_CONSENSUS_VERSION_HEIGHTS.last().unwrap();
             let mut current_height = self.ledger.latest_height();
-            println!("Skipping {height} blocks to reach {version}");
-
-            while current_height < *height && result.len() < num_blocks {
-                let options = GenerateBlockOptions {
-                    skip_votes: options.skip_votes,
-                    skip_nodes: options.skip_nodes.clone(),
-                    ..Default::default()
-                };
-
-                let block = self.generate_block_with_opts(options, rng)?;
-                current_height = block.height();
-                result.push(block);
-            }
 
-            // Subtract the number of placeholder blocks from the number of blocks to generate.
-            assert!(result.len() <= num_blocks);
-            num_blocks -= result.len();
+            let diff = target_height.saturating_sub(current_height);
+
+            if diff > 0 {
+                println!("Skipping {diff} blocks to reach {version}");
+
+                while current_height < *target_height && result.len() < num_blocks {
+                    let options = GenerateBlockOptions {
+                        skip_votes: options.skip_votes,
+                        skip_nodes: options.skip_nodes.clone(),
+                        ..Default::default()
+                    };
 
-            println!("Advanced to the current consensus version");
+                    let block = self.generate_block_with_opts(options, rng)?;
+                    current_height = block.height();
+                    result.push(block);
+                }
+
+                println!("Advanced to the current consensus version at height {target_height}");
+            } else {
+                debug!("Already at the current consensus version. No blocks to skip.");
+            }
         }
 
-        for _ in 0..num_blocks {
-            let num_txs = cmp::min(
-                BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH * options.num_validators,
-                options.transactions.len(),
-            );
+        while result.len() < num_blocks {
+            let num_txs = (BatchHeader::<N>::MAX_TRANSMISSIONS_PER_BATCH * options.num_validators)
+                .min(options.transactions.len());
 
             let options = GenerateBlockOptions {
                 skip_votes: options.skip_votes,

ledger/testchain-generator/src/main.rs

@@ -32,12 +32,7 @@ use snarkvm_ledger::{
 use aleo_std::StorageMode;
 use anyhow::{Context, Result, bail};
 use clap::{Parser, builder::PossibleValuesParser};
-use std::{
-    fs::{self, File},
-    io::Read,
-    path::Path,
-    str::FromStr,
-};
+use std::{fs, path::Path, str::FromStr};
 use tracing::debug;
 
 #[derive(Parser)]
@@ -120,24 +115,25 @@ fn generate_testchain<N: Network>(args: Args) -> Result<()> {
 
     let mut txs = if let Some(path) = args.txs_path {
         let path = Path::new(&path);
-        println!("Attempting to load txs from {}", path.display());
 
-        let mut txs = Vec::new();
         if path.is_dir() {
-            let mut buffer = String::new();
-            for entry in fs::read_dir(path)? {
-                let entry = entry?;
-                let path = entry.path();
-
-                let mut file = File::open(path)?;
-                let _ = file.read_to_string(&mut buffer)?;
-                let tx = Transaction::<N>::from_str(&buffer)?;
-                txs.push(tx);
-                buffer.clear();
-            }
+            println!("Attempting to load transactions from \"{}\"", path.display());
+        } else {
+            bail!("Cannot load transactions from \"{}\": not a valid directory", path.display());
+        }
+
+        let mut txs = Vec::new();
+        for entry in fs::read_dir(path)? {
+            let entry = entry?;
+            let path = entry.path();
+
+            let buffer = fs::read_to_string(path)?;
+            let tx = Transaction::<N>::from_str(&buffer)?;
+
+            txs.push(tx);
         }
 
-        println!("Loaded {} txs from {}", txs.len(), path.display());
+        println!("Loaded {} tranactionss from \"{}\"", txs.len(), path.display());
         txs
     } else {
         Default::default()