{-# LANGUAGE DeriveDataTypeable #-}

-- |
-- Module: Data.Knob
-- Copyright: 2011 John Millikin
-- License: MIT
--
-- Maintainer: n@monade.li
-- Portability: GHC only
--
-- Create memory-backed 'IO.Handle's, referencing virtual files. This is
-- mostly useful for testing 'IO.Handle'-based APIs without having to
-- interact with the filesystem.
--
-- > import Data.ByteString (pack)
-- > import Data.Knob
-- > import System.IO
-- >
-- > main = do
-- >     knob <- newKnob (pack [])
-- >     h <- newFileHandle knob "test.txt" WriteMode
-- >     hPutStrLn h "Hello world!"
-- >     hClose h
-- >     bytes <- Data.Knob.getContents knob
-- >     putStrLn ("Wrote bytes: " ++ show bytes)
module Data.Knob
  ( Knob
  , newKnob
  , Data.Knob.getContents
  , setContents

  , newFileHandle
  , withFileHandle
  ) where

import qualified Control.Concurrent.MVar as MVar
import           Control.Exception (bracket, throwIO)
import           Control.Monad (when)
import           Control.Monad.IO.Class (MonadIO, liftIO)
import qualified Data.ByteString
import           Data.ByteString (ByteString)
import           Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import           Data.Typeable (Typeable)
import qualified Foreign
import qualified GHC.IO.Buffer as IO
import qualified GHC.IO.BufferedIO as IO
import qualified GHC.IO.Device as IO
import qualified GHC.IO.Exception as IO
import qualified GHC.IO.Handle as IO
import qualified System.IO as IO

-- | A knob is a basic virtual file, which contains a byte buffer. A knob can
-- have multiple 'IO.Handle's open to it, each of which behaves like a standard
-- file handle.
--
-- Use 'Data.Knob.getContents' and 'setContents' to inspect and modify the knob's
-- byte buffer.
newtype Knob = Knob (MVar.MVar ByteString)

data Device = Device IO.IOMode (MVar.MVar ByteString) (MVar.MVar Int)
  deriving (Typeable)

instance IO.IODevice Device where
  ready :: Device -> Bool -> Int -> IO Bool
ready Device
_ Bool
_ Int
_ = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
  close :: Device -> IO ()
close Device
_ = forall (m :: * -> *) a. Monad m => a -> m a
return ()
  isTerminal :: Device -> IO Bool
isTerminal Device
_ = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  isSeekable :: Device -> IO Bool
isSeekable Device
_ = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

  seek :: Device -> SeekMode -> Integer -> IO Integer
seek (Device IOMode
_ MVar ByteString
_ MVar Int
var) SeekMode
IO.AbsoluteSeek Integer
off = do
    Integer -> IO ()
checkOffset Integer
off
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
var (\Int
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Num a => Integer -> a
fromInteger Integer
off, Integer
off))

  seek (Device IOMode
_ MVar ByteString
_ MVar Int
var) SeekMode
IO.RelativeSeek Integer
off = do
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
var (\Int
old_off -> do
      let new_off :: Integer
new_off = forall a. Integral a => a -> Integer
toInteger Int
old_off forall a. Num a => a -> a -> a
+ Integer
off
      Integer -> IO ()
checkOffset Integer
new_off
      forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Num a => Integer -> a
fromInteger Integer
new_off, Integer
new_off))

  seek dev :: Device
dev@(Device IOMode
_ MVar ByteString
_ MVar Int
off_var) SeekMode
IO.SeekFromEnd Integer
off = do
    forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
off_var (\Int
_ -> do
      Integer
size <- forall a. IODevice a => a -> IO Integer
IO.getSize Device
dev
      let new_off :: Integer
new_off = Integer
size forall a. Num a => a -> a -> a
+ Integer
off
      Integer -> IO ()
checkOffset Integer
new_off
      forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Num a => Integer -> a
fromInteger Integer
new_off, Integer
new_off))

  tell :: Device -> IO Integer
tell (Device IOMode
_ MVar ByteString
_ MVar Int
var) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Integral a => a -> Integer
toInteger (forall a. MVar a -> IO a
MVar.readMVar MVar Int
var)
  getSize :: Device -> IO Integer
getSize (Device IOMode
_ MVar ByteString
var MVar Int
_) = do
    ByteString
