b-tree (empty) → 0.1.0.0
raw patch · 13 files changed
+928/−0 lines, 13 filesdep +basedep +binarydep +bytestringsetup-changed
Dependencies added: base, binary, bytestring, containers, criterion, directory, errors, filepath, lens, mmap, mtl, pipes, pipes-interleave, transformers, vector
Files
- BTree.hs +31/−0
- BTree/BinaryFile.hs +87/−0
- BTree/BinaryList.hs +83/−0
- BTree/BuildUnordered.hs +109/−0
- BTree/Builder.hs +231/−0
- BTree/Lookup.hs +47/−0
- BTree/Merge.hs +61/−0
- BTree/Types.hs +97/−0
- BTree/Walk.hs +49/−0
- Benchmark.hs +28/−0
- LICENSE +30/−0
- Setup.hs +2/−0
- b-tree.cabal +73/−0
+ BTree.hs view
@@ -0,0 +1,31 @@+module BTree ( -- * Basic types+ BLeaf(..)+ , Size+ , Order+ -- * Building trees+ , fromOrderedToFile+ , fromOrderedToByteString+ , fromUnorderedToFile+ -- * Looking up in trees+ , LookupTree+ , open+ , fromByteString+ , lookup+ -- * Merging trees+ , mergeTrees+ , mergeLeaves+ , sizedProducerForTree+ -- * Iterating over leaves+ , walkLeaves+ ) where++import Prelude hiding (lookup)+import BTree.Types+import BTree.Merge+import BTree.Builder+import BTree.Lookup+import BTree.Walk+import BTree.BuildUnordered++-- | This package provides immutable B* trees targetting large data+-- sets requiring secondary storage.
+ BTree/BinaryFile.hs view
@@ -0,0 +1,87 @@+module BTree.BinaryFile+ ( writeWithHeader+ , readWithHeader+ ) where++import Control.Monad (when)+import Control.Error+import Control.Monad.Trans.Class+import Control.Applicative+import Data.Word+import System.IO++import qualified Data.ByteString.Lazy as LBS+import qualified Data.Binary as B+import qualified Data.Binary.Get as B+import qualified Data.Binary.Put as B+import Pipes++-- | This module provides helpers for emitting and reading binary files with+-- a trailing "header".++-- | An internal data structure placed at the very end of the file which+-- describes the header and provides a magic number for sanity checking.+data Epilogue = Epilogue { magic :: Word64+ , headerLen :: Word64+ }+ deriving (Show)++epiLength :: Integer+epiLength = 16++magicNumber :: Word64+magicNumber = 0xdeadbeef++instance B.Binary Epilogue where+ get = Epilogue <$> B.getWord64le <*> B.getWord64le+ put (Epilogue m l) = B.putWord64le m >> B.putWord64le l++-- | Write the produced bytestrings to the file followed by the+-- returned header+writeWithHeader :: (MonadIO m, B.Binary hdr)+ => FilePath+ -> Producer LBS.ByteString m (hdr, r)+ -> m r+writeWithHeader fname prod = do+ h <- liftIO $ openFile fname WriteMode+ (hdr, r) <- runEffect $ for prod (liftIO . LBS.hPut h)+ let encoded = B.encode hdr+ liftIO $ LBS.hPut h encoded+ let epi = Epilogue { magic = magicNumber+ , headerLen = fromIntegral $ LBS.length encoded }+ liftIO $ LBS.hPut h (B.encode epi)+ liftIO $ hClose h+ return r++annotate :: Monad m => String -> EitherT String m a -> EitherT String m a+annotate ann = fmapLT ((ann++": ")++)++runGetT :: Monad m => B.Get a -> LBS.ByteString -> EitherT String m a+runGetT _get bs = do+ case B.runGetOrFail _get bs of+ Left (_, _, e) -> left e+ Right (_, _, a) -> right a++-- | Read and verify the header from the file, then pass it along with the+-- file's handle to an action. The file handle sits at the beginning of the+-- written content when passed to the action.+readWithHeader :: (MonadIO m, B.Binary hdr)+ => FilePath+ -> (hdr -> Handle -> m a)+ -> EitherT String m a+readWithHeader fname action = do+ h <- liftIO $ openFile fname ReadMode+ -- read epilogue+ liftIO $ hSeek h SeekFromEnd (-epiLength)+ epiBytes <- liftIO (LBS.hGet h $ fromIntegral epiLength)+ epi <- annotate "Error reading epilogue" (runGetT B.get epiBytes)+ when (magic epi /= magicNumber) $+ left "BinaryFile.readWithHeader: Bad magic number"+ -- read header+ let offset = fromIntegral epiLength + fromIntegral (headerLen epi)+ liftIO $ hSeek h SeekFromEnd (negate offset)+ hdrBytes <- liftIO (LBS.hGet h $ fromIntegral $ headerLen epi)+ hdr <- annotate "Error reading header" (runGetT B.get hdrBytes)+ -- pass control to action+ liftIO $ hSeek h AbsoluteSeek 0+ lift $ action hdr h
+ BTree/BinaryList.hs view
@@ -0,0 +1,83 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}++module BTree.BinaryList+ ( BinaryList+ -- * Construction+ , toBinaryList+ -- * Fetching contents+ , stream+ -- * Other queries+ , length+ , filePath+ ) where++import Prelude hiding (length)+import Control.Applicative+import Control.Monad.Trans.Class+import Control.Error+import Data.Word+import System.IO++import qualified Data.ByteString.Lazy as LBS+import qualified Data.Binary as B+import qualified Data.Binary.Get as B+import qualified Data.Binary.Put as B+import Pipes++import BTree.BinaryFile++-- | A file containing a finite list of binary encoded items+newtype BinaryList a = BinList FilePath+ deriving (Show)++-- | Get the path to the @BinaryList@ file+filePath :: BinaryList a -> FilePath+filePath (BinList f) = f++data Header = Header { hdrLength :: Word64 }+ deriving (Show)++instance B.Binary Header where+ get = Header <$> B.getWord64le+ put (Header l) = B.putWord64le l++-- | Encode the items of the given producer+toBinaryList :: forall m a r. (MonadIO m, B.Binary a)+ => FilePath -> Producer a m r -> m (BinaryList a, r)+toBinaryList fname producer = do+ writeWithHeader fname (go 0 producer)+ where+ go :: Int -> Producer a m r+ -> Producer LBS.ByteString m (Header, (BinaryList a, r))+ go !n prod = do+ result <- lift $ next prod+ case result of+ Left r ->+ let hdr = Header (fromIntegral n)+ in return (hdr, (BinList fname, r))+ Right (a, prod') -> do+ yield (B.encode a)+ go (n+1) prod'++withHeader :: MonadIO m+ => BinaryList a -> (Header -> Handle -> m b) -> EitherT String m b+withHeader (BinList fname) action = readWithHeader fname action++length :: MonadIO m => BinaryList a -> EitherT String m Word64+length bl = withHeader bl $ \hdr _ -> return $ hdrLength hdr++-- | Stream the items out of a @BinaryList@+stream :: forall m a. (B.Binary a, MonadIO m)+ => BinaryList a -> EitherT String m (Producer a m (Either String ()))+stream bl = withHeader bl readContents+ where+ readContents :: Header -> Handle -> m (Producer a m (Either String ()))+ readContents hdr h = return $ liftIO (LBS.hGetContents h) >>= go (hdrLength hdr)++ go :: Word64 -> LBS.ByteString -> Producer a m (Either String ())+ go 0 _ = return $ Right ()+ go !n bs =+ case B.decodeOrFail bs of+ Left (_, _, e) -> return $ Left e+ Right (bs', _, a) -> yield a >> go (n-1) bs'
+ BTree/BuildUnordered.hs view
@@ -0,0 +1,109 @@+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE FlexibleContexts #-}++module BTree.BuildUnordered+ ( fromUnorderedToFile ) where++import Control.Monad.Trans.State+import Control.Error+import Data.Traversable (forM)++import qualified Data.Binary as B+import qualified Data.Set as S+import System.IO+import System.Directory (removeFile)++import Pipes+import Pipes.Interleave+import qualified BTree.BinaryList as BL+import BTree.Types+import BTree.Builder++-- | Maximum number of leaf lists to attempt to merge at once.+-- This is bounded by the maximum file handle count.+maxChunkMerge :: Int+maxChunkMerge = 100++tempFilePath :: FilePath -> String -> IO FilePath+tempFilePath dir template = do+ (fname, h) <- liftIO $ openTempFile dir template+ hClose h+ return fname++-- | Build a B-tree into the given file.+--+-- This does not assume that the leaves are produced in order. Instead,+-- the sorting is handled internally through a simple merge sort. Chunks of+-- leaves are collected, sorted in memory, and then written to intermediate+-- trees. At the end these trees are then merged.+fromUnorderedToFile :: forall m e k r.+ (MonadIO m, B.Binary (BLeaf k e), B.Binary k, B.Binary e, Ord k)+ => FilePath -- ^ Path to scratch directory+ -> Int -- ^ Maximum chunk size+ -> Order -- ^ Order of tree+ -> FilePath -- ^ Output file+ -> Producer (BLeaf k e) m r -- ^ 'Producer' of elements+ -> EitherT String m ()+fromUnorderedToFile scratch maxChunk order dest producer = do+ bList <- lift (execStateT (fillLists producer) []) >>= goMerge+ size <- BL.length bList+ stream <- BL.stream bList+ lift $ fromOrderedToFile order size dest stream+ liftIO $ removeFile $ BL.filePath bList+ where+ fillLists :: Producer (BLeaf k e) m r -> StateT [BL.BinaryList (BLeaf k e)] m r+ fillLists prod = do+ fname <- liftIO $ tempFilePath scratch "chunk.list"+ (leaves, rest) <- lift $ takeChunk maxChunk prod+ (bList, ()) <- lift $ BL.toBinaryList fname $ each leaves+ modify (bList:)+ case rest of+ Left r -> return r+ Right nextProd -> fillLists nextProd++ goMerge :: [BL.BinaryList (BLeaf k e)] -> EitherT String m (BL.BinaryList (BLeaf k e))+ goMerge [l] = return l+ goMerge ls = do+ ls'' <- forM (splitChunks maxChunkMerge ls) $ \ls'->do+ fname <- liftIO $ tempFilePath scratch "merged.list"+ list <- mergeLists fname ls'+ liftIO $ mapM_ (removeFile . BL.filePath) ls'+ return list+ goMerge ls''++-- | Split the list into chunks of bounded size and run each through a function+splitChunks :: Int -> [a] -> [[a]]+splitChunks chunkSize = go+ where+ go [] = []+ go xs = let (prefix,suffix) = splitAt chunkSize xs+ in prefix : go suffix++throwLeft :: Monad m => m (Either String r) -> m r+throwLeft action = action >>= either error return++mergeLists :: (Ord a, B.Binary a, MonadIO m)+ => FilePath -> [BL.BinaryList a] -> EitherT String m (BL.BinaryList a)+mergeLists dest lists = do+ streams <- mapM BL.stream lists+ let prod = interleave compare (map throwLeft streams)+ (bList, ()) <- lift $ BL.toBinaryList dest prod+ return bList++-- | Take the first 'n' elements and collect them in a 'Set'. Return+-- a 'Producer' which will emit the remaining elements (or the return+-- value).+takeChunk :: forall m a r. (Monad m, Ord a)+ => Int+ -> Producer a m r+ -> m (S.Set a, Either r (Producer a m r))+takeChunk n = go n S.empty+ where+ go :: Int -> S.Set a -> Producer a m r -> m (S.Set a, Either r (Producer a m r))+ go 0 s prod = return (s, Right prod)+ go i s prod = do+ result <- next prod+ case result of+ Left r -> return (s, Left r)+ Right (a, prod') -> go (i-1) (S.insert a s) prod'
+ BTree/Builder.hs view
@@ -0,0 +1,231 @@+{-# LANGUAGE TemplateHaskell #-}++module BTree.Builder+ ( buildNodes, putBS+ , fromOrderedToFile+ , fromOrderedToByteString+ ) where++import Control.Monad.Trans.State.Strict+import Control.Monad.IO.Class+import Control.Monad++import Data.Foldable as F+import qualified Data.Sequence as Seq+import Data.Sequence (Seq)++import Data.Ratio+import Control.Lens+import System.IO++import qualified Data.Binary as B+import Data.Binary (Binary)+import qualified Data.ByteString.Lazy as LBS++import Pipes+import Pipes.Core+import qualified Pipes.Internal as PI++import BTree.Types++-- | A Producer which accepts offsets for the yielded objects in return+type DiskProducer a = Proxy X () (OnDisk a) a++putBS :: (Binary a, Monad m) => a -> Proxy (OnDisk a) a () LBS.ByteString m r+putBS a0 = evalStateT (go a0) 0+ where+ go a = do+ s <- get+ let bs = B.encode a+ put $! s + fromIntegral (LBS.length bs)+ lift $ yield bs+ a' <- lift $ request (OnDisk s)+ go a'++data DepthState k e = DepthS { -- | nodes to be included in the active node+ _dNodes :: !(Seq (k, OnDisk (BTree k OnDisk e)))+ -- | the length of @dNodes@+ , _dNodeCount :: !Int+ -- | the desired number of elements to fill the active node+ , _dMinFill :: [Int]+ }+makeLenses ''DepthState++next' :: (Monad m) => Proxy X () a' a m r -> m (Either r (a, a' -> Proxy X () a' a m r))+next' = go+ where+ go p = case p of+ PI.Request _ fu -> go (fu ())+ PI.Respond a fu -> return (Right (a, fu))+ PI.M m -> m >>= go+ PI.Pure r -> return (Left r)++-- | Compute the optimal node sizes for each stratum of a tree of+-- given size and order+optimalFill :: Order -> Size -> [[Int]]+optimalFill order size = go (fromIntegral size)+ where+ go :: Int -> [[Int]]+ go 0 = error "BTree.Builder.optimalFill: zero size"+ go n =+ let nNodes = ceiling (n % order')+ order' = fromIntegral order :: Int+ nodes = let (nPerNode, leftover) = n `divMod` nNodes+ in zipWith (+) (replicate nNodes nPerNode)+ (replicate leftover 1 ++ repeat 0)+ rest = case nNodes of+ 1 -> []+ _ -> go nNodes+ in nodes : rest++-- | Given a producer of a known number of leaves, produces an optimal B-tree.+-- Technically the size is only an upper bound: the producer may+-- terminate before providing the given number of leaves although the resulting+-- tree will break the minimal fill invariant.+buildNodes :: Monad m+ => Order -> Size+ -> DiskProducer (BLeaf k e) m r+ -> DiskProducer (BTree k OnDisk e) m (BTreeHeader k e)+buildNodes order size =+ flip evalStateT initialState . loop size+ where+ initialState = map (DepthS Seq.empty 0) $ optimalFill order size+ -- depth=0 denotes the bottom (leaves) of the tree.+ loop :: Monad m+ => Size -> DiskProducer (BLeaf k e) m r+ -> StateT [DepthState k e] (DiskProducer (BTree k OnDisk e) m)+ (BTreeHeader k e)+ loop n producer = do+ _next <- lift $ lift $ next' producer+ case _next of+ Left _ -> do+ flushAll (size - n)+ Right _ | n == 0 -> do+ flushAll (size - n)+ Right (leaf@(BLeaf k _), producer') -> do+ -- TODO: Is there a way to check this coercion with the type system?+ OnDisk offset <- processNode k $ Leaf leaf+ loop (n-1) $ producer' (OnDisk offset)++ isFilled :: Monad m+ => StateT [DepthState k e] m Bool+ isFilled = zoom (singular _head) $ do+ nodeCount <- use dNodeCount+ minFill:_ <- use dMinFill+ return $ nodeCount >= minFill++ emitNode :: Monad m+ => StateT [DepthState k e] (DiskProducer (BTree k OnDisk e) m)+ (OnDisk (BTree k OnDisk e))+ emitNode = do+ (k0,node0):nodes <- zoom (singular _head) $ do+ nodes <- uses dNodes F.toList+ dNodes .= Seq.empty+ dNodeCount .= 0+ dMinFill %= tail+ return nodes++ --when (null nodes)+ -- $ error "BTree.Builder.buildNodes: Internal invariant broken: unexpected empty node"+ let newNode = Node node0 nodes+ s <- get+ case s of+ [_] -> lift $ respond newNode+ _ -> zoom (singular _tail) $ processNode k0 newNode++ processNode :: Monad m+ => k -> BTree k OnDisk e+ -> StateT [DepthState k e]+ (DiskProducer (BTree k OnDisk e) m)+ (OnDisk (BTree k OnDisk e))+ processNode startKey tree = do+ filled <- isFilled+ when filled $ void $ emitNode+ offset <- lift $ respond tree+ zoom _head $ do+ dNodes %= (Seq.|> (startKey, offset))+ dNodeCount += 1+ return offset++ flushAll :: Monad m+ => Size+ -> StateT [DepthState k e]+ (DiskProducer (BTree k OnDisk e) m)+ (BTreeHeader k e)+ flushAll realSize = do+ s <- get+ case s of+ [] -> error "BTree.Builder.flushAll: We should never get here"+ [_] -> do -- We are at the top node, this shouldn't be flushed yet+ root <- emitNode+ return $ BTreeHeader magic 1 order realSize root+ d:_ -> do when (not $ Seq.null $ d^.dNodes) $ void $ emitNode+ zoom (singular _tail) $ flushAll realSize++-- | Produce a bytestring representing the nodes and leaves of the+-- B-tree and return a suitable header+buildTree :: (Monad m, Binary e, Binary k)+ => Order -> Size+ -> Producer (BLeaf k e) m r+ -> Producer LBS.ByteString m (BTreeHeader k e)+buildTree order size producer =+ dropUpstream $ buildNodes order size (dropUpstream producer) >>~ putBS++dropUpstream :: Monad m => Proxy X () () b m r -> Proxy X () b' b m r+dropUpstream = go+ where+ go producer = do+ n <- lift $ next producer+ case n of+ Left r -> return r+ Right (a, producer') -> respond a >> go producer'++-- | Build a B-tree into the given file.+--+-- As the name suggests, this requires that the @Producer@ emits+-- leaves in ascending key order.+fromOrderedToFile :: (MonadIO m, Binary e, Binary k)+ => Order -- ^ Order of tree+ -> Size -- ^ Maximum tree size+ -> FilePath -- ^ Output file+ -> Producer (BLeaf k e) m r -- ^ 'Producer' of elements+ -> m ()+fromOrderedToFile order size fname producer = do+ h <- liftIO $ openFile fname WriteMode+ liftIO $ LBS.hPut h $ B.encode invalidHeader+ hdr <- runEffect $ for (buildTree order size producer) $ liftIO . LBS.hPut h+ liftIO $ hSeek h AbsoluteSeek 0+ liftIO $ LBS.hPut h $ B.encode hdr+ liftIO $ hClose h+ return ()+ where+ invalidHeader = BTreeHeader 0 0 0 0 (OnDisk 0)++-- | Build a B-tree into @ByteString@+--+-- As the name suggests, this requires that the @Producer@ emits+-- leaves in ascending key order.+--+-- This is primarily used for testing. In particular, note that+-- this is a bad idea for large trees as the entire contents of the+-- tree will need to be kept in memory until all leaves have been+-- added so that the header can be prepended.+fromOrderedToByteString :: (Monad m, Binary e, Binary k)+ => Order -- ^ Order of tree+ -> Size -- ^ Maximum tree size+ -> Producer (BLeaf k e) m r -- ^ 'Producer' of elements+ -> m LBS.ByteString+fromOrderedToByteString order size producer = do+ (bs, hdr) <- foldR LBS.append LBS.empty id $ buildTree order size producer+ return $ B.encode hdr `LBS.append` bs++-- | Like @Pipes.Prelude.fold@ but provides returns producer result+-- in addition to accumulator+foldR :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Producer a m r -> m (b, r)+foldR step begin done p0 = loop p0 begin+ where+ loop p x = case p of+ PI.Request _ fu -> loop (fu ()) x+ PI.Respond a fu -> loop (fu ()) $! step x a+ PI.M m -> m >>= \p' -> loop p' x+ PI.Pure r -> return (done x, r)
+ BTree/Lookup.hs view
@@ -0,0 +1,47 @@+module BTree.Lookup ( LookupTree+ , open+ , fromByteString+ , lookup+ ) where++import Prelude hiding (lookup)+import Control.Error+import Control.Lens hiding (children)+import qualified Data.ByteString as BS+import qualified Data.ByteString.Lazy as LBS+import Data.Binary+import System.IO.MMap+import BTree.Types++fetch :: (Binary a) => LookupTree k e -> OnDisk a -> a+fetch lt (OnDisk offset) =+ decode $ LBS.fromStrict $ BS.drop (fromIntegral offset) (lt^.ltData)++-- | Read a B-tree from a 'ByteString' produced by 'BTree.Builder'+fromByteString :: LBS.ByteString -> Either String (LookupTree k e)+fromByteString bs = do+ (rest, _, hdr) <- fmapL (\(_,_,e)->e) $ decodeOrFail bs+ validateHeader hdr+ return $ LookupTree (LBS.toStrict rest) hdr++-- | Open a B-tree file.+open :: FilePath -> IO (Either String (LookupTree k e))+open fname = runEitherT $ do+ d <- fmapLT show $ tryIO $ mmapFileByteString fname Nothing+ EitherT $ return $ fromByteString (LBS.fromStrict d)+ +-- | Lookup a key in a B-tree.+lookup :: (Binary k, Binary e, Ord k)+ => LookupTree k e -> k -> Maybe e+lookup lt k = go $ fetch lt (lt ^. ltHeader . btRoot)+ where+ go (Leaf (BLeaf k' e))+ | k' == k = Just e+ | otherwise = Nothing+ go (Node c0 []) = go $ fetch lt c0 -- is this case necessary?+ go (Node c0 children@((k0,_):_))+ | k < k0 = go $ fetch lt c0+ | otherwise =+ case takeWhile (\(k',_)->k' <= k) children of+ [] -> Nothing+ xs -> go $ fetch lt $ snd $ last xs
+ BTree/Merge.hs view
@@ -0,0 +1,61 @@+{-# LANGUAGE TemplateHaskell, BangPatterns, GeneralizedNewtypeDeriving #-}++module BTree.Merge ( mergeTrees+ , mergeLeaves+ , sizedProducerForTree+ ) where++import Prelude hiding (sum)+import Control.Applicative+import Data.Foldable+import Data.Function (on)+import Control.Monad.State hiding (forM_)+import Data.Binary+import Control.Lens+import Pipes+import Pipes.Interleave++import BTree.Types+import BTree.Builder+import BTree.Walk++-- | Merge trees' leaves taking ordered leaves from a set of producers.+--+-- Each producer must be annotated with the number of leaves it is+-- expected to produce. The size of the resulting tree will be at most+-- the sum of these sizes.+mergeLeaves :: (MonadIO m, Functor m, Binary k, Binary e, Ord k)+ => (e -> e -> m e) -- ^ merge operation on elements+ -> Order -- ^ order of merged tree+ -> FilePath -- ^ name of output file+ -> [(Size, Producer (BLeaf k e) m ())] -- ^ producers of leaves to merge+ -> m ()+mergeLeaves append destOrder destFile producers = do+ let size = sum $ map fst producers+ fromOrderedToFile destOrder size destFile $+ mergeM (compare `on` key) doAppend (map snd producers)+ where+ doAppend (BLeaf k e) (BLeaf _ e') = BLeaf k <$> append e e'+ key (BLeaf k _) = k++-- | Merge several 'LookupTrees'+--+-- This is a convenience function for merging several trees already on+-- disk. For a more flexible interface, see 'mergeLeaves'.+mergeTrees :: (MonadIO m, Functor m, Binary k, Binary e, Ord k)+ => (e -> e -> m e) -- ^ merge operation on elements+ -> Order -- ^ order of merged tree+ -> FilePath -- ^ name of output file+ -> [LookupTree k e] -- ^ trees to merge+ -> m ()+mergeTrees append destOrder destFile trees = do+ mergeLeaves append destOrder destFile+ $ map sizedProducerForTree trees++-- | Get a sized 'Producer' suitable for 'mergeLeaves' from a 'LookupTree'+sizedProducerForTree :: (Monad m, Binary k, Binary e)+ => LookupTree k e -- ^ a tree+ -> (Size, Producer (BLeaf k e) m ())+ -- ^ a sized 'Producer' suitable for passing+ -- to 'mergeLeaves'+sizedProducerForTree lt = (lt ^. ltHeader . btSize, void $ walkLeaves lt)
+ BTree/Types.hs view
@@ -0,0 +1,97 @@+{-# LANGUAGE DeriveGeneric, FlexibleContexts, TemplateHaskell, UndecidableInstances, StandaloneDeriving #-}++module BTree.Types where++import Data.Binary+import GHC.Generics+import Control.Monad (when)+import Control.Applicative+import Control.Lens+import Data.Int+import qualified Data.ByteString as BS++-- | An offset within the stream+type Offset = Int64++-- | The number of entries in a B-tree+type Size = Word64++-- | The maximum number of children of a B-tree inner node+type Order = Word64++-- | 'OnDisk a' is a reference to an object of type 'a' on disk.+-- The offset does not include the header; e.g. the first object after+-- the header is located at offset 0.+newtype OnDisk a = OnDisk Offset+ deriving (Show, Eq, Ord)++instance Binary (OnDisk a) where+ get = OnDisk <$> get+ put (OnDisk off) = put off++-- | A tree leaf (e.g. key/value pair)+data BLeaf k e = BLeaf !k !e+ deriving (Generic)++deriving instance (Show k, Show e) => Show (BLeaf k e)++-- | This only compares on the keys+instance (Eq k) => Eq (BLeaf k e) where+ BLeaf a _ == BLeaf b _ = a == b++-- | This only compares on the keys+instance Ord k => Ord (BLeaf k e) where+ compare (BLeaf a _) (BLeaf b _) = compare a b++instance (Binary k, Binary e) => Binary (BLeaf k e) where+ get = BLeaf <$> get <*> get+ put (BLeaf k e) = put k >> put e++-- | 'BTree k f e' is a B* tree of key type 'k' with elements of type 'e'.+-- Subtree references are contained within a type 'f'+--+-- The Node constructor contains a left child, and a list of key/child pairs+-- where each child's keys are greater than or equal to the given key.+data BTree k f e = Node (f (BTree k f e)) [(k, f (BTree k f e))]+ | Leaf !(BLeaf k e)+ deriving (Generic)++deriving instance (Show e, Show k, Show (f (BTree k f e))) => Show (BTree k f e)+deriving instance (Eq e, Eq k, Eq (f (BTree k f e))) => Eq (BTree k f e)++instance (Binary k, Binary (f (BTree k f e)), Binary e)+ => Binary (BTree k f e) where+ get = do typ <- getWord8+ case typ of+ 0 -> Node <$> get <*> get+ 1 -> bleaf <$> get <*> get+ _ -> fail "BTree.Types/get: Unknown node type"+ where bleaf k v = Leaf (BLeaf k v)+ put (Node e0 es) = putWord8 0 >> put e0 >> put es+ put (Leaf (BLeaf k0 e)) = putWord8 1 >> put k0 >> put e++magic :: Word64+magic = 0xdeadbeefbbbbcccc++-- | B-tree file header+data BTreeHeader k e = BTreeHeader { _btMagic :: !Word64+ , _btVersion :: !Word64+ , _btOrder :: !Order+ , _btSize :: !Size+ , _btRoot :: !(OnDisk (BTree k OnDisk e))+ }+ deriving (Show, Eq, Generic)+makeLenses ''BTreeHeader++instance Binary (BTreeHeader k e)++validateHeader :: BTreeHeader k e -> Either String ()+validateHeader hdr = do+ when (hdr^.btMagic /= magic) $ Left "Invalid magic number"+ when (hdr^.btVersion > 1) $ Left "Invalid version"++-- | A read-only B-tree for lookups+data LookupTree k e = LookupTree { _ltData :: !BS.ByteString+ , _ltHeader :: !(BTreeHeader k e)+ }+makeLenses ''LookupTree
+ BTree/Walk.hs view
@@ -0,0 +1,49 @@+module BTree.Walk ( walkLeaves+ , walkNodes+ , walkNodesWithOffset+ ) where++import BTree.Types+import qualified Data.ByteString.Lazy as LBS+import Pipes+import qualified Pipes.Prelude as PP+import Data.Binary+import Data.Binary.Get (runGetOrFail)+import Control.Lens++-- If we only look at leaves keys will increase monotonically as we+-- progress through the file.++filterLeaves :: Monad m => Pipe (BTree k OnDisk v) (BLeaf k v) m r+filterLeaves = do+ a <- await+ case a of+ Leaf leaf -> yield leaf+ _ -> return ()+ filterLeaves++-- | Iterate over the leaves of the given tree in ascending key order.+walkLeaves :: (Binary k, Binary v, Monad m)+ => LookupTree k v+ -> Producer (BLeaf k v) m (LBS.ByteString, Maybe String)+walkLeaves b = walkNodes b >-> filterLeaves++-- | Iterate over the nodes and leaves of the given tree. These aren't+-- necessarily sorted.+walkNodes :: (Binary k, Binary v, Monad m)+ => LookupTree k v+ -> Producer (BTree k OnDisk v) m (LBS.ByteString, Maybe String)+walkNodes b = walkNodesWithOffset b >-> PP.map snd++walkNodesWithOffset :: (Binary k, Binary v, Monad m)+ => LookupTree k v+ -> Producer (Offset, BTree k OnDisk v) m (LBS.ByteString, Maybe String)+walkNodesWithOffset = go 0 . view (ltData . to LBS.fromStrict)+ where go offset bs =+ case runGetOrFail get bs of+ Left (rest,_,err) -> return (rest, Just err)+ Right (rest,o,a) -> do+ yield (offset, a)+ if LBS.null rest+ then return (rest, Nothing)+ else go (offset+o) rest
+ Benchmark.hs view
@@ -0,0 +1,28 @@+import BTree as BT+import Criterion.Main+import Pipes++main = do+ buildTree "hello.btree" 100 100000+ Right lt <- BT.open "hello.btree"+ :: IO (Either String (LookupTree Int Int))++ defaultMain $ benchmarks lt++benchmarks lt =+ [ bench "build tree (order 10, 100000 elements)"+ $ nfIO $ buildTree "test.btree" 10 100000+ , bench "build tree (order 100, 100000 elements)"+ $ nfIO $ buildTree "test.btree" 100 100000+ , bench "lookup (5000 lookups)"+ $ nf (\lt->map (lookupBench lt) [0..5000]) lt+ ]++buildTree :: FilePath -> Order -> Int -> IO ()+buildTree fname order n = + BT.fromOrderedToFile 4 (fromIntegral n) fname (each things)+ where things :: [BLeaf Int Int]+ things = [BLeaf (2*i) i | i <- [0..n-1]]+ +lookupBench :: LookupTree Int Int -> Int -> Maybe Int+lookupBench lt = BT.lookup lt
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright (c) 2013, Ben Gamari++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++ * Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++ * Redistributions in binary form must reproduce the above+ copyright notice, this list of conditions and the following+ disclaimer in the documentation and/or other materials provided+ with the distribution.++ * Neither the name of Ben Gamari nor the names of other+ contributors may be used to endorse or promote products derived+ from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ b-tree.cabal view
@@ -0,0 +1,73 @@+name: b-tree+version: 0.1.0.0+synopsis: Immutable disk-based B* trees+description: Immutable disk-based B* trees+homepage: http://github.com/bgamari/b-tree+license: BSD3+license-file: LICENSE+author: Ben Gamari+maintainer: bgamari.foss@gmail.com+-- copyright: +category: Data+build-type: Simple+-- extra-source-files: +cabal-version: >=1.10++source-repository head+ type: git+ location: git@github.com:bgamari/b-tree++library+ exposed-modules: BTree+ other-modules: BTree.Types,+ BTree.Merge,+ BTree.Builder,+ BTree.Lookup,+ BTree.Walk,+ BTree.BinaryFile,+ BTree.BinaryList,+ BTree.BuildUnordered++ other-extensions: DeriveGeneric,+ FlexibleContexts,+ TemplateHaskell,+ UndecidableInstances,+ StandaloneDeriving,+ BangPatterns,+ GeneralizedNewtypeDeriving++ build-depends: base >=4.6 && <4.8,+ mtl >=2.0 && <3.0,+ pipes >=4.1 && <4.2,+ pipes-interleave >= 0.1 && <1.0,+ bytestring >=0.10 && <0.11,+ binary >=0.7 && <0.8,+ transformers >=0.3 && <0.5,+ lens >=3.10 && <4.3,+ containers >=0.5 && <0.6,+ vector >=0.10 && <0.11,+ errors >=1.4 && <1.5,+ filepath >=1.3 && <1.4,+ directory >=1.2 && <1.3,+ mmap >=0.5 && <0.6+ default-language: Haskell2010+ ghc-options: -Wall++benchmark btree-benchmark+ type: exitcode-stdio-1.0+ main-is: Benchmark.hs+ build-depends: base >=4.6 && <4.8,+ mtl >=2.0 && <3.0,+ criterion >=0.8 && <0.10,+ pipes >=4.1 && <4.2,+ pipes-interleave >= 0.1 && <1.0,+ bytestring >=0.10 && <0.11,+ binary >=0.7 && <0.8,+ transformers >=0.3 && <0.5,+ lens >=3.10 && <4.3,+ containers >=0.5 && <0.6,+ vector >=0.10 && <0.11,+ errors >=1.4 && <1.5,+ mmap >=0.5 && <0.6+ default-language: Haskell2010+