1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

#![cfg(feature = "arb")]

use std::convert::TryFrom;

use multihash::{Code, MultihashDigest, MultihashGeneric};
use quickcheck::Gen;
use rand::{
    distributions::{weighted::WeightedIndex, Distribution},
    Rng,
};

use arbitrary::{size_hint, Unstructured};
use rand::SeedableRng;

use crate::{CidGeneric, Version};

impl quickcheck::Arbitrary for Version {
    fn arbitrary(g: &mut Gen) -> Self {
        let version = u64::from(bool::arbitrary(g));
        Version::try_from(version).unwrap()
    }
}

impl<const S: usize> quickcheck::Arbitrary for CidGeneric<S> {
    fn arbitrary(g: &mut Gen) -> Self {
        if S >= 32 && Version::arbitrary(g) == Version::V0 {
            let data: Vec<u8> = Vec::arbitrary(g);
            let hash = Code::Sha2_256
                .digest(&data)
                .resize()
                .expect("digest too large");
            CidGeneric::new_v0(hash).expect("sha2_256 is a valid hash for cid v0")
        } else {
            // In real world lower IPLD Codec codes more likely to happen, hence distribute them
            // with bias towards smaller values.
            let weights = [128, 32, 4, 4, 2, 2, 1, 1];
            let dist = WeightedIndex::new(weights.iter()).unwrap();
            let mut rng = rand::rngs::SmallRng::seed_from_u64(u64::arbitrary(g));
            let codec = match dist.sample(&mut rng) {
                0 => rng.gen_range(0..u64::pow(2, 7)),
                1 => rng.gen_range(u64::pow(2, 7)..u64::pow(2, 14)),
                2 => rng.gen_range(u64::pow(2, 14)..u64::pow(2, 21)),
                3 => rng.gen_range(u64::pow(2, 21)..u64::pow(2, 28)),
                4 => rng.gen_range(u64::pow(2, 28)..u64::pow(2, 35)),
                5 => rng.gen_range(u64::pow(2, 35)..u64::pow(2, 42)),
                6 => rng.gen_range(u64::pow(2, 42)..u64::pow(2, 49)),
                7 => rng.gen_range(u64::pow(2, 56)..u64::pow(2, 63)),
                _ => unreachable!(),
            };

            let hash: MultihashGeneric<S> = quickcheck::Arbitrary::arbitrary(g);
            CidGeneric::new_v1(codec, hash)
        }
    }
}

impl<'a, const S: usize> arbitrary::Arbitrary<'a> for CidGeneric<S> {
    fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
        if S >= 32 && u.ratio(1, 10)? {
            let mh = MultihashGeneric::wrap(Code::Sha2_256.into(), u.bytes(32)?).unwrap();
            return Ok(CidGeneric::new_v0(mh).expect("32 bytes is correct for v0"));
        }

        let mut codec = 0u64;
        let mut len_choice = u.arbitrary::<u8>()? | 1;

        while len_choice & 1 == 1 {
            len_choice >>= 1;

            let x = u.arbitrary::<u8>();
            let next = codec
                .checked_shl(8)
                .zip(x.ok())
                .map(|(next, x)| next.saturating_add(x as u64));

            match next {
                None => break,
                Some(next) => codec = next,
            }
        }

        Ok(CidGeneric::new_v1(codec, u.arbitrary()?))
    }

    fn size_hint(depth: usize) -> (usize, Option<usize>) {
        let v1 = size_hint::and_all(&[
            <[u8; 2]>::size_hint(depth),
            (0, Some(8)),
            <MultihashGeneric<S> as arbitrary::Arbitrary>::size_hint(depth),
        ]);
        if S >= 32 {
            size_hint::and(<u8>::size_hint(depth), size_hint::or((32, Some(32)), v1))
        } else {
            v1
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::CidGeneric;
    use arbitrary::{Arbitrary, Unstructured};
    use multihash::MultihashGeneric;

    #[test]
    fn arbitrary() {
        let mut u = Unstructured::new(&[
            1, 22, 41, 13, 5, 6, 7, 8, 9, 6, 10, 243, 43, 231, 123, 43, 153, 127, 67, 76, 24, 91,
            23, 32, 32, 23, 65, 98, 193, 108, 3,
        ]);
        let c = <CidGeneric<16> as Arbitrary>::arbitrary(&mut u).unwrap();
        let c2 =
            CidGeneric::<16>::new_v1(22, MultihashGeneric::wrap(13, &[6, 7, 8, 9, 6]).unwrap());
        assert_eq!(c.hash(), c2.hash());
        assert_eq!(c.codec(), c2.codec());
        assert_eq!(c, c2)
    }
}