From 1904869b5a0ec86751196c38387f48cb2cc22b9f Mon Sep 17 00:00:00 2001
From: Preetham Gujjula <gitcommit@mail.preetham.io>
Date: Thu, 9 May 2024 01:12:20 -0700
Subject: [PATCH] Add Data.List.NonEmpty.ApplyMerge

---
 apply-merge.cabal                    |  3 +++
 package.yaml                         |  1 +
 src/ApplyMerge/DoublyLinkedList.hs   |  2 +-
 src/ApplyMerge/IntMap.hs             |  2 +-
 src/ApplyMerge/IntSet.hs             |  2 +-
 src/Data/List/NonEmpty/ApplyMerge.hs | 29 ++++++++++++++++++++++++++++
 6 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 src/Data/List/NonEmpty/ApplyMerge.hs

diff --git a/apply-merge.cabal b/apply-merge.cabal
index c1f6c4e..492c71d 100644
--- a/apply-merge.cabal
+++ b/apply-merge.cabal
@@ -35,6 +35,7 @@ source-repository head
 library
   exposed-modules:
       Data.List.ApplyMerge
+      Data.List.NonEmpty.ApplyMerge
   other-modules:
       ApplyMerge.IntSet
   hs-source-dirs:
@@ -55,6 +56,7 @@ test-suite apply-merge-tests
       ApplyMerge.IntSet
       Data.DoublyLinkedList.STRef
       Data.List.ApplyMerge
+      Data.List.NonEmpty.ApplyMerge
       Data.PQueue.Prio.Min.Mutable
       Test.ApplyMerge.Common
       Test.ApplyMerge.DoublyLinkedList
@@ -88,6 +90,7 @@ benchmark apply-merge-benchmarks
       ApplyMerge.IntSet
       Data.DoublyLinkedList.STRef
       Data.List.ApplyMerge
+      Data.List.NonEmpty.ApplyMerge
       Data.PQueue.Prio.Min.Mutable
       Bench.Data.DoublyLinkedList.STRef
       Bench.PriorityQueue.MinPQueue
diff --git a/package.yaml b/package.yaml
index 6aae25c..99aaf1d 100644
--- a/package.yaml
+++ b/package.yaml
@@ -38,6 +38,7 @@ library:
   source-dirs: src
   exposed-modules:
   - Data.List.ApplyMerge
+  - Data.List.NonEmpty.ApplyMerge
   other-modules:
   - ApplyMerge.IntSet
   dependencies:
diff --git a/src/ApplyMerge/DoublyLinkedList.hs b/src/ApplyMerge/DoublyLinkedList.hs
index 8ac8758..14fed78 100644
--- a/src/ApplyMerge/DoublyLinkedList.hs
+++ b/src/ApplyMerge/DoublyLinkedList.hs
@@ -3,7 +3,7 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
-module ApplyMerge.DoublyLinkedList (applyMerge) where
+module ApplyMerge.DoublyLinkedList (applyMerge, applyMergeNonEmpty) where
 
 import Control.Monad (guard, (>=>))
 import Control.Monad.ST qualified as Strict
diff --git a/src/ApplyMerge/IntMap.hs b/src/ApplyMerge/IntMap.hs
index 384d4eb..a328258 100644
--- a/src/ApplyMerge/IntMap.hs
+++ b/src/ApplyMerge/IntMap.hs
@@ -3,7 +3,7 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
-module ApplyMerge.IntMap (applyMerge) where
+module ApplyMerge.IntMap (applyMerge, applyMergeNonEmpty) where
 
 import Control.Arrow ((>>>))
 import Control.Monad (guard)
diff --git a/src/ApplyMerge/IntSet.hs b/src/ApplyMerge/IntSet.hs
index 2fde91a..b63bee2 100644
--- a/src/ApplyMerge/IntSet.hs
+++ b/src/ApplyMerge/IntSet.hs
@@ -3,7 +3,7 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
-module ApplyMerge.IntSet (applyMerge) where
+module ApplyMerge.IntSet (applyMerge, applyMergeNonEmpty) where
 
 import Control.Arrow ((>>>))
 import Control.Monad (guard)
diff --git a/src/Data/List/NonEmpty/ApplyMerge.hs b/src/Data/List/NonEmpty/ApplyMerge.hs
new file mode 100644
index 0000000..41fe838
--- /dev/null
+++ b/src/Data/List/NonEmpty/ApplyMerge.hs
@@ -0,0 +1,29 @@
+-- SPDX-FileCopyrightText: Copyright Preetham Gujjula
+-- SPDX-License-Identifier: BSD-3-Clause
+
+-- |
+-- Module: Data.List.NonEmpty.ApplyMerge
+-- License: BSD-3-Clause
+-- Maintainer: Preetham Gujjula <libraries@mail.preetham.io>
+-- Stability: experimental
+module Data.List.NonEmpty.ApplyMerge (applyMerge, applyMergeOn) where
+
+import ApplyMerge.IntSet qualified
+import Data.List.NonEmpty (NonEmpty)
+import Data.List.NonEmpty qualified as NonEmpty
+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 applies a custom projection function before
+--   performing comparisons.
+applyMergeOn ::
+  (Ord d) => (c -> d) -> (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
+applyMergeOn p f as bs =
+  let f' a b =
+        let c = f a b
+         in Arg (p c) c
+   in NonEmpty.map (\(Arg _ c) -> c) (applyMerge f' as bs)