@@ -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
11−-- handling.
12−
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"
19−
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.
30−
31−{-# LANGUAGE CPP #-}
32−
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)
42−
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
48−
49−-------------------------------------------------------------------------------
50−-- Parallel char, line and word counting
51−-------------------------------------------------------------------------------
52−
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−-------------------------------------------------------------------------------
59−
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−-------------------------------------------------------------------------------
89−
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−-------------------------------------------------------------------------------
114−
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.
117−
118−-------------------------------------------------------------------------------
119−-- Counting state
120−-------------------------------------------------------------------------------
121−
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.
131−
132−-- XXX we can also count the number of decoding errors separately
133−data 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)
173−
174−-------------------------------------------------------------------------------
175−-- Default/initial state of the block
176−-------------------------------------------------------------------------------
177−
178−readField :: V.IOVector Int -> Field -> IO Int
179−readField v fld = V.read v (fromEnum fld)
180−
181−writeField :: V.IOVector Int -> Field -> Int -> IO ()
182−writeField v fld val = V.write v (fromEnum fld) val
183−
184−modifyField :: V.IOVector Int -> Field -> (Int -> Int) -> IO ()
185−modifyField v fld f = V.modify v f (fromEnum fld)
186−
187−newCounts :: IO (V.IOVector Int)
188−newCounts = 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
199−
200−-------------------------------------------------------------------------------
201−-- Counting chars
202−-------------------------------------------------------------------------------
203−
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
211−
212−-------------------------------------------------------------------------------
213−-- Manipulating the header bytes
214−-------------------------------------------------------------------------------
215−
216−addToHeader :: V.IOVector Int -> Int -> IO Bool
217−addToHeader 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
234−
235−resetHeaderOnNewChar :: V.IOVector Int -> IO ()
236−resetHeaderOnNewChar counts = do
237− hdone <- readField counts HeaderDone
238− when (hdone == 0) $ writeField counts HeaderDone 1
239−
240−-------------------------------------------------------------------------------
241−-- Manipulating the trailer
242−-------------------------------------------------------------------------------
243−
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
249−
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
256−
257−-------------------------------------------------------------------------------
258−-- Counting the stream
259−-------------------------------------------------------------------------------
260−
261−{-# INLINE countChar #-}
262−countChar :: V.IOVector Int -> Either DecodeError Char -> IO ()
263−countChar 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
269−
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
298−
299−printCounts :: V.IOVector Int -> IO ()
300−printCounts 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
305−
306−-------------------------------------------------------------------------------
307−-- Serial counting using parallel version of countChar
308−-------------------------------------------------------------------------------
309−
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
317−
318−-------------------------------------------------------------------------------
319−-- Serial word counting with UTF-8 handling
320−-------------------------------------------------------------------------------
321−
322−data Counts = Counts !Int !Int !Int !Bool deriving Show
323−
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
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)
333−
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)
340−
341−-------------------------------------------------------------------------------
342−-- Parallel counting
343−-------------------------------------------------------------------------------
344−
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
350−reconstructChar :: Int
351− -> V.IOVector Int
352− -> V.IOVector 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"
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
374−
375− state <- readField v1 TrailerState
376− cp <- readField v1 TrailerCodePoint
377− return $ resumeDecodeUtf8Either (fromIntegral state) cp stream3
378−
379−getHdrChar :: V.IOVector Int -> IO (Maybe Int)
380−getHdrChar 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"
402−
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
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 ()
417−
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)
423−addCounts 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
446−
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
453−
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)
467−
468− when (c1 == 0) $ do
469− if c2 == 0 && hdrCnt2 /= 0
470− then writeField v1 FirstIsSpace 1
471− else writeField v1 FirstIsSpace firstIsSpace2
472−
473− if c2 == 0 && hdrCnt2 /= 0
474− then writeField v1 WasSpace 1
475− else writeField v1 WasSpace wasSpace2
476−
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)
491−
492− when (c1 == 0) $ writeField v1 FirstIsSpace 1
493− if (c2 == 0)
494− then writeField v1 WasSpace 1
495− else writeField v1 WasSpace wasSpace2
496−
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)
532−
533− when (c1 == 0) $
534− writeField v1 FirstIsSpace
535− (if firstSpace then 1 else 0)
536−
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
543−
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)
564−
565− when (c1 == 0) $
566− writeField v1 FirstIsSpace 1
567−
568− writeField v1 WasSpace 1
569−
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)
578−
579− when (c1 == 0) $
580− writeField v1 FirstIsSpace 1
581−
582− if c2 == 0
583− then writeField v1 WasSpace 1
584− else writeField v1 WasSpace wasSpace2
585−
586− writeField v1 TrailerPresent trailerPresent2
587− writeField v1 TrailerState trailerState2
588− writeField v1 TrailerCodePoint trailerCodePoint2
589− return v1
590−
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)
595−countArray src = do
596− counts <- newCounts
597− S.mapM_ (countChar counts)
598− $ decodeUtf8Either
599− $ S.unfold A.read src
600− return counts
601−
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
606− $ S.aheadly
607− $ S.maxThreads numCapabilities
608− $ S.mapM (countArray)
609− $ S.unfold FH.readChunksWithBufferOf (n, src)
610−
611−-------------------------------------------------------------------------------
612−-- Main
613−-------------------------------------------------------------------------------
614−
615−main :: IO ()
616−main = 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