chia-consensus 0.36.0

use super::conditions::{NewCoin, SpendBundleConditions, SpendConditions};
use super::run_block_generator::{
    get_coinspends_for_trusted_block, get_coinspends_with_conditions_for_trusted_block,
    run_block_generator, run_block_generator2,
};
use crate::allocator::make_allocator;
use crate::consensus_constants::TEST_CONSTANTS;
use crate::flags::{COST_CONDITIONS, DONT_VALIDATE_SIGNATURE, MEMPOOL_MODE, SIMPLE_GENERATOR};
use crate::run_block_generator::check_generator_node;
use crate::validation_error::ErrorCode;
use chia_bls::Signature;
use chia_protocol::Program;
use chia_protocol::{Bytes, Bytes48};
use chia_puzzles::CHIALISP_DESERIALISATION;
use clvmr::Allocator;
use clvmr::allocator::NodePtr;
use clvmr::serde::{node_from_bytes, node_from_bytes_backrefs};
use std::iter::zip;
use text_diff::Difference;
use text_diff::diff;

use rstest::rstest;

pub(crate) fn print_conditions(a: &Allocator, c: &SpendBundleConditions, a2: &Allocator) -> String {
    let mut ret = String::new();
    if c.reserve_fee > 0 {
        ret += &format!("RESERVE_FEE: {}\n", c.reserve_fee);
    }

    if c.height_absolute > 0 {
        ret += &format!("ASSERT_HEIGHT_ABSOLUTE {}\n", c.height_absolute);
    }
    if c.seconds_absolute > 0 {
        ret += &format!("ASSERT_SECONDS_ABSOLUTE {}\n", c.seconds_absolute);
    }
    if let Some(val) = c.before_seconds_absolute {
        ret += &format!("ASSERT_BEFORE_SECONDS_ABSOLUTE {val}\n");
    }
    if let Some(val) = c.before_height_absolute {
        ret += &format!("ASSERT_BEFORE_HEIGHT_ABSOLUTE {val}\n");
    }
    let mut agg_sigs = Vec::<(Bytes48, Bytes)>::new();
    for (pk, msg) in &c.agg_sig_unsafe {
        agg_sigs.push((pk.to_bytes().into(), a.atom(*msg).as_ref().into()));
    }
    agg_sigs.sort();
    for (pk, msg) in agg_sigs {
        ret += &format!(
            "AGG_SIG_UNSAFE pk: {} msg: {}\n",
            hex::encode(pk),
            hex::encode(msg)
        );
    }
    ret += "SPENDS:\n";

    let mut spends: Vec<SpendConditions> = c.spends.clone();
    spends.sort_by_key(|s| *s.coin_id);
    for s in spends {
        ret += &format!(
            "- coin id: {} ph: {} exe-cost: {} cond-cost: {}\n",
            hex::encode(*s.coin_id),
            hex::encode(a.atom(s.puzzle_hash)),
            s.execution_cost,
            s.condition_cost,
        );

        if let Some(val) = s.height_relative {
            ret += &format!("  ASSERT_HEIGHT_RELATIVE {val}\n");
        }
        if let Some(val) = s.seconds_relative {
            ret += &format!("  ASSERT_SECONDS_RELATIVE {val}\n");
        }
        if let Some(val) = s.before_height_relative {
            ret += &format!("  ASSERT_BEFORE_HEIGHT_RELATIVE {val}\n");
        }
        if let Some(val) = s.before_seconds_relative {
            ret += &format!("  ASSERT_BEFORE_SECONDS_RELATIVE {val}\n");
        }
        let mut create_coin: Vec<&NewCoin> = s.create_coin.iter().collect();
        create_coin.sort_by_key(|cc| (cc.puzzle_hash, cc.amount));
        for cc in create_coin {
            if cc.hint == NodePtr::NIL {
                ret += &format!(
                    "  CREATE_COIN: ph: {} amount: {}\n",
                    hex::encode(cc.puzzle_hash),
                    cc.amount
                );
            } else {
                ret += &format!(
                    "  CREATE_COIN: ph: {} amount: {} hint: {}\n",
                    hex::encode(cc.puzzle_hash),
                    cc.amount,
                    hex::encode(a.atom(cc.hint))
                );
            }
        }

        for sig_conds in [
            (&s.agg_sig_me, "AGG_SIG_ME"),
            (&s.agg_sig_parent, "AGG_SIG_PARENT"),
            (&s.agg_sig_puzzle, "AGG_SIG_PUZZLE"),
            (&s.agg_sig_amount, "AGG_SIG_AMOUNT"),
            (&s.agg_sig_puzzle_amount, "AGG_SIG_PUZZLE_AMOUNT"),
            (&s.agg_sig_parent_amount, "AGG_SIG_PARENT_AMOUNT"),
            (&s.agg_sig_parent_puzzle, "AGG_SIG_PARENT_PUZZLE"),
        ] {
            let mut agg_sigs = Vec::<(Bytes48, Bytes)>::new();
            for (pk, msg) in sig_conds.0 {
                agg_sigs.push((pk.to_bytes().into(), a.atom(*msg).as_ref().into()));
            }
            agg_sigs.sort();
            for (pk, msg) in &agg_sigs {
                ret += &format!(
                    "  {} pk: {} msg: {}\n",
                    sig_conds.1,
                    hex::encode(pk),
                    hex::encode(msg)
                );
            }
        }
    }

    ret += &format!("cost: {}\n", c.cost);
    ret += &format!("execution-cost: {}\n", c.execution_cost);
    ret += &format!("condition-cost: {}\n", c.condition_cost);
    ret += &format!("removal_amount: {}\n", c.removal_amount);
    ret += &format!("addition_amount: {}\n", c.addition_amount);
    ret += &format!("atoms: {}\n", a2.atom_count() + a2.small_atom_count());
    ret += &format!("pairs: {}\n", a2.pair_count());
    ret += &format!("heap: {}\n", a2.heap_size());
    ret
}

