module Data.Conduit.VFS.InMemory
( InMemoryVFS
, InMemoryVFSRoot
, runInMemoryVFS
, runInMemoryVFS'
, mkInMemoryVFSRoot
) where
import ClassyPrelude hiding (ByteString, handle)
import Control.Monad.Extra (ifM)
import Data.Conduit.VFS.Import
import qualified Data.ByteString.Lazy as LBS
import qualified Data.HashMap.Lazy as HashMap
import Data.HashMap.Lazy (HashMap)
import System.FilePath (splitPath)
import Control.Monad.Fail (MonadFail)
import qualified Data.Text as Text
data IMNode
= IMNodeDir IMDirectory
| IMNodeFile IMFile
data IMFile
= Resident ByteString
| EmptyFile
newtype IMDirectory = IMDirectory
{ imdNodes :: HashMap Text IMNode
}
instance Semigroup IMDirectory where
(<>) imdLeft imdRight =
IMDirectory { imdNodes = HashMap.unionWith mergeImpl nodeLeft nodeRight }
where
nodeLeft = imdNodes imdLeft
nodeRight = imdNodes imdRight
mergeImpl (IMNodeDir nodeDirLeft) (IMNodeDir nodeDirRight) = IMNodeDir $ nodeDirLeft <> nodeDirRight
mergeImpl _ right = right
{-# INLINE (<>) #-}
instance Monoid IMDirectory where
mempty = IMDirectory { imdNodes = mempty }
{-# INLINE mempty #-}
type instance Element IMDirectory = (Text, IMNode)
instance MonoFunctor IMDirectory where
omap f oldImd = IMDirectory{ imdNodes = HashMap.fromList (f <$> HashMap.toList (imdNodes oldImd)) }
{-# INLINE omap #-}
instance MonoPointed IMDirectory where
opoint (k,v) = IMDirectory { imdNodes = HashMap.singleton k v }
{-# INLINE opoint #-}
newtype InMemoryVFSRoot = InMemoryVFSRoot { imvfsStore :: MVar IMDirectory }
newtype InMemoryVFS m a = InMemoryVFS
{ unIMVFS :: ReaderT InMemoryVFSRoot m a }
deriving (Functor, Applicative, Monad, MonadTrans, MonadIO, MonadFail, MonadReader InMemoryVFSRoot)
mkInMemoryVFSRoot :: (MonadIO m) => m InMemoryVFSRoot
mkInMemoryVFSRoot = do
mvar <- newMVar mempty
return $ InMemoryVFSRoot { imvfsStore = mvar }
{-# INLINE mkInMemoryVFSRoot #-}
runInMemoryVFS :: (MonadUnliftIO m) => InMemoryVFS m a -> m a
runInMemoryVFS imvfs = mkInMemoryVFSRoot >>= flip runInMemoryVFS' imvfs
{-# INLINE runInMemoryVFS #-}
runInMemoryVFS' :: InMemoryVFSRoot -> InMemoryVFS m a -> m a
runInMemoryVFS' root imvfs =
let monad = unIMVFS imvfs in
runReaderT monad root
{-# INLINE runInMemoryVFS' #-}
withIMVFSRootDir :: (MonadUnliftIO m) => ( IMDirectory -> m a ) -> InMemoryVFS m a
withIMVFSRootDir f = do
mvar <- imvfsStore <$> ask
rootDir <- readMVar mvar
lift $ f rootDir
{-# INLINE withIMVFSRootDir #-}
modifyIMVFSRootDir :: (MonadUnliftIO m) => ( IMDirectory -> m IMDirectory ) -> InMemoryVFS m ()
modifyIMVFSRootDir f = do
mvar <- imvfsStore <$> ask
lift $ modifyMVar_ mvar f
{-# INLINE modifyIMVFSRootDir #-}
instance (MonadUnliftIO m) => ReadVFSC (InMemoryVFS m) where
vfsTypeC = awaitForever $ \filepath -> do
result <- lift $ withIMVFSRootDir $ return . loop (splitPath filepath)
yield (filepath, result)
where
loop :: [FilePath] -> IMDirectory -> Maybe VFileType
loop [] (IMDirectory _) = Just VDirectory
loop ("/":rest) imd = loop rest imd
loop (nextDirPath:rest) IMDirectory{imdNodes} =
case HashMap.lookup (Text.pack nextDirPath) imdNodes of
Nothing -> Nothing
(Just imnode) ->
case imnode of
(IMNodeDir imd) -> loop rest imd
(IMNodeFile _)
| null rest -> Just VFile
| otherwise -> Nothing
{-# INLINEABLE vfsTypeC #-}
vfsContentsC = awaitForever $ \filepath -> do
result <- lift $ withIMVFSRootDir $ loop (splitPath filepath)
case result of
Nothing -> return ()
(Just resultLBS) -> yield (filepath, resultLBS)
where
loop :: [FilePath] -> IMDirectory -> m (Maybe ByteString)
loop [] _ = return Nothing
loop [filename] IMDirectory{imdNodes} =
case HashMap.lookup (Text.pack filename) imdNodes of
Nothing -> return Nothing
(Just (IMNodeDir _)) -> return Nothing
(Just (IMNodeFile fileData)) ->
case fileData of
(Resident bytes) -> return (Just bytes)
EmptyFile -> return (Just mempty)
loop (dirname:rest) IMDirectory{imdNodes} =
case HashMap.lookup (Text.pack dirname) imdNodes of
(Just (IMNodeDir imd)) -> loop rest imd
_ -> return Nothing
{-# INLINEABLE vfsContentsC #-}
vfsContentsEitherC = awaitForever $ \filepath -> do
maybeResult <- lift $ withIMVFSRootDir $ loop (splitPath filepath)
case maybeResult of
Nothing -> return ()
(Just bytes) -> do
yield $ Left filepath
yield $ Right bytes
where
loop :: [FilePath] -> IMDirectory -> m (Maybe ByteString)
loop [] _ = return Nothing
loop [filename] IMDirectory{imdNodes} =
case HashMap.lookup (Text.pack filename) imdNodes of
Nothing -> return Nothing
(Just (IMNodeDir _)) -> return Nothing
(Just (IMNodeFile filedata)) ->
case filedata of
(Resident bytes) -> return $ Just bytes
EmptyFile -> return $ Just mempty
loop (dirname:rest) IMDirectory{imdNodes} =
case HashMap.lookup (Text.pack dirname) imdNodes of
(Just (IMNodeDir imd)) -> loop rest imd
_ -> return Nothing
{-# INLINEABLE vfsContentsEitherC #-}
vfsChildrenC = awaitForever $ \filepath ->
lift ( withIMVFSRootDir $ return . loop filepath (splitPath filepath) ) >>= yieldMany
where
loop :: FilePath -> [FilePath] -> IMDirectory -> [FilePath]
loop _ [] IMDirectory{imdNodes} = Text.unpack <$> HashMap.keys imdNodes
loop filepath (foo:rest) IMDirectory{imdNodes} =
case HashMap.lookup (Text.pack foo) imdNodes of
(Just (IMNodeDir imd@(IMDirectory dir)))
| null rest -> (filepath </>) . Text.unpack <$> HashMap.keys dir
| otherwise -> loop filepath rest imd
_ -> mempty
{-# INLINEABLE vfsChildrenC #-}
instance (MonadUnliftIO m) => WriteVFSC (InMemoryVFS m) where
vfsWriteEitherSink = awaitForever $ \case
(Right _) -> return ()
(Left filepath) -> awaitBytes >>= \bytes ->
let imfile =
if null bytes then
EmptyFile
else
Resident bytes
in
lift $ modifyIMVFSRootDir $ return . loop (IMNodeFile imfile) (splitPath filepath)
where
hasMoreBytes = peekC >>= \case
(Just (Right _)) -> return True
_ -> return False
awaitBytes =
flip (ifM hasMoreBytes) (return mempty) $
await >>= \case
(Just (Right bytes)) -> LBS.append bytes <$> awaitBytes
_ -> fail "We should have more bytes, but we don't."
loop _ [] imd = imd
loop node [filename] imd@IMDirectory{imdNodes} = imd { imdNodes = HashMap.insert (Text.pack filename) node imdNodes }
loop node (name:rest) imd@IMDirectory{imdNodes} = imd { imdNodes = HashMap.alter
(\case
Nothing -> Just . IMNodeDir . loop node rest $ IMDirectory { imdNodes=mempty }
(Just (IMNodeDir childImd)) -> Just . IMNodeDir $ loop node rest childImd
whatever -> whatever
) (Text.pack name) imdNodes }
{-# INLINEABLE vfsWriteEitherSink #-}
vfsRemoveSink = awaitForever $ \filepath -> lift . modifyIMVFSRootDir $ return . loop (splitPath filepath)
where
loop [] imd = imd
loop [filename] imd@IMDirectory{imdNodes} = imd { imdNodes = HashMap.delete (Text.pack filename) imdNodes }
loop (name:rest) imd@IMDirectory{imdNodes} = imd
{ imdNodes = HashMap.adjust
(\case
(IMNodeDir childImd) -> IMNodeDir $ loop rest childImd
whatever -> whatever
) (Text.pack name) imdNodes
}
{-# INLINE vfsRemoveSink #-}
instance (MonadUnliftIO m) => VFSC (InMemoryVFS m)