module Database.CDB.Write (
  CDBMake(),
  cdbMake,
  cdbAdd,
  cdbAddMany
) where

import Control.Applicative
import Control.Monad
import Control.Monad.State
import Data.Array.IO
import Data.Array.Unboxed
import Data.Array.Unsafe (unsafeFreeze)
import qualified Data.ByteString as ByteString
import Data.ByteString (ByteString)
import Data.IORef
import Data.List
import Data.Maybe
import Data.Word
import Database.CDB.Packable
import Database.CDB.Util
import System.Directory
import System.FilePath
import System.IO

-- |Construct a CDB as described inside the supplied CDBMake computation.
--  During construction, it will be written to a temporary file and then
--  moved over top of the given file atomically.
cdbMake :: FilePath -> CDBMake -> IO ()
cdbMake :: FilePath -> CDBMake -> IO ()
cdbMake FilePath
fileName CDBMake
f = do
  let tmp :: FilePath
tmp = FilePath
fileName FilePath -> FilePath -> FilePath
<.> FilePath
"tmp"
  Handle
h <- FilePath -> IOMode -> IO Handle
openBinaryFile FilePath
tmp IOMode
WriteMode
  Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
h SeekMode
AbsoluteSeek (Integer
256forall a. Num a => a -> a -> a
*Integer
8)
  CDBMakeState
initState <- Handle -> IO CDBMakeState
initialMakeState Handle
h
  CDBMakeState
cdb <- forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m s
execStateT CDBMake
f CDBMakeState
initState
  Array Word8 [CDBSlot]
tablesArrays <- forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze (CDBMakeState -> IOArray Word8 [CDBSlot]
cdbstTables CDBMakeState
cdb) :: IO (Array Word8 [CDBSlot])
  let tables :: [[CDBSlot]]
tables = forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
elems Array Word8 [CDBSlot]
tablesArrays
  Handle -> [[CDBSlot]] -> IO ()
writeHashTables Handle
h [[CDBSlot]]
tables
  Handle -> IO ()
hClose Handle
h
  FilePath -> FilePath -> IO ()
renameFile FilePath
tmp FilePath
fileName

-- |Adds a given key-value pair to the CDB being built.
cdbAdd :: (Packable k, Packable v) => k -> v -> CDBMake
cdbAdd :: forall k v. (Packable k, Packable v) => k -> v -> CDBMake
cdbAdd k
k v
v = do
  let (ByteString
pk, ByteString
pv) = (forall k. Packable k => k -> ByteString
pack k
k, forall k. Packable k => k -> ByteString
pack v
v)
  ByteString -> ByteString -> CDBMake
cdbAddSlot ByteString
pk ByteString
pv
  ByteString -> ByteString -> CDBMake
cdbWriteRecord ByteString
pk ByteString
pv

-- |Add a list of key-value pairs to the CDB being built.
cdbAddMany :: (Packable k, Packable v) => [(k,v)] -> CDBMake
cdbAddMany :: forall k v. (Packable k, Packable v) => [(k, v)] -> CDBMake
cdbAddMany = forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall k v. (Packable k, Packable v) => k -> v -> CDBMake
cdbAdd)

-----------------------
-- write implementation
-----------------------

type CDBMake = StateT CDBMakeState IO ()

data CDBMakeState = CDBMakeState {
  CDBMakeState -> Handle
cdbstHandle        :: Handle,
  CDBMakeState -> Word32
cdbstRecordsEnd    :: Word32,
  CDBMakeState -> IOArray Word8 [CDBSlot]
cdbstTables        :: IOArray Word8 [CDBSlot] 
}

type CDBSlot = (Word32, Word32)

initialMakeState :: Handle -> IO CDBMakeState
initialMakeState :: Handle -> IO CDBMakeState
initialMakeState Handle
h = do
  IOArray Word8 [CDBSlot]
tables <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Word8
0, Word8
255) []
  forall (m :: * -> *) a. Monad m => a -> m a
return CDBMakeState {
    cdbstTables :: IOArray Word8 [CDBSlot]
cdbstTables     = IOArray Word8 [CDBSlot]
tables,
    cdbstRecordsEnd :: Word32
cdbstRecordsEnd = Word32
256forall a. Num a => a -> a -> a
*Word32
8,
    cdbstHandle :: Handle
cdbstHandle     = Handle
h
  }

