From d0e2f0d9ae7e133e5edb6ab703e3e19e2ce30d7f Mon Sep 17 00:00:00 2001
From: Preetham Gujjula <gitcommit@mail.preetham.io>
Date: Wed, 8 May 2024 16:19:05 -0700
Subject: [PATCH] [#2] Add DoublyLinkedList peekInsertChildren

---
 src/ApplyMerge/DoublyLinkedList.hs | 33 ++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/ApplyMerge/DoublyLinkedList.hs b/src/ApplyMerge/DoublyLinkedList.hs
index 7e3700a..5db97a9 100644
--- a/src/ApplyMerge/DoublyLinkedList.hs
+++ b/src/ApplyMerge/DoublyLinkedList.hs
@@ -5,7 +5,7 @@
 
 module ApplyMerge.DoublyLinkedList (applyMerge) 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)
@@ -37,7 +37,7 @@ applyMergeNonEmpty ::
   (Ord c) => (a -> b -> c) -> NonEmpty a -> NonEmpty b -> [c]
 applyMergeNonEmpty f as bs = Lazy.runST $ do
   frontier <- Lazy.strictToLazyST (initialFrontier f as bs)
-  unfoldrM (Lazy.strictToLazyST . step f) frontier
+  unfoldrM (Lazy.strictToLazyST . runMaybeT . step f) frontier
 
 unfoldrM :: (Monad m) => (b -> m (Maybe (a, b))) -> b -> m [a]
 unfoldrM f seed = do
@@ -58,17 +58,12 @@ 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 +78,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) ->