From 2251812717e75ba59630494fd76cd979840bcc3d Mon Sep 17 00:00:00 2001
From: Josh Stone <cuviper@gmail.com>
Date: Wed, 16 Nov 2022 18:22:31 -0800
Subject: [PATCH] impl Arbitrary for IndexMap and IndexSet

This implements both `arbitrary::Arbitrary` and `quickcheck::Arbitrary`
behind optional dependencies on their respective crates.

(cherry picked from commit 3e78c62b964b0cb1e19506156511df7384b708f1)
---
 .github/workflows/ci.yml |  6 +++-
 Cargo.toml               |  4 ++-
 src/arbitrary.rs         | 75 ++++++++++++++++++++++++++++++++++++++++
 src/lib.rs               |  1 +
 4 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 src/arbitrary.rs

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 8c6c10cb..dcde8f4b 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -19,11 +19,15 @@ jobs:
           - rust: 1.56.0 # MSRV
             features:
           - rust: stable
-            features: serde
+            features: arbitrary
+          - rust: stable
+            features: quickcheck
           - rust: stable
             features: rayon
           - rust: stable
             features: rustc-rayon
+          - rust: stable
+            features: serde
           - rust: stable
             features: std
           - rust: beta
diff --git a/Cargo.toml b/Cargo.toml
index 1ecd8ba2..7758759f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,6 +17,8 @@ bench = false
 autocfg = "1"
 
 [dependencies]
+arbitrary = { version = "1.0", optional = true, default-features = false }
+quickcheck = { version = "1.0", optional = true, default-features = false }
 serde = { version = "1.0", optional = true, default-features = false }
 rayon = { version = "1.4.1", optional = true }
 
@@ -57,7 +59,7 @@ no-dev-version = true
 tag-name = "{{version}}"
 
 [package.metadata.docs.rs]
-features = ["serde-1", "rayon"]
+features = ["arbitrary", "quickcheck", "serde-1", "rayon"]
 
 [workspace]
 members = ["test-nostd", "test-serde"]
diff --git a/src/arbitrary.rs b/src/arbitrary.rs
new file mode 100644
index 00000000..1347c8b5
--- /dev/null
+++ b/src/arbitrary.rs
@@ -0,0 +1,75 @@
+#[cfg(feature = "arbitrary")]
+mod impl_arbitrary {
+    use crate::{IndexMap, IndexSet};
+    use arbitrary::{Arbitrary, Result, Unstructured};
+    use core::hash::{BuildHasher, Hash};
+
+    impl<'a, K, V, S> Arbitrary<'a> for IndexMap<K, V, S>
+    where
+        K: Arbitrary<'a> + Hash + Eq,
+        V: Arbitrary<'a>,
+        S: BuildHasher + Default,
+    {
+        fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
+            u.arbitrary_iter()?.collect()
+        }
+
+        fn arbitrary_take_rest(u: Unstructured<'a>) -> Result<Self> {
+            u.arbitrary_take_rest_iter()?.collect()
+        }
+    }
+
+    impl<'a, T, S> Arbitrary<'a> for IndexSet<T, S>
+    where
+        T: Arbitrary<'a> + Hash + Eq,
+        S: BuildHasher + Default,
+    {
+        fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
+            u.arbitrary_iter()?.collect()
+        }
+
+        fn arbitrary_take_rest(u: Unstructured<'a>) -> Result<Self> {
+            u.arbitrary_take_rest_iter()?.collect()
+        }
+    }
+}
+
+#[cfg(feature = "quickcheck")]
+mod impl_quickcheck {
+    use crate::{IndexMap, IndexSet};
+    use alloc::boxed::Box;
+    use alloc::vec::Vec;
+    use core::hash::{BuildHasher, Hash};
+    use quickcheck::{Arbitrary, Gen};
+
+    impl<K, V, S> Arbitrary for IndexMap<K, V, S>
+    where
+        K: Arbitrary + Hash + Eq,
+        V: Arbitrary,
+        S: BuildHasher + Default + Clone + 'static,
+    {
+        fn arbitrary(g: &mut Gen) -> Self {
+            Self::from_iter(Vec::arbitrary(g))
+        }
+
+        fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
+            let vec = Vec::from_iter(self.clone());
+            Box::new(vec.shrink().map(Self::from_iter))
+        }
+    }
+
+    impl<T, S> Arbitrary for IndexSet<T, S>
+    where
+        T: Arbitrary + Hash + Eq,
+        S: BuildHasher + Default + Clone + 'static,
+    {
+        fn arbitrary(g: &mut Gen) -> Self {
+            Self::from_iter(Vec::arbitrary(g))
+        }
+
+        fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
+            let vec = Vec::from_iter(self.clone());
+            Box::new(vec.shrink().map(Self::from_iter))
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index ba613342..6e949361 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -85,6 +85,7 @@ extern crate std;
 
 use alloc::vec::{self, Vec};
 
+mod arbitrary;
 #[macro_use]
 mod macros;
 mod equivalent;