neptune-consensus 0.14.0

use std::sync::OnceLock;

use get_size2::GetSize;
use itertools::Itertools;
use neptune_mutator_set::addition_record::AdditionRecord;
use neptune_mutator_set::mutator_set_accumulator::MutatorSetAccumulator;
use neptune_primitives::block_height::BlockHeight;
use neptune_primitives::mast_hash::HasDiscriminant;
use neptune_primitives::mast_hash::MastHash;
use serde::Deserialize;
use serde::Serialize;
use strum::EnumCount;
use tasm_lib::prelude::TasmObject;
use tasm_lib::prelude::Tip5;
use tasm_lib::triton_vm::prelude::Digest;
use tasm_lib::twenty_first::math::b_field_element::BFieldElement;
use tasm_lib::twenty_first::math::bfield_codec::BFieldCodec;
use tasm_lib::twenty_first::prelude::MerkleTree;
use tasm_lib::twenty_first::prelude::Mmr;
use tasm_lib::twenty_first::util_types::mmr::mmr_accumulator::MmrAccumulator;

use crate::block::block_validation_error::BlockValidationError;
use crate::block::mutator_set_update::MutatorSetUpdate;
use crate::prelude::twenty_first;
use crate::transaction::transaction_kernel::TransactionKernel;
use crate::type_scripts::native_currency_amount::NativeCurrencyAmount;

pub(crate) const NUM_GUESSER_FEE_OUTPUTS: u64 = 2;

#[derive(Debug, Copy, Clone, EnumCount)]
pub enum BlockBodyField {
    TransactionKernel,
    MutatorSetAccumulator,
    LockFreeMmrAccumulator,
    BlockMmrAccumulator,
}

impl HasDiscriminant for BlockBodyField {
    fn discriminant(&self) -> usize {
        *self as usize
    }
}

/// Public fields of `BlockBody` are read-only, enforced by #[readonly::make].
/// Modifications are possible only through `BlockBody` methods.
///
// ## About the private `mast_hash` field:
//
// The `mast_hash` field represents the `BlockBody` MAST hash.  It is an
// optimization so that the hash can be lazily computed at most once (per
// modification). Without it, the PoW hash rate depends on the number of inputs
// and outputs in a transaction. This caching of a hash value is similar to that
// of `Block`.
//
// The field must be reset whenever the block body is modified.  As such, we
// should not permit direct modification of internal fields.
//
// Therefore `[readonly::make]` is used to make public `BlockBody` fields read-
// only (not mutable) outside of this module.  All methods that modify BlockBody
// must reset this field.
//
// We manually implement `PartialEq` and `Eq` so that digest field will not be
// compared.  Otherwise, we could have identical blocks except one has
// initialized digest field and the other has not.
//
// The field should not be serialized, so it has the `#[serde(skip)]` attribute.
// Upon deserialization, the field will have Digest::default() which is desired
// so that the digest will be recomputed if/when hash() is called.
//
// We likewise skip the field for `BFieldCodec`, and `GetSize` because there
// exist no impls for `OnceLock<_>` so derive fails.
#[derive(Clone, Debug, Serialize, Deserialize, BFieldCodec, GetSize, TasmObject)]
pub struct BlockBody {
    /// Every block contains exactly one transaction, which represents the merger of all
    /// broadcasted transactions that the miner decided to confirm. The inputs
    /// to this transaction kernel must be packed if the consensus rule dictate
    /// that.
    pub transaction_kernel: TransactionKernel,

    /// The mutator set accumulator represents the UTXO set. It is simultaneously an
    /// accumulator (=> compact representation and membership proofs) and an anonymity
    /// construction (=> outputs from one transaction do not look like inputs to another).
    ///
    /// This field represents the state of the MS *after* applying the update
    /// induced by the transaction, but *before* applying the update induced by
    /// guesser fees (and perhaps later composer fees).
    ///
    /// For the final post-block state, refer to
    /// [`Self::mutator_set_accumulator_after`].
    pub(crate) mutator_set_accumulator: MutatorSetAccumulator,

