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