module Data.BTree.BTree
(
makeParam
, execTree
, Key
, Value
, Interval(..)
, TreeBackend
, TreeResult
, insert
, delete
, lookup
, modify
, save
, rebalanceProcess
, toList
, foldli
, foldri
, search
, search_
, findMin
, findMax
, height
) where
import Data.BTree.Types
import Data.Maybe
import Control.Concurrent
import Control.Concurrent.STM
import Control.Exception
import Control.Monad.Reader
import Control.Monad.Trans
import System.Random (randomIO)
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.ByteString as B
import Data.Serialize (Serialize)
import qualified Data.BTree.Cache.Class as C
import qualified Data.BTree.Cache.STM as Cstm
import qualified Data.BTree.KVBackend.Class as KV
import Debug.Trace
import Prelude hiding (lookup, catch)
type CacheSTM m k v = Cstm.CacheSTM m (Ref (Node k v)) (Node k v)
type CacheSTMP m k v = Cstm.Param m (Ref (Node k v)) (Node k v)
class (Ord k, Serialize k, Interval k) => Key k
instance (Ord k, Serialize k, Interval k) => Key k
class (Eq v, Serialize v) => Value v
instance (Eq v, Serialize v) => Value v
class Interval k where
between :: k -> k -> k
between = const
instance Interval B.ByteString where
between a b = if n < B.length b then B.take n b
else a
where
n = 1 + (length $ takeWhile (uncurry equals) $ B.zip a b)
type TreeBackend mc k v = KV.KVBackend mc (Ref (Node k v)) B.ByteString
type TreeResult m mc k v a = BTreeM m (Cstm.Param mc (Ref (Node k v)) (Node k v)) k v a
ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM m t f | true <- m = t
| otherwise = f
equals a b = a `compare` b == EQ
makeParam :: (MonadIO mIO, C.Cache m p (Ref a) (Node k v))
=> Int
-> Maybe (Ref (Node k1 v1))
-> p
-> mIO (Param p k1 v1)
makeParam ord mroot cacheP = do
un <- liftIO $ newTVarIO []
tv <- liftIO $ newTVarIO $ fromMaybe 0 mroot
when (isNothing mroot) $
liftIO $ C.eval cacheP $ C.store Nothing (Ref 0) emptyTree
rebCh <- liftIO $ newTChanIO
return $ Param ord tv cacheP rebCh un
where
emptyTree = Leaf M.empty
execTree :: Param st k v
-> BTreeM m st k v a
-> m a
execTree p m =
runReaderT (runBTreeM m) p
newRef p = do
ls <- Cstm.liftSTM $ readTVar $ unused p
case ls of
(r:rs) -> do Cstm.liftSTM $ writeTVar (unused p) rs
return r
_ -> Cstm.fail newNode
where
newNode = do
refids <- map Ref `fmap` replicateM 8 randomIO
forM_ refids $ \ref -> C.eval (state p) $ do
n <- C.fetch ref
case n of
Nothing -> do Cstm.liftSTM $ do ls <- readTVar $ unused p
writeTVar (unused p) $ ref : ls
Just _ -> return ()
addMarked p k = do
Cstm.liftSTM $ writeTChan (marked p) k
root' p = do
liftIO $ atomically $ readTVar $ root p
fetch' ref = do
n <- C.fetch ref
return $ fromMaybe (error $ "invalid tree: " ++ show ref) n
storeNew p parent ns = do
refs <- replicateM (length ns) $ newRef p
mapM_ (\(r, n) -> C.store (Just parent) r n) $ zip refs ns
return refs
split p parent r (Leaf ks) = do
[refL, refR] <- storeNew p r [leafL, leafR]
when (keyL > keyB || keyR <= keyB) $ error "Between has failed!"
let branch = Branch [keyB] [refL, refR]
C.store (Just parent) r $! branch
return branch
where
(keyL, keyR) = ((fst $ last keysL), (fst $ head keysR))
keyB = keyL `between` keyR
keys = M.toList ks
(keysL, keysR) = L.splitAt (order p) keys
leafL = Leaf $! M.fromAscList keysL
leafR = Leaf $! M.fromAscList keysR
split p parent r (Branch keys rs) =
do [refL, refR] <- storeNew p r [ Branch keysL rsL
, Branch keysR rsR ]
let branch = Branch [b] [refL, refR]
mapM_ (Cstm.updateTag $ Just r) rs
C.store (Just parent) r $! branch
return branch
where
i = order p
(keysL, (b : keysR)) = L.splitAt i keys
(rsL, rsR) = L.splitAt (i+1) rs
insert :: ( MonadIO m, TreeBackend mc k v
, Key k, Value v)
=> k
-> v
-> TreeResult m mc k v ()
insert k v = do _ <- modify const k v
return ()
delete :: ( MonadIO m, TreeBackend mc k v
, Key k, Value v)
=> k
-> TreeResult m mc k v (Maybe v)
delete k = modifyLeaf (findChild k) $ \p parent r ks -> do
let (vOld, ks') = M.updateLookupWithKey (\_ _ -> Nothing) k ks
C.store (Just parent) r $! Leaf ks'
return vOld
lookup :: ( MonadIO m, TreeBackend mc k v
, Key k, Value v)
=> k
-> TreeResult m mc k v (Maybe v)
lookup k = modifyLeaf (findChild k) $ \_ _ _ -> return . M.lookup k
modify :: ( MonadIO m, TreeBackend mc k v
, Key k, Value v)
=> (v -> v -> v)
-> k
-> v
-> TreeResult m mc k v (Maybe v)
modify f k v = do
modifyLeaf (findChild k) $ \p parent r ks -> do
let (vOld, ks') = M.insertLookupWithKey (const f) k v ks
lf' = Leaf ks'
if M.size ks' <= 2 * order p
then C.store (Just parent) r $! lf'
else do Branch [bk] _ <- split p parent r lf'
addMarked p bk
return vOld
findChild k (Branch ks rs) =
let idx = L.length $ L.takeWhile (<k) ks in rs !! idx
height :: (MonadIO m, C.Cache m1 p (Ref (Node k v)) (Node k v)) => BTreeM m p k v Int
height = do p <- ask
r <- root' p
liftIO $ go p r
where
go p r = do
n <- C.eval (state p) $ C.fetch r
case n of
Nothing -> error $ "height: invalid tree: " ++ show r
Just (Leaf ks) -> return 1
Just (Branch _ rs) -> do h <- foldM (\h r -> do h' <- go p r
return $ max h h') 0 rs
return $ 1 + h
modifyLeaf :: (Ord k, MonadIO m, C.Cache mc p (Ref (Node k v)) (Node k v)) =>
(Node k v -> Ref (Node k v)) ->
(Param p k v -> Ref (Node k v) -> Ref (Node k v) -> M.Map k v -> mc a) ->
BTreeM m p k v a
modifyLeaf pick f = retry 100 $ do p <- ask
r <- root' p
liftIO $ try $ go p r r
where
retry 0 m = do x <- m
either throw return x
retry n m = do x <- m
case x of
Left (_::ErrorCall) -> retry (n1) m
Right a -> return a
go p parent r = do
res <- C.eval (state p) transaction
either (go p r) return res
where
transaction =
do n <- C.fetch r
case n of
Just (Leaf ks) ->
do v <- f p parent r ks
return $! Right $! v
Just n -> return $! Left $! pick n
Nothing ->
error $ "modifyLeaf: invalid tree: " ++ show r
rebalanceProcess :: (MonadIO m, TreeBackend m2 k v, Key k, Value v)
=> Param (CacheSTMP m2 k v) k v
-> m (MVar ThreadId)
rebalanceProcess p = liftIO $ do
mv <- newEmptyMVar
pid <- forkIO $ forever $ execTree p $ do
p <- ask
r <- root' p
k <- liftIO $ atomically $ readTChan $ marked p
x <- liftIO $ takeMVar mv
rebalanceKey p r r k
liftIO $ putMVar mv x
putMVar mv pid
return mv
rebalanceAll :: (MonadIO m, Key k, Value v,
C.Cache (CacheSTM m2 k v) (CacheSTMP m2 k v) (Ref (Node k v)) (Node k v),
TreeBackend m2 k v) =>
BTreeM m (CacheSTMP m2 k v) k v ()
rebalanceAll = ifM rebalance rebalanceAll $ return ()
rebalance :: (Key k, Value v, MonadIO m,
C.Cache (CacheSTM m2 k v) (CacheSTMP m2 k v) (Ref (Node k v)) (Node k v),
TreeBackend m2 k v) =>
BTreeM m (CacheSTMP m2 k v) k v Bool
rebalance = do p <- ask
r <- root' p
let ch = marked p
mk <- liftIO $ atomically $
ifM (isEmptyTChan ch)
(return Nothing)
(Just `fmap` readTChan ch)
case mk of
Just k -> do rebalanceKey p r r k
return True
Nothing -> return False
rebalance' :: (Key k, Value v, MonadIO m,
C.Cache (CacheSTM m2 k v) (CacheSTMP m2 k v) (Ref (Node k v)) (Node k v),
TreeBackend m2 k v) =>
BTreeM m (CacheSTMP m2 k v) k v ()
rebalance' = do p <- ask
r <- root' p
let ch = marked p
k <- liftIO $ atomically $ readTChan ch
rebalanceKey p r r k
rebalanceKey p parent r k = do
ord <- asks order
mr <- liftIO $ C.eval (state p) $ trans ord
case mr of
Nothing -> return ()
Just r' -> rebalanceKey p r r' k
where
trans ord = do
mn <- C.fetch r
case mn of
Nothing -> error $ "rebalanceKey: invalid tree: " ++ show r
Just (Leaf _) -> return Nothing
Just (n@(Branch ks rs)) -> do
let cr = findChild k n
mcn <- C.fetch cr
case mcn of
Nothing -> error $ "rebalanceKey: invalid tree: cr = " ++ show cr
Just (Leaf _) -> return Nothing
Just (cn@(Branch [k'] [r0, r1])) -> do
let ks' = L.insert k' ks
Just idx = L.elemIndex k' ks'
(a, _ : b) = L.splitAt idx rs
rs' = a ++ (r0 : r1 : b)
branch = Branch ks' rs'
if length rs' <= 2 * ord
then do C.store (Just parent) r branch
C.remove Nothing cr
else do branch'@(Branch _ frs) <- split p parent r branch
C.store (Just parent) r branch'
C.remove Nothing cr
addMarked p k'
if k `equals` k' then return Nothing
else return $ Just r
Just (Branch cks crs) -> return $ Just cr
deleteByFst k = L.deleteBy cmpfst (k, undefined)
where
cmpfst (a, _) (b, _) = a `equals` b
size :: (Ord k, MonadIO m, C.Cache mc p (Ref (Node k v)) (Node k v)) =>
BTreeM m p k v Int
size = do p <- ask
r <- root' p
go p r
where
go p r = do
n <- liftIO $ C.eval (state p) $ fetch' r
case n of
Leaf ks -> return $ M.size ks
Branch ks rs -> do
tot <- foldM (\n r ->
do m <- go p r
return $ n + m
) 0 rs
return tot
toList :: ( MonadIO m, TreeBackend mc k v, Key k, Value v)
=> TreeResult m mc k v [(k, v)]
toList = do p <- ask
r <- root' p
Cstm.withGeneration (state p) $ \n -> go n p r
where
go gen p r = do
n <- liftIO $ C.eval (state p) $ Cstm.fetchGen gen r
case n of
Nothing -> error $ "toList: invalid tree: " ++ show r
Just (Leaf ks) -> return $ M.toList ks
Just (Branch ks rs) -> do
ls <- mapM (go gen p) rs
return $ concat ls
save :: ( MonadIO m, TreeBackend mc k v, Key k, Value v)
=> TreeResult m mc k v (Ref (Node k v))
save = do
p <- ask
liftIO $ C.sync $ state p
root' p
findMin :: (Ord k, MonadIO m, C.Cache mc p (Ref (Node k v)) (Node k v)) =>
BTreeM m p k v (k, v)
findMin = modifyLeaf (\(Branch _ rs) -> head rs) $ \_ _ _ -> return . M.findMin
findMax :: (Ord k, MonadIO m, C.Cache mc p (Ref (Node k v)) (Node k v)) =>
BTreeM m p k v (k, v)
findMax = modifyLeaf (\(Branch _ rs) -> last rs) $ \_ _ _ -> return . M.findMax
foldli :: (MonadIO m, TreeBackend mc k v, Key k, Value v) =>
(a -> k -> v -> a) -> a -> TreeResult m mc k v a
foldli f a = do p <- ask
r <- root' p
go p a r
where
go p a r = do
n <- liftIO $ C.eval (state p) $ fetch' r
case n of
Leaf ks -> return $ M.foldlWithKey f a ks
Branch ks rs -> foldM (go p) a rs
foldri :: (MonadIO m, TreeBackend mc k v, Key k, Value v) =>
(k -> v -> a -> a) -> a -> TreeResult m mc k v a
foldri f a = do p <- ask
r <- root' p
go p a r
where
go p a r = do
n <- liftIO $ C.eval (state p) $ fetch' r
case n of
Leaf ks -> return $ M.foldrWithKey f a ks
Branch ks rs -> foldM (go p) a $ reverse rs
search :: (MonadIO m, TreeBackend mc k v, Key k, Value v)
=> ((k, k) -> Bool)
-> TreeResult m mc k v [(k, v)]
search f = do p <- ask
r <- root' p
lb <- fst `fmap'` findMin
ub <- fst `fmap'` findMax
go p r lb ub
where
fmap' f m = m >>= return . f
go p r lb ub = do
n <- liftIO $ C.eval (state p) $ fetch' r
case n of
Leaf ks -> return $ filter (\(k, _) -> f (k, k)) $ M.toList ks
Branch ks rs -> do
let ints = findInts 0 (lb:ks ++ [ub])
ls <- mapM (\(n, lb, ub) -> go p (rs!!n) lb ub) ints
return $ concat ls
where
findInts (!n) (lb:ub:xs) =
if f (lb, ub) then (n, lb, ub) : rest else rest
where
rest = findInts (n+1) (ub:xs)
findInts _ _ = []
search_ :: (MonadIO m, TreeBackend mc k v, Key k, Value v)
=> ((k, k) -> Bool)
-> TreeResult m mc k v ()
search_ f = do p <- ask
r <- root' p
lb <- fst `fmap'` findMin
ub <- fst `fmap'` findMax
tv <- liftIO $ newEmptyMVar
liftIO $ go p r lb ub tv
liftIO $ takeMVar tv
where
fmap' f m = m >>= return . f
go p r lb ub tv = do
n <- C.eval (state p) $ fetch' r
case n of
Leaf ks -> putMVar tv ()
Branch ks rs -> do
let ints = findInts 0 (lb:ks ++ [ub])
tv' <- newEmptyMVar
mapM_ (\(n, lb, ub) -> forkIO $ go p (rs!!n) lb ub tv') ints
replicateM_ (length ints) $ takeMVar tv'
putMVar tv ()
where
findInts (!n) (lb:ub:xs) =
if f (lb, ub) then (n, lb, ub) : rest else rest
where
rest = findInts (n+1) (ub:xs)
findInts _ _ = []