bytes <- forall a. MVar a -> IO a
MVar.readMVar MVar ByteString
var
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Integral a => a -> Integer
toInteger (ByteString -> Int
Data.ByteString.length ByteString
bytes))
  setSize :: Device -> Integer -> IO ()
setSize Device
dev Integer
size = Device -> Integer -> IO ()
setDeviceSize Device
dev Integer
size
  devType :: Device -> IO IODeviceType
devType Device
_ = forall (m :: * -> *) a. Monad m => a -> m a
return IODeviceType
IO.RegularFile

checkOffset :: Integer -> IO ()
checkOffset :: Integer -> IO ()
checkOffset Integer
off = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Integral a => a -> Integer
toInteger (forall a. Bounded a => a
maxBound :: Int) forall a. Ord a => a -> a -> Bool
< Integer
off) (forall e a. Exception e => e -> IO a
throwIO IOException
err) where
  err :: IOException
err = Maybe Handle
-> IOErrorType
-> String
-> String
-> Maybe CInt
-> Maybe String
-> IOException
IO.IOError forall a. Maybe a
Nothing IOErrorType
IO.InvalidArgument String
"" String
"offset > (maxBound :: Int)" forall a. Maybe a
Nothing forall a. Maybe a
Nothing

setDeviceSize :: Device -> Integer -> IO ()
setDeviceSize :: Device -> Integer -> IO ()
setDeviceSize (Device IOMode
mode MVar ByteString
bytes_var MVar Int
_) Integer
size = IO ()
checkSize forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
setBytes where
  intSize :: Int
  intSize :: Int
intSize = forall a. Num a => Integer -> a
fromInteger Integer
size

  checkSize :: IO ()
checkSize = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Integer
size forall a. Ord a => a -> a -> Bool
> forall a. Integral a => a -> Integer
toInteger (forall a. Bounded a => a
maxBound :: Int)) forall a b. (a -> b) -> a -> b
$ do
    forall e a. Exception e => e -> IO a
throwIO (Maybe Handle
-> IOErrorType
-> String
-> String
-> Maybe CInt
-> Maybe String
-> IOException
IO.IOError forall a. Maybe a
Nothing IOErrorType
IO.InvalidArgument String
"" String
"size > (maxBound :: Int)" forall a. Maybe a
Nothing forall a. Maybe a
Nothing)

  setBytes :: IO ()
setBytes = forall a. MVar a -> (a -> IO a) -> IO ()
MVar.modifyMVar_ MVar ByteString
bytes_var forall a b. (a -> b) -> a -> b
$ \ByteString
bytes -> case IOMode
mode of
    IOMode
IO.ReadMode -> forall e a. Exception e => e -> IO a
throwIO (Maybe Handle
-> IOErrorType
-> String
-> String
-> Maybe CInt
-> Maybe String
-> IOException
IO.IOError forall a. Maybe a
Nothing IOErrorType
IO.IllegalOperation String
"" String
"handle in ReadMode" forall a. Maybe a
Nothing forall a. Maybe a
Nothing)
    IOMode
IO.WriteMode -> forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Word8 -> ByteString
Data.ByteString.replicate Int
intSize Word8
0)
    IOMode
IO.ReadWriteMode -> forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> ByteString
clip ByteString
bytes)
    IOMode
IO.AppendMode -> forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> ByteString
clip ByteString
bytes)

  clip :: ByteString -> ByteString
clip ByteString
bytes = case Int
intSize forall a. Num a => a -> a -> a
- ByteString -> Int
Data.ByteString.length ByteString
bytes of
    Int
padLen | Int
padLen forall a. Ord a => a -> a -> Bool
> Int
0 -> ByteString -> ByteString -> ByteString
Data.ByteString.append ByteString
bytes (Int -> Word8 -> ByteString
Data.ByteString.replicate Int
padLen Word8
0)
    Int
_ -> Int -> ByteString -> ByteString
Data.ByteString.take Int
intSize ByteString
bytes

{- What about non-POSIX environment? -}
instance IO.RawIO Device where
  read :: Device -> Ptr Word8 -> Word64 -> Int -> IO Int
read (Device IOMode
_ MVar ByteString
bytes_var MVar Int
pos_var) Ptr Word8
ptr Word64
_ Int
bufSize = do
    forall a b. MVar a -> (a -> IO b) -> IO b