    /// Lock-free UTXOs do not come with lock scripts and do not live in the mutator set.
    pub lock_free_mmr_accumulator: MmrAccumulator,

    /// All blocks live in an MMR, so that we can efficiently prove that a given block
    /// lives on the line between the tip and genesis. This MMRA does not contain the
    /// current block.
    pub block_mmr_accumulator: MmrAccumulator,

    // This caching ensures that the hash rate is independent of the size of
    // the block's transaction.
    #[serde(skip)]
    #[bfield_codec(ignore)]
    #[get_size(ignore)]
    #[tasm_object(ignore)]
    merkle_tree: OnceLock<MerkleTree>,
}

impl PartialEq for BlockBody {
    fn eq(&self, other: &Self) -> bool {
        self.mast_hash() == other.mast_hash()
    }
}
impl Eq for BlockBody {}

impl BlockBody {
    /// Caller must pack the removal records if required.
    pub fn new(
        transaction_kernel: TransactionKernel,
        mutator_set_accumulator: MutatorSetAccumulator,
        lock_free_mmr_accumulator: MmrAccumulator,
        block_mmr_accumulator: MmrAccumulator,
    ) -> Self {
        Self {
            transaction_kernel,
            mutator_set_accumulator,
            lock_free_mmr_accumulator,
            block_mmr_accumulator,
            merkle_tree: OnceLock::default(), // calc'd in merkle_tree()
        }
    }

    /// The kernel of the transaction contained in the block.
    pub fn transaction_kernel(&self) -> &TransactionKernel {
        &self.transaction_kernel
    }

    /// The mutator set accumulator *before* guesser-fee UTXOs are applied.
    pub fn mutator_set_accumulator(&self) -> &MutatorSetAccumulator {
        &self.mutator_set_accumulator
    }

    /// Mutable access to the mutator set accumulator. Test-support only.
    #[cfg(any(test, feature = "test-helpers"))]
    pub fn mutator_set_accumulator_mut(&mut self) -> &mut MutatorSetAccumulator {
        &mut self.mutator_set_accumulator
    }

    /// Return the mutator set as it looks after the application of this block.
    ///
    /// Includes the guesser-fee UTXOs which are not included by the
    /// `mutator_set_accumulator` field on the block body.
    pub fn mutator_set_accumulator_after(
        &self,
        guesser_fee_addition_records: Vec<AdditionRecord>,
    ) -> MutatorSetAccumulator {
        let mutator_set_update = MutatorSetUpdate::new(vec![], guesser_fee_addition_records);
        let mut msa = self.mutator_set_accumulator.clone();
        mutator_set_update.apply_to_accumulator(&mut msa).expect("mutator set update derived from guesser fees should be applicable to mutator set accumulator contained in body");

        msa
    }

    /// Return the number of AOCL leafs in the mutator set *prior* to the
    /// application of any outputs from this block.
    pub fn num_aocl_leafs_prior(&self) -> u64 {
        let num_outputs: u64 = self
            .transaction_kernel
            .outputs
            .len()
            .try_into()
            .expect("Can't contain more than u64::MAX outputs");

        // The mutator set field on the block body does not contain guesser
        // outputs. So those should not be subtracted here either.
        self.mutator_set_accumulator.aocl.num_leafs() - num_outputs
    }

    /// Return the mutator set as it looks after the application of the
    /// transaction in this block but before the guesser-fee UTXOs are applied.
    ///
    /// The returned mutator set accumulator is the same as the field
    /// [`Self::mutator_set_accumulator`].
    pub fn mutator_set_accumulator_without_guesser_fees(&self) -> MutatorSetAccumulator {
        self.mutator_set_accumulator.clone()
    }

