zebra_chain/transparent/
arbitrary.rs

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