io-streams + beautiful folding

foldIOStream.hs

{-# LANGUAGE OverloadedStrings, BangPatterns #-}

import System.IO.Streams hiding (mapM_, stderr)
import qualified System.IO.Streams as Streams
import qualified  System.IO as IO

import Prelude hiding (writeFile, splitAt, read)
import Control.Applicative
import Control.Monad ((>=>), liftM)
import Control.Foldl (FoldM(..), Fold(..),impurely)
import qualified Control.Foldl as L

import qualified Data.Text.Encoding as T
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.Text (Text)
import Data.Vector.Unboxed (Vector)

import Control.Lens

-- ------------------------------------
-- proposed new io-stream folds, for use with 
--   `purely` and `impurely` in the `foldl` library.
--   Note that they involve no dependency.
--    The existing `Streams.fold` and `Streams.foldM` can be 
--    defined as:  fold op seed  = fold_ op seed id
--                 foldM op seed = foldM_ op seed return
-- ------------------------------------

fold_ :: (x -> a -> x)    -- ^ accumulator update function
      -> x                -- ^ initial seed
      -> (x -> s)         -- ^ recover folded value
      -> InputStream a    -- ^ input stream
      -> IO s
fold_ op seed done stream = liftM done (go seed)
   where 
     go !s = Streams.read stream >>= maybe (return s) (go . op s)

-- > :t L.purely fold_
-- L.purely fold_ :: L.Fold a b -> InputStream a -> IO b

foldM_ :: (x -> a -> IO x)   -- ^ accumulator update action
       -> IO x               -- ^ initial seed
       -> (x -> IO s)        -- ^ recover folded value
       -> InputStream a      -- ^ input stream
       -> IO s
foldM_ f seed done stream = seed >>= go 
  where
    go !x = Streams.read stream >>= maybe (done x) ((go =<<) . f x)

    -- > :t L.impurely foldM_
    -- L.impurely foldM_ :: FoldM IO a b -> InputStream a -> IO b


-- ------------------------------------
-- streaming calculation of average: 
--   (the textbook illustration of composable folds)
--   here we just use `purely` and the variant io-streams fold
--   defined above. Here we don't use `pretraverse\M` of `premap\M`.
-- ------------------------------------
stream1 :: IO ()
stream1 = do 
  input <- Streams.fromList [1..20::Int] 
  m <- L.purely fold_ (div <$> L.sum <*>  L.length) input 
  print m
  
-- >>> stream1
-- 10

-- ------------------------------------
-- similar to average, but we add an 
--     impure fold building a vector
--     thus 'generalizing' the pure folds
-- ------------------------------------
stream2 = do 
  input <- Streams.fromList [1..20::Int] 
  let stats =  (,,,) <$> (div <$> L.sum <*> L.length)  
                     <*> L.last  
                     <*> L.head  
                     <*> L.all even
  (v,out) <- L.impurely foldM_ ((,) <$> L.vector <*> L.generalize stats) input
  printStats out 
  putStr "Whole vector:   "
  print (v :: Vector Int)
 where 
  printStats (av, lst, h, e) = mapM_ (putStrLn . uncurry (++)) 
                                [("Vector average: ", show av) 
                                ,("Last element:   ", show lst)
                                ,("First element:  ", show h)  
                                ,("All are even?:  ", show e)]
-- >>> stream2
-- Vector average: 10
-- Last element:   Just 20
-- First element:  Just 1
-- All are even?:  False
-- Whole vector:   [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20]



-- ------------------------------------
-- a simple example using the stock _Left _Right traversals
-- ------------------------------------

sheep_from_goats :: IO.Handle -> IO.Handle -> FoldM IO (Either Text Text) ()
sheep_from_goats h1 h2 = 
         L.handlesM _Left (L.sink (T.hPutStrLn h1)) 
     <*  L.handlesM _Right (L.sink (T.hPutStrLn h2))

sheep_and_goats :: [Either Text Text]
sheep_and_goats = [Left "Goat", Right "Sheep", Right "Sheep", Right "Sheep", Left "Goat"]

stream3 = IO.withFile "sheep.txt" IO.WriteMode $ \h -> 
       do input <- Streams.fromList sheep_and_goats
          L.impurely foldM_ (sheep_from_goats IO.stderr h) input
          
-- >>> stream3
-- Goat
-- Goat
-- *Main
-- >>> :! cat sheep.txt
-- Sheep
-- Sheep
-- Sheep