MVar.withMVar MVar ByteString
bytes_var forall a b. (a -> b) -> a -> b
$ \ByteString
bytes -> do
      forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
pos_var forall a b. (a -> b) -> a -> b
$ \Int
pos -> do
        if Int
pos forall a. Ord a => a -> a -> Bool
>= ByteString -> Int
Data.ByteString.length ByteString
bytes
          then forall (m :: * -> *) a. Monad m => a -> m a
return (Int
pos, Int
0)
          else do
            let chunk :: ByteString
chunk = Int -> ByteString -> ByteString
Data.ByteString.take Int
bufSize (Int -> ByteString -> ByteString
Data.ByteString.drop Int
pos ByteString
bytes)
            forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
chunk forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
chunkPtr, Int
chunkLen) -> do
              forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
Foreign.copyArray Ptr Word8
ptr (forall a b. Ptr a -> Ptr b
Foreign.castPtr Ptr CChar
chunkPtr) Int
chunkLen
              forall (m :: * -> *) a. Monad m => a -> m a
return (Int
pos forall a. Num a => a -> a -> a
+ Int
chunkLen, Int
chunkLen)

  write :: Device -> Ptr Word8 -> Word64 -> Int -> IO ()
write (Device IOMode
_ MVar ByteString
bytes_var MVar Int
pos_var) Ptr Word8
ptr Word64
_ Int
bufSize = do
    forall a. MVar a -> (a -> IO a) -> IO ()
MVar.modifyMVar_ MVar ByteString
bytes_var forall a b. (a -> b) -> a -> b
$ \ByteString
bytes -> do
      forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
pos_var forall a b. (a -> b) -> a -> b
$ \Int
pos -> do
        let (ByteString
before, ByteString
after) = Int -> ByteString -> (ByteString, ByteString)
Data.ByteString.splitAt Int
pos ByteString
bytes
        let padding :: ByteString
padding = Int -> Word8 -> ByteString
Data.ByteString.replicate (Int
pos forall a. Num a => a -> a -> a
- ByteString -> Int
Data.ByteString.length ByteString
before) Word8
0

        ByteString
bufBytes <- CStringLen -> IO ByteString
Data.ByteString.packCStringLen (forall a b. Ptr a -> Ptr b
Foreign.castPtr Ptr Word8
ptr, Int
bufSize)
        let newBytes :: ByteString
newBytes = [ByteString] -> ByteString
Data.ByteString.concat [ByteString
before, ByteString
padding, ByteString
bufBytes, Int -> ByteString -> ByteString
Data.ByteString.drop Int
bufSize ByteString
after]
        forall (m :: * -> *) a. Monad m => a -> m a
return (Int
pos forall a. Num a => a -> a -> a
+ Int
bufSize, ByteString
newBytes)
    forall (m :: * -> *) a. Monad m => a -> m a
return ()

  readNonBlocking :: Device -> Ptr Word8 -> Word64 -> Int -> IO (Maybe Int)
readNonBlocking Device
dev Ptr Word8
buf Word64
off Int
size = forall a. RawIO a => a -> Ptr Word8 -> Word64 -> Int -> IO Int
IO.read Device
dev Ptr Word8
buf Word64
off Int
size forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Int
cnt -> if Int
cnt forall a. Eq a => a -> a -> Bool
== Int
0
    then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Int
cnt
  writeNonBlocking :: Device -> Ptr Word8 -> Word64 -> Int -> IO Int
writeNonBlocking Device
dev Ptr Word8
buf Word64
off Int
cnt = forall a. RawIO a => a -> Ptr Word8 -> Word64 -> Int -> IO ()
IO.write Device
dev Ptr Word8
buf Word64
off Int
cnt forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return Int
cnt

instance IO.BufferedIO Device where
  newBuffer :: Device -> BufferState -> IO (Buffer Word8)
newBuffer Device
_ = Int -> BufferState -> IO (Buffer Word8)
IO.newByteBuffer Int
4096

  fillReadBuffer :: Device -> Buffer Word8 -> IO (Int, Buffer Word8)
fillReadBuffer Device
dev Buffer Word8
buf = do
    (Maybe Int
numRead, Buffer Word8
newBuf) <- forall dev.
BufferedIO dev =>
dev -> Buffer Word8 -> IO (Maybe Int, Buffer Word8)
IO.fillReadBuffer0 Device
dev Buffer Word8
buf
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 forall a. a -> a
id Maybe Int
numRead, Buffer Word8
newBuf)

  fillReadBuffer0 :: Device -> Buffer Word8 -> IO (Maybe Int, Buffer Word8)
