[#3] Add applyMergeBy

apply-merge.cabal

@@ -45,6 +45,7 @@ library
       base >=4.16 && <4.17 || >=4.17 && <4.18 || >=4.18 && <4.19 || >=4.19 && <4.20
     , containers ==0.6.*
     , pqueue >=1.4 && <1.5 || >=1.5 && <1.6
+    , reflection ==2.1.*
   default-language: GHC2021
 
 test-suite apply-merge-tests
@@ -73,6 +74,7 @@ test-suite apply-merge-tests
       base >=4.16 && <4.17 || >=4.17 && <4.18 || >=4.18 && <4.19 || >=4.19 && <4.20
     , containers ==0.6.*
     , pqueue >=1.4 && <1.5 || >=1.5 && <1.6
+    , reflection ==2.1.*
     , tasty >=1.4 && <1.5 || >=1.5 && <1.6
     , tasty-expected-failure ==0.12.*
     , tasty-hunit ==0.10.*
@@ -103,6 +105,7 @@ benchmark apply-merge-benchmarks
       base >=4.16 && <4.17 || >=4.17 && <4.18 || >=4.18 && <4.19 || >=4.19 && <4.20
     , containers ==0.6.*
     , pqueue >=1.4 && <1.5 || >=1.5 && <1.6
+    , reflection ==2.1.*
     , tasty-bench ==0.3.*
     , transformers >=0.5 && <0.6 || >=0.6 && <0.7
     , vector >=0.12 && <0.13 || >=0.13 && <0.14

package.yaml

@@ -44,6 +44,7 @@ library:
   dependencies:
   - containers ^>= {0.6}
   - pqueue ^>= {1.4, 1.5}
+  - reflection ^>= {2.1}
 
 # The library implements apply-merge in a few different ways, with slightly
 # different performance characteristics, but only exports the current best
@@ -64,6 +65,7 @@ tests:
     dependencies:
     - containers ^>= {0.6}
     - pqueue ^>= {1.4, 1.5}
+    - reflection ^>= {2.1}
     - tasty ^>= {1.4, 1.5}
     - tasty-expected-failure ^>= {0.12}
     - tasty-hunit ^>= {0.10}
@@ -80,6 +82,7 @@ benchmarks:
     dependencies:
     - containers ^>= {0.6}
     - pqueue ^>= {1.4, 1.5}
+    - reflection ^>= {2.1}
     - tasty-bench ^>= {0.3}
     - transformers ^>= {0.5, 0.6}
     - vector ^>= {0.12, 0.13}

src/Data/List/ApplyMerge.hs

@@ -1,14 +1,17 @@
 -- SPDX-FileCopyrightText: Copyright Preetham Gujjula
 -- SPDX-License-Identifier: BSD-3-Clause
+{-# LANGUAGE UndecidableInstances #-}
 
 -- |
 -- Module: Data.List.ApplyMerge
 -- License: BSD-3-Clause
 -- Maintainer: Preetham Gujjula <[email protected]>
 -- Stability: experimental
-module Data.List.ApplyMerge (applyMerge, applyMergeOn) where
+module Data.List.ApplyMerge (applyMerge, applyMergeBy, applyMergeOn) where
 
 import ApplyMerge.IntSet qualified
+import Data.Proxy (Proxy (..))
+import Data.Reflection (Reifies, reflect, reify)
 import Data.Semigroup (Arg (..))
 
 -- | If given a binary function @f@ that is non-decreasing in both arguments,
@@ -31,6 +34,35 @@ import Data.Semigroup (Arg (..))
 applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
 applyMerge = ApplyMerge.IntSet.applyMerge
 
+-- | Like 'applyMerge', but uses a custom comparison function.
+applyMergeBy :: (c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
+applyMergeBy = applyMergeBy_
+
+-- Reflection logic in applyMerge_ is based on "All about reflection: a
+-- tutorial" [1] by Arnaud Spiwack, licensed under CC BY 4.0 [2].
+--
+-- [1]: https://www.tweag.io/blog/2017-12-21-reflection-tutorial/
+-- [2]: https://creativecommons.org/licenses/by/4.0/
+applyMergeBy_ ::
+  forall a b c. (c -> c -> Ordering) -> (a -> b -> c) -> [a] -> [b] -> [c]
+applyMergeBy_ cmp f as bs =
+  reify cmp $ \(_ :: Proxy s) ->
+    let f' :: a -> b -> ReflectedOrd s c
+        f' a b = ReflectedOrd (f a b)
+     in map unReflectedOrd (applyMerge f' as bs)
+
+newtype ReflectedOrd s a = ReflectedOrd {unReflectedOrd :: a}
+
+instance (Reifies s (a -> a -> Ordering)) => Eq (ReflectedOrd s a) where
+  (==) (ReflectedOrd x) (ReflectedOrd y) =
+    let cmp = reflect (Proxy :: Proxy s)
+     in cmp x y == EQ
+
+instance (Reifies s (a -> a -> Ordering)) => Ord (ReflectedOrd s a) where
+  compare (ReflectedOrd x) (ReflectedOrd y) =
+    let cmp = reflect (Proxy :: Proxy s)
+     in cmp x y
+
 -- | Like 'applyMerge', but applies a custom projection function before
 --   performing comparisons.
 --

src/Data/List/NonEmpty/ApplyMerge.hs

@@ -1,23 +1,70 @@
 -- SPDX-FileCopyrightText: Copyright Preetham Gujjula
 -- SPDX-License-Identifier: BSD-3-Clause
+{-# LANGUAGE UndecidableInstances #-}
 
 -- |
 -- Module: Data.List.NonEmpty.ApplyMerge
 -- License: BSD-3-Clause
 -- Maintainer: Preetham Gujjula <[email protected]>
 -- Stability: experimental
-module Data.List.NonEmpty.ApplyMerge (applyMerge, applyMergeOn) where
+module Data.List.NonEmpty.ApplyMerge
+  ( applyMerge,
+    applyMergeBy,
+    applyMergeOn,
+  )
+where
 
 import ApplyMerge.IntSet qualified
 import Data.List.NonEmpty (NonEmpty)
 import Data.List.NonEmpty qualified as NonEmpty
+import Data.Proxy (Proxy (..))
+import Data.Reflection (Reifies, reflect, reify)
 import Data.Semigroup (Arg (..))
 
 -- | Like 'Data.List.ApplyMerge.applyMerge', but operates on 'NonEmpty's instead
 --   of lists.
 applyMerge :: (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
 applyMerge = ApplyMerge.IntSet.applyMergeNonEmpty
 
+-- | Like 'applyMerge', but uses a custom comparison function.
+applyMergeBy ::
+  (c -> c -> Ordering) ->
+  (a -> b -> c) ->
+  NonEmpty a ->
+  NonEmpty b ->
+  NonEmpty c
+applyMergeBy = applyMergeBy_
+
+-- Reflection logic in applyMerge_ is based on "All about reflection: a
+-- tutorial" [1] by Arnaud Spiwack, licensed under CC BY 4.0 [2].
+--
+-- [1]: https://www.tweag.io/blog/2017-12-21-reflection-tutorial/
+-- [2]: https://creativecommons.org/licenses/by/4.0/
+applyMergeBy_ ::
+  forall a b c.
+  (c -> c -> Ordering) ->
+  (a -> b -> c) ->
+  NonEmpty a ->
+  NonEmpty b ->
+  NonEmpty c
+applyMergeBy_ cmp f as bs =
+  reify cmp $ \(_ :: Proxy s) ->
+    let f' :: a -> b -> ReflectedOrd s c
+        f' a b = ReflectedOrd (f a b)
+     in NonEmpty.map unReflectedOrd (applyMerge f' as bs)
+
+newtype ReflectedOrd s a = ReflectedOrd {unReflectedOrd :: a}
+
+instance (Reifies s (a -> a -> Ordering)) => Eq (ReflectedOrd s a) where
+  (==) (ReflectedOrd x) (ReflectedOrd y) =
+    let cmp = reflect (Proxy :: Proxy s)
+     in cmp x y == EQ
+
+instance (Reifies s (a -> a -> Ordering)) => Ord (ReflectedOrd s a) where
+  compare (ReflectedOrd x) (ReflectedOrd y) =
+    let cmp = reflect (Proxy :: Proxy s)
+     in cmp x y
+
 -- | Like 'applyMerge', but applies a custom projection function before
 --   performing comparisons.
 applyMergeOn ::

test/Test/Data/List/ApplyMerge.hs

@@ -9,7 +9,7 @@ import Control.Applicative (liftA2)
 #endif
 import Data.Complex (Complex ((:+)))
 import Data.List (sortOn)
-import Data.List.ApplyMerge (applyMergeOn)
+import Data.List.ApplyMerge (applyMergeBy, applyMergeOn)
 import Data.Ratio ((%))
 import Test.ApplyMerge.Common (basicTest, blockTest, maxTest, skewedTest)
 import Test.Tasty (TestTree, testGroup)
@@ -18,12 +18,22 @@ 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
+    "Data.List.ApplyMerge"
+    [ testGroup
+        "applyMergeOn"
+        [ basicTest (applyMergeOn id),
+          skewedTest (applyMergeOn id),
+          blockTest (applyMergeOn id),
+          maxTest (applyMergeOn id),
+          gaussianIntegerTest
+        ],
+      testGroup
+        "applyMergeBy"
+        [ basicTest (applyMergeBy compare),
+          skewedTest (applyMergeBy compare),
+          blockTest (applyMergeBy compare),
+          maxTest (applyMergeBy compare)
+        ]
     ]
 
 gaussianIntegerTest :: TestTree