module Database.HongoDB.HashMem (
HashMem,
runHashMem,
) where
import Database.HongoDB.Base
import Control.Applicative
import Control.Monad.IO.Control
import Control.Monad.Reader
import qualified Data.ByteString.Char8 as B
import qualified Data.Enumerator as E
import qualified Data.HashMap.Strict as HM
import Data.IORef
type Table = HM.HashMap B.ByteString B.ByteString
newtype HashMem m a =
HashMem { unHashMem :: ReaderT (IORef Table) m a }
deriving (Monad, MonadIO, MonadTrans, Functor, Applicative, MonadControlIO)
instance (MonadControlIO m) => DB (HashMem m) where
accept key f = modifyTable $ \t -> do
(act, r) <- f (HM.lookup key t)
let nt = case act of
Replace val -> HM.insert key val t
Remove -> HM.delete key t
Nop -> t
return (nt, r)
count = withTable $ \t ->
return $ HM.size t
clear = modifyTable $ \_ ->
return (HM.empty, ())
enum = withTable $ \t ->
return $ go (HM.toList t)
where
go [] (E.Continue k) = E.continue k
go xs (E.Continue k) =
let (as, bs) = splitAt 256 xs in
k (E.Chunks as) E.>>== go bs
go _ step = E.returnI step
withTable :: MonadIO m => (Table -> HashMem m a) -> HashMem m a
withTable f = do
r <- HashMem ask
f =<< liftIO (readIORef r)
modifyTable :: MonadIO m => (Table -> HashMem m (Table, a)) -> HashMem m a
modifyTable f = do
r <- HashMem ask
(t, v) <- f =<< liftIO (readIORef r)
liftIO $ writeIORef r t
return v
runHashMem :: MonadIO m => HashMem m a -> m a
runHashMem hdb = do
ref <- liftIO $ newIORef HM.empty
runReaderT (unHashMem hdb) ref