86 lines
2.5 KiB
Haskell
86 lines
2.5 KiB
Haskell
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)
|
|
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 (using, parList, rseq, parBuffer, rdeepseq, parMap)
|
|
import qualified Data.List as L
|
|
import Control.Concurrent (getNumCapabilities)
|
|
import Debug.Trace (traceShow)
|
|
|
|
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"
|
|
xs <- LBS.getContents >>= stream
|
|
let x = F.fold (parMap rseq F.fold (chunks 200 xs))
|
|
runResourceT $
|
|
runConduit (CL.sourceList (Set.elems x) .| C.map (<> "\n") .| CB.sinkFileCautious (datDir </> ".addressbook.dat"))
|
|
where
|
|
separate = \case
|
|
From x -> [x]
|
|
To xs -> F.toList xs
|
|
stream :: LBS.ByteString -> IO ([Set ByteString])
|
|
stream = 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
|
|
|
|
run_ :: IO ()
|
|
run_ = do
|
|
datDir <- fromMaybe "./" <$> lookupEnv "HOME"
|
|
runResourceT $ do
|
|
x <- runConduit stream
|
|
runConduit (CL.sourceList (Set.elems x) .| C.map (<> "\n") .| CB.sinkFileCautious (datDir </> ".addressbook.dat"))
|
|
where
|
|
separate = \case
|
|
From x -> [x]
|
|
To xs -> F.toList xs
|
|
stream =
|
|
CB.sourceHandle stdin
|
|
.| CT.decode CT.utf8
|
|
.| CT.lines
|
|
.| C.map T.unpack
|
|
.| combine
|
|
.| C.concatMap separate
|
|
.| C.foldMap (Set.singleton)
|