Make the analysis parallel (#2)

Co-authored-by: Mats Rauhala <mats.rauhala@iki.fi>
Reviewed-on: #2
Co-authored-by: Mats Rauhala <masse@rauhala.info>
Co-committed-by: Mats Rauhala <masse@rauhala.info>

src/Control/Addressbook/Streaming.hs

@@ -1,23 +1,18 @@
 module Control.Addressbook.Streaming where
 
-import qualified Data.Text as T
 
 import Conduit
 import qualified Data.Conduit.Binary as CB
 import qualified Data.Conduit.Combinators as C
 import qualified Data.Conduit.List as CL
-import qualified Data.Conduit.Text as CT
 
 import Data.Email
 import Data.Email.Header
        (Header(..))
 
-import System.IO
-       (stdin)
 
 import qualified Data.Foldable as F
 
-import qualified Data.Map.Strict as Map
 
 import Data.Maybe
        (fromMaybe)
@@ -25,28 +20,47 @@ import System.Environment
        (lookupEnv)
 import System.FilePath
        ((</>))
+import Data.Set (Set)
+import Data.ByteString (ByteString)
+import qualified Data.Set as Set
+import qualified Data.ByteString.Lazy as LBS
+import Data.Char (ord)
+import qualified Data.ByteString.Lazy.Char8 as LBC
+import System.IO.Unsafe (unsafeInterleaveIO)
+import Control.Parallel.Strategies (rseq, parMap)
+import qualified Data.List as L
 
 combine :: (MonadUnliftIO m, MonadResource m, MonadThrow m, MonadIO m) => ConduitM FilePath Header m ()
 combine = await >>= \case
   Nothing -> pure ()
   Just path -> (CB.sourceFile path .| parseEmail) >> combine
 
+chunks :: Int -> [a] -> [[a]]
+chunks n = L.unfoldr $ \case
+  [] -> Nothing
+  xs -> Just (splitAt n xs)
+
 run :: IO ()
 run = do
   datDir <- fromMaybe "./" <$> lookupEnv "HOME"
-  runResourceT $ do
-    x <- runConduit stream
-    runConduit (CL.sourceList (Map.keys x) .| C.map (<> "\n") .| CB.sinkFileCautious (datDir </> ".addressbook.dat"))
+  xs <- LBS.getContents >>= stream
+  let set = F.fold (parMap rseq F.fold (chunks 20 xs))
+  runResourceT $
+    runConduit $
+         CL.sourceList (Set.elems set)
+      .| C.map (<> "\n")
+      .| CB.sinkFileCautious (datDir </> ".addressbook.dat")
   where
     separate = \case
       From x -> [x]
       To xs -> F.toList xs
+    -- A set of (locally) unique addresses. Composes with parMap
+    stream :: LBS.ByteString -> IO [Set ByteString]
     stream =
-      CB.sourceHandle stdin
-        .| CT.decode CT.utf8
-        .| CT.lines
-        .| C.map T.unpack
-        .| combine
-        .| C.concatMap separate
-        .| CT.encode CT.utf8
-        .| C.foldMap (`Map.singleton` ())
+        traverse (unsafeInterleaveIO . parse . LBC.unpack)
+      . filter (not . LBS.null)
+      . LBS.split (fromIntegral $ ord '\n')
+    parse path =
+      runResourceT $
+        runConduit $
+          CB.sourceFile path .| parseEmail .| C.concatMap separate .| C.foldMap Set.singleton