zebra_chain/transparent/
arbitrary.rs

1use proptest::{collection::vec, prelude::*};
2
3use crate::{
4    block,
5    parameters::NetworkKind,
6    serialization::TrustedPreallocate,
7    transparent::{Output, OutputIndex},
8    LedgerState,
9};
10
11use super::{Address, CoinbaseData, Input, OutPoint, Script, GENESIS_COINBASE_DATA};
12
13impl Input {
14    /// Construct a strategy for creating valid-ish vecs of Inputs.
15    pub fn vec_strategy(ledger_state: &LedgerState, max_size: usize) -> BoxedStrategy<Vec<Self>> {
16        if ledger_state.has_coinbase {
17            Self::arbitrary_with(Some(ledger_state.height))
18                .prop_map(|input| vec![input])
19                .boxed()
20        } else {
21            vec(Self::arbitrary_with(None), 1..=max_size).boxed()
22        }
23    }
24}
25
26impl Arbitrary for Input {
27    type Parameters = Option<block::Height>;
28
29    fn arbitrary_with(height: Self::Parameters) -> Self::Strategy {
30        if let Some(height) = height {
31            (vec(any::<u8>(), 0..95), any::<u32>())
32                .prop_map(move |(data, sequence)| Input::Coinbase {
33                    height,
34                    data: if height == block::Height(0) {
35                        CoinbaseData(GENESIS_COINBASE_DATA.to_vec())
36                    } else {
37                        CoinbaseData(data)
38                    },
39                    sequence,
40                })
41                .boxed()
42        } else {
43            (any::<OutPoint>(), any::<Script>(), any::<u32>())
44                .prop_map(|(outpoint, unlock_script, sequence)| Input::PrevOut {
45                    outpoint,
46                    unlock_script,
47                    sequence,
48                })
49                .boxed()
50        }
51    }
52
53    type Strategy = BoxedStrategy<Self>;
54}
55
56impl Arbitrary for Address {
57    type Parameters = ();
58
59    fn arbitrary_with(_args: ()) -> Self::Strategy {
60        any::<(bool, bool, [u8; 20])>()
61            .prop_map(|(is_mainnet, is_p2pkh, hash_bytes)| {
62                let network = if is_mainnet {
63                    NetworkKind::Mainnet
64                } else {
65                    NetworkKind::Testnet
66                };
67
68                if is_p2pkh {
69                    Address::from_pub_key_hash(network, hash_bytes)
70                } else {
71                    Address::from_script_hash(network, hash_bytes)
72                }
73            })
74            .boxed()
75    }
76
77    type Strategy = BoxedStrategy<Self>;
78}
79
80impl Arbitrary for OutputIndex {
81    type Parameters = ();
82
83    fn arbitrary_with(_args: ()) -> Self::Strategy {
84        (0..=Output::max_allocation())
85            .prop_map(OutputIndex::from_u64)
86            .boxed()
87    }
88
89    type Strategy = BoxedStrategy<Self>;
90}