@@ -1,4 +1,4 @@
1−module WordCount (foldWith, foldCountElements, printSelecttedOutput) where
1+module WordCount (countElements, printSelecttedOutput) where
22
33import Streamly.Data.Fold (Fold)
44import qualified Streamly.Data.Stream as Stream
@@ -6,11 +6,18 @@ import qualified Streamly.Data.Fold as Fold
66import qualified Streamly.FileSystem.File as File
77import Data.Function ((&))
88import Data.Char (ord)
9+import Data.Word (Word8)
910import qualified Streamly.Unicode.Stream as Stream
11+import Streamly.Data.Array (Array)
12+import qualified Streamly.Data.Stream.Prelude as Stream
13+import GHC.Conc (numCapabilities)
14+import qualified Streamly.Data.Array as Array
15+import Data.Char (chr)
1016
1117isSpace :: Char -> Bool
12−isSpace c = uc == 0x20 || uc - 0x9 <= 4
13− where uc = fromIntegral (ord c) :: Word
18+isSpace c = uc == 32 || uc - 0x9 <= 4 || uc == 0xa0
19+ where
20+ uc = fromIntegral (ord c) :: Word
1421
1522data WCOutputs =
1623 WCOutputs {
@@ -21,21 +28,43 @@ data WCOutputs =
2128 }
2229 deriving (Show)
2330
24−count :: WCOutputs -> Char -> WCOutputs
25−count (WCOutputs charsAmount wordsAmount linesAmount wasSpace) char =
31+processChar :: WCOutputs -> Char -> WCOutputs
32+processChar (WCOutputs charsAmount wordsAmount linesAmount wasSpace) char =
2633 let newLinesAmount = if char == '\n' then linesAmount + 1 else linesAmount
27− (newWordsAmount, isSpaceChar) = if isSpace char then (wordsAmount, True)
28− else (if wasSpace then wordsAmount + 1 else wordsAmount, False)
29− in WCOutputs (charsAmount + 1) newWordsAmount newLinesAmount isSpaceChar
30−
31−foldWith :: Fold IO Char a -> String -> IO a
32−foldWith f file =
33− File.read file
34− & Stream.decodeUtf8
35− & Stream.fold f
36−
37−foldCountElements :: Fold IO Char WCOutputs
38−foldCountElements = Fold.foldl' count (WCOutputs 0 0 0 True)
39−
40−printSelecttedOutput :: WCOutputs -> String
41−printSelecttedOutput output = show (lineCount output) ++ " " ++ show (wordsCount output) ++ " " ++ show (charCount output)
34+ (newWordsAmount, isSpaceChar') = if isSpace char then (wordsAmount, True)
35+ else (if wasSpace then wordsAmount + 1 else wordsAmount, False)
36+ in WCOutputs (charsAmount + 1) newWordsAmount newLinesAmount isSpaceChar'
37+
38+processChunk :: Array Word8 -> IO (Bool, WCOutputs)
39+processChunk chunk = do
40+ let firstChar = Array.getIndex 0 chunk
41+ case firstChar of
42+ Just x -> do
43+ output <- countChunkElements chunk
44+ return (isSpace (chr (fromIntegral x)), output)
45+ Nothing -> return (False, WCOutputs 0 0 0 True)
46+ where
47+ countChunkElements chunk' =
48+ Array.read chunk'
49+ & Stream.decodeUtf8
50+ & Stream.fold (Fold.foldl' processChar (WCOutputs 0 0 0 True))
51+
52+combineChunkOutputs :: (Bool, WCOutputs) -> (Bool, WCOutputs) -> (Bool, WCOutputs)
53+combineChunkOutputs (fstBeginsWithSpace, WCOutputs fstLines fstWords fstChars fstEndsWithSpace) (sndBeginsWithSpace, WCOutputs sndLines sndWords sndChars sndEndsWithSpace) =
54+ -- For correctly processing multiple following spaces and first/last words
55+ let combinedWordCount = if not fstEndsWithSpace && not sndBeginsWithSpace
56+ then fstWords + sndWords - 1
57+ else fstWords + sndWords
58+ in (fstBeginsWithSpace, WCOutputs (fstLines + sndLines) combinedWordCount (fstChars + sndChars) sndEndsWithSpace)
59+
60+foldChunks :: Fold IO (Bool, WCOutputs) (Bool, WCOutputs)
61+foldChunks = Fold.foldl' combineChunkOutputs (False, WCOutputs 0 0 0 True)
62+
63+countElements :: String -> IO (Bool, WCOutputs)
64+countElements file = do
65+ File.readChunks file
66+ & Stream.parMapM (Stream.maxThreads numCapabilities . Stream.ordered True) processChunk
67+ & Stream.fold foldChunks
68+
69+printSelecttedOutput :: (Bool, WCOutputs) -> String
70+printSelecttedOutput output = show (lineCount . snd $ output) ++ " " ++ show (wordsCount . snd $ output) ++ " " ++ show (charCount . snd $ output)