Add applyMergeOn, with tests

apply-merge.cabal

@@ -61,6 +61,7 @@ test-suite apply-merge-tests
       Test.ApplyMerge.IntMap
       Test.ApplyMerge.IntSet
       Test.Data.DoublyLinkedList.STRef
+      Test.Data.List.ApplyMerge
       Test.Data.PQueue.Prio.Min.Mutable
   hs-source-dirs:
       src

src/Data/List/ApplyMerge.hs

@@ -6,9 +6,10 @@
 -- License: BSD-3-Clause
 -- Maintainer: Preetham Gujjula <[email protected]>
 -- Stability: experimental
-module Data.List.ApplyMerge (applyMerge) where
+module Data.List.ApplyMerge (applyMerge, applyMergeOn) where
 
 import ApplyMerge.IntSet qualified
+import Data.Semigroup (Arg (..))
 
 -- | If given a binary function @f@ that is non-decreasing in both arguments,
 --   and two (potentially infinite) ordered lists @xs@ and @ys@, then
@@ -29,3 +30,21 @@ import ApplyMerge.IntSet qualified
 --   [README#examples](https://github.com/pgujjula/apply-merge/#examples).
 applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
 applyMerge = ApplyMerge.IntSet.applyMerge
+
+-- | Like 'applyMerge', but applies a custom projection function before
+--   performing comparisons.
+--
+--   For example, to compute the Gaussian integers, ordered by norm:
+--
+--   > zs :: [Integer]
+--   > zs = 0 : concatMap (\i -> [i, -i]) [1..]
+--   >
+--   > gaussianIntegers :: [GaussianInteger]      -- `GaussianInteger` from arithmoi
+--   > gaussianIntegers = applyMergeOn norm (:+) zs zs
+applyMergeOn ::
+  (Ord d) => (c -> d) -> (a -> b -> c) -> [a] -> [b] -> [c]
+applyMergeOn p f as bs =
+  let f' a b =
+        let c = f a b
+         in Arg (p c) c
+   in map (\(Arg _ c) -> c) (applyMerge f' as bs)

test/Main.hs

@@ -7,6 +7,7 @@ import Test.ApplyMerge.DoublyLinkedList qualified (tests)
 import Test.ApplyMerge.IntMap qualified (tests)
 import Test.ApplyMerge.IntSet qualified (tests)
 import Test.Data.DoublyLinkedList.STRef qualified (tests)
+import Test.Data.List.ApplyMerge qualified (tests)
 import Test.Data.PQueue.Prio.Min.Mutable qualified (tests)
 import Test.Tasty (TestTree, defaultMain, testGroup)
 
@@ -20,6 +21,7 @@ tests =
     [ Test.ApplyMerge.DoublyLinkedList.tests,
       Test.ApplyMerge.IntMap.tests,
       Test.ApplyMerge.IntSet.tests,
+      Test.Data.List.ApplyMerge.tests,
       Test.Data.DoublyLinkedList.STRef.tests,
       Test.Data.PQueue.Prio.Min.Mutable.tests
     ]

test/Test/Data/List/ApplyMerge.hs

@@ -0,0 +1,57 @@
+-- SPDX-FileCopyrightText: Copyright Preetham Gujjula
+-- SPDX-License-Identifier: BSD-3-Clause
+{-# LANGUAGE CPP #-}
+
+module Test.Data.List.ApplyMerge (tests) where
+
+#if !MIN_VERSION_base(4,18,0)
+import Control.Applicative (liftA2)
+#endif
+import Data.Complex (Complex ((:+)))
+import Data.List (sortOn)
+import Data.List.ApplyMerge (applyMergeOn)
+import Data.Ratio ((%))
+import Test.ApplyMerge.Common (basicTest, blockTest, maxTest, skewedTest)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (testCase, (@?=))
+
+tests :: TestTree
+tests =
+  testGroup
+    "Data.List.ApplyMerge.applyMergeOn"
+    [ basicTest (applyMergeOn id),
+      skewedTest (applyMergeOn id),
+      blockTest (applyMergeOn id),
+      maxTest (applyMergeOn id),
+      gaussianIntegerTest
+    ]
+
+gaussianIntegerTest :: TestTree
+gaussianIntegerTest =
+  testCase "gaussian integers x + yi, with 0 <= x <= y, ordered by norm" $ do
+    let actual =
+          take 100 $
+            applyMergeOn
+              (\x -> (norm x, slope x))
+              (\x k -> x :+ (x + k))
+              [0 ..]
+              [0 ..]
+        expected =
+          take 100 $
+            filter (\(x :+ y) -> x <= y) $
+              sortOn
+                (\x -> (norm x, slope x))
+                (liftA2 (:+) [0 .. 100] [0 .. 100])
+     in actual @?= expected
+
+norm :: Complex Integer -> Integer
+norm (a :+ b) = a * a + b * b
+
+data Slope = Finite Rational | Infinity
+  deriving (Eq, Show, Ord)
+
+slope :: Complex Integer -> Slope
+slope (x :+ y) =
+  if x == 0
+    then Infinity
+    else Finite (y % x)