Ticket #2148: Bug.hs

{-# OPTIONS_GHC -XOverloadedStrings #-}
module Main where

import System
import qualified System.IO as IO
import qualified System.Directory as IO (removeFile)

import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BS (w2c,c2w)
import qualified Data.String as S

instance S.IsString BS.ByteString where
    fromString = BS.pack . map BS.c2w

--------------------------------------------------------------------------------
-- Pipeable Streams (based on Get/Put continuations)
--
-- lift f       -- turn f into a stream by passing through f
--
-- liftEither f -- wrap f with Either, passing Left through f and and Right
--                 arround f (typically used for Left being normal operation
--                 and Right being out-of-stream operation -- e.g., errors)
--
-- liftMaybe f  -- wrap f with Maybe, passing Just through f and Nothing around
--                 f (typically used for end-of-stream operation)
--
data Stream a b = Get (a -> Stream a b)
                | Put b (Stream a b)

pipe :: Stream a b -> Stream b c -> Stream a c
pipe f (Put y g)       = Put y $ pipe f g
pipe (Put x f) (Get g) = pipe f $ g x
pipe (Get f) g         = Get (\x -> pipe (f x) g)

stack :: Int -> Stream a a -> Stream a a
stack n f | n>1 = f `pipe` stack (n-1) f
stack n f       = f

lift :: (a -> b) -> Stream a b
lift f = Get get
    where --get :: a -> Stream a b
          get x = Put (f x) $ Get get

liftEither :: Stream a b -> Stream (Either a e) (Either b e)
liftEither (Get f)   = Get get
    where -- get :: Either a e -> Stream (Either a e) (Either b e)
          get (Left x)  = liftEither $ f x
          get (Right e) = Put (Right e) $ Get get
liftEither (Put x f) = Put (Left x) $ liftEither f

liftMaybe :: Stream a b -> Stream (Maybe a) (Maybe b)
liftMaybe (Get f)   = Get get
    where -- get :: Maybe a -> Stream (Maybe a) (Maybe b)
          get (Just x) = liftMaybe $ f x
          get Nothing  = Put Nothing $ Get get
liftMaybe (Put x f) = Put (Just x) $ liftMaybe f


--------------------------------------------------------------------------------
-- record
--
-- Input  ">Header 1\nABCDDFF\nAHDHFDSHE\n>Header 2\nAHSDHFF\nSDHJFD"...
-- Output Record "Header 1" "ABCDDFFAHDHFDSHE",
--        Record "Header 2" "AHSDHFFSDHJFD"...
--
type Header = BS.ByteString
type Sequence = BS.ByteString
data Record = Record !Header !Sequence deriving Show

record :: Stream (Maybe BS.ByteString) (Either (Maybe Record) BS.ByteString)
record = Get recordHeader
    where recordHeader :: Maybe BS.ByteString -> Stream (Maybe BS.ByteString) (Either (Maybe Record) BS.ByteString)
          recordHeader (Just l) | BS.head l == BS.c2w '>' = Get $ recordBody (BS.tail l) []
          recordHeader (Just l)                           = Put (Right $ BS.append "bad header line: " l) $ Get recordHeader
          recordHeader Nothing                            = Put (Left Nothing) $ Get recordHeader
          recordBody :: Header -> [Sequence] -> Maybe BS.ByteString -> Stream (Maybe BS.ByteString) (Either (Maybe Record) BS.ByteString)
          recordBody h ls (Just l) | BS.head l /= BS.c2w '>' = Get $ recordBody h (l:ls)
          recordBody h [] (Just l)                           = Put (Right $ BS.append "no body for: " h) $ recordHeader $ Just l
          recordBody h ls (Just l)                           = Put (Left $ Just $ Record h $ BS.concat $ reverse ls) $ recordHeader $ Just l
          recordBody h [] Nothing                            = Put (Right $ BS.append "no body for: " h) $ Put (Left Nothing) $ Get recordHeader
          recordBody h ls Nothing                            = Put (Left $ Just $ Record h $ BS.concat $ reverse ls) $ Put (Left Nothing) $ Get recordHeader

--------------------------------------------------------------------------------
-- divide n
--
-- Input  Record "Header 1" "ABCDEF", ...
-- Output "ABC","BCD, "CDE", "DEF", ...
--
-- Where n is the length of the bits (n=3 above)
--
divide :: Int -> Stream Record BS.ByteString
divide n = Get extract
    where extract :: Record -> Stream Record BS.ByteString
          extract (Record h l) | BS.length l >= n = Put (BS.take n l) $ extract $ Record h $ BS.tail l
          extract (Record _ l)                    = Get extract

--------------------------------------------------------------------------------
-- group n
--
-- Input  "ABC", "DEF", "BCD", "ABC", ....
-- Output ["ABC","DEF"], ["BCD","ABC"], ...
--
-- Where n is the group size (n=2 above)
--
group :: Int -> Stream (Maybe a) (Maybe [a])
group n = Get $ group [] n
    where group :: [a] -> Int -> Maybe a -> Stream (Maybe a) (Maybe [a])
          group xs n (Just x) | n > 1 = Get $ group (x:xs) (n-1)
          group xs _ (Just x)         = Put (Just $ reverse $ x:xs) $ Get $ group [] n
          group [] _ Nothing          = Put Nothing $ Get $ group [] n
          group xs _ Nothing          = Put (Just $ reverse xs) $ Put Nothing $ Get $ group [] n

--------------------------------------------------------------------------------
-- process s d
--
-- Input  ">Header 1\nABCDDFF\nAHDHFDSHE\n>Header 2\nAHSDHFF\nSDHJFD"...
-- Output [("ABC,1),("BCD",1),("CDD",1),("DDF",1)],
--        [("AHD,1),("DFF",1),("DHF",1),("HDH",1)],...
--
-- Where s is length of the bits (s=3 above) and 2^d is the blocking size
-- (n=2 above)
--
process :: Int -> Int -> Stream (Maybe BS.ByteString) (Either (Maybe [BS.ByteString]) BS.ByteString)
process s d = record `pipe` (liftEither $ (liftMaybe $ divide s) `pipe` group d)
--    where wrap :: Stream BS.ByteString [BS.ByteString]
--          wrap = lift (\x -> [x])


--------------------------------------------------------------------------------
-- main loop
--
-- Setup the above process stream according to command-line paramaters, pass
-- the contents of the file through it (a line at a time), and write each
-- result (a collection of processed elements) out to seperate temporary file.
---
main = do IO.hSetBuffering IO.stdout IO.LineBuffering
          args <- getArgs :: IO [String]
          if length args /= 3 then do prog <- getProgName
                                      fail $ "command line: " ++ prog ++ " input size depth"
                              else return ()
          input <- return $ args !! 0 :: IO String
          size <- catch (read $ args !! 1) $ \_ -> fail $ "size should be an integer not " ++ (show $ args !! 1) :: IO Int
          depth <- catch (read $ args !! 2) $ \_ -> fail $ "depth should be an integer not " ++ (show $ args !! 2) :: IO Int
          file <- IO.openFile input IO.ReadMode :: IO IO.Handle
          hdls <- chew file [] $ process size depth :: IO [IO.Handle]
          mapM_ IO.hClose hdls
          IO.hClose file
    where read :: (Read a) => String -> IO a
          read xs = case [ x | (x,xs') <- reads xs, ("","") <- lex xs' ] of
                    [x] -> return x
                    _   -> fail "read: ambiguous parse"
          chew :: IO.Handle -> [IO.Handle] -> Stream (Maybe BS.ByteString) (Either (Maybe [BS.ByteString]) BS.ByteString) -> IO [IO.Handle]
          chew h hs (Put (Left (Just xs)) f) = do IO.hTell h >>= \p -> putStr $ "Starting new sorted block at " ++ show p ++ " ("
                                                  h' <- IO.openTempFile "/tmp" "PChunk.tmp" >>= \(nm,h) -> putStrLn (nm ++ ")") >> return h
                                                  mapM_ (BS.hPutStrLn h') xs
                                                  chew h (h':hs) f
          chew h hs (Put (Left Nothing) f)   = mapM (\h -> IO.hSeek h IO.AbsoluteSeek 0 >> return h) hs
          chew h hs (Put (Right m) f)        = do BS.hPutStrLn IO.stderr $ BS.append "**WARNING** " m
                                                  chew h hs f
          chew h hs (Get f)                  = do eof <- IO.hIsEOF h
                                                  if eof
                                                     then chew h hs $ f Nothing
                                                     else do l <- BS.hGetLine h
                                                             chew h hs $ f $ Just l