-- add a slot to the set of slots
cdbAddSlot :: ByteString -> ByteString -> CDBMake
cdbAddSlot :: ByteString -> ByteString -> CDBMake
cdbAddSlot ByteString
k ByteString
v = do
  let hash :: Word32
hash     = ByteString -> Word32
cdbHash ByteString
k
  let tableNum :: Word8
tableNum = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Word32
hash forall a. Integral a => a -> a -> a
`mod` Word32
256
  CDBMakeState
cdb <- forall s (m :: * -> *). MonadState s m => m s
get
  let pointer :: Word32
pointer = CDBMakeState -> Word32
cdbstRecordsEnd CDBMakeState
cdb
  let tables :: IOArray Word8 [CDBSlot]
tables = CDBMakeState -> IOArray Word8 [CDBSlot]
cdbstTables CDBMakeState
cdb
  [CDBSlot]
oldTable <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray IOArray Word8 [CDBSlot]
tables Word8
tableNum
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOArray Word8 [CDBSlot]
tables Word8
tableNum forall a b. (a -> b) -> a -> b
$ (Word32
hash, Word32
pointer)forall a. a -> [a] -> [a]
:[CDBSlot]
oldTable


cdbWriteRecord :: ByteString -> ByteString -> CDBMake
cdbWriteRecord :: ByteString -> ByteString -> CDBMake
cdbWriteRecord ByteString
k ByteString
v = 
  let lk :: Word32
lk  = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
k
      lv :: Word32
lv  = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ ByteString -> Int
ByteString.length ByteString
v
      record :: ByteString
record = [ByteString] -> ByteString
ByteString.concat [forall k. Packable k => k -> ByteString
pack Word32
lk, forall k. Packable k => k -> ByteString
pack Word32
lv, ByteString
k, ByteString
v]
  in do
    CDBMakeState
cdb <- forall s (m :: * -> *). MonadState s m => m s
get
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
ByteString.hPut (CDBMakeState -> Handle
cdbstHandle CDBMakeState
cdb) ByteString
record
    forall s (m :: * -> *). MonadState s m => s -> m ()
put forall a b. (a -> b) -> a -> b
$ CDBMakeState
cdb { cdbstRecordsEnd :: Word32
cdbstRecordsEnd = CDBMakeState -> Word32
cdbstRecordsEnd CDBMakeState
cdb forall a. Num a => a -> a -> a
+ Word32
lk forall a. Num a => a -> a -> a
+ Word32
lv forall a. Num a => a -> a -> a
+ Word32
8 }

-- assumes the Handle is pointing to right after the last record written
writeHashTables :: Handle -> [[CDBSlot]] -> IO ()
writeHashTables :: Handle -> [[CDBSlot]] -> IO ()
writeHashTables Handle
h [[CDBSlot]]
tables = do
  Word32
tableBase <- forall a b. (Integral a, Num b) => a -> b
fromIntegral forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle -> IO Integer
hTell Handle
h
  let bufSize :: Word32
bufSize = forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ (forall a. Num a => a -> a -> a
*Int
4)  forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b. (a -> b) -> [a] -> [b]
map forall (t :: * -> *) a. Foldable t => t a -> Int
length [[CDBSlot]]
tables)
  IOUArray Word32 Word32
buf <- forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Word32
0, Word32
bufSizeforall a. Num a => a -> a -> a
-Word32
1) Word32
0 
  IORef Word32
bufOffset <- forall a. a -> IO (IORef a)
newIORef Word32
0
  [CDBSlot]
pointers <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (IOUArray Word32 Word32
-> IORef Word32 -> Word32 -> [CDBSlot] -> IO CDBSlot
writeTable IOUArray Word32 Word32
buf IORef Word32
bufOffset Word32
tableBase) [[CDBSlot]]
tables
  UArray Word32 Word32
ibuf <- forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze IOUArray Word32 Word32
buf :: IO (UArray Word32 Word32)
  Handle -> ByteString -> IO ()
ByteString.hPut Handle
h (forall k. Packable k => k -> ByteString
pack UArray Word32 Word32
ibuf)
  Handle -> [CDBSlot] -> IO ()
writePointers Handle
h [CDBSlot]
pointers

writeTable :: IOUArray Word32 Word32 ->
              IORef Word32 ->
              Word32 ->
              [CDBSlot] ->
              IO (Word32, Word32)
