From 9efa9745a22382af89ee4d9d38908c1795834cf6 Mon Sep 17 00:00:00 2001
From: Preetham Gujjula <gitcommit@mail.preetham.io>
Date: Fri, 13 Dec 2024 08:38:16 -0800
Subject: [PATCH] wip: Add more tests

---
 test/Test/Data/List/ApplyMerge/New.hs | 70 +++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/test/Test/Data/List/ApplyMerge/New.hs b/test/Test/Data/List/ApplyMerge/New.hs
index 9dc1d19..2c79eff 100644
--- a/test/Test/Data/List/ApplyMerge/New.hs
+++ b/test/Test/Data/List/ApplyMerge/New.hs
@@ -171,6 +171,76 @@ testNonEmptyFunctions label am funcs op =
             expected = NE.sort $ on (liftA2 f) (take1 limit) xs ys
          in on (===) (NE.take limit) actual expected
 
+data family Ordered (f :: Type -> Type)
+
+data instance Ordered []
+  = OrderedList
+  { olOrigin :: Integer,
+    olSequence :: Either [Natural] (InfiniteList Natural)
+  }
+  deriving (Generic)
+
+instance Arbitrary (Ordered []) where
+  arbitrary = liftA2 OrderedList arbitrary arbitrary
+  shrink = genericShrink
+
+data instance Ordered NonEmpty
+  = OrderedNonEmpty
+  { onOrigin :: Integer,
+    onSequence :: (Natural, Either [Natural] (InfiniteList Natural))
+  }
+  deriving (Generic)
+
+instance Arbitrary (Ordered NonEmpty) where
+  arbitrary = liftA2 OrderedNonEmpty arbitrary arbitrary
+  shrink = genericShrink
+
+class (Functor f) => ApplyMergeable (f :: Type -> Type) where
+  getIncreasing :: Ordered f -> f Natural
+  getDecreasing :: Ordered f -> f Integer
+  toList :: f a -> [a]
+
+instance ApplyMergeable [] where
+  getIncreasing =
+    olSequence
+      >>> either id getInfiniteList
+      >>> scanl1 (+)
+
+  getDecreasing ol =
+    let x0 = olOrigin ol
+        xs = map (negate . toInteger) (getIncreasing ol)
+     in map (+ x0) xs
+
+  toList = id
+
+instance ApplyMergeable NonEmpty where
+  getIncreasing =
+    onSequence
+      >>> second (either id getInfiniteList)
+      >>> uncurry (:|)
+      >>> NonEmpty.scanl1 (+)
+
+  getDecreasing one =
+    let x0 = onOrigin one
+        xs = fmap (negate . toInteger) (getIncreasing one)
+     in fmap (+ x0) xs
+
+  toList = NonEmpty.toList
+
+testGenericFunctions :: (Arbitrary (Ordered f), Show (Ordered f)) => (ApplyMergeable f) => TestFunctions f
+testGenericFunctions label am _ _ =
+  QC.testProperty label $ do
+    (fName, f) <- QC.elements increasingNaturalFuncs
+    let limit = 100
+    pure . QC.counterexample fName $
+      \(getIncreasing -> xs) (getIncreasing -> ys) ->
+        let actual :: [Natural]
+            actual = toList (am f xs ys)
+
+            expected :: [Natural]
+            expected = sort $ on (liftA2 f) (take limit) (toList xs) (toList ys)
+         in on (===) (take limit) actual expected
+
 -- Utilities
 take1 :: Int -> NonEmpty a -> NonEmpty a
 take1 n (x :| xs) = x :| take n xs