module Data.FilePlow
( PlowHandle(..), Handle, SeekMode(..)
, seekUntil, seekUntilRev
, MultiHandle, withMultiHandle
)
where
import Control.Monad
import Data.IORef
import GHC.IO.Exception
import GHC.IO.Handle
import System.IO
import qualified Data.ByteString as BS
import qualified Data.Vector as V
class PlowHandle hdl where
pSeek :: hdl -> SeekMode -> Integer -> IO ()
pTell :: hdl -> IO Integer
pGetChar :: hdl -> IO Char
pGetLine :: hdl -> IO BS.ByteString
pIsEOF :: hdl -> IO Bool
pFileSize :: hdl -> IO Integer
instance PlowHandle Handle where
pSeek = hSeek
pTell = hTell
pGetChar = hGetChar
pGetLine = BS.hGetLine
pIsEOF = hIsEOF
pFileSize = hFileSize
data MultiHandle h
= MultiHandle
{ mh_handles :: !(V.Vector h)
, mh_currentIndex :: !(IORef Int)
}
seekUntilRev :: PlowHandle hdl => hdl -> (Char -> Bool) -> IO Bool
seekUntilRev hdl pp =
do let getLoop =
do p <- pTell hdl
if p == 0
then pure False
else do x <- pGetChar hdl
if pp x
then pure True
else do pSeek hdl RelativeSeek (2)
getLoop
getLoop
seekUntil :: PlowHandle hdl => hdl -> (Char -> Bool) -> IO Bool
seekUntil hdl pp =
do let getLoop =
do eof <- pIsEOF hdl
if eof
then pure False
else do x <- pGetChar hdl
if pp x then pure True else getLoop
getLoop
instance (PlowHandle h) => PlowHandle (MultiHandle h) where
pFileSize mh =
V.foldM' (\total hdl -> (+ total) <$> pFileSize hdl) 0 (mh_handles mh)
pIsEOF mh =
do idx <- readIORef (mh_currentIndex mh)
let ct = V.length (mh_handles mh)
if idx >= ct 1
then pIsEOF $ mh_handles mh V.! (ct 1)
else pure False
pGetChar mh =
do h <- getCurrentHandle "pGetChar" mh
pGetChar h
pGetLine mh =
do h <- getCurrentHandle "pGetLine" mh
pGetLine h
pSeek mh sm val =
case sm of
AbsoluteSeek -> absoluteSeek mh val
RelativeSeek -> relativeSeek mh val
SeekFromEnd -> seekFromEnd mh val
pTell mh =
do idx <- readIORef (mh_currentIndex mh)
let h = mh_handles mh V.! idx
localPos <- pTell h
foldM
(\total i -> (+ total) <$> pFileSize (mh_handles mh V.! i))
localPos [0 .. (idx 1)]
seekFromEnd :: (PlowHandle h) => MultiHandle h -> Integer -> IO ()
seekFromEnd mh tval =
let ct = V.length (mh_handles mh)
seekLoop !tgtVal !idx
| idx >= ct || idx < 0 =
ioException (IOError Nothing EOF "seekFromEnd" "No more file handles" Nothing Nothing)
| otherwise =
do let h = mh_handles mh V.! idx
size <- pFileSize h
if abs tgtVal >= size
then seekLoop ((1) * (abs tgtVal size)) (idx 1)
else do writeIORef (mh_currentIndex mh) idx
pSeek h SeekFromEnd tgtVal
in seekLoop tval (ct 1)
relativeSeek :: (PlowHandle h) => MultiHandle h -> Integer -> IO ()
relativeSeek mh tval =
do myIdx <- readIORef (mh_currentIndex mh)
let localH = mh_handles mh V.! myIdx
localPos <- pTell localH
pSeek localH AbsoluteSeek 0
let seekVal = tval + localPos
seekHelp mh seekVal myIdx
absoluteSeek :: PlowHandle h => MultiHandle h -> Integer -> IO ()
absoluteSeek mh tval =
seekHelp mh tval 0
seekHelp :: PlowHandle h => MultiHandle h -> Integer -> Int -> IO ()
seekHelp mh t i =
let ct = V.length (mh_handles mh)
seekLoop !tgtVal !idx
| idx >= ct =
let msg = "No more file handles, desired pos: " ++ show t
in ioException (IOError Nothing EOF "seekHelp" msg Nothing Nothing)
| tgtVal < 0 && idx > 0 =
do let h = mh_handles mh V.! (idx 1)
size <- pFileSize h
pSeek h AbsoluteSeek 0
seekLoop (tgtVal + size) (idx 1)
| otherwise =
do let h = mh_handles mh V.! idx
size <- pFileSize h
if tgtVal >= size
then seekLoop (tgtVal size) (idx + 1)
else do writeIORef (mh_currentIndex mh) idx
pSeek h AbsoluteSeek tgtVal
in seekLoop t i
getCurrentHandle :: PlowHandle h => String -> MultiHandle h -> IO h
getCurrentHandle helpTxt mh =
do startIdx <- readIORef (mh_currentIndex mh)
let ct = V.length (mh_handles mh)
gotoUsefulHdl !i !nh
| i >= ct =
ioException (IOError Nothing EOF ("currentHandle/" ++ helpTxt) "No more file handles" Nothing Nothing)
| otherwise =
do let h = mh_handles mh V.! i
when nh $
pSeek h AbsoluteSeek 0
eof <- pIsEOF h
if eof
then gotoUsefulHdl (i + 1) True
else do writeIORef (mh_currentIndex mh) i
pure h
gotoUsefulHdl startIdx False
withMultiHandle :: [FilePath] -> (MultiHandle Handle -> IO a) -> IO a
withMultiHandle files go =
loop [] files
where
loop accum xs =
case xs of
[] ->
do r <- newIORef 0
go (MultiHandle (V.fromList (reverse accum)) r)
(fp : more) ->
withFile fp ReadMode $ \hdl ->
loop (hdl : accum) more