test: improve the data generation for the consensus simulation (#510)

* test: adjust some log messages for the simulation

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: fix a possible deadlock when running the simulation with too many downstream peers

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: only reverse heap entries when added to the heap

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: parametrize the number of peers for the data generation

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* chore: rebase on main

Signed-off-by: etorreborre <etorreborre@yahoo.com>

| Conflicts:
|	simulation/amaru-sim/src/simulator/run.rs

* test: use the test generation from headers tree for the simulation

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* feat: don't terminate on validation errors, just log

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: make the simulation env. variables a bit more lenient

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: adjust the simulation oracle for generated data

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: remove the unused data generation files

Signed-off-by: etorreborre <etorreborre@yahoo.com>

| Conflicts:
|	simulation/amaru-sim/src/simulator/args.rs
|	simulation/amaru-sim/src/simulator/ledger.rs
|	simulation/amaru-sim/src/simulator/node_config.rs

* test: only check the messages sent to the first downstream peer

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: add an argument to vary the depth of generated chains

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: add generation arguments

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* chore: implement some coderabbit suggestions

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* chore: rebase on main

Signed-off-by: etorreborre <etorreborre@yahoo.com>

* test: pass around the best chain for a generated tree of headers to write assertions

Signed-off-by: etorreborre <etorreborre@yahoo.com>

---------

Signed-off-by: etorreborre <etorreborre@yahoo.com>

Author: etorreborre

Cargo.lock

@@ -23,7 +23,7 @@
 use crate::consensus::errors::ConsensusError;
 use crate::consensus::headers_tree::Tracker::{Me, SomePeer};
 use crate::consensus::headers_tree::data_generation::SelectionResult::{Back, Forward};
-use crate::consensus::headers_tree::data_generation::any_tree_of_headers;
+use crate::consensus::headers_tree::data_generation::any_tree_of_headers_and_best_chain;
 use crate::consensus::headers_tree::tree::Tree;
 use crate::consensus::headers_tree::{HeadersTree, Tracker};
 use crate::consensus::stages::select_chain::RollbackChainSelection::RollbackBeyondLimit;
@@ -44,6 +44,7 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use std::cmp::Reverse;
 use std::collections::BTreeMap;
 use std::fmt::{Debug, Display, Formatter};
+use std::str::FromStr;
 use std::sync::Arc;
 
 /// This data type models the events sent by the ChainSync mini-protocol with simplified data for the tests.
@@ -315,6 +316,40 @@ pub fn random_walk(
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Ratio(pub u32, pub u32);
 
+impl Display for Ratio {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}/{}", self.0, self.1)
+    }
+}
+
+impl FromStr for Ratio {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let parts: Vec<&str> = s.split('/').collect();
+        if parts.len() != 2 {
+            return Err("Ratio must be in the form 'numerator/denominator'".to_string());
+        }
+        let numerator = parts[0]
+            .parse::<u32>()
+            .map_err(|_| format!("Invalid numerator in ratio: {}", parts[0]))?;
+        let denominator = parts[1]
+            .parse::<u32>()
+            .map_err(|_| format!("Invalid denominator in ratio: {}", parts[1]))?;
+
+        if denominator == 0 {
+            return Err("Ratio denominator must be > 0".to_string());
+        }
+        if numerator > denominator {
+            return Err(format!(
+                "Ratio must be <= 1.0, got {}/{}",
+                numerator, denominator
+            ));
+        }
+        Ok(Ratio(numerator, denominator))
+    }
+}
+
 /// Generate random walks for a fixed number of peers on a given tree of headers.
 ///
 /// The returned list of actions is transposed so that the actions from different peers are interleaved.
@@ -342,25 +377,34 @@ pub fn generate_random_walks(
         .collect()
 }
 
-/// Generate a random list of actions, for a fixed number of peers, with:
+/// Generate a random list of actions, for a number of peers.
 ///
-///  - A tree of headers of depth `depth`.
-///  - A `max_length` for how far rollback actions can go in the past.
+/// We first generate a tree of headers of depth `depth` with a given `branching_ratio`
+/// (the probability that a branch is created from an original spine of headers).
 ///
