Wcs2.hs

{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell   #-}

module Wcs2 where

import qualified Streaming.Prelude               as S

import           Control.Lens                            ( (+~)
                                                         , (^.)
                                                         , makeLenses
                                                         )
import           Control.Monad.Trans.Resource            ( runResourceT )
import qualified Data.ByteString.Char8           as C8
import qualified Data.ByteString.Lazy            as Lazy
import qualified Data.ByteString.Streaming.Char8 as Q
import           Data.Foldable                           ( traverse_ )
import           Data.Text                               ( Text )
import           Data.Text.Encoding                      ( encodeUtf8 )
import           System.Environment                      ( getArgs )
import           Text.Printf                             ( printf )

data CountState = CountState
  { _lineCount :: !Int
  , _wordCount :: !Int
  , _charCount :: !Int
  , _lastChar  :: !(Maybe Char)
  } deriving (Show)

makeLenses ''CountState

emptyCountState :: CountState
emptyCountState = CountState 0 0 0 Nothing

countStep :: Lazy.ByteString -> CountState -> CountState
countStep s = (lineCount +~ 1)
  . (wordCount +~ length (C8.words . Lazy.toStrict $ s))
  . (charCount +~ fromIntegral (Lazy.length s) + (1 {- for the newline -}))

countStreamingByteString :: Monad m => Q.ByteString m r -> m CountState
countStreamingByteString = id
  . S.fold_ (flip countStep) emptyCountState id
  . S.mapped Q.toLazy
  . Q.lines

wcFile :: FilePath -> IO CountState
wcFile = runResourceT . countStreamingByteString . Q.readFile

wcText :: Text -> IO CountState
wcText = countStreamingByteString . Q.fromStrict . encodeUtf8

display :: (Int, Int, Int) -> String
display (l, w, c) = printf "%8d %8d %8d" l w c

main :: IO ()
main = traverse_ ((putStrLn =<<) . fmap (display . extract) . wcFile) =<< getArgs

extract :: CountState -> (Int, Int, Int)
extract cs =
  ( cs ^. lineCount
  , cs ^. wordCount
  , cs ^. charCount
  )

as :: Text
as = "The quick brown fox jumped over the lazy\n log. And\n Stuff and things\n And they're going to be mad at us."