writeTable :: IOUArray Word32 Word32
-> IORef Word32 -> Word32 -> [CDBSlot] -> IO CDBSlot
writeTable IOUArray Word32 Word32
buf IORef Word32
bufOffset Word32
tableBase [CDBSlot]
table = do
  -- compute the number of slots
  -- twice the number of actual entries to help prevent collision
  let tableLength :: Int
tableLength = forall (t :: * -> *) a. Foldable t => t a -> Int
length [CDBSlot]
table forall a. Num a => a -> a -> a
* Int
2
  Word32
pointer <- forall a. IORef a -> IO a
readIORef IORef Word32
bufOffset 
  -- write the slots in the order they came in
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IOUArray Word32 Word32 -> Word32 -> Int -> CDBSlot -> IO ()
writeSlot IOUArray Word32 Word32
buf Word32
pointer Int
tableLength) (forall a. [a] -> [a]
reverse [CDBSlot]
table)
  forall a. IORef a -> a -> IO ()
writeIORef IORef Word32
bufOffset forall a b. (a -> b) -> a -> b
$ Word32
pointer forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
tableLength forall a. Num a => a -> a -> a
* Word32
2
  forall (m :: * -> *) a. Monad m => a -> m a
return (Word32
pointer forall a. Num a => a -> a -> a
* Word32
4 forall a. Num a => a -> a -> a
+ Word32
tableBase, forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
tableLength)

writeSlot :: IOUArray Word32 Word32 -> Word32 -> Int -> CDBSlot -> IO ()
writeSlot :: IOUArray Word32 Word32 -> Word32 -> Int -> CDBSlot -> IO ()
writeSlot IOUArray Word32 Word32
buf Word32
bufOffset Int
tableLength (Word32
hash, Word32
pointer) = do
  UArray Word32 Word32
ibuf <- forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze IOUArray Word32 Word32
buf
  let slot :: Word32
slot = UArray Word32 Word32 -> Word32 -> Int -> Word32 -> Word32
findEmptySlot UArray Word32 Word32
ibuf Word32
bufOffset Int
tableLength Word32
hash
  forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Word32 Word32
buf Word32
slot Word32
hash
  forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray IOUArray Word32 Word32
buf (Word32
slotforall a. Num a => a -> a -> a
+Word32
1) Word32
pointer

findEmptySlot :: UArray Word32 Word32 -> Word32 -> Int -> Word32 -> Word32
findEmptySlot :: UArray Word32 Word32 -> Word32 -> Int -> Word32 -> Word32
findEmptySlot UArray Word32 Word32
buf Word32
bufOffset Int
tl Word32
hash =
  let tl' :: Word32
tl' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
tl
      searchStart :: Word32
searchStart = (Word32
hash forall a. Integral a => a -> a -> a
`div` Word32
256 forall a. Integral a => a -> a -> a
`mod` Word32
tl') forall a. Num a => a -> a -> a
* Word32
2
      linearSearch :: Word32 -> Word32
linearSearch Word32
i = if UArray Word32 Word32
buf forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! (Word32
bufOffsetforall a. Num a => a -> a -> a
+Word32
iforall a. Num a => a -> a -> a
+Word32
1) forall a. Eq a => a -> a -> Bool
== Word32
0
                         then Word32
bufOffset forall a. Num a => a -> a -> a
+ Word32
i
                         else Word32 -> Word32
linearSearch forall a b. (a -> b) -> a -> b
$ (Word32
i forall a. Num a => a -> a -> a
+ Word32
2) forall a. Integral a => a -> a -> a
`mod` (Word32
tl' forall a. Num a => a -> a -> a
* Word32
2)
  in
  Word32 -> Word32
linearSearch Word32
searchStart

writePointers :: Handle -> [(Word32, Word32)] -> IO ()
writePointers :: Handle -> [CDBSlot] -> IO ()
writePointers Handle
h [CDBSlot]
pointers = do
  Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
h SeekMode
AbsoluteSeek Integer
0
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(Word32
pointer, Word32
tableLength) -> do Handle -> ByteString -> IO ()
ByteString.hPut Handle
h (forall k. Packable k => k -> ByteString
pack Word32
pointer)
                                       Handle -> ByteString -> IO ()
ByteString.hPut Handle
h (forall k. Packable k => k -> ByteString
pack Word32
tableLength))
        [CDBSlot]
pointers