@@ -1,628 +0,0 @@
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(..))
39−import Streamly.Internal.Unicode.Stream
40− (DecodeState, DecodeError(..), CodePoint, decodeUtf8Either,
41− resumeDecodeUtf8Either)
43−import qualified Streamly.Internal.Data.Stream.IsStream as S
44−import qualified Streamly.Unicode.Stream as S
45−import qualified Streamly.FileSystem.Handle as FH
46−import qualified Streamly.Data.Array.Storable.Foreign as A
47−import 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.
56−-------------------------------------------------------------------------------
58−-------------------------------------------------------------------------------
60−-- To count chars each block needs the following:
62−-- -- | header | char counts | trailer |
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.
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.
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:
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.
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.
86−-------------------------------------------------------------------------------
88−-------------------------------------------------------------------------------
90−-- For word counting we need the following in each block:
92−-- -- | header | startsWithSpace | word counts | endsWithSpace | trailer |
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
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.
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.
111−-------------------------------------------------------------------------------
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−-------------------------------------------------------------------------------
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.
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
134− -- The number of "\n" characters found in the block.
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
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
148− -- whether the last counted char in this block was a space char
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.
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.
159− -- XXX This is ugly to manipulate, we can implement a statically max sized
160− -- mutable ring structure within this record.
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.
172− deriving (Show, Enum, Bounded)
174−-------------------------------------------------------------------------------
175−-- Default/initial state of the block
176−-------------------------------------------------------------------------------
178−readField :: V.IOVector Int -> Field -> IO Int
179−readField v fld = V.read v (fromEnum fld)
181−writeField :: V.IOVector Int -> Field -> Int -> IO ()
182−writeField v fld val = V.write v (fromEnum fld) val
184−modifyField :: V.IOVector Int -> Field -> (Int -> Int) -> IO ()
185−modifyField v fld f = V.modify v f (fromEnum fld)
187−newCounts :: IO (V.IOVector Int)
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
200−-------------------------------------------------------------------------------
202−-------------------------------------------------------------------------------
204−accountChar :: V.IOVector Int -> Bool -> IO ()
205−accountChar 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−-------------------------------------------------------------------------------
216−addToHeader :: V.IOVector Int -> Int -> IO Bool
217−addToHeader counts cp = do
218− cnt <- readField counts HeaderWordCount
221− writeField counts HeaderWord1 cp
222− writeField counts HeaderWordCount 1
225− writeField counts HeaderWord2 cp
226− writeField counts HeaderWordCount 2
229− writeField counts HeaderWord3 cp
230− writeField counts HeaderWordCount 3
231− writeField counts HeaderDone 1
235−resetHeaderOnNewChar :: V.IOVector Int -> IO ()
236−resetHeaderOnNewChar counts = do
237− hdone <- readField counts HeaderDone
238− when (hdone == 0) $ writeField counts HeaderDone 1
240−-------------------------------------------------------------------------------
241−-- Manipulating the trailer
242−-------------------------------------------------------------------------------
244−setTrailer :: V.IOVector Int -> DecodeState -> CodePoint -> IO ()
245−setTrailer counts st cp = do
246− writeField counts TrailerState (fromIntegral st)
247− writeField counts TrailerCodePoint cp
248− writeField counts TrailerPresent 1
250−resetTrailerOnNewChar :: V.IOVector Int -> IO ()
251−resetTrailerOnNewChar 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 #-}
262−countChar :: V.IOVector Int -> Either DecodeError Char -> IO ()
263−countChar counts inp =
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)
272− then accountChar counts True
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
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"
288− -- We got an error in a non-initial decoder state, it may
289− -- be an input underflow error, keep it as incomplete in
291− writeField counts HeaderDone 1
292− setTrailer counts st cp
294− resetTrailerOnNewChar counts
296− then accountChar counts True
297− else setTrailer counts st cp
299−printCounts :: V.IOVector Int -> IO ()
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
313− S.mapM_ (countChar counts)
315− $ S.unfold FH.read src
318−-------------------------------------------------------------------------------
319−-- Serial word counting with UTF-8 handling
320−-------------------------------------------------------------------------------
322−data Counts = Counts !Int !Int !Int !Bool deriving Show
324−{-# INLINE countCharSerial #-}
325−countCharSerial :: Counts -> Char -> Counts
326−countCharSerial (Counts l w c wasSpace) ch =
327− let l1 = if (ch == '\n') then l + 1 else l
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)
339− $ S.unfold FH.read src)
341−-------------------------------------------------------------------------------
343−-------------------------------------------------------------------------------
345−-- XXX we need a better data structure to store the header bytes to make these
348−-- combine trailing bytes in preceding block with leading bytes in the next
349−-- block and decode them into a codepoint
350−reconstructChar :: Int
353− -> IO (S.SerialT IO (Either DecodeError Char))
354−reconstructChar hdrCnt v1 v2 = do
355− when (hdrCnt > 3 || hdrCnt < 0) $ error "reconstructChar: hdrCnt > 3"
359− x <- readField v2 HeaderWord3
360− return $ (fromIntegral x :: Word8) `S.cons` S.nil
365− x <- readField v2 HeaderWord2
366− return $ fromIntegral x `S.cons` stream1
371− x <- readField v2 HeaderWord1
372− return $ fromIntegral x `S.cons` stream2
375− state <- readField v1 TrailerState
376− cp <- readField v1 TrailerCodePoint
377− return $ resumeDecodeUtf8Either (fromIntegral state) cp stream3
379−getHdrChar :: V.IOVector Int -> IO (Maybe Int)
381− hdrCnt <- readField v HeaderWordCount
385− writeField v HeaderWordCount 0
386− fmap Just $ readField v HeaderWord1
388− x1 <- readField v HeaderWord1
389− x2 <- readField v HeaderWord2
390− writeField v HeaderWord1 x2
391− writeField v HeaderWordCount 1
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
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.
405−combineHeaders :: V.IOVector Int -> V.IOVector Int -> IO ()
406−combineHeaders v1 v2 = do
407− hdone1 <- readField v1 HeaderDone
410− res <- getHdrChar v2
412− Nothing -> return ()
414− r <- addToHeader v1 x
415− when (not r) $ error "combineHeaders: Bug, addToHeader failed"
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.
422−addCounts :: V.IOVector Int -> V.IOVector Int -> IO (V.IOVector Int)
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)
429− combineHeaders v1 v2
433− writeField v1 HeaderDone 1
436− trailerPresent2 <- readField v2 TrailerPresent
437− trailerState2 <- readField v2 TrailerState
438− trailerCodePoint2 <- readField v2 TrailerCodePoint
440− then error "addCounts: Bug, trying to add completely empty second block"
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
457− -- header2, if any, complete or incomplete, is just invalid
458− -- bytes, count them as whitespace
459− let firstIsSpace2' = firstIsSpace2 /= 0 || hdrCnt2 /= 0
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)
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
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
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
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
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
510− Nothing -> error "addCounts: Bug. empty reconstructed char"
512− tlength <- S.length t
515− -- If we have an error case after this
516− -- char then that would be treated
518− let lcount = l1 + l2
519− lcount1 = if (ch == '\n') then lcount + 1 else lcount
521− firstSpace = isSpace ch
522− wasSpace = firstSpace || tlength /= 0
526− else if w2 == 0 || firstIsSpace2 /= 0
529− writeField v1 LineCount lcount1
530− writeField v1 WordCount wcount1
531− writeField v1 CharCount (c1 + c2 + 1 + tlength)
534− writeField v1 FirstIsSpace
535− (if firstSpace then 1 else 0)
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
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.
557− Just lst -> case lst of
558− Right _ -> error "addCounts: Bug"
559− Left (DecodeError st1 cp1) -> (st1, cp1)
560− if hdone2 == 0 && st' /= 12
562− -- all elements before the last one must be errors
563− writeField v1 CharCount (c1 + tlength)
566− writeField v1 FirstIsSpace 1
568− writeField v1 WasSpace 1
570− writeField v1 TrailerState (fromIntegral st')
571− writeField v1 TrailerCodePoint cp'
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)
580− writeField v1 FirstIsSpace 1
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
591−-- Individual array processing is an isolated loop, fusing it with the bigger
592−-- loop may be counter productive.
593−{-# NOINLINE countArray #-}
594−countArray :: A.Array Word8 -> IO (V.IOVector Int)
597− S.mapM_ (countChar counts)
599− $ S.unfold A.read src
602−{-# INLINE wc_mwl_parallel #-}
603−wc_mwl_parallel :: Handle -> Int -> IO (V.IOVector Int)
604−wc_mwl_parallel src n = do
605− S.foldlM' addCounts newCounts
607− $ S.maxThreads numCapabilities
608− $ S.mapM (countArray)
609− $ S.unfold FH.readChunksWithBufferOf (n, src)
611−-------------------------------------------------------------------------------
613−-------------------------------------------------------------------------------
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.
625− let chunkSize = read $ args !! 1
627− let chunkSize = 32 * 1024
628− printCounts =<< wc_mwl_parallel src chunkSize -- Unix wc -l program