add instance for NonEmpty

nick8325 · Mar 21, 2024 · cbe44c4 · cbe44c4
1 parent f1e767c
commit cbe44c4
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/src/Test/QuickCheck/Arbitrary.hs b/src/Test/QuickCheck/Arbitrary.hs
@@ -147,6 +147,12 @@ import System.IO
 #endif
 #endif
 
+#if defined(MIN_VERSION_base)
+#if MIN_VERSION_base(4,9,0)
+import Data.List.NonEmpty (NonEmpty(..))
+#endif
+#endif
+
 import Control.Monad
   ( liftM
   , liftM2
@@ -1073,6 +1079,17 @@ instance Arbitrary a => Arbitrary (Down a) where
 #endif
 #endif
 
+#if defined(MIN_VERSION_base)
+#if MIN_VERSION_base(4,9,0)
+instance Arbitrary a => Arbitrary (NonEmpty a) where
+  arbitrary = (:|) <$> arbitrary <*> arbitrary
+  shrink (a :| bs) =
+    [ a' :| bs | a' <- shrink a ] ++
+    [ b :| bs' | b : bs' <- [bs] ] ++
+    [ a :| bs' | bs' <- shrink bs ]
+#endif
+#endif
+
 -- ** Helper functions for implementing arbitrary
 
 -- | Apply a binary function to random arguments.