-/// If we use `max_length` < depth we will generate test cases where the `HeadersTree` will
-/// have to be truncated after a number of roll forwards.
+/// Then we execute a random walk on that tree for `peers_nb` peers, with a given `rollback_ratio`,
+/// which is how often a given peer will rollback.
 ///
 /// Important note: this function returns a prop test strategy but the random walks are generated
 /// using a `StdRng` generator. This makes the generator reproducible, because the `StdGenerator`
 /// is given a seed controlled by `proptest` but this makes the resulting list of actions non-shrinkable.
 ///
 pub fn any_select_chains(
+    peers_nb: usize,
     depth: usize,
     rollback_ratio: Ratio,
-) -> impl Strategy<Value = Vec<Action>> {
-    any_tree_of_headers(depth, Ratio(1, 2)).prop_flat_map(move |tree| {
-        (1..u64::MAX).prop_map(move |seed| generate_random_walks(&tree, 5, rollback_ratio, seed))
-    })
+    branching_ratio: Ratio,
+) -> impl Strategy<Value = (Vec<Action>, Chain)> {
+    any_tree_of_headers_and_best_chain(depth, branching_ratio).prop_flat_map(
+        move |(tree, best_chain)| {
+            (1..u64::MAX).prop_map(move |seed| {
+                (
+                    generate_random_walks(&tree, peers_nb, rollback_ratio, seed),
+                    best_chain.clone(),
+                )
+            })
+        },
+    )
 }
 
 /// Create an empty `HeadersTree` handling chains of maximum length `max_length` and
@@ -476,12 +520,15 @@ pub fn execute_json_actions(
 }
 
 /// Type alias for a chain of headers tracked by a peer
-type Chain = Vec<BlockHeader>;
+pub type Chain = Vec<BlockHeader>;
 
 /// This function computes the chains sent by each peer from a list of actions.
 /// Once all the actions have been executed it returns the chains that are the longest.
 ///