    /// Return the highest AOCL leaf index of outputs defined by this block.
    /// Includes the guesser reward outputs in this number.
    pub fn max_aocl_leaf_index(&self) -> u64 {
        const NUMBER_TO_INDEX_OFFSET: u64 = 1;
        let height: BlockHeight = self.block_mmr_accumulator.num_leafs().into();
        let num_guesser_fee_outputs = if height.is_genesis() {
            0
        } else {
            NUM_GUESSER_FEE_OUTPUTS
        };
        self.mutator_set_accumulator.aocl.num_leafs() + num_guesser_fee_outputs
            - NUMBER_TO_INDEX_OFFSET
    }

    /// The amount rewarded to the guesser who finds a valid nonce for this
    /// block.
    pub fn total_guesser_reward(&self) -> Result<NativeCurrencyAmount, BlockValidationError> {
        let r = self.transaction_kernel.fee;
        if r.is_negative() {
            Err(BlockValidationError::NegativeFee)
        } else {
            Ok(r)
        }
    }
}

impl MastHash for BlockBody {
    type FieldEnum = BlockBodyField;

    fn mast_sequences(&self) -> Vec<Vec<BFieldElement>> {
        vec![
            self.transaction_kernel.mast_hash().encode(),
            self.mutator_set_accumulator.encode(),
            self.lock_free_mmr_accumulator.encode(),
            self.block_mmr_accumulator.encode(),
        ]
    }

    fn merkle_tree(&self) -> MerkleTree {
        self.merkle_tree
            .get_or_init(|| {
                let mut digests = self
                    .mast_sequences()
                    .iter()
                    .map(|seq| Tip5::hash_varlen(seq))
                    .collect_vec();

                // pad until length is a power of two
                while digests.len() & (digests.len() - 1) != 0 {
                    digests.push(Digest::default());
                }

                twenty_first::prelude::MerkleTree::par_new(&digests).unwrap()
            })
            .clone()
    }
}

#[cfg(any(test, feature = "arbitrary-impls"))]
impl<'a> arbitrary::Arbitrary<'a> for BlockBody {
    fn arbitrary(u: &mut arbitrary::Unstructured<'a>) -> arbitrary::Result<Self> {
        Ok(Self {
            transaction_kernel: u.arbitrary()?,
            mutator_set_accumulator: u.arbitrary()?,
            lock_free_mmr_accumulator: u.arbitrary()?,
            block_mmr_accumulator: u.arbitrary()?,
            merkle_tree: OnceLock::new(), // always empty in fuzzing
        })
    }
}

#[cfg(any(test, feature = "test-helpers"))]
mod test_support {
    use neptune_mutator_set::removal_record::removal_record_list::RemovalRecordList;
    use proptest::prelude::BoxedStrategy;
    use proptest::strategy::Strategy;
    use proptest_arbitrary_interop::arb;

    use super::*;
    use crate::transaction::transaction_kernel::TransactionKernelModifier;

    impl rand::distr::Distribution<BlockBody> for rand::distr::StandardUniform {
        fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> BlockBody {
            let random_mmr_accumulator = |seed| {
                use rand::rngs::StdRng;
                use rand::Rng;
                use rand::RngCore;
                use rand::SeedableRng;

                let mut inner_rng = StdRng::from_seed(seed);
                let leaf_count = inner_rng.next_u64();
                let num_peaks = leaf_count.count_ones();
                MmrAccumulator::init(
                    (0..num_peaks).map(|_| inner_rng.random()).collect_vec(),
                    leaf_count,
                )
            };

            BlockBody {
                transaction_kernel: rng.random(),
                mutator_set_accumulator: rng.random(),
                lock_free_mmr_accumulator: random_mmr_accumulator(rng.random()),
                block_mmr_accumulator: random_mmr_accumulator(rng.random()),
                merkle_tree: OnceLock::new(),
            }
        }
    }

