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