1 // This Source Code Form is subject to the terms of the Mozilla Public
2 // License, v. 2.0. If a copy of the MPL was not distributed with this
3 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
4 
5 use std::hash::{BuildHasher, Hash};
6 use std::iter;
7 
8 use ::arbitrary::{size_hint, Arbitrary, Result, Unstructured};
9 
10 use crate::{HashMap, HashSet, OrdMap, OrdSet, Vector};
11 
empty<T: 'static>() -> Box<dyn Iterator<Item = T>>12 fn empty<T: 'static>() -> Box<dyn Iterator<Item = T>> {
13     Box::new(iter::empty())
14 }
15 
shrink_collection<T: Clone, A: Clone + Arbitrary>( entries: impl Iterator<Item = T>, f: impl Fn(&T) -> Box<dyn Iterator<Item = A>>, ) -> Box<dyn Iterator<Item = Vec<A>>>16 fn shrink_collection<T: Clone, A: Clone + Arbitrary>(
17     entries: impl Iterator<Item = T>,
18     f: impl Fn(&T) -> Box<dyn Iterator<Item = A>>,
19 ) -> Box<dyn Iterator<Item = Vec<A>>> {
20     let entries: Vec<_> = entries.collect();
21     if entries.is_empty() {
22         return empty();
23     }
24 
25     let mut shrinkers: Vec<Vec<_>> = vec![];
26     let mut i = entries.len();
27     loop {
28         shrinkers.push(entries.iter().take(i).map(&f).collect());
29         i /= 2;
30         if i == 0 {
31             break;
32         }
33     }
34     Box::new(iter::once(Vec::new()).chain(iter::from_fn(move || loop {
35         let mut shrinker = shrinkers.pop()?;
36         let x: Option<Vec<A>> = shrinker.iter_mut().map(|s| s.next()).collect();
37         if x.is_none() {
38             continue;
39         }
40         shrinkers.push(shrinker);
41         return x;
42     })))
43 }
44 
45 impl<A: Arbitrary + Clone> Arbitrary for Vector<A> {
arbitrary(u: &mut Unstructured<'_>) -> Result<Self>46     fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
47         u.arbitrary_iter()?.collect()
48     }
49 
arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self>50     fn arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self> {
51         u.arbitrary_take_rest_iter()?.collect()
52     }
53 
size_hint(depth: usize) -> (usize, Option<usize>)54     fn size_hint(depth: usize) -> (usize, Option<usize>) {
55         size_hint::recursion_guard(depth, |depth| {
56             size_hint::and(<usize as Arbitrary>::size_hint(depth), (0, None))
57         })
58     }
59 
shrink(&self) -> Box<dyn Iterator<Item = Self>>60     fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
61         let collections = shrink_collection(self.iter(), |x| x.shrink());
62         Box::new(collections.map(|entries| entries.into_iter().collect()))
63     }
64 }
65 
66 impl<K: Arbitrary + Ord + Clone, V: Arbitrary + Clone> Arbitrary for OrdMap<K, V> {
arbitrary(u: &mut Unstructured<'_>) -> Result<Self>67     fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
68         u.arbitrary_iter()?.collect()
69     }
70 
arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self>71     fn arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self> {
72         u.arbitrary_take_rest_iter()?.collect()
73     }
74 
size_hint(depth: usize) -> (usize, Option<usize>)75     fn size_hint(depth: usize) -> (usize, Option<usize>) {
76         size_hint::recursion_guard(depth, |depth| {
77             size_hint::and(<usize as Arbitrary>::size_hint(depth), (0, None))
78         })
79     }
80 
shrink(&self) -> Box<dyn Iterator<Item = Self>>81     fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
82         let collections =
83             shrink_collection(self.iter(), |(k, v)| Box::new(k.shrink().zip(v.shrink())));
84         Box::new(collections.map(|entries| entries.into_iter().collect()))
85     }
86 }
87 
88 impl<A: Arbitrary + Ord + Clone> Arbitrary for OrdSet<A> {
arbitrary(u: &mut Unstructured<'_>) -> Result<Self>89     fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
90         u.arbitrary_iter()?.collect()
91     }
92 
arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self>93     fn arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self> {
94         u.arbitrary_take_rest_iter()?.collect()
95     }
96 
size_hint(depth: usize) -> (usize, Option<usize>)97     fn size_hint(depth: usize) -> (usize, Option<usize>) {
98         size_hint::recursion_guard(depth, |depth| {
99             size_hint::and(<usize as Arbitrary>::size_hint(depth), (0, None))
100         })
101     }
102 
shrink(&self) -> Box<dyn Iterator<Item = Self>>103     fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
104         let collections = shrink_collection(self.iter(), |v| v.shrink());
105         Box::new(collections.map(|entries| entries.into_iter().collect()))
106     }
107 }
108 
109 impl<K, V, S> Arbitrary for HashMap<K, V, S>
110 where
111     K: Arbitrary + Hash + Eq + Clone,
112     V: Arbitrary + Clone,
113     S: BuildHasher + Default + 'static,
114 {
arbitrary(u: &mut Unstructured<'_>) -> Result<Self>115     fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
116         u.arbitrary_iter()?.collect()
117     }
118 
arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self>119     fn arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self> {
120         u.arbitrary_take_rest_iter()?.collect()
121     }
122 
size_hint(depth: usize) -> (usize, Option<usize>)123     fn size_hint(depth: usize) -> (usize, Option<usize>) {
124         size_hint::recursion_guard(depth, |depth| {
125             size_hint::and(<usize as Arbitrary>::size_hint(depth), (0, None))
126         })
127     }
128 
shrink(&self) -> Box<dyn Iterator<Item = Self>>129     fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
130         let collections =
131             shrink_collection(self.iter(), |(k, v)| Box::new(k.shrink().zip(v.shrink())));
132         Box::new(collections.map(|entries| entries.into_iter().collect()))
133     }
134 }
135 
136 impl<A, S> Arbitrary for HashSet<A, S>
137 where
138     A: Arbitrary + Hash + Eq + Clone,
139     S: BuildHasher + Default + 'static,
140 {
arbitrary(u: &mut Unstructured<'_>) -> Result<Self>141     fn arbitrary(u: &mut Unstructured<'_>) -> Result<Self> {
142         u.arbitrary_iter()?.collect()
143     }
144 
arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self>145     fn arbitrary_take_rest(u: Unstructured<'_>) -> Result<Self> {
146         u.arbitrary_take_rest_iter()?.collect()
147     }
148 
size_hint(depth: usize) -> (usize, Option<usize>)149     fn size_hint(depth: usize) -> (usize, Option<usize>) {
150         size_hint::recursion_guard(depth, |depth| {
151             size_hint::and(<usize as Arbitrary>::size_hint(depth), (0, None))
152         })
153     }
154 
shrink(&self) -> Box<dyn Iterator<Item = Self>>155     fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
156         let collections = shrink_collection(self.iter(), |v| v.shrink());
157         Box::new(collections.map(|entries| entries.into_iter().collect()))
158     }
159 }
160