module Database.CDB.Read (
  CDB(),
  cdbInit,
  cdbGet,
  cdbGetAll,
  cdbHasKey,
  cdbCount
) where

import Control.Monad
import Data.Bits
import qualified Data.ByteString as ByteString
import Data.ByteString (ByteString)
import Data.Word
import Database.CDB.Packable
import Database.CDB.Util
import System.IO.MMap

-- |Internal representation of a CDB file on disk.
data CDB = CDB { CDB -> ByteString
cdbMem :: ByteString }

-- |Loads a CDB from a file.
cdbInit :: FilePath -> IO CDB
cdbInit :: FilePath -> IO CDB
cdbInit FilePath
f = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ByteString -> CDB
CDB forall a b. (a -> b) -> a -> b
$ FilePath -> Maybe (Int64, Int) -> IO ByteString
mmapFileByteString FilePath
f forall a. Maybe a
Nothing

-- |Finds the first entry associated with a key in a CDB.
cdbGet :: (Packable k, Unpackable v) => CDB -> k -> Maybe v
cdbGet :: forall k v. (Packable k, Unpackable v) => CDB -> k -> Maybe v
cdbGet CDB
cdb k
key = case CDB -> ByteString -> [Word32]
cdbFind CDB
cdb (forall k. Packable k => k -> ByteString
pack k
key) of
  []    -> forall a. Maybe a
Nothing
  (Word32
x:[Word32]
_) -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall v. Unpackable v => ByteString -> v
unpack forall a b. (a -> b) -> a -> b
$ CDB -> Word32 -> ByteString
readData CDB
cdb Word32
x 

-- |Finds all entries associated with a key in a CDB.
cdbGetAll :: (Packable k, Unpackable v) => CDB -> k -> [v]
cdbGetAll :: forall k v. (Packable k, Unpackable v) => CDB -> k -> [v]
cdbGetAll CDB
cdb k
key = forall a b. (a -> b) -> [a] -> [b]
map (forall v. Unpackable v => ByteString -> v
unpack forall b c a. (b -> c) -> (a -> b) -> a -> c
. CDB -> Word32 -> ByteString
readData CDB
cdb) (CDB -> ByteString -> [Word32]
cdbFind CDB
cdb (forall k. Packable k => k -> ByteString
pack k
key))

-- |Returns True if the CDB has a value associated with the given key.
cdbHasKey :: (Packable k) => CDB -> k -> Bool
cdbHasKey :: forall k. Packable k => CDB -> k -> Bool
cdbHasKey CDB
cdb k
key = case CDB -> ByteString -> [Word32]
cdbFind CDB
cdb (forall k. Packable k => k -> ByteString
pack k
key) of
  [] -> Bool
False
  [Word32]
_  -> Bool
True

-- |Returns the number of values a CDB has for a given key.
cdbCount :: (Packable k) => CDB -> k -> Int
cdbCount :: forall k. Packable k => CDB -> k -> Int
cdbCount CDB
cdb k
key = forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ CDB -> ByteString -> [Word32]
cdbFind CDB
cdb (forall k. Packable k => k -> ByteString
pack k
key)

substr :: ByteString -> Int -> Int -> ByteString
substr :: ByteString -> Int -> Int -> ByteString
substr ByteString
bs Int
i Int
n = Int -> ByteString -> ByteString
ByteString.take Int
n (forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> (ByteString, ByteString)
ByteString.splitAt Int
i ByteString
bs)

cdbRead32 :: CDB -> Word32 -> Word32
cdbRead32 :: CDB -> Word32 -> Word32
cdbRead32 CDB
cdb Word32
i =
  [Word8] -> Word32
bytesToWord forall a b. (a -> b) -> a -> b
$ ByteString -> [Word8]
ByteString.unpack forall a b. (a -> b) -> a -> b
$ ByteString -> Int -> Int -> ByteString
substr (CDB -> ByteString
cdbMem CDB
cdb) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
i) Int
4

bytesToWord :: [Word8] -> Word32
bytesToWord :: [Word8] -> Word32
bytesToWord = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Word8
x Word32
y -> (Word32
y forall a. Bits a => a -> Int -> a
`shiftL` Int
8) forall a. Bits a => a -> a -> a
.|. forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
x) Word32
0

