concurrent-merge: refactoring

Author: vaibhavsagar

concurrent-merge/Main.hs

@@ -6,38 +6,27 @@ module Main where
 
 import Control.Concurrent.MVar (MVar, newEmptyMVar, readMVar, takeMVar, putMVar)
 import Control.Concurrent      (forkIO)
-import Data.Maybe              (isJust)
-import Data.List               (sortOn)
-
+import Data.Maybe              (isJust, fromJust)
+import Data.List               (minimumBy)
+import Data.Ord                (comparing)
 
 input1 = [ 1, 1, 1, 2 ]
 input2 = [ 2, 3, 4, 4 ]
 
 worker :: [Int] -> MVar (Maybe Int) -> IO ()
 worker []     slot = putMVar slot Nothing
-worker (x:xs) slot = do
-    putMVar slot (Just x)
-    worker xs slot
+worker (x:xs) slot = putMVar slot (Just x) >> worker xs slot
 
 merge :: [MVar (Maybe Int)] -> [Int] -> IO [Int]
-merge slots output = do
-    values <- traverse readMVar slots
-    let values' = zip [0..] values
-    let filled = filter (\(_, i) -> isJust i) values'
-    case filled of
+merge slots output = traverse readMVar slots >>= \values ->
+    case filter (isJust . snd) $ zip [0..] values of
         [] -> return output
-        _  -> do
-            let sorted = sortOn (\(_, Just i) -> i) filled
-            let (idx, Just el) = head sorted
-            let output' = output ++ [el]
-            takeMVar (slots !! idx)
-            merge slots output'
+        ls -> let
+            (ix, Just el) = minimumBy (comparing (fromJust . snd)) ls
+            in takeMVar (slots !! ix) >> merge slots (output ++ [el])
 
 main :: IO ()
 main = do
-    slot1 <- newEmptyMVar
-    slot2 <- newEmptyMVar
-    forkIO $ worker input1 slot1
-    forkIO $ worker input2 slot2
-    result <- merge [slot1, slot2] []
-    print result
+    slots <- sequenceA $ replicate 2 newEmptyMVar
+    traverse (forkIO . uncurry worker) $ zip [input1, input2] slots
+    print =<< merge slots []