fillReadBuffer0 (Device IOMode
_ MVar ByteString
bytes_var MVar Int
pos_var) Buffer Word8
buf = do
    forall a b. MVar a -> (a -> IO b) -> IO b
MVar.withMVar MVar ByteString
bytes_var forall a b. (a -> b) -> a -> b
$ \ByteString
bytes -> do
      forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
pos_var forall a b. (a -> b) -> a -> b
$ \Int
pos -> do
        if Int
pos forall a. Ord a => a -> a -> Bool
>= ByteString -> Int
Data.ByteString.length ByteString
bytes
          then forall (m :: * -> *) a. Monad m => a -> m a
return (Int
pos, (forall a. Maybe a
Nothing, Buffer Word8
buf))
          else do
            let chunk :: ByteString
chunk = Int -> ByteString -> ByteString
Data.ByteString.take (forall e. Buffer e -> Int
IO.bufSize Buffer Word8
buf) (Int -> ByteString -> ByteString
Data.ByteString.drop Int
pos ByteString
bytes)
            forall a. ByteString -> (CStringLen -> IO a) -> IO a
unsafeUseAsCStringLen ByteString
chunk forall a b. (a -> b) -> a -> b
$ \(Ptr CChar
chunkPtr, Int
chunkLen) -> do
              forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
Foreign.withForeignPtr (forall e. Buffer e -> RawBuffer e
IO.bufRaw Buffer Word8
buf) forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> do
                forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
Foreign.copyArray Ptr Word8
ptr (forall a b. Ptr a -> Ptr b
Foreign.castPtr Ptr CChar
chunkPtr) Int
chunkLen
              forall (m :: * -> *) a. Monad m => a -> m a
return (Int
pos forall a. Num a => a -> a -> a
+ Int
chunkLen, (forall a. a -> Maybe a
Just Int
chunkLen, (Buffer Word8
buf { bufL :: Int
IO.bufL = Int
0, bufR :: Int
IO.bufR = Int
chunkLen })))

  flushWriteBuffer :: Device -> Buffer Word8 -> IO (Buffer Word8)
flushWriteBuffer (Device IOMode
_ MVar ByteString
bytes_var MVar Int
pos_var) Buffer Word8
buf = do
    forall a. MVar a -> (a -> IO a) -> IO ()
MVar.modifyMVar_ MVar ByteString
bytes_var forall a b. (a -> b) -> a -> b
$ \ByteString
bytes -> do
      forall a b. MVar a -> (a -> IO (a, b)) -> IO b
MVar.modifyMVar MVar Int
pos_var forall a b. (a -> b) -> a -> b
$ \Int
pos -> do
        let (ByteString
before, ByteString
after) = Int -> ByteString -> (ByteString, ByteString)
Data.ByteString.splitAt Int
pos ByteString
bytes
        let padding :: ByteString
padding = Int -> Word8 -> ByteString
Data.ByteString.replicate (Int
pos forall a. Num a => a -> a -> a
- ByteString -> Int
Data.ByteString.length ByteString
before) Word8
0

        let bufStart :: Ptr a -> Ptr b
bufStart Ptr a
ptr = forall a b. Ptr a -> Ptr b
Foreign.castPtr (forall a b. Ptr a -> Int -> Ptr b
Foreign.plusPtr Ptr a
ptr (forall e. Buffer e -> Int
IO.bufL Buffer Word8
buf))
        let bufLen :: Int
bufLen = forall e. Buffer e -> Int
IO.bufR Buffer Word8
buf forall a. Num a => a -> a -> a
- forall e. Buffer e -> Int
IO.bufL Buffer Word8
buf
        ByteString
bufBytes <- forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
Foreign.withForeignPtr (forall e. Buffer e -> RawBuffer e
IO.bufRaw Buffer Word8
buf) (\Ptr Word8
ptr ->
          CStringLen -> IO ByteString
Data.ByteString.packCStringLen (forall a b. Ptr a -> Ptr b
bufStart Ptr Word8
ptr, Int
bufLen))
        let newBytes :: ByteString
