Add applyMerge for NonEmpty

.github/workflows/ci.yml

@@ -3,8 +3,10 @@
 
 name: CI
 on:
-  - push
-  - pull_request
+  push:
+    branches:
+    - main
+  pull_request:
 jobs:
   generate-matrix:
     name: Generate matrix from cabal

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

package.yaml

@@ -38,6 +38,7 @@ library:
   source-dirs: src
   exposed-modules:
   - Data.List.ApplyMerge
+  - Data.List.NonEmpty.ApplyMerge
   other-modules:
   - ApplyMerge.IntSet
   dependencies:

src/ApplyMerge/DoublyLinkedList.hs

@@ -3,15 +3,15 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
-module ApplyMerge.DoublyLinkedList (applyMerge) where
+module ApplyMerge.DoublyLinkedList (applyMerge, applyMergeNonEmpty) where
 
-import Control.Monad (guard)
+import Control.Monad (guard, (>=>))
 import Control.Monad.ST qualified as Strict
 import Control.Monad.ST.Lazy qualified as Lazy
 import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
 import Data.DoublyLinkedList.STRef qualified as DoublyLinked
-import Data.List.NonEmpty (NonEmpty, nonEmpty)
+import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty)
 import Data.List.NonEmpty qualified as NonEmpty
 import Data.Maybe (fromMaybe)
 import Data.PQueue.Prio.Min (MinPQueue)
@@ -29,15 +29,16 @@ newtype Frontier s a b c = Frontier
   }
 
 applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
-applyMerge f as bs =
-  fromMaybe [] $
-    applyMergeNonEmpty f <$> nonEmpty as <*> nonEmpty bs
+applyMerge f as bs = fromMaybe [] $ do
+  as' <- nonEmpty as
+  bs' <- nonEmpty bs
+  pure (NonEmpty.toList (applyMergeNonEmpty f as' bs'))
 
 applyMergeNonEmpty ::
