@@ -0,0 +1,627 @@
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
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+-------------------------------------------------------------------------------
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.
33+import Control.Monad (when)
34+import Data.Char (isSpace)
35+import Data.Word (Word8)
36+import GHC.Conc (numCapabilities)
37+import System.Environment (getArgs)
38+import System.IO (Handle, openFile, IOMode(..))
40+import qualified Streamly as S
41+import qualified Streamly.Data.String as S
42+import qualified Streamly.FileSystem.Handle as FH
43+import qualified Streamly.Memory.Array as A
44+import qualified Streamly.Prelude as S
45+import qualified Data.Vector.Storable.Mutable as V
47+-------------------------------------------------------------------------------
48+-- Parallel char, line and word counting
49+-------------------------------------------------------------------------------
51+-- We process individual chunks in the stream independently and parallely and
52+-- the combine the chunks to combine what they have counted.
54+-------------------------------------------------------------------------------
56+-------------------------------------------------------------------------------
58+-- To count chars each block needs the following:
60+-- -- | header | char counts | trailer |
62+-- header and trailer are incomplete utf8 byte sequences that may be combined
63+-- with the previous or the next block to complete them later.
65+-- The trailer may have one or more bytes in a valid utf8 sequence and is
66+-- expecting more bytes to complete the sequence. The header stores any
67+-- possible continuation from the previous block. It contains a maximum of 3
68+-- bytes which all must be non-starter bytes.
70+-- When two blocks are combined, the trailer of the first block is combined
71+-- with the header of the next block and then utf8 decoded. The combined
72+-- header+trailer may yield:
74+-- * Nothing - when there is no trailer and header
75+-- * All errors - when there is no trailer in the previous block, and there is
76+-- a header in the next block. In this case there is no starting char which
77+-- means all header bytes are errors.
78+-- * It can yield at most one valid character followed by 0, 1 or 2 errors.
80+-- We count an incomplete utf8 sequence of 2 or more bytes starting with a
81+-- valid starter byte as a single codepoint. Bytes not following a valid
82+-- starter byte are treated as individual codepoints for counting.
84+-------------------------------------------------------------------------------
86+-------------------------------------------------------------------------------
88+-- For word counting we need the following in each block:
90+-- -- | header | startsWithSpace | word counts | endsWithSpace | trailer |
92+-- The word counts in individual blocks are performed assuming that the
93+-- previous char before the block is a space.
94+-- When combining two blocks, after combining the trailer of previous blocks
95+-- with the header of the next we determine if the resulting char is a space or
98+-- 1) If there is no new char joining the two blocks then we use endsWithSpace
99+-- of the previous block and startsWithSpace of the next block to determine if
100+-- the word counts are to be adjusted. If the previous block ends with
101+-- non-space and the next block starts with non-space we need to decrement the
102+-- word count by one if it is non-zero in the next block.
104+-- 2) If the new joining char is a space then we combine it with
105+-- startsWithSpace and endsWithSpace to determine the
106+-- startsWithSpace/endsWithSpace of the combined block and adjust the word
107+-- counts appropriately.
109+-------------------------------------------------------------------------------
111+-------------------------------------------------------------------------------
113+-- Line counting is performed by counting "\n" in the stream. No new "\n" can
114+-- result from patching the trailer and header as it is always a single byte.
116+-------------------------------------------------------------------------------
118+-------------------------------------------------------------------------------
120+-- We use a mutable vector for the counting state. A fold using an immutable
121+-- structure for such a large state does not perform well. However, mutability
122+-- is confined to just the accumulator.
124+-- XXX we need convenient mutable records (like C structs) to handle things
125+-- like this. It may be possible to achieve the same performance with an
126+-- immutable accumulator, but that will require more research. Since we are
127+-- always discarding the previous state, we can perhaps make use of that memory
128+-- using safe in-place modifications, without having to allocate new memory.
130+-- XXX we can also count the number of decoding errors separately
132+ -- The number of "\n" characters found in the block.
134+ -- Number of full words found in the block, words are counted on a
135+ -- transition from space char to a non-space char. We always assume the
136+ -- char before the first starter char in a block is a space. If this is
137+ -- found to be incorrect when joining two blocks then we fix the counts
140+ -- The number of successfully decoded characters plus the number of
141+ -- decoding failures in the block. Each byte or sequence of bytes on which
142+ -- decoding fails is also counted as one char. The header and trailer bytes
143+ -- are not accounted in this, they are accounted only when we join two
146+ -- whether the last counted char in this block was a space char
148+ -- whether the first successfully decoded char in this block is a space. A
149+ -- decoding failure, after the trailing bytes from the previous block are
150+ -- accounted, is also considered as space.
152+ -- If no starter byte is found in the first three bytes in the block then
153+ -- store those bytes to possibly combine them with the trailing incomplete
154+ -- byte sequence in the previous block. We mark it done when either we have
155+ -- stored three bytes or we have found a starter byte.
157+ -- XXX This is ugly to manipulate, we can implement a statically max sized
158+ -- mutable ring structure within this record.
164+ -- If a byte sequence at the end of the block is not complete then store
165+ -- the current state of the utf8 decoder to continue it later using the
166+ -- incomplete leading byte sequence in the next block.
170+ deriving (Show, Enum, Bounded)
172+-------------------------------------------------------------------------------
173+-- Default/initial state of the block
174+-------------------------------------------------------------------------------
176+readField :: V.IOVector Int -> Field -> IO Int
177+readField v fld = V.read v (fromEnum fld)
179+writeField :: V.IOVector Int -> Field -> Int -> IO ()
180+writeField v fld val = V.write v (fromEnum fld) val
182+modifyField :: V.IOVector Int -> Field -> (Int -> Int) -> IO ()
183+modifyField v fld f = V.modify v f (fromEnum fld)
185+newCounts :: IO (V.IOVector Int)
187+ counts <- V.new (fromEnum (maxBound :: Field) + 1)
188+ writeField counts LineCount 0
189+ writeField counts WordCount 0
190+ writeField counts CharCount 0
191+ writeField counts WasSpace 1
192+ writeField counts FirstIsSpace 0
193+ writeField counts HeaderDone 0
194+ writeField counts HeaderWordCount 0
195+ writeField counts TrailerPresent 0
198+-------------------------------------------------------------------------------
200+-------------------------------------------------------------------------------
202+accountChar :: V.IOVector Int -> Bool -> IO ()
203+accountChar counts isSp = do
204+ c <- readField counts CharCount
205+ let space = if isSp then 1 else 0
206+ when (c == 0) $ writeField counts FirstIsSpace space
207+ writeField counts CharCount (c + 1)
208+ writeField counts WasSpace space
210+-------------------------------------------------------------------------------
211+-- Manipulating the header bytes
212+-------------------------------------------------------------------------------
214+addToHeader :: V.IOVector Int -> Int -> IO Bool
215+addToHeader counts cp = do
216+ cnt <- readField counts HeaderWordCount
219+ writeField counts HeaderWord1 cp
220+ writeField counts HeaderWordCount 1
223+ writeField counts HeaderWord2 cp
224+ writeField counts HeaderWordCount 2
227+ writeField counts HeaderWord3 cp
228+ writeField counts HeaderWordCount 3
229+ writeField counts HeaderDone 1
233+resetHeaderOnNewChar :: V.IOVector Int -> IO ()
234+resetHeaderOnNewChar counts = do
235+ hdone <- readField counts HeaderDone
236+ when (hdone == 0) $ writeField counts HeaderDone 1
238+-------------------------------------------------------------------------------
239+-- Manipulating the trailer
240+-------------------------------------------------------------------------------
242+setTrailer :: V.IOVector Int -> S.DecodeState -> S.CodePoint -> IO ()
243+setTrailer counts st cp = do
244+ writeField counts TrailerState (fromIntegral st)
245+ writeField counts TrailerCodePoint cp
246+ writeField counts TrailerPresent 1
248+resetTrailerOnNewChar :: V.IOVector Int -> IO ()
249+resetTrailerOnNewChar counts = do
250+ trailer <- readField counts TrailerPresent
251+ when (trailer /= 0) $ do
252+ writeField counts TrailerPresent 0
253+ accountChar counts True
255+-------------------------------------------------------------------------------
256+-- Counting the stream
257+-------------------------------------------------------------------------------
259+{-# INLINE countChar #-}
260+countChar :: V.IOVector Int -> Either S.DecodeError Char -> IO ()
261+countChar counts inp =
264+ resetHeaderOnNewChar counts
265+ -- account the last stored error as whitespace and clear it
266+ resetTrailerOnNewChar counts
268+ when (ch == '\n') $ modifyField counts LineCount (+ 1)
270+ then accountChar counts True
272+ wasSpace <- readField counts WasSpace
273+ when (wasSpace /= 0) $ modifyField counts WordCount (+ 1)
274+ accountChar counts False
275+ Left (S.DecodeError st cp) -> do
276+ hdone <- readField counts HeaderDone
281+ -- We got a non-starter in initial decoder state, there may
282+ -- be something that comes before this to complete it.
283+ r <- addToHeader counts cp
284+ when (not r) $ error "countChar: Bug addToHeader failed"
286+ -- We got an error in a non-initial decoder state, it may
287+ -- be an input underflow error, keep it as incomplete in
289+ writeField counts HeaderDone 1
290+ setTrailer counts st cp
292+ resetTrailerOnNewChar counts
294+ then accountChar counts True
295+ else setTrailer counts st cp
297+printCounts :: V.IOVector Int -> IO ()
299+ l <- readField v LineCount
300+ w <- readField v WordCount
301+ c <- readField v CharCount
302+ putStrLn $ show l ++ " " ++ show w ++ " " ++ show c
304+-------------------------------------------------------------------------------
305+-- Serial counting using parallel version of countChar
306+-------------------------------------------------------------------------------
308+_wc_mwl_parserial :: Handle -> IO (V.IOVector Int)
309+_wc_mwl_parserial src = do
311+ S.mapM_ (countChar counts)
312+ $ S.decodeUtf8Either
313+ $ S.unfold FH.read src
316+-------------------------------------------------------------------------------
317+-- Serial word counting with UTF-8 handling
318+-------------------------------------------------------------------------------
320+data Counts = Counts !Int !Int !Int !Bool deriving Show
322+{-# INLINE countCharSerial #-}
323+countCharSerial :: Counts -> Char -> Counts
324+countCharSerial (Counts l w c wasSpace) ch =
325+ let l1 = if (ch == '\n') then l + 1 else l
329+ else (if wasSpace then w + 1 else w, False)
330+ in (Counts l1 w1 (c + 1) wasSpace1)
332+-- Note: This counts invalid byte sequences are non-space chars
333+_wc_mwl_serial :: Handle -> IO ()
334+_wc_mwl_serial src = print =<< (
335+ S.foldl' countCharSerial (Counts 0 0 0 True)
336+ $ S.decodeUtf8Lenient
337+ $ S.unfold FH.read src)
339+-------------------------------------------------------------------------------
341+-------------------------------------------------------------------------------
343+-- XXX we need a better data structure to store the header bytes to make these
346+-- combine trailing bytes in preceding block with leading bytes in the next
347+-- block and decode them into a codepoint
348+reconstructChar :: Int
351+ -> IO (S.SerialT IO (Either S.DecodeError Char))
352+reconstructChar hdrCnt v1 v2 = do
353+ when (hdrCnt > 3 || hdrCnt < 0) $ error "reconstructChar: hdrCnt > 3"
357+ x <- readField v2 HeaderWord3
358+ return $ (fromIntegral x :: Word8) `S.cons` S.nil
363+ x <- readField v2 HeaderWord2
364+ return $ fromIntegral x `S.cons` stream1
369+ x <- readField v2 HeaderWord1
370+ return $ fromIntegral x `S.cons` stream2
373+ state <- readField v1 TrailerState
374+ cp <- readField v1 TrailerCodePoint
375+ return $ S.resumeDecodeUtf8Either (fromIntegral state) cp stream3
377+getHdrChar :: V.IOVector Int -> IO (Maybe Int)
379+ hdrCnt <- readField v HeaderWordCount
383+ writeField v HeaderWordCount 0
384+ fmap Just $ readField v HeaderWord1
386+ x1 <- readField v HeaderWord1
387+ x2 <- readField v HeaderWord2
388+ writeField v HeaderWord1 x2
389+ writeField v HeaderWordCount 1
392+ x1 <- readField v HeaderWord1
393+ x2 <- readField v HeaderWord2
394+ x3 <- readField v HeaderWord3
395+ writeField v HeaderWord1 x2
396+ writeField v HeaderWord2 x3
397+ writeField v HeaderWordCount 2
399+ _ -> error "getHdrChar: Bug, hdrCnt not in range 0-3"
401+-- If the header of the first block is not done then combine the header
402+-- with the header of the next block.
403+combineHeaders :: V.IOVector Int -> V.IOVector Int -> IO ()
404+combineHeaders v1 v2 = do
405+ hdone1 <- readField v1 HeaderDone
408+ res <- getHdrChar v2
410+ Nothing -> return ()
412+ r <- addToHeader v1 x
413+ when (not r) $ error "combineHeaders: Bug, addToHeader failed"
416+-- We combine the contents of the second vector into the first vector, mutating
417+-- the first vector and returning it.
418+-- XXX This is a quick hack and can be refactored to reduce the size
419+-- and understandability considerably.
420+addCounts :: V.IOVector Int -> V.IOVector Int -> IO (V.IOVector Int)
422+ hdone1 <- readField v1 HeaderDone
423+ hdone2 <- readField v2 HeaderDone
424+ hdrCnt2_0 <- readField v2 HeaderWordCount
425+ if hdone1 == 0 && (hdrCnt2_0 /= 0 || hdone2 /= 0)
427+ combineHeaders v1 v2
431+ writeField v1 HeaderDone 1
434+ trailerPresent2 <- readField v2 TrailerPresent
435+ trailerState2 <- readField v2 TrailerState
436+ trailerCodePoint2 <- readField v2 TrailerCodePoint
438+ then error "addCounts: Bug, trying to add completely empty second block"
440+ l1 <- readField v1 LineCount
441+ w1 <- readField v1 WordCount
442+ c1 <- readField v1 CharCount
443+ wasSpace1 <- readField v1 WasSpace
445+ l2 <- readField v2 LineCount
446+ w2 <- readField v2 WordCount
447+ c2 <- readField v2 CharCount
448+ wasSpace2 <- readField v2 WasSpace
449+ firstIsSpace2 <- readField v2 FirstIsSpace
450+ hdrCnt2 <- readField v2 HeaderWordCount
452+ trailer1 <- readField v1 TrailerPresent
453+ if trailer1 == 0 -- no trailer in the first block
455+ -- header2, if any, complete or incomplete, is just invalid
456+ -- bytes, count them as whitespace
457+ let firstIsSpace2' = firstIsSpace2 /= 0 || hdrCnt2 /= 0
459+ if w2 /= 0 && wasSpace1 == 0 && not firstIsSpace2'
460+ then return $ w1 + w2 - 1
461+ else return $ w1 + w2
462+ writeField v1 LineCount (l1 + l2)
463+ writeField v1 WordCount w
464+ writeField v1 CharCount (c1 + c2 + hdrCnt2)
467+ if c2 == 0 && hdrCnt2 /= 0
468+ then writeField v1 FirstIsSpace 1
469+ else writeField v1 FirstIsSpace firstIsSpace2
471+ if c2 == 0 && hdrCnt2 /= 0
472+ then writeField v1 WasSpace 1
473+ else writeField v1 WasSpace wasSpace2
475+ writeField v1 TrailerPresent trailerPresent2
476+ writeField v1 TrailerState trailerState2
477+ writeField v1 TrailerCodePoint trailerCodePoint2
482+ when (hdone2 /= 0) $ do -- empty and Done header
483+ -- count trailer as whitespace, its counted as one char
484+ -- Note: hdrCnt2 == 0 means either header is not done
486+ writeField v1 LineCount (l1 + l2)
487+ writeField v1 WordCount (w1 + w2)
488+ writeField v1 CharCount (c1 + c2 + 1)
490+ when (c1 == 0) $ writeField v1 FirstIsSpace 1
492+ then writeField v1 WasSpace 1
493+ else writeField v1 WasSpace wasSpace2
495+ writeField v1 TrailerPresent trailerPresent2
496+ writeField v1 TrailerState trailerState2
497+ writeField v1 TrailerCodePoint trailerCodePoint2
498+ -- If header of the next block is not done we continue the
499+ -- trailer from the previous block instead of treating it
503+ -- join the trailing part of the first block with the
504+ -- header of the next
505+ decoded <- reconstructChar hdrCnt2 v1 v2
506+ res <- S.uncons decoded
508+ Nothing -> error "addCounts: Bug. empty reconstructed char"
510+ tlength <- S.length t
513+ -- If we have an error case after this
514+ -- char then that would be treated
516+ let lcount = l1 + l2
517+ lcount1 = if (ch == '\n') then lcount + 1 else lcount
519+ firstSpace = isSpace ch
520+ wasSpace = firstSpace || tlength /= 0
524+ else if w2 == 0 || firstIsSpace2 /= 0
527+ writeField v1 LineCount lcount1
528+ writeField v1 WordCount wcount1
529+ writeField v1 CharCount (c1 + c2 + 1 + tlength)
532+ writeField v1 FirstIsSpace
533+ (if firstSpace then 1 else 0)
538+ then writeField v1 WasSpace 1
539+ else writeField v1 WasSpace 0
540+ else writeField v1 WasSpace wasSpace2
542+ writeField v1 TrailerPresent trailerPresent2
543+ writeField v1 TrailerState trailerState2
544+ writeField v1 TrailerCodePoint trailerCodePoint2
546+ Left (S.DecodeError st cp) -> do
547+ -- if header was incomplete it may result
548+ -- in partially decoded char to be written
549+ -- as trailer. Check if the last error is
550+ -- an incomplete decode.
555+ Just lst -> case lst of
556+ Right _ -> error "addCounts: Bug"
557+ Left (S.DecodeError st1 cp1) -> (st1, cp1)
558+ if hdone2 == 0 && st' /= 12
560+ -- all elements before the last one must be errors
561+ writeField v1 CharCount (c1 + tlength)
564+ writeField v1 FirstIsSpace 1
566+ writeField v1 WasSpace 1
568+ writeField v1 TrailerState (fromIntegral st')
569+ writeField v1 TrailerCodePoint cp'
571+ -- all elements must be errors
572+ -- treat them as whitespace
573+ writeField v1 LineCount (l1 + l2)
574+ writeField v1 WordCount (w1 + w2)
575+ writeField v1 CharCount (c1 + c2 + tlength + 1)
578+ writeField v1 FirstIsSpace 1
581+ then writeField v1 WasSpace 1
582+ else writeField v1 WasSpace wasSpace2
584+ writeField v1 TrailerPresent trailerPresent2
585+ writeField v1 TrailerState trailerState2
586+ writeField v1 TrailerCodePoint trailerCodePoint2
589+-- Individual array processing is an isolated loop, fusing it with the bigger
590+-- loop may be counter productive.
591+{-# NOINLINE countArray #-}
592+countArray :: A.Array Word8 -> IO (V.IOVector Int)
595+ S.mapM_ (countChar counts)
596+ $ S.decodeUtf8Either
597+ $ S.unfold A.read src
600+{-# INLINE wc_mwl_parallel #-}
601+wc_mwl_parallel :: Handle -> Int -> IO (V.IOVector Int)
602+wc_mwl_parallel src n = do
604+ S.foldlM' addCounts counts
606+ $ S.maxThreads numCapabilities
607+ $ S.mapM (countArray)
608+ $ S.unfold FH.readArraysOf (n, src)
610+-------------------------------------------------------------------------------
612+-------------------------------------------------------------------------------
616+ name <- fmap head getArgs
617+ src <- openFile name ReadMode
618+ -- _wc_mwl_serial src -- Unix wc -l program
619+ -- printCounts =<< _wc_mwl_parserial src -- Unix wc -l program
620+ -- Using different sizes of chunks (1,2,3,4,5,10,128,256) is a good testing
621+ -- mechanism for parallel counting code.
624+ let chunkSize = read $ args !! 1
626+ let chunkSize = 32 * 1024
627+ printCounts =<< wc_mwl_parallel src chunkSize -- Unix wc -l program