newBytes = [ByteString] -> ByteString
Data.ByteString.concat [ByteString
before, ByteString
padding, ByteString
bufBytes, Int -> ByteString -> ByteString
Data.ByteString.drop Int
bufLen ByteString
after]
        forall (m :: * -> *) a. Monad m => a -> m a
return (Int
pos forall a. Num a => a -> a -> a
+ Int
bufLen, ByteString
newBytes)
    forall (m :: * -> *) a. Monad m => a -> m a
return (Buffer Word8
buf { bufL :: Int
IO.bufL = Int
0, bufR :: Int
IO.bufR = Int
0 })

  flushWriteBuffer0 :: Device -> Buffer Word8 -> IO (Int, Buffer Word8)
flushWriteBuffer0 Device
dev Buffer Word8
buf = do
    Buffer Word8
newBuf <- forall dev.
BufferedIO dev =>
dev -> Buffer Word8 -> IO (Buffer Word8)
IO.flushWriteBuffer Device
dev Buffer Word8
buf
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall e. Buffer e -> Int
IO.bufR Buffer Word8
buf forall a. Num a => a -> a -> a
- forall e. Buffer e -> Int
IO.bufL Buffer Word8
buf, Buffer Word8
newBuf)

newKnob :: MonadIO m => ByteString -> m Knob
newKnob :: forall (m :: * -> *). MonadIO m => ByteString -> m Knob
newKnob ByteString
bytes = do
  MVar ByteString
var <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. a -> IO (MVar a)
MVar.newMVar ByteString
bytes)
  forall (m :: * -> *) a. Monad m => a -> m a
return (MVar ByteString -> Knob
Knob MVar ByteString
var)

getContents :: MonadIO m => Knob -> m ByteString
getContents :: forall (m :: * -> *). MonadIO m => Knob -> m ByteString
getContents (Knob MVar ByteString
var) = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. MVar a -> IO a
MVar.readMVar MVar ByteString
var)

setContents :: MonadIO m => Knob -> ByteString -> m ()
setContents :: forall (m :: * -> *). MonadIO m => Knob -> ByteString -> m ()
setContents (Knob MVar ByteString
var) ByteString
bytes = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. MVar a -> (a -> IO a) -> IO ()
MVar.modifyMVar_ MVar ByteString
var (\ByteString
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bytes))

-- | Create a new 'IO.Handle' pointing to a 'Knob'. This handle behaves like
-- a file-backed handle for most purposes.
newFileHandle :: MonadIO m
              => Knob
              -> String -- ^ Filename shown in error messages
              -> IO.IOMode -> m IO.Handle
newFileHandle :: forall (m :: * -> *).
MonadIO m =>
Knob -> String -> IOMode -> m Handle
newFileHandle (Knob MVar ByteString
var) String
name IOMode
mode = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
  Int
startPosition <- forall a b. MVar a -> (a -> IO b) -> IO b
MVar.withMVar MVar ByteString
var forall a b. (a -> b) -> a -> b
$ \ByteString
bytes -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case IOMode
mode of
    IOMode
IO.AppendMode -> ByteString -> Int
Data.ByteString.length ByteString
bytes
    IOMode
_ -> Int
0
  MVar Int
posVar <- forall a. a -> IO (MVar a)
MVar.newMVar Int
startPosition
  forall dev.
(RawIO dev, IODevice dev, BufferedIO dev, Typeable dev) =>
dev
-> String
-> IOMode
-> Maybe TextEncoding
-> NewlineMode
-> IO Handle
IO.mkFileHandle (IOMode -> MVar ByteString -> MVar Int -> Device
Device IOMode
mode MVar ByteString
var MVar Int
posVar) String
name IOMode
mode forall a. Maybe a
Nothing NewlineMode
IO.noNewlineTranslation

-- | See 'newFileHandle'.
withFileHandle :: MonadIO m
               => Knob
               -> String -- ^ Filename shown in error messages.
               -> IO.IOMode -> (IO.Handle -> IO a) -> m a
withFileHandle :: forall (m :: * -> *) a.
MonadIO m =>
Knob -> String -> IOMode -> (Handle -> IO a) -> m a
withFileHandle Knob
knob String
name IOMode
mode Handle -> IO a
io = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (forall (m :: * -> *).
MonadIO m =>
Knob -> String -> IOMode -> m Handle
newFileHandle Knob
knob String
name IOMode
mode) Handle -> IO ()
IO.hClose Handle -> IO a
io)