From 7675571a00341b9b3932daa00d7bcf5f5bded7ae Mon Sep 17 00:00:00 2001
From: Preetham Gujjula <gitcommit@mail.preetham.io>
Date: Wed, 8 May 2024 22:12:15 -0700
Subject: [PATCH] [#2] Add IntMap applyMergeNonEmpty

---
 src/ApplyMerge/IntMap.hs | 35 ++++++++++++++++++++++++++---------
 1 file changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/ApplyMerge/IntMap.hs b/src/ApplyMerge/IntMap.hs
index a63b599..384d4eb 100644
--- a/src/ApplyMerge/IntMap.hs
+++ b/src/ApplyMerge/IntMap.hs
@@ -10,7 +10,7 @@ import Control.Monad (guard)
 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,15 +32,32 @@ 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 = fmap (uncurry (peekInsertChildren f)) . deleteMinNode