tableLength :: CDB -> Word8 -> Word32
tableLength :: CDB -> Word8 -> Word32
tableLength CDB
cdb Word8
n = CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` ((forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
n forall a. Num a => a -> a -> a
* Word32
8) forall a. Num a => a -> a -> a
+ Word32
4)

tableOffset :: CDB -> Word8 -> Word32
tableOffset :: CDB -> Word8 -> Word32
tableOffset CDB
cdb Word8
n = CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
n forall a. Num a => a -> a -> a
* Word32
8)

-- finds the indices of hash table entries for a given key
cdbFind :: CDB -> ByteString -> [Word32]
cdbFind :: CDB -> ByteString -> [Word32]
cdbFind CDB
cdb ByteString
key =
  let hash :: Word32
hash     = ByteString -> Word32
cdbHash ByteString
key
      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
      tl :: Word32
tl       = CDB -> Word8 -> Word32
tableLength CDB
cdb Word8
tableNum 
      in
      if Word32
tl forall a. Eq a => a -> a -> Bool
== Word32
0
        then []
        else
          let slotNum :: Word32
slotNum = Word32
hash forall a. Integral a => a -> a -> a
`div` Word32
256 forall a. Integral a => a -> a -> a
`mod` Word32
tl
              linearSearch :: Word32 -> [Word32]
linearSearch Word32
slotNum = case CDB -> Word8 -> Word32 -> Maybe (Word32, Word32)
probe CDB
cdb Word8
tableNum Word32
slotNum of
                Just (Word32
recordOffset, Word32
hash') ->
                  let nextSlot :: Word32
nextSlot = (Word32
slotNum forall a. Num a => a -> a -> a
+ Word32
1) forall a. Integral a => a -> a -> a
`mod` Word32
tl in
                  if Word32
hash forall a. Eq a => a -> a -> Bool
== Word32
hash' Bool -> Bool -> Bool
&& ByteString
key forall a. Eq a => a -> a -> Bool
== CDB -> Word32 -> ByteString
readKey CDB
cdb Word32
recordOffset
                    then Word32
recordOffset forall a. a -> [a] -> [a]
: Word32 -> [Word32]
linearSearch Word32
nextSlot
                    else Word32 -> [Word32]
linearSearch Word32
nextSlot
                Maybe (Word32, Word32)
Nothing -> []
          in
          Word32 -> [Word32]
linearSearch Word32
slotNum

-- returns a tuple (offset, hash) if the slot contains anything
probe :: CDB -> Word8 -> Word32 -> Maybe (Word32, Word32)
probe :: CDB -> Word8 -> Word32 -> Maybe (Word32, Word32)
probe CDB
cdb Word8
tableNum Word32
slotNum =
  let offset :: Word32
offset       = CDB -> Word8 -> Word32
tableOffset CDB
cdb Word8
tableNum forall a. Num a => a -> a -> a
+ (Word32
slotNum forall a. Num a => a -> a -> a
* Word32
8)
      recordOffset :: Word32
recordOffset = CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` (Word32
offset forall a. Num a => a -> a -> a
+ Word32
4)
  in
  if Word32
recordOffset forall a. Eq a => a -> a -> Bool
== Word32
0 then forall a. Maybe a
Nothing
                       else forall (m :: * -> *) a. Monad m => a -> m a
return (Word32
recordOffset, CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` Word32
offset)

readKey :: CDB -> Word32 -> ByteString
readKey :: CDB -> Word32 -> ByteString
readKey CDB
cdb Word32
offset =
  let len :: Word32
len = CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` Word32
offset in
  ByteString -> Int -> Int -> ByteString
substr (CDB -> ByteString
cdbMem CDB
cdb) (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Word32
offset forall a. Num a => a -> a -> a
+ Word32
8) (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)

readData :: CDB -> Word32 -> ByteString
readData :: CDB -> Word32 -> ByteString
readData CDB
cdb Word32
offset = 
  let keyLen :: Word32
keyLen = CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` Word32
offset
      dataLen :: Word32
dataLen = CDB
cdb CDB -> Word32 -> Word32
`cdbRead32` (Word32
offset forall a. Num a => a -> a -> a
+ Word32
4)
  in
  ByteString -> Int -> Int -> ByteString
substr (CDB -> ByteString
cdbMem CDB
cdb) (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ Word32
offset forall a. Num a => a -> a -> a
+ Word32
8 forall a. Num a => a -> a -> a
+ Word32
keyLen)
    (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
dataLen)