git.iko.soy/streamly
on 511fe85 help sign in
streamly/examples/WordCount.hs
630 lines 28 kB history raw
1-------------------------------------------------------------------------------
2-- Fast, streaming and parallel word counting (wc) program.
3-------------------------------------------------------------------------------
4-- 1) On utf8 inputs the serial version is around 3x faster than MacOS wc
5-- 2) It can run parallely on multiple cores providing further speedup
6-- 3) Parallel version works efficiently on stdin/streaming input as well
7-- 4) Parallel version handles utf8 input correctly (including multi-byte space
8-- chars) and gives the same output as the serial version on all inputs.
9-- 5) There may be differences in word/char counts when there are invalid utf8
10-- byte sequences present in the input because of different styles of error
11-- handling.
13-------------------------------------------------------------------------------
14-- Build with the following options:
15-------------------------------------------------------------------------------
16-- streamly optimization plugin is required for best performance
17-- ghc -O2 -fplugin Plugin -fspec-constr-recursive=10 -fmax-worker-args=16
18-- For concurrent version add: -threaded -with-rtsopts "-N"
20-------------------------------------------------------------------------------
21-- Comparing with "wc -mwl" command:
22-------------------------------------------------------------------------------
23--
24-- 1) To enable UTF8 with wc: export LANG=en_US.UTF-8; export LC_ALL=$LANG
25-- 2) To test whether it is acutally using utf8, copy and paste this string
26-- "U+1680 U+2000 U+2001 U+2002" and run "wc -mwl" on this. Without proper UTF8
27-- handling word count would be 1, with proper UTF8 handling word count would
28-- be 4. Note that the spaces in this string are not regular space chars they
29-- are different unicode space chars.
31{-# LANGUAGE CPP #-}
33import Control.Monad (when)
34import Data.Char (isSpace)
35import Data.Word (Word8)
36import GHC.Conc (numCapabilities)
37import System.Environment (getArgs)
38import System.IO (Handle, openFile, IOMode(..))
39import Streamly.Internal.Data.Unicode.Stream
40 (DecodeState, DecodeError(..), CodePoint, decodeUtf8Either,
41 resumeDecodeUtf8Either)
43import qualified Streamly as S
44import qualified Streamly.Data.Unicode.Stream as S
45import qualified Streamly.FileSystem.Handle as FH
46import qualified Streamly.Memory.Array as A
47import qualified Streamly.Prelude as S
48import qualified Data.Vector.Storable.Mutable as V
50-------------------------------------------------------------------------------
51-- Parallel char, line and word counting
52-------------------------------------------------------------------------------
54-- We process individual chunks in the stream independently and parallely and
55-- the combine the chunks to combine what they have counted.
56--
57-------------------------------------------------------------------------------
58-- Char counting
59-------------------------------------------------------------------------------
61-- To count chars each block needs the following:
62--
63-- -- | header | char counts | trailer |
64--
65-- header and trailer are incomplete utf8 byte sequences that may be combined
66-- with the previous or the next block to complete them later.
67--
68-- The trailer may have one or more bytes in a valid utf8 sequence and is
69-- expecting more bytes to complete the sequence. The header stores any
70-- possible continuation from the previous block. It contains a maximum of 3
71-- bytes which all must be non-starter bytes.
72--
73-- When two blocks are combined, the trailer of the first block is combined
74-- with the header of the next block and then utf8 decoded. The combined
75-- header+trailer may yield:
76--
77-- * Nothing - when there is no trailer and header
78-- * All errors - when there is no trailer in the previous block, and there is
79-- a header in the next block. In this case there is no starting char which
80-- means all header bytes are errors.
81-- * It can yield at most one valid character followed by 0, 1 or 2 errors.
82--
83-- We count an incomplete utf8 sequence of 2 or more bytes starting with a
84-- valid starter byte as a single codepoint. Bytes not following a valid
85-- starter byte are treated as individual codepoints for counting.
86--
87-------------------------------------------------------------------------------
88-- Word counting
89-------------------------------------------------------------------------------
91-- For word counting we need the following in each block:
92--
93-- -- | header | startsWithSpace | word counts | endsWithSpace | trailer |
94--
95-- The word counts in individual blocks are performed assuming that the
96-- previous char before the block is a space.
97-- When combining two blocks, after combining the trailer of previous blocks
98-- with the header of the next we determine if the resulting char is a space or
99-- not.
100--
101-- 1) If there is no new char joining the two blocks then we use endsWithSpace
102-- of the previous block and startsWithSpace of the next block to determine if
103-- the word counts are to be adjusted. If the previous block ends with
104-- non-space and the next block starts with non-space we need to decrement the
105-- word count by one if it is non-zero in the next block.
106--
107-- 2) If the new joining char is a space then we combine it with
108-- startsWithSpace and endsWithSpace to determine the
109-- startsWithSpace/endsWithSpace of the combined block and adjust the word
110-- counts appropriately.
111--
112-------------------------------------------------------------------------------
113-- Line counting
114-------------------------------------------------------------------------------
116-- Line counting is performed by counting "\n" in the stream. No new "\n" can
117-- result from patching the trailer and header as it is always a single byte.
119-------------------------------------------------------------------------------
120-- Counting state
121-------------------------------------------------------------------------------
123-- We use a mutable vector for the counting state. A fold using an immutable
124-- structure for such a large state does not perform well. However, mutability
125-- is confined to just the accumulator.
126--
127-- XXX we need convenient mutable records (like C structs) to handle things
128-- like this. It may be possible to achieve the same performance with an
129-- immutable accumulator, but that will require more research. Since we are
130-- always discarding the previous state, we can perhaps make use of that memory
131-- using safe in-place modifications, without having to allocate new memory.
133-- XXX we can also count the number of decoding errors separately
134data Field =
135 -- The number of "\n" characters found in the block.
136 LineCount
137 -- Number of full words found in the block, words are counted on a
138 -- transition from space char to a non-space char. We always assume the
139 -- char before the first starter char in a block is a space. If this is
140 -- found to be incorrect when joining two blocks then we fix the counts
141 -- accordingly.
142 | WordCount
143 -- The number of successfully decoded characters plus the number of
144 -- decoding failures in the block. Each byte or sequence of bytes on which
145 -- decoding fails is also counted as one char. The header and trailer bytes
146 -- are not accounted in this, they are accounted only when we join two
147 -- blocks.
148 | CharCount
149 -- whether the last counted char in this block was a space char
150 | WasSpace
151 -- whether the first successfully decoded char in this block is a space. A
152 -- decoding failure, after the trailing bytes from the previous block are
153 -- accounted, is also considered as space.
154 | FirstIsSpace
155 -- If no starter byte is found in the first three bytes in the block then
156 -- store those bytes to possibly combine them with the trailing incomplete
157 -- byte sequence in the previous block. We mark it done when either we have
158 -- stored three bytes or we have found a starter byte.
159 --
160 -- XXX This is ugly to manipulate, we can implement a statically max sized
161 -- mutable ring structure within this record.
162 | HeaderDone
163 | HeaderWordCount
164 | HeaderWord1
165 | HeaderWord2
166 | HeaderWord3
167 -- If a byte sequence at the end of the block is not complete then store
168 -- the current state of the utf8 decoder to continue it later using the
169 -- incomplete leading byte sequence in the next block.
170 | TrailerPresent
171 | TrailerState
172 | TrailerCodePoint
173 deriving (Show, Enum, Bounded)
175-------------------------------------------------------------------------------
176-- Default/initial state of the block
177-------------------------------------------------------------------------------
179readField :: V.IOVector Int -> Field -> IO Int
180readField v fld = V.read v (fromEnum fld)
182writeField :: V.IOVector Int -> Field -> Int -> IO ()
183writeField v fld val = V.write v (fromEnum fld) val
185modifyField :: V.IOVector Int -> Field -> (Int -> Int) -> IO ()
186modifyField v fld f = V.modify v f (fromEnum fld)
188newCounts :: IO (V.IOVector Int)
189newCounts = do
190 counts <- V.new (fromEnum (maxBound :: Field) + 1)
191 writeField counts LineCount 0
192 writeField counts WordCount 0
193 writeField counts CharCount 0
194 writeField counts WasSpace 1
195 writeField counts FirstIsSpace 0
196 writeField counts HeaderDone 0
197 writeField counts HeaderWordCount 0
198 writeField counts TrailerPresent 0
199 return counts
201-------------------------------------------------------------------------------
202-- Counting chars
203-------------------------------------------------------------------------------
205accountChar :: V.IOVector Int -> Bool -> IO ()
206accountChar counts isSp = do
207 c <- readField counts CharCount
208 let space = if isSp then 1 else 0
209 when (c == 0) $ writeField counts FirstIsSpace space
210 writeField counts CharCount (c + 1)
211 writeField counts WasSpace space
213-------------------------------------------------------------------------------
214-- Manipulating the header bytes
215-------------------------------------------------------------------------------
217addToHeader :: V.IOVector Int -> Int -> IO Bool
218addToHeader counts cp = do
219 cnt <- readField counts HeaderWordCount
220 case cnt of
221 0 -> do
222 writeField counts HeaderWord1 cp
223 writeField counts HeaderWordCount 1
224 return True
225 1 -> do
226 writeField counts HeaderWord2 cp
227 writeField counts HeaderWordCount 2
228 return True
229 2 -> do
230 writeField counts HeaderWord3 cp
231 writeField counts HeaderWordCount 3
232 writeField counts HeaderDone 1
233 return True
234 _ -> return False
236resetHeaderOnNewChar :: V.IOVector Int -> IO ()
237resetHeaderOnNewChar counts = do
238 hdone <- readField counts HeaderDone
239 when (hdone == 0) $ writeField counts HeaderDone 1
241-------------------------------------------------------------------------------
242-- Manipulating the trailer
243-------------------------------------------------------------------------------
245setTrailer :: V.IOVector Int -> DecodeState -> CodePoint -> IO ()
246setTrailer counts st cp = do
247 writeField counts TrailerState (fromIntegral st)
248 writeField counts TrailerCodePoint cp
249 writeField counts TrailerPresent 1
251resetTrailerOnNewChar :: V.IOVector Int -> IO ()
252resetTrailerOnNewChar counts = do
253 trailer <- readField counts TrailerPresent
254 when (trailer /= 0) $ do
255 writeField counts TrailerPresent 0
256 accountChar counts True
258-------------------------------------------------------------------------------
259-- Counting the stream
260-------------------------------------------------------------------------------
262{-# INLINE countChar #-}
263countChar :: V.IOVector Int -> Either DecodeError Char -> IO ()
264countChar counts inp =
265 case inp of
266 Right ch -> do
267 resetHeaderOnNewChar counts
268 -- account the last stored error as whitespace and clear it
269 resetTrailerOnNewChar counts
271 when (ch == '\n') $ modifyField counts LineCount (+ 1)
272 if isSpace ch
273 then accountChar counts True
274 else do
275 wasSpace <- readField counts WasSpace
276 when (wasSpace /= 0) $ modifyField counts WordCount (+ 1)
277 accountChar counts False
278 Left (DecodeError st cp) -> do
279 hdone <- readField counts HeaderDone
280 if hdone == 0
281 then do
282 if st == 0
283 then do
284 -- We got a non-starter in initial decoder state, there may
285 -- be something that comes before this to complete it.
286 r <- addToHeader counts cp
287 when (not r) $ error "countChar: Bug addToHeader failed"
288 else do
289 -- We got an error in a non-initial decoder state, it may
290 -- be an input underflow error, keep it as incomplete in
291 -- the trailer.
292 writeField counts HeaderDone 1
293 setTrailer counts st cp
294 else do
295 resetTrailerOnNewChar counts
296 if st == 0
297 then accountChar counts True
298 else setTrailer counts st cp
300printCounts :: V.IOVector Int -> IO ()
301printCounts v = do
302 l <- readField v LineCount
303 w <- readField v WordCount
304 c <- readField v CharCount
305 putStrLn $ show l ++ " " ++ show w ++ " " ++ show c
307-------------------------------------------------------------------------------
308-- Serial counting using parallel version of countChar
309-------------------------------------------------------------------------------
311_wc_mwl_parserial :: Handle -> IO (V.IOVector Int)
312_wc_mwl_parserial src = do
313 counts <- newCounts
314 S.mapM_ (countChar counts)
315 $ decodeUtf8Either
316 $ S.unfold FH.read src
317 return counts
319-------------------------------------------------------------------------------
320-- Serial word counting with UTF-8 handling
321-------------------------------------------------------------------------------
323data Counts = Counts !Int !Int !Int !Bool deriving Show
325{-# INLINE countCharSerial #-}
326countCharSerial :: Counts -> Char -> Counts
327countCharSerial (Counts l w c wasSpace) ch =
328 let l1 = if (ch == '\n') then l + 1 else l
329 (w1, wasSpace1) =
330 if (isSpace ch)
331 then (w, True)
332 else (if wasSpace then w + 1 else w, False)
333 in (Counts l1 w1 (c + 1) wasSpace1)
335-- Note: This counts invalid byte sequences are non-space chars
336_wc_mwl_serial :: Handle -> IO ()
337_wc_mwl_serial src = print =<< (
338 S.foldl' countCharSerial (Counts 0 0 0 True)
339 $ S.decodeUtf8Lax
340 $ S.unfold FH.read src)
342-------------------------------------------------------------------------------
343-- Parallel counting
344-------------------------------------------------------------------------------
346-- XXX we need a better data structure to store the header bytes to make these
347-- routines simpler.
348--
349-- combine trailing bytes in preceding block with leading bytes in the next
350-- block and decode them into a codepoint
351reconstructChar :: Int
352 -> V.IOVector Int
353 -> V.IOVector Int
354 -> IO (S.SerialT IO (Either DecodeError Char))
355reconstructChar hdrCnt v1 v2 = do
356 when (hdrCnt > 3 || hdrCnt < 0) $ error "reconstructChar: hdrCnt > 3"
357 stream1 <-
358 if (hdrCnt > 2)
359 then do
360 x <- readField v2 HeaderWord3
361 return $ (fromIntegral x :: Word8) `S.cons` S.nil
362 else return S.nil
363 stream2 <-
364 if (hdrCnt > 1)
365 then do
366 x <- readField v2 HeaderWord2
367 return $ fromIntegral x `S.cons` stream1
368 else return stream1
369 stream3 <-
370 if (hdrCnt > 0)
371 then do
372 x <- readField v2 HeaderWord1
373 return $ fromIntegral x `S.cons` stream2
374 else return stream2
376 state <- readField v1 TrailerState
377 cp <- readField v1 TrailerCodePoint
378 return $ resumeDecodeUtf8Either (fromIntegral state) cp stream3
380getHdrChar :: V.IOVector Int -> IO (Maybe Int)
381getHdrChar v = do
382 hdrCnt <- readField v HeaderWordCount
383 case hdrCnt of
384 0 -> return Nothing
385 1 -> do
386 writeField v HeaderWordCount 0
387 fmap Just $ readField v HeaderWord1
388 2 -> do
389 x1 <- readField v HeaderWord1
390 x2 <- readField v HeaderWord2
391 writeField v HeaderWord1 x2
392 writeField v HeaderWordCount 1
393 return $ Just x1
394 3 -> do
395 x1 <- readField v HeaderWord1
396 x2 <- readField v HeaderWord2
397 x3 <- readField v HeaderWord3
398 writeField v HeaderWord1 x2
399 writeField v HeaderWord2 x3
400 writeField v HeaderWordCount 2
401 return $ Just x1
402 _ -> error "getHdrChar: Bug, hdrCnt not in range 0-3"
404-- If the header of the first block is not done then combine the header
405-- with the header of the next block.
406combineHeaders :: V.IOVector Int -> V.IOVector Int -> IO ()
407combineHeaders v1 v2 = do
408 hdone1 <- readField v1 HeaderDone
409 if hdone1 == 0
410 then do
411 res <- getHdrChar v2
412 case res of
413 Nothing -> return ()
414 Just x -> do
415 r <- addToHeader v1 x
416 when (not r) $ error "combineHeaders: Bug, addToHeader failed"
417 else return ()
419-- We combine the contents of the second vector into the first vector, mutating
420-- the first vector and returning it.
421-- XXX This is a quick hack and can be refactored to reduce the size
422-- and understandability considerably.
423addCounts :: V.IOVector Int -> V.IOVector Int -> IO (V.IOVector Int)
424addCounts v1 v2 = do
425 hdone1 <- readField v1 HeaderDone
426 hdone2 <- readField v2 HeaderDone
427 hdrCnt2_0 <- readField v2 HeaderWordCount
428 if hdone1 == 0 && (hdrCnt2_0 /= 0 || hdone2 /= 0)
429 then do
430 combineHeaders v1 v2
431 if hdone2 == 0
432 then return v1
433 else do
434 writeField v1 HeaderDone 1
435 addCounts v1 v2
436 else do
437 trailerPresent2 <- readField v2 TrailerPresent
438 trailerState2 <- readField v2 TrailerState
439 trailerCodePoint2 <- readField v2 TrailerCodePoint
440 if hdone1 == 0
441 then error "addCounts: Bug, trying to add completely empty second block"
442 else do
443 l1 <- readField v1 LineCount
444 w1 <- readField v1 WordCount
445 c1 <- readField v1 CharCount
446 wasSpace1 <- readField v1 WasSpace
448 l2 <- readField v2 LineCount
449 w2 <- readField v2 WordCount
450 c2 <- readField v2 CharCount
451 wasSpace2 <- readField v2 WasSpace
452 firstIsSpace2 <- readField v2 FirstIsSpace
453 hdrCnt2 <- readField v2 HeaderWordCount
455 trailer1 <- readField v1 TrailerPresent
456 if trailer1 == 0 -- no trailer in the first block
457 then do
458 -- header2, if any, complete or incomplete, is just invalid
459 -- bytes, count them as whitespace
460 let firstIsSpace2' = firstIsSpace2 /= 0 || hdrCnt2 /= 0
461 w <-
462 if w2 /= 0 && wasSpace1 == 0 && not firstIsSpace2'
463 then return $ w1 + w2 - 1
464 else return $ w1 + w2
465 writeField v1 LineCount (l1 + l2)
466 writeField v1 WordCount w
467 writeField v1 CharCount (c1 + c2 + hdrCnt2)
469 when (c1 == 0) $ do
470 if c2 == 0 && hdrCnt2 /= 0
471 then writeField v1 FirstIsSpace 1
472 else writeField v1 FirstIsSpace firstIsSpace2
474 if c2 == 0 && hdrCnt2 /= 0
475 then writeField v1 WasSpace 1
476 else writeField v1 WasSpace wasSpace2
478 writeField v1 TrailerPresent trailerPresent2
479 writeField v1 TrailerState trailerState2
480 writeField v1 TrailerCodePoint trailerCodePoint2
481 return v1
482 else do
483 if hdrCnt2 == 0
484 then do
485 when (hdone2 /= 0) $ do -- empty and Done header
486 -- count trailer as whitespace, its counted as one char
487 -- Note: hdrCnt2 == 0 means either header is not done
488 -- or c2 /= 0
489 writeField v1 LineCount (l1 + l2)
490 writeField v1 WordCount (w1 + w2)
491 writeField v1 CharCount (c1 + c2 + 1)
493 when (c1 == 0) $ writeField v1 FirstIsSpace 1
494 if (c2 == 0)
495 then writeField v1 WasSpace 1
496 else writeField v1 WasSpace wasSpace2
498 writeField v1 TrailerPresent trailerPresent2
499 writeField v1 TrailerState trailerState2
500 writeField v1 TrailerCodePoint trailerCodePoint2
501 -- If header of the next block is not done we continue the
502 -- trailer from the previous block instead of treating it
503 -- as whitespace
504 return v1
505 else do
506 -- join the trailing part of the first block with the
507 -- header of the next
508 decoded <- reconstructChar hdrCnt2 v1 v2
509 res <- S.uncons decoded
510 case res of
511 Nothing -> error "addCounts: Bug. empty reconstructed char"
512 Just (h, t) -> do
513 tlength <- S.length t
514 case h of
515 Right ch -> do
516 -- If we have an error case after this
517 -- char then that would be treated
518 -- as whitespace
519 let lcount = l1 + l2
520 lcount1 = if (ch == '\n') then lcount + 1 else lcount
521 wcount = w1 + w2
522 firstSpace = isSpace ch
523 wasSpace = firstSpace || tlength /= 0
524 wcount1 =
525 if wasSpace
526 then wcount
527 else if w2 == 0 || firstIsSpace2 /= 0
528 then wcount
529 else wcount - 1
530 writeField v1 LineCount lcount1
531 writeField v1 WordCount wcount1
532 writeField v1 CharCount (c1 + c2 + 1 + tlength)
534 when (c1 == 0) $
535 writeField v1 FirstIsSpace
536 (if firstSpace then 1 else 0)
538 if c2 == 0
539 then do
540 if wasSpace
541 then writeField v1 WasSpace 1
542 else writeField v1 WasSpace 0
543 else writeField v1 WasSpace wasSpace2
545 writeField v1 TrailerPresent trailerPresent2
546 writeField v1 TrailerState trailerState2
547 writeField v1 TrailerCodePoint trailerCodePoint2
548 return v1
549 Left (DecodeError st cp) -> do
550 -- if header was incomplete it may result
551 -- in partially decoded char to be written
552 -- as trailer. Check if the last error is
553 -- an incomplete decode.
554 r <- S.last t
555 let (st', cp') =
556 case r of
557 Nothing -> (st, cp)
558 Just lst -> case lst of
559 Right _ -> error "addCounts: Bug"
560 Left (DecodeError st1 cp1) -> (st1, cp1)
561 if hdone2 == 0 && st' /= 12
562 then do
563 -- all elements before the last one must be errors
564 writeField v1 CharCount (c1 + tlength)
566 when (c1 == 0) $
567 writeField v1 FirstIsSpace 1
569 writeField v1 WasSpace 1
571 writeField v1 TrailerState (fromIntegral st')
572 writeField v1 TrailerCodePoint cp'
573 else do
574 -- all elements must be errors
575 -- treat them as whitespace
576 writeField v1 LineCount (l1 + l2)
577 writeField v1 WordCount (w1 + w2)
578 writeField v1 CharCount (c1 + c2 + tlength + 1)
580 when (c1 == 0) $
581 writeField v1 FirstIsSpace 1
583 if c2 == 0
584 then writeField v1 WasSpace 1
585 else writeField v1 WasSpace wasSpace2
587 writeField v1 TrailerPresent trailerPresent2
588 writeField v1 TrailerState trailerState2
589 writeField v1 TrailerCodePoint trailerCodePoint2
590 return v1
592-- Individual array processing is an isolated loop, fusing it with the bigger
593-- loop may be counter productive.
594{-# NOINLINE countArray #-}
595countArray :: A.Array Word8 -> IO (V.IOVector Int)
596countArray src = do
597 counts <- newCounts
598 S.mapM_ (countChar counts)
599 $ decodeUtf8Either
600 $ S.unfold A.read src
601 return counts
603{-# INLINE wc_mwl_parallel #-}
604wc_mwl_parallel :: Handle -> Int -> IO (V.IOVector Int)
605wc_mwl_parallel src n = do
606 counts <- newCounts
607 S.foldlM' addCounts counts
608 $ S.aheadly
609 $ S.maxThreads numCapabilities
610 $ S.mapM (countArray)
611 $ S.unfold FH.readChunksWithBufferOf (n, src)
613-------------------------------------------------------------------------------
614-- Main
615-------------------------------------------------------------------------------
617main :: IO ()
618main = do
619 name <- fmap head getArgs
620 src <- openFile name ReadMode
621 -- _wc_mwl_serial src -- Unix wc -l program
622 -- printCounts =<< _wc_mwl_parserial src -- Unix wc -l program
623 -- Using different sizes of chunks (1,2,3,4,5,10,128,256) is a good testing
624 -- mechanism for parallel counting code.
625 {-
626 args <- getArgs
627 let chunkSize = read $ args !! 1
628 -}
629 let chunkSize = 32 * 1024
630 printCounts =<< wc_mwl_parallel src chunkSize -- Unix wc -l program