    #[cfg(any(test, feature = "test-helpers"))]
    impl BlockBody {
        pub fn arbitrary_with_mutator_set_accumulator(
            mutator_set_accumulator: MutatorSetAccumulator,
        ) -> BoxedStrategy<BlockBody> {
            (NativeCurrencyAmount::arbitrary_non_negative())
                .prop_flat_map(move |fee| {
                    let transaction_kernel_strategy = TransactionKernel::strategy_with_fee(fee);
                    let lock_free_mmr_accumulator_strategy = arb::<MmrAccumulator>();
                    let block_mmr_accumulator_strategy = arb::<MmrAccumulator>();
                    let mutator_set_accumulator = mutator_set_accumulator.clone();
                    (
                        transaction_kernel_strategy,
                        lock_free_mmr_accumulator_strategy,
                        block_mmr_accumulator_strategy,
                    )
                        .prop_map(
                            move |(
                                transaction_kernel,
                                lock_free_mmr_accumulator,
                                block_mmr_accumulator,
                            )| {
                                let inputs =
                                    RemovalRecordList::pack(transaction_kernel.inputs.clone());
                                let transaction_kernel = TransactionKernelModifier::default()
                                    .inputs(inputs)
                                    .modify(transaction_kernel);
                                BlockBody {
                                    transaction_kernel,
                                    mutator_set_accumulator: mutator_set_accumulator.clone(),
                                    lock_free_mmr_accumulator,
                                    block_mmr_accumulator,
                                    merkle_tree: OnceLock::default(),
                                }
                            },
                        )
                })
                .boxed()
        }
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
    use neptune_mutator_set::msa_and_records::MsaAndRecords;
    use neptune_primitives::network::Network;
    use num_traits::Zero;
    use proptest::prop_assert_eq;
    use test_strategy::proptest;

    use super::*;
    use crate::block::test_helpers::invalid_empty_block;
    use crate::block::Block;

    #[test]
    fn max_aocl_leaf_index_genesis() {
        let genesis = Block::genesis(Network::Main);
        assert_eq!(
            genesis.body().mutator_set_accumulator.aocl.num_leafs() - 1,
            genesis.body().max_aocl_leaf_index(),
            "Genesis block does not contain guesser rewards"
        );
    }

    #[proptest(cases = 4)]
    fn aocl_leaf_count_and_index_consistency(
        #[strategy(0..=5_000_000_000u64)] _num_leafs_aocl: u64,
        #[strategy(MsaAndRecords::arbitrary_with((vec![], #_num_leafs_aocl)))]
        _msa_and_records: MsaAndRecords,
        #[strategy(BlockBody::arbitrary_with_mutator_set_accumulator(#_msa_and_records.mutator_set_accumulator))]
        body: BlockBody,
    ) {
        let guesser_outputs =
            [AdditionRecord::new(Digest::default()); NUM_GUESSER_FEE_OUTPUTS as usize];
        let expected_max_index = body
            .mutator_set_accumulator_after(guesser_outputs.to_vec())
            .aocl
            .num_leafs()
            - 1;
        prop_assert_eq!(expected_max_index, body.max_aocl_leaf_index());

        let num_outputs: u64 = body.transaction_kernel.outputs.len().try_into().unwrap();
        let num_outputs = num_outputs + NUM_GUESSER_FEE_OUTPUTS;
        prop_assert_eq!(
            expected_max_index + 1,
            body.num_aocl_leafs_prior() + num_outputs
        );
    }

    #[test]
    fn num_aocl_leafs_prior_genesis() {
        assert!(Block::genesis(Network::Main)
            .kernel
            .body
            .num_aocl_leafs_prior()
            .is_zero());
    }

    #[test]
    fn num_aocl_leafs_prior_block1() {
        let network = Network::Main;
        let num_premine_outputs: u64 = Block::premine_distribution().len().try_into().unwrap();
        assert_eq!(
            num_premine_outputs,
            invalid_empty_block(&Block::genesis(network), network)
                .kernel
                .body
                .num_aocl_leafs_prior()
        );
    }
}