-/// The return value is a list of lists of chains, because it returns one list of chains per action.
+/// The return value is a list of lists of chains:
+///
+///  - For each action we produce the resulting best chains.
+///  - So that the last list of chains is the best chains after executing all the actions.
 ///
 pub fn make_best_chains_from_actions(actions: &Vec<Action>) -> Vec<Vec<Chain>> {
     let mut all_best_chains: Vec<Vec<Chain>> = vec![];
@@ -584,6 +631,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use super::*;
+
     #[test]
     fn transpose_works() {
         let rows = vec![
@@ -600,7 +649,37 @@ mod tests {
             vec![3, 8],
             vec![9],
         ];
-        let result = super::transpose(rows);
+        let result = transpose(rows);
         assert_eq!(result, expected);
     }
+
+    #[test]
+    fn parse_ratio_works() {
+        let ratio: Ratio = "3/4".parse().unwrap();
+        assert_eq!(ratio, Ratio(3, 4));
+
+        let ratio: super::Ratio = "0/1".parse().unwrap();
+        assert_eq!(ratio, Ratio(0, 1));
+
+        let ratio: super::Ratio = "1/1".parse().unwrap();
+        assert_eq!(ratio, Ratio(1, 1));
+    }
+
+    #[test]
+    fn parse_ratio_with_errors() {
+        let err = "3".parse::<Ratio>().unwrap_err();
+        assert_eq!(err, "Ratio must be in the form 'numerator/denominator'");
+
+        let err = "a/2".parse::<Ratio>().unwrap_err();
+        assert_eq!(err, "Invalid numerator in ratio: a");
+
+        let err = "1/b".parse::<Ratio>().unwrap_err();
+        assert_eq!(err, "Invalid denominator in ratio: b");
+
+        let err = "1/0".parse::<Ratio>().unwrap_err();
+        assert_eq!(err, "Ratio denominator must be > 0");
+
+        let err = "5/4".parse::<Ratio>().unwrap_err();
+        assert_eq!(err, "Ratio must be <= 1.0, got 5/4");
+    }
 }

crates/amaru-consensus/src/consensus/headers_tree/data_generation/actions.rs

@@ -21,7 +21,7 @@
 //!
 
 use crate::consensus::headers_tree::HeadersTree;
-use crate::consensus::headers_tree::data_generation::Ratio;
+use crate::consensus::headers_tree::data_generation::{Chain, Ratio};
 use crate::consensus::headers_tree::tree::Tree;
 use amaru_kernel::peer::Peer;
 use amaru_kernel::{Bytes, HEADER_HASH_SIZE, Header, HeaderHash};
@@ -33,23 +33,46 @@ use rand::prelude::StdRng;
 use rand::{Rng, RngCore, SeedableRng};
 use std::sync::Arc;
 
+/// Return a `proptest` Strategy producing a random `Tree<BlockHeader>` of a given depth
+/// Additionally, we return the best chain corresponding to the spine of the tree.
+pub fn any_tree_of_headers_and_best_chain(
+    depth: usize,
+    branching_ratio: Ratio,
+) -> impl Strategy<Value = (Tree<BlockHeader>, Chain)> {
+    (0..u64::MAX)
+        .prop_map(move |seed| generate_header_tree_and_best_chain(depth, seed, branching_ratio))
+}
+
 /// Return a `proptest` Strategy producing a random `Tree<BlockHeader>` of a given depth
 pub fn any_tree_of_headers(
     depth: usize,
     branching_ratio: Ratio,
 ) -> impl Strategy<Value = Tree<BlockHeader>> {
-    (0..u64::MAX).prop_map(move |seed| generate_header_tree(depth, seed, branching_ratio))
+    any_tree_of_headers_and_best_chain(depth, branching_ratio).prop_map(|(tree, _)| tree)
 }
 
 /// Generate a tree of headers of a given depth.
 /// A seed is used to control the random generation of subtrees on top of a spine of length `depth`.
-pub fn generate_header_tree(depth: usize, seed: u64, branching_ratio: Ratio) -> Tree<BlockHeader> {
+///
+/// We also return the chain corresponding to the spine of the tree. This is the best chain.
+pub fn generate_header_tree_and_best_chain(
+    depth: usize,
+    seed: u64,
+    branching_ratio: Ratio,
+) -> (Tree<BlockHeader>, Chain) {
     let mut rng = StdRng::seed_from_u64(seed);
 
     let root = generate_header(1, 1, None, &mut rng);
     let mut root_tree = Tree::make_leaf(&root);
-    generate_header_subtree(&mut rng, &mut root_tree, depth - 1, branching_ratio);
-    root_tree
+    let mut spine = generate_header_subtree(&mut rng, &mut root_tree, depth - 1, branching_ratio);
+    spine.insert(0, root);
+    (root_tree, spine)
+}
+
+/// Generate a tree of headers of a given depth.
+/// A seed is used to control the random generation of subtrees on top of a spine of length `depth`.
+pub fn generate_header_tree(depth: usize, seed: u64, branching_ratio: Ratio) -> Tree<BlockHeader> {
+    generate_header_tree_and_best_chain(depth, seed, branching_ratio).0
 }
 
 /// Given a random generator and a tree:
@@ -66,7 +89,7 @@ fn generate_header_subtree(
     tree: &mut Tree<BlockHeader>,
     depth: usize,
     branching_ratio: Ratio,
-) {
+) -> Chain {
     let header_body = tree.value.header_body().clone();
     let mut spine = generate_headers(
         depth,
@@ -90,6 +113,7 @@ fn generate_header_subtree(
             generate_header_subtree(rng, current, other_branch_depth, branching_ratio);
         }
     }
+    spine
 }
 
 /// Generate a chain of headers anchored at a given header.

crates/amaru-consensus/src/consensus/headers_tree/data_generation/data_generation.rs

@@ -1284,11 +1284,13 @@ mod tests {
     const MAX_LENGTH: usize = 5;
     const TEST_CASES_NB: u32 = 1000;
     const ROLLBACK_RATIO: Ratio = Ratio(1, 2);
+    const BRANCHING_RATIO: Ratio = Ratio(1, 2);
 
     proptest! {
         #![proptest_config(config_begin().no_shrink().with_cases(TEST_CASES_NB).end())]
         #[test]
-        fn run_chain_selection(actions in any_select_chains(DEPTH, ROLLBACK_RATIO)) {
+        fn run_chain_selection(generated in any_select_chains(DEPTH, 5, ROLLBACK_RATIO, BRANCHING_RATIO)) {
+            let actions = generated.0;
             let results = execute_actions(MAX_LENGTH, &actions, false).unwrap();
             let actual_chains = make_best_chains_from_results(&results);
             let expected_chains = make_best_chains_from_actions(&actions);

crates/amaru-consensus/src/consensus/headers_tree/headers_tree.rs

@@ -44,6 +44,8 @@ pub async fn stage(
             raw_header,
             span,
         } => {
+            // TODO: check the point vs the deserialized header point and invalidate if they don't match
+            // then simplify and don't pass the point separately
             let header = match decode_header(&point, raw_header.as_slice()) {
                 Ok(header) => {
                     if header.point() != point {
@@ -282,6 +284,15 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn decode_header_on_generated_header() {
+        let header = run(any_header());
+        let raw_header = cbor::to_vec(header.clone()).unwrap();
+        let point = header.point();
+        let decoded = decode_header(&point, &raw_header).unwrap();
+        assert_eq!(header, decoded);
+    }
+
     // HELPERS
 
     fn make_state() -> State {

crates/amaru-consensus/src/consensus/stages/receive_header.rs

@@ -134,6 +134,10 @@ impl BlockHeader {
         &self.header.header_body
     }
 
+    pub fn parent_hash(&self) -> Option<HeaderHash> {
+        self.header.header_body.prev_hash
+    }
+
     fn recompute_hash(&mut self) {
         self.hash = Hasher::<{ HEADER_HASH_SIZE * 8 }>::hash_cbor(&self.header);
     }

crates/amaru-ouroboros-traits/src/is_header/mod.rs

@@ -15,7 +15,7 @@
 use super::*;
 use proptest::prelude::*;
 use proptest::strategy::ValueTree;
-use proptest::test_runner::TestRunner;
+use proptest::test_runner::{Config, RngSeed, TestRunner};
 
 /// Make a mostly empty Header with the given block_number, slot and previous hash
 pub fn make_header(block_number: u64, slot: u64, prev_hash: Option<HeaderHash>) -> Header {
@@ -121,3 +121,15 @@ pub fn run<T>(s: impl Strategy<Value = T>) -> T {
     let mut runner = TestRunner::default();
     s.new_tree(&mut runner).unwrap().current()
 }
+
+/// Run a strategy with a seed provided by a random generator
+/// and return the generated value, panicking if generation fails.
+#[expect(clippy::unwrap_used)]
+pub fn run_with_rng<T, RNG: Rng>(rng: &mut RNG, s: impl Strategy<Value = T>) -> T {
+    let config = Config {
+        rng_seed: RngSeed::Fixed(rng.random()),
+        ..Default::default()
+    };
+    let mut runner = TestRunner::new(config);
+    s.new_tree(&mut runner).unwrap().current()
+}

crates/amaru-ouroboros-traits/src/is_header/tests.rs

@@ -59,11 +59,9 @@ pub fn build_stage_graph(
     // TODO: currently only validate_header errors, will need to grow into all error handling
     let validation_errors_stage = network.stage(
         "validation_errors",
-        async |_, error: ValidationFailed, eff| {
+        async |_, error: ValidationFailed, _eff| {
             tracing::error!(%error, "stage error");
-            // TODO: implement specific actions once we have an upstream network
-            // termination here will tear down the entire stage graph
-            eff.terminate().await
+            // TODO: disconnect bad behaving peers
         },
     );
 

crates/amaru/src/stages/build_stage_graph.rs

@@ -35,7 +35,8 @@ tracing.workspace = true
 # Internal dependencies ───────────────────────────────────────────────────────┐
 amaru.workspace = true
 amaru-kernel = { workspace = true, features = ["test-utils"] }
-amaru-consensus.workspace = true
+amaru-consensus = { workspace = true, features = ["test-utils"] }
+amaru-ouroboros-traits = { workspace = true, features = ["test-utils"] }
 amaru-ouroboros.workspace = true
 amaru-slot-arithmetic.workspace = true