Iteratee attempt

gistfile1.hs

{-# LANGUAGE OverloadedStrings #-}

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Enumerator (Iteratee, Enumeratee, ($=), (=$), (=$=), (==<<))
import qualified Data.Enumerator as E
import qualified Data.Enumerator.Binary as EB
import qualified Data.Enumerator.Text as ET
import qualified Data.Enumerator.List as EL
import Data.Text
import System.Environment

-- http://okmij.org/ftp/Haskell/Iteratee/describe.pdf


------------------------------------------------------------------------------
countCharBS :: (Monad m) => Char -> Iteratee ByteString m Integer
countCharBS needle = loop 0
  where loop n = EL.head >>= check n
        check n Nothing = return n
        check n (Just t) = loop (n + toInteger (B.count needle t))


------------------------------------------------------------------------------
countLines :: (Monad m) => Iteratee Text m Integer
countLines = E.length


------------------------------------------------------------------------------
convertUtf8 :: (Monad m) => Enumeratee ByteString Text m b
convertUtf8 = ET.decode ET.utf8


------------------------------------------------------------------------------
main :: IO ()
main = do
    args <- getArgs
    case args of
      (fname:needle:[]) ->
          case needle of
            (c:"") -> do
                print "Switching to single char scan..."
                cnt <- E.run $ EB.enumFile fname ==<< countCharBS c
                print cnt
            _ -> do
                cnt <- E.run $ (EB.enumFile fname $= convertUtf8) ==<<
                      countLines
                print cnt
      _ -> print "Error in accessing file."