-- | Simple, stupid implementation of a RAM-based 'Disk' for testing.
module System.Mem.Disk.Memory where

import Data.ByteString
    ( ByteString )
import Data.Foldable
    ( foldl' )
import Data.Int
    ( Int64 )
import Data.IORef

import qualified Control.Concurrent.STM as STM
import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map
import qualified System.Mem.Disk.DiskApi as DiskApi

{-----------------------------------------------------------------------------
    Disk
------------------------------------------------------------------------------}
data Disk = Disk
    { Disk -> IORef (Map Int64 ByteString)
db :: IORef (Map.Map Int64 ByteString)
    , Disk -> IORef Int64
counter :: IORef Int64
    }

-- | Create a 'Disk' in memory for the purpose of testing and profiling
-- — by swapping v'System.Mem.Disk.withDiskMemory'
-- for  v'System.Mem.Disk.withDiskSqlite' and looking at the
-- heap profile of your program, you can quickly find out whether the
-- use of t'System.Mem.Disk.DiskBytes' really helps.
--
-- Ignores the 'FilePath' argument.
withDiskMemory :: FilePath -> (DiskApi.Disk -> IO a) -> IO a
withDiskMemory :: forall a. FilePath -> (Disk -> IO a) -> IO a
withDiskMemory FilePath
_ Disk -> IO a
action = do
    IORef (Map Int64 ByteString)
db <- forall a. a -> IO (IORef a)
newIORef forall k a. Map k a
Map.empty
    IORef Int64
counter <- forall a. a -> IO (IORef a)
newIORef Int64
0
    Disk -> IO a
action forall a b. (a -> b) -> a -> b
$ Disk -> Disk
mkDiskApi forall a b. (a -> b) -> a -> b
$ Disk{IORef (Map Int64 ByteString)
db :: IORef (Map Int64 ByteString)
db :: IORef (Map Int64 ByteString)
db,IORef Int64
counter :: IORef Int64
counter :: IORef Int64
counter}

getDiskSize_ :: Disk -> IO Integer
getDiskSize_ :: Disk -> IO Integer
getDiskSize_ Disk{IORef (Map Int64 ByteString)
db :: IORef (Map Int64 ByteString)
db :: Disk -> IORef (Map Int64 ByteString)
db} = Map Int64 ByteString -> Integer
sumBytes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef (Map Int64 ByteString)
db
  where
    sumBytes :: Map Int64 ByteString -> Integer
sumBytes = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Integer
n ByteString
bs -> Integer
n forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
bs)) Integer
0

mkDiskApi :: Disk -> DiskApi.Disk
mkDiskApi :: Disk -> Disk
mkDiskApi Disk
disk = DiskApi.Disk
    { put :: ByteString -> IO Int64
DiskApi.put = Disk -> ByteString -> IO Int64
put_ Disk
disk
    , get :: Int64 -> IO ByteString
DiskApi.get = Disk -> Int64 -> IO ByteString
get_ Disk
disk
    , delete :: Int64 -> IO ()
DiskApi.delete = Disk -> Int64 -> IO ()
delete_ Disk
disk
    , getDiskSize :: IO Integer
DiskApi.getDiskSize = Disk -> IO Integer
getDiskSize_ Disk
disk
    }

{-----------------------------------------------------------------------------
    Disk operations
------------------------------------------------------------------------------}
put_ :: Disk -> ByteString -> IO Int64
put_ :: Disk -> ByteString -> IO Int64
put_ Disk{IORef (Map Int64 ByteString)
db :: IORef (Map Int64 ByteString)
db :: Disk -> IORef (Map Int64 ByteString)
db,IORef Int64
counter :: IORef Int64
counter :: Disk -> IORef Int64
counter} ByteString
bytes = do
    Int64
index <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int64
counter forall a b. (a -> b) -> a -> b
$ \Int64
x -> (Int64
xforall a. Num a => a -> a -> a
+Int64
1,Int64
x)
    forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Map Int64 ByteString)
db forall a b. (a -> b) -> a -> b
$ \Map Int64 ByteString
db_ -> (forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Int64
index ByteString
bytes Map Int64 ByteString
db_, ())
    forall (f :: * -> *) a. Applicative f => a -> f a
pure Int64
index

get_ :: Disk -> Int64 -> IO ByteString
get_ :: Disk -> Int64 -> IO ByteString
get_ Disk{IORef (Map Int64 ByteString)
db :: IORef (Map Int64 ByteString)
db :: Disk -> IORef (Map Int64 ByteString)
db} Int64
index = do
    Just ByteString
bytes <- forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int64
index forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef (Map Int64 ByteString)
db
    forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
bytes

delete_ :: Disk -> Int64 -> IO ()
delete_ :: Disk -> Int64 -> IO ()
delete_ Disk{IORef (Map Int64 ByteString)
db :: IORef (Map Int64 ByteString)
db :: Disk -> IORef (Map Int64 ByteString)
db} Int64
index =
    forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Map Int64 ByteString)
db forall a b. (a -> b) -> a -> b
$ \Map Int64 ByteString
db_ -> (forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Int64
index Map Int64 ByteString
db_, ())