pub(crate) fn print_diff(output: &str, expected: &str) {
    println!("\x1b[102m \x1b[0m - output from test");
    println!("\x1b[101m \x1b[0m - expected output");
    for diff in diff(expected, output, "\n").1 {
        match diff {
            Difference::Same(s) => {
                let lines: Vec<&str> = s.split('\n').collect();
                if lines.len() <= 6 {
                    for l in &lines {
                        println!(" {l}");
                    }
                } else {
                    for l in &lines[0..3] {
                        println!(" {l}");
                    }
                    println!(" ...");
                    for l in &lines[lines.len() - 3..] {
                        println!(" {l}");
                    }
                }
            }
            Difference::Rem(s) => {
                println!("\x1b[91m");
                for l in s.split('\n') {
                    println!("-{l}");
                }
                println!("\x1b[0m");
            }
            Difference::Add(s) => {
                println!("\x1b[92m");
                for l in s.split('\n') {
                    println!("+{l}");
                }
                println!("\x1b[0m");
            }
        }
    }
}

#[rstest]
// in CI we run with the clvmr/debug-allocator feature enabled, which makes this
// test use too much RAM (about 6.8 GB)
//#[case("aa-million-messages")]
#[case("single-coin-only-garbage")]
#[case("many-coins-announcement-cap")]
#[case("3000000-conditions-single-coin")]
#[case("100000-remarks-prefab")]
#[case("29500-remarks-procedural")]
#[case("aa-million-message-spends")]
#[case("puzzle-hash-stress-test")]
#[case("puzzle-hash-stress-tree")]
#[case("new-agg-sigs")]
#[case("infinity-g1")]
#[case("block-1ee588dc")]
#[case("block-6fe59b24")]
#[case("block-b45268ac")]
#[case("block-c2a8df0d")]
#[case("block-e5002df2")]
#[case("block-4671894")]
#[case("block-225758")]
#[case("assert-puzzle-announce-fail")]
#[case("block-834752")]
#[case("block-834752-compressed")]
#[case("block-834760")]
#[case("block-834761")]
#[case("block-834765")]
#[case("block-834766")]
#[case("block-834768")]
#[case("create-coin-different-amounts")]
#[case("create-coin-hint-duplicate-outputs")]
#[case("create-coin-hint")]
#[case("create-coin-hint2")]
#[case("deep-recursion-plus")]
#[case("double-spend")]
#[case("duplicate-coin-announce")]
#[case("duplicate-create-coin")]
#[case("duplicate-height-absolute-div")]
#[case("duplicate-height-absolute-substr-tail")]
#[case("duplicate-height-absolute-substr")]
#[case("duplicate-height-absolute")]
#[case("duplicate-height-relative")]
#[case("duplicate-outputs")]
#[case("duplicate-reserve-fee")]
#[case("duplicate-seconds-absolute")]
#[case("duplicate-seconds-relative")]
#[case("height-absolute-ladder")]
#[case("infinite-recursion1")]
#[case("infinite-recursion2")]
#[case("infinite-recursion3")]
#[case("infinite-recursion4")]
#[case("invalid-conditions")]
#[case("just-puzzle-announce")]
#[case("many-create-coin")]
#[case("many-large-ints-negative")]
#[case("many-large-ints")]
#[case("max-height")]
#[case("multiple-reserve-fee")]
#[case("negative-reserve-fee")]
#[case("recursion-pairs")]
#[case("unknown-condition")]
#[case("duplicate-messages")]
#[case("non-quote-0001-start")]
fn run_generator(#[case] name: &str) {
    use std::fs::{read_to_string, write};

    // When making changes to print_conditions() enabling this will update the
    // test cases to match. Make sure to carefully review the diff before
    // landing an automatic update of the test case.
    const UPDATE_TESTS: bool = false;

    let run_generator_one: bool = ![
        "single-coin-only-garbage",
        "many-coins-announcement-cap",
        "puzzle-hash-stress-test",
        "puzzle-hash-stress-tree",
        "aa-million-message-spends",
        "29500-remarks-procedural",
        "100000-remarks-prefab",
        "3000000-conditions-single-coin",
    ]
    .contains(&name); // these could be stored in the case, but there's so few it's easier to look at here

    let filename = format!("../../generator-tests/{name}.txt");
    println!("file: {filename}");
    let test_file = read_to_string(&filename).expect("test file not found");
    let (generator, expected) = test_file.split_once('\n').expect("invalid test file");
    let generator = hex::decode(generator).expect("invalid hex encoded generator");

    let expected = match expected.split_once("STRICT:\n") {
        Some((c, m)) => [c, m],
        None => [expected, expected],
    };

    let mut block_refs = Vec::<Vec<u8>>::new();

    let filename2 = format!("../../generator-tests/{name}.env");
    if let Ok(env_hex) = read_to_string(&filename2) {
        println!("block-ref file: {filename2}");
        block_refs.push(hex::decode(env_hex).expect("hex decode env-file"));
    }

    let mut write_back = format!("{}\n", hex::encode(&generator));
    let mut last_output = String::new();

    for (flags, expected) in zip(&[0, MEMPOOL_MODE], expected) {
        let mut flags = *flags;
        if name == "aa-million-messages" || name == "aa-million-message-spends" {
            // this test requires running after hard fork 2, where the COST_CONDITIONS
            // flag is set
            flags |= COST_CONDITIONS;
        }

        // These are generators that are programs, not a simple list of spends
        // when the SIMPLE_GENERATOR flag is set, these should fail
        if [
            "single-coin-only-garbage",
            "many-coins-announcement-cap",
            "puzzle-hash-stress-test",
            "puzzle-hash-stress-tree",
            "aa-million-message-spends",
            "aa-million-messages",
            "29500-remarks-procedural",
            "100000-remarks-prefab",
            "3000000-conditions-single-coin",
            "block-4671894",
            "block-225758",
            "infinite-recursion1",
            "infinite-recursion2",
            "infinite-recursion3",
            "infinite-recursion4",
            "recursion-pairs",
            "non-quote-0001-start", // this will fail the generator format check before it fails at runtime
        ]
        .contains(&name)
        {
            // lets test that the procedural generators are filtered with the flag
            let mut a = make_allocator(flags);
            let test_conds = run_block_generator2(
                &mut a,
                &generator,
                &block_refs, // we're not allowed to pass in block references when SIMPLE_GENERATOR is set
                11_000_000_000,
                flags | DONT_VALIDATE_SIGNATURE | SIMPLE_GENERATOR,
                &Signature::default(),
                None,
                &TEST_CONSTANTS,
            );
            assert_eq!(
                test_conds.unwrap_err().1,
                ErrorCode::ComplexGeneratorReceived
            );

            // now lets specifically check the node generator check
            let program = node_from_bytes_backrefs(&mut a, generator.as_ref())
                .expect("node_from_bytes_backref");
            let res = check_generator_node(&a, program, flags | SIMPLE_GENERATOR);
            assert_eq!(res.unwrap_err().1, ErrorCode::ComplexGeneratorReceived);
        } else {
            flags |= SIMPLE_GENERATOR;
            // ensure SIMPLE_GENERATOR fails if there are any block references
            // passed in. We pass in a dummy block reference
            let mut a = make_allocator(flags);
            let test_conds = run_block_generator2(
                &mut a,
                &generator,
                &[&[0_u8, 1, 2, 3]],
                11_000_000_000,
                flags | DONT_VALIDATE_SIGNATURE,
                &Signature::default(),
                None,
                &TEST_CONSTANTS,
            );
            assert_eq!(test_conds.unwrap_err().1, ErrorCode::TooManyGeneratorRefs);

            // now lets specifically check the node generator check
            let program = node_from_bytes_backrefs(&mut a, generator.as_ref())
                .expect("node_from_bytes_backref");
            let res = check_generator_node(&a, program, flags | SIMPLE_GENERATOR);
            assert!(res.is_ok());
        }

        println!("flags: {flags:x}");
        let mut a2 = make_allocator(flags);
        let conds2 = run_block_generator2(
            &mut a2,
            &generator,
            &block_refs,
            11_000_000_000,
            flags | DONT_VALIDATE_SIGNATURE,
            &Signature::default(),
            None,
            &TEST_CONSTANTS,
        );

        let (expected_cost, output) = match conds2 {
            Ok(ref conditions) => {
                let cond_cost: u64 = conditions.spends.iter().map(|v| v.condition_cost).sum();
                assert_eq!(cond_cost, conditions.condition_cost);
                let exe_cost: u64 = conditions.spends.iter().map(|v| v.execution_cost).sum();
                // the generator itself has execution cost. At least the cost of
                // a quote
                assert!(exe_cost <= conditions.execution_cost);

                if (flags & SIMPLE_GENERATOR) != 0 {
                    // when running generators with the SIMPLE_GENERATOR flag
                    // set, we don't pass in the CLVM deserializer program. This
                    // causes the atoms and pairs counters to be lower than
                    // before (when the test cases were created). In order to
                    // match the test cases, we increment those counters in
                    // order to print compatible output.
                    // these are the atoms and pairs that would have been
                    // allocated by the deserializer program
                    let _ =
                        node_from_bytes(&mut a2, &CHIALISP_DESERIALISATION).expect("deserializer");
                    a2.add_ghost_pair(2).expect("add_ghost_pair");
                }
                (conditions.cost, print_conditions(&a2, &conditions, &a2))
            }
            Err(code) => (0, format!("FAILED: {}\n", u32::from(code.1))),
        };

        if UPDATE_TESTS {
            if (flags & MEMPOOL_MODE) != 0 {
                if output != last_output {
                    write_back.push_str(&format!("STRICT:\n{output}"));
                }
            } else {
                last_output = output.clone();
                write_back.push_str(&format!("{output}"));
            }
        }
        if run_generator_one {
            let mut a1 = make_allocator(flags);
            let conds1 = run_block_generator(
                &mut a1,
                &generator,
                &block_refs,
                11_000_000_000,
                flags | DONT_VALIDATE_SIGNATURE,
                &Signature::default(),
                None,
                &TEST_CONSTANTS,
            );
            let output_pre_hard_fork = match conds1 {
                Ok(mut conditions) => {
                    // before the hard fork, the cost of running the genrator +
                    // puzzles should never be lower than after the hard-fork
                    // but it's likely higher.
                    assert!(conditions.cost >= expected_cost);
                    // pre-hard fork, we don't have access to per-puzzle costs, so
                    // set those to whatever run_block_generator2() produced, to
                    // make the check pass
                    if let Ok(ref conds2) = conds2 {
                        // update the cost we print here, just to be compatible with
                        // the test cases we have. We've already ensured the cost is
                        // lower
                        conditions.cost = conds2.cost;
                        conditions.execution_cost = conds2.execution_cost;
                        for s in &conds2.spends {
                            for ms in conditions.spends.iter_mut() {
                                if ms.coin_id == s.coin_id {
                                    ms.execution_cost = s.execution_cost;
                                    break;
                                }
                            }
                        }
                    }

                    print_conditions(&a1, &conditions, &a2)
                }
                Err(code) => {
                    format!("FAILED: {}\n", u32::from(code.1))
                }
            };
            if output != output_pre_hard_fork {
                print_diff(&output, &output_pre_hard_fork);
                if !UPDATE_TESTS {
                    panic!("run_block_generator 1 and 2 produced a different result!");
                }
            }
        }

        if output != expected {
            print_diff(&output, expected);
            if !UPDATE_TESTS {
                panic!("mismatching generator output");
            }
        }

        let vec_of_slices: Vec<&[u8]> = block_refs.iter().map(std::vec::Vec::as_slice).collect();
        let result = get_coinspends_for_trusted_block(
            &TEST_CONSTANTS,
            &Program::new(generator.clone().into()),
            &vec_of_slices,
            flags,
        );

        let result2 = get_coinspends_with_conditions_for_trusted_block(
            &TEST_CONSTANTS,
            &Program::new(generator.clone().into()),
            &vec_of_slices,
            flags,
        );

        // now lets check get_coinspends_for_trusted_block
        // but only if we trust it not to do shenanigans
        if let Ok(ref conds) = conds2 {
            // if run_block_generator2 is OK then check we're equal
            let coinspends = result.expect("get_coinspends");

            // check that we're getting the same info from get_coinspends_with_conditions_for_trusted_block
            let coinspends2 = result2.expect("get_coinspends_with_conds");

            // get_coinspends_for_trusted_block skips certain attacking spends as it is required to
            // return serialized puzzles, which may not be possible. We can ignore the comparison
            // if it has skipped a spend for this reason
            if coinspends.len() == conds.spends.len() {
                for (i, spend) in conds.spends.iter().enumerate() {
                    let runnable = {
                        let mut a = make_allocator(flags);
                        coinspends[i]
                            .puzzle_reveal
                            .run(
                                &mut a,
                                flags,
                                TEST_CONSTANTS.max_block_cost_clvm,
                                &coinspends[i].solution,
                            )
                            .is_ok()
                    };
                    assert!(runnable);
                    let parent_id = a2.atom(spend.parent_id);
                    assert_eq!(
                        parent_id.as_ref(),
                        coinspends[i].coin.parent_coin_info.as_slice()
                    );
                    let puzhash = a2.atom(spend.puzzle_hash);
                    assert_eq!(puzhash.as_ref(), coinspends[i].coin.puzzle_hash.as_slice());
                    assert_eq!(spend.coin_amount, coinspends[i].coin.amount);

                    // test get_coinspends_with_conditions_for_trusted_block()
                    assert_eq!(
                        parent_id.as_ref(),
                        coinspends2[i].0.coin.parent_coin_info.as_slice()
                    );
                    assert_eq!(
                        puzhash.as_ref(),
                        coinspends2[i].0.coin.puzzle_hash.as_slice()
                    );
                    assert_eq!(spend.coin_amount, coinspends2[i].0.coin.amount);
                }
            }
        }
    }

    if UPDATE_TESTS {
        write(&filename, write_back.into_bytes()).expect("write file");
    }
}