-  (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> [c]
+  (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
 applyMergeNonEmpty f as bs = Lazy.runST $ do
-  frontier <- Lazy.strictToLazyST (initialFrontier f as bs)
-  unfoldrM (Lazy.strictToLazyST . step f) frontier
+  (c, frontier) <- Lazy.strictToLazyST (initialState f as bs)
+  (c :|) <$> unfoldrM (Lazy.strictToLazyST . runMaybeT . step f) frontier
 
 unfoldrM :: (Monad m) => (b -> m (Maybe (a, b))) -> b -> m [a]
 unfoldrM f seed = do
@@ -46,29 +47,29 @@ unfoldrM f seed = do
     Nothing -> pure []
     Just (x, newSeed) -> (x :) <$> unfoldrM f newSeed
 
-initialFrontier ::
-  (a -> b -> c) -> NonEmpty a -> NonEmpty b -> Strict.ST s (Frontier s a b c)
-initialFrontier f as bs = do
+initialState ::
+  (Ord c) =>
+  (a -> b -> c) ->
+  NonEmpty a ->
+  NonEmpty b ->
+  Strict.ST s (c, Frontier s a b c)
+initialState f as bs = do
   list <- DoublyLinked.empty
-  position <- DoublyLinked.cons list (0 :: Int, 0 :: Int)
-  let node = mkNode f position as bs
-  pure $ Frontier $ MinPQueue.singleton node.value node
+  rootPosition <- DoublyLinked.cons list (0 :: Int, 0 :: Int)
+  let rootNode = mkNode f rootPosition as bs
+  let emptyFrontier = Frontier MinPQueue.empty
+  peekInsertChildren f rootNode emptyFrontier
 
 step ::
   (Ord c) =>
   (a -> b -> c) ->
   Frontier s a b c ->
-  Strict.ST s (Maybe (c, Frontier s a b c))
-step f frontier = runMaybeT $ do
-  (node, frontier') <- MaybeT (deleteMinNode frontier)
-  frontier'' <- lift $ insertChildA f node frontier'
-  frontier''' <- lift $ insertChildB f node frontier''
-  lift (DoublyLinked.delete node.position)
-  pure (node.value, frontier''')
+  MaybeT (Strict.ST s) (c, Frontier s a b c)
+step f = deleteMinNode >=> lift . uncurry (peekInsertChildren f)
 
 deleteMinNode ::
-  (Ord c) => Frontier s a b c -> Strict.ST s (Maybe (Node s a b c, Frontier s a b c))
-deleteMinNode frontier = runMaybeT $ do
+  (Ord c) => Frontier s a b c -> MaybeT (Strict.ST s) (Node s a b c, Frontier s a b c)
+deleteMinNode frontier = do
   (node, queue') <- hoistMaybe (MinPQueue.minView frontier.queue)
   let frontier' = Frontier queue'
   pure (node, frontier')
@@ -83,6 +84,22 @@ prevNodeValue valueNode = runMaybeT $ do
   valueNode' <- MaybeT $ DoublyLinked.prev valueNode
   pure (DoublyLinked.value valueNode')
 
+-- Take a node not in the frontier but whose position is still in the position
+-- list, add its children to the frontier, and remove the node from the
+-- position list.
+peekInsertChildren ::
+  (Ord c) =>
+  (a -> b -> c) ->
+  Node s a b c ->
+  Frontier s a b c ->
+  Strict.ST s (c, Frontier s a b c)
+peekInsertChildren f node frontier = do
+  frontier' <-
+    insertChildA f node frontier
+      >>= insertChildB f node
+  DoublyLinked.delete node.position
+  pure (node.value, frontier')
+
 insertChildA ::
   (Ord c) =>
   (a -> b -> c) ->

src/ApplyMerge/IntMap.hs

@@ -3,14 +3,14 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
-module ApplyMerge.IntMap (applyMerge) where
+module ApplyMerge.IntMap (applyMerge, applyMergeNonEmpty) where
 
+import Control.Arrow ((>>>))
 import Control.Monad (guard)
-import Data.Function ((&))
 import Data.IntMap.Strict (IntMap)
 import Data.IntMap.Strict qualified as IntMap
 import Data.List (unfoldr)
-import Data.List.NonEmpty (NonEmpty, nonEmpty)
+import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty)
 import Data.List.NonEmpty qualified as NonEmpty
 import Data.Maybe (fromMaybe)
 import Data.PQueue.Prio.Min (MinPQueue)
@@ -32,24 +32,35 @@ applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
 applyMerge f as bs = fromMaybe [] $ do
   as' <- nonEmpty as
   bs' <- nonEmpty bs
-  pure (unfoldr (step f) (initialFrontier f as' bs'))
+  pure (NonEmpty.toList (applyMergeNonEmpty f as' bs'))
 
-initialFrontier :: (a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
-initialFrontier f as bs =
-  let node = mkNode f (0, 0) as bs
-   in Frontier
-        { queue = MinPQueue.singleton node.value node,
-          indexMap = IntMap.singleton 0 0
-        }
+applyMergeNonEmpty ::
+  (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
+applyMergeNonEmpty f as bs =
+  let (c, frontier) = initialState f as bs
+   in c :| unfoldr (step f) frontier
+
+initialState ::
+  forall a b c.
+  (Ord c) =>
+  (a -> b -> c) ->
+  NonEmpty a ->
+  NonEmpty b ->
+  (c, Frontier a b c)
+initialState f as bs =
+  let initialNode :: Node a b c
+      initialNode = mkNode f (0, 0) as bs
+
+      emptyFrontier :: Frontier a b c
+      emptyFrontier =
+        Frontier
+          { queue = MinPQueue.empty,
+            indexMap = IntMap.empty
+          }
+   in peekInsertChildren f initialNode emptyFrontier
 
 step :: (Ord c) => (a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
-step f frontier = do
-  (node, frontier') <- deleteMinNode frontier
-  let frontier'' =
-        frontier'
-          & insertChildA f node
-          & insertChildB f node
-  pure (node.value, frontier'')
+step f = fmap (uncurry (peekInsertChildren f)) . deleteMinNode
 
 deleteMinNode :: (Ord c) => Frontier a b c -> Maybe (Node a b c, Frontier a b c)
 deleteMinNode frontier = do
@@ -62,6 +73,17 @@ deleteMinNode frontier = do
           }
   pure (node, frontier')
 
+peekInsertChildren ::
+  (Ord c) =>
+  (a -> b -> c) ->
+  Node a b c ->
+  Frontier a b c ->
+  (c, Frontier a b c)
+peekInsertChildren f node =
+  insertChildA f node
+    >>> insertChildB f node
+    >>> (node.value,)
+
 insertChildA ::
   (Ord c) => (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
 insertChildA f (Node (ia, ib) _ as bs) frontier = fromMaybe frontier $ do

src/ApplyMerge/IntSet.hs

@@ -3,14 +3,14 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
-module ApplyMerge.IntSet (applyMerge) where
+module ApplyMerge.IntSet (applyMerge, applyMergeNonEmpty) where
 
+import Control.Arrow ((>>>))
 import Control.Monad (guard)
-import Data.Function ((&))
 import Data.IntSet (IntSet)
 import Data.IntSet qualified as IntSet
 import Data.List (unfoldr)
-import Data.List.NonEmpty (NonEmpty, nonEmpty)
+import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty)
 import Data.List.NonEmpty qualified as NonEmpty
 import Data.Maybe (fromMaybe)
 import Data.PQueue.Prio.Min (MinPQueue)
@@ -33,25 +33,36 @@ applyMerge :: (Ord c) => (a -> b -> c) -> [a] -> [b] -> [c]
 applyMerge f as bs = fromMaybe [] $ do
   as' <- nonEmpty as
   bs' <- nonEmpty bs
-  pure (unfoldr (step f) (initialFrontier f as' bs'))
+  pure (NonEmpty.toList (applyMergeNonEmpty f as' bs'))
 
-initialFrontier :: (a -> b -> c) -> NonEmpty a -> NonEmpty b -> Frontier a b c
-initialFrontier f as bs =
-  let node = mkNode f (0, 0) as bs
-   in Frontier
-        { queue = MinPQueue.singleton node.value node,
-          indexSetA = IntSet.singleton 0,
-          indexSetB = IntSet.singleton 0
-        }
+applyMergeNonEmpty ::
+  (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
+applyMergeNonEmpty f as bs =
+  let (c, frontier) = initialState f as bs
+   in c :| unfoldr (step f) frontier
+
+initialState ::
+  forall a b c.
+  (Ord c) =>
+  (a -> b -> c) ->
+  NonEmpty a ->
+  NonEmpty b ->
+  (c, Frontier a b c)
+initialState f as bs =
+  let initialNode :: Node a b c
+      initialNode = mkNode f (0, 0) as bs
+
+      emptyFrontier :: Frontier a b c
+      emptyFrontier =
+        Frontier
+          { queue = MinPQueue.empty,
+            indexSetA = IntSet.empty,
+            indexSetB = IntSet.empty
+          }
+   in peekInsertChildren f initialNode emptyFrontier
 
 step :: (Ord c) => (a -> b -> c) -> Frontier a b c -> Maybe (c, Frontier a b c)
-step f frontier = do
-  (node, frontier') <- deleteMinNode frontier
-  let frontier'' =
-        frontier'
-          & insertChildA f node
-          & insertChildB f node
-  pure (node.value, frontier'')
+step f = fmap (uncurry (peekInsertChildren f)) . deleteMinNode
 
 deleteMinNode :: (Ord c) => Frontier a b c -> Maybe (Node a b c, Frontier a b c)
 deleteMinNode frontier = do
@@ -65,6 +76,17 @@ deleteMinNode frontier = do
           }
   pure (node, frontier')
 
+peekInsertChildren ::
+  (Ord c) =>
+  (a -> b -> c) ->
+  Node a b c ->
+  Frontier a b c ->
+  (c, Frontier a b c)
+peekInsertChildren f node =
+  insertChildA f node
+    >>> insertChildB f node
+    >>> (node.value,)
+
 insertChildA ::
   (Ord c) => (a -> b -> c) -> Node a b c -> Frontier a b c -> Frontier a b c
 insertChildA f (Node (ia, ib) _ as bs) frontier = fromMaybe frontier $ do

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 <[email protected]>
+-- 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)