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 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
cdbMake :: FilePath -> CDBMake -> IO ()
cdbMake fileName f = do
let tmp = fileName <.> "tmp"
h <- openBinaryFile tmp WriteMode
hSeek h AbsoluteSeek (256*8)
initState <- initialMakeState h
cdb <- execStateT f initState
tablesArrays <- unsafeFreeze (cdbstTables cdb) :: IO (Array Word8 [CDBSlot])
let tables = elems tablesArrays
writeHashTables h tables
hClose h
renameFile tmp fileName
cdbAdd :: (Packable k, Packable v) => k -> v -> CDBMake
cdbAdd k v = do
let (pk, pv) = (pack k, pack v)
cdbAddSlot pk pv
cdbWriteRecord pk pv
cdbAddMany :: (Packable k, Packable v) => [(k,v)] -> CDBMake
cdbAddMany = mapM_ (uncurry cdbAdd)
type CDBMake = StateT CDBMakeState IO ()
data CDBMakeState = CDBMakeState {
cdbstHandle :: Handle,
cdbstRecordsEnd :: Word32,
cdbstTables :: IOArray Word8 [CDBSlot]
}
type CDBSlot = (Word32, Word32)
initialMakeState :: Handle -> IO CDBMakeState
initialMakeState h = do
tables <- newArray (0, 255) []
return CDBMakeState {
cdbstTables = tables,
cdbstRecordsEnd = 256*8,
cdbstHandle = h
}
cdbAddSlot :: ByteString -> ByteString -> CDBMake
cdbAddSlot k v = do
let hash = cdbHash k
let tableNum = fromIntegral $ hash `mod` 256
cdb <- get
let pointer = cdbstRecordsEnd cdb
let tables = cdbstTables cdb
oldTable <- liftIO $ readArray tables tableNum
liftIO $ writeArray tables tableNum $ (hash, pointer):oldTable
cdbWriteRecord :: ByteString -> ByteString -> CDBMake
cdbWriteRecord k v =
let lk = fromIntegral $ ByteString.length k
lv = fromIntegral $ ByteString.length v
record = ByteString.concat [pack lk, pack lv, k, v]
in do
cdb <- get
liftIO $ ByteString.hPut (cdbstHandle cdb) record
put $ cdb { cdbstRecordsEnd = cdbstRecordsEnd cdb + lk + lv + 8 }
writeHashTables :: Handle -> [[CDBSlot]] -> IO ()
writeHashTables h tables = do
tableBase <- fromIntegral <$> hTell h
let bufSize = fromIntegral $ (*4) $ sum (map length tables)
buf <- newArray (0, bufSize1) 0
bufOffset <- newIORef 0
pointers <- mapM (writeTable buf bufOffset tableBase) tables
ibuf <- unsafeFreeze buf :: IO (UArray Word32 Word32)
ByteString.hPut h (pack ibuf)
writePointers h pointers
writeTable :: IOUArray Word32 Word32 ->
IORef Word32 ->
Word32 ->
[CDBSlot] ->
IO (Word32, Word32)
writeTable buf bufOffset tableBase table = do
let tableLength = length table * 2
pointer <- readIORef bufOffset
mapM_ (writeSlot buf pointer tableLength) (reverse table)
writeIORef bufOffset $ pointer + fromIntegral tableLength * 2
return (pointer * 4 + tableBase, fromIntegral tableLength)
writeSlot :: IOUArray Word32 Word32 -> Word32 -> Int -> CDBSlot -> IO ()
writeSlot buf bufOffset tableLength (hash, pointer) = do
ibuf <- unsafeFreeze buf
let slot = findEmptySlot ibuf bufOffset tableLength hash
writeArray buf slot hash
writeArray buf (slot+1) pointer
findEmptySlot :: UArray Word32 Word32 -> Word32 -> Int -> Word32 -> Word32
findEmptySlot buf bufOffset tl hash =
let tl' = fromIntegral tl
searchStart = (hash `div` 256 `mod` tl') * 2
linearSearch i = if buf ! (bufOffset+i+1) == 0
then bufOffset + i
else linearSearch $ (i + 2) `mod` (tl' * 2)
in
linearSearch searchStart
writePointers :: Handle -> [(Word32, Word32)] -> IO ()
writePointers h pointers = do
hSeek h AbsoluteSeek 0
mapM_ (\(pointer, tableLength) -> do ByteString.hPut h (pack pointer)
ByteString.hPut h (pack tableLength))
pointers