b-tree 0.1.1 → 0.1.2
raw patch · 21 files changed
+871/−845 lines, 21 filesdep ~basedep ~binarydep ~containers
Dependency ranges changed: base, binary, containers, criterion, errors, filepath, lens, pipes, vector
Files
- BTree.hs +0/−31
- BTree/BinaryFile.hs +0/−87
- BTree/BinaryList.hs +0/−83
- BTree/BuildUnordered.hs +0/−109
- BTree/Builder.hs +0/−231
- BTree/Lookup.hs +0/−47
- BTree/Merge.hs +0/−61
- BTree/Types.hs +0/−97
- BTree/Walk.hs +0/−49
- Benchmark.hs +0/−28
- b-tree.cabal +16/−22
- benchmarks/Benchmark.hs +28/−0
- src/BTree.hs +31/−0
- src/BTree/BinaryFile.hs +88/−0
- src/BTree/BinaryList.hs +90/−0
- src/BTree/BuildUnordered.hs +113/−0
- src/BTree/Builder.hs +237/−0
- src/BTree/Lookup.hs +47/−0
- src/BTree/Merge.hs +63/−0
- src/BTree/Types.hs +105/−0
- src/BTree/Walk.hs +53/−0
− BTree.hs
@@ -1,31 +0,0 @@-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
@@ -1,87 +0,0 @@-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
@@ -1,83 +0,0 @@-{-# 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
@@ -1,109 +0,0 @@-{-# 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 $ S.toAscList 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
@@ -1,231 +0,0 @@-{-# 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
@@ -1,47 +0,0 @@-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
@@ -1,61 +0,0 @@-{-# 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
@@ -1,97 +0,0 @@-{-# 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
@@ -1,49 +0,0 @@-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
@@ -1,28 +0,0 @@-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
b-tree.cabal view
@@ -1,5 +1,5 @@ name: b-tree-version: 0.1.1+version: 0.1.2 synopsis: Immutable disk-based B* trees description: Immutable disk-based B* trees homepage: http://github.com/bgamari/b-tree@@ -7,10 +7,9 @@ license-file: LICENSE author: Ben Gamari maintainer: bgamari.foss@gmail.com--- copyright: +copyright: (c) 2014 Ben Gamari category: Data build-type: Simple--- extra-source-files: cabal-version: >=1.10 source-repository head@@ -36,18 +35,19 @@ BangPatterns, GeneralizedNewtypeDeriving - build-depends: base >=4.6 && <4.8,+ hs-source-dirs: src+ build-depends: base >=4.6 && <4.9, 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.8,+ lens >=3.10 && <4.13, containers >=0.5 && <0.6,- vector >=0.10 && <0.11,- errors >=1.4 && <1.5,- filepath >=1.3 && <1.4,+ vector >=0.10 && <0.12,+ errors >=2.0 && <2.1,+ filepath >=1.3 && <1.5, directory >=1.2 && <1.3, mmap >=0.5 && <0.6 default-language: Haskell2010@@ -57,7 +57,8 @@ type: exitcode-stdio-1.0 main-is: QuickCheck.hs hs-source-dirs: tests- build-depends: base >=4.6 && <4.8,+ default-language: Haskell2010+ build-depends: base >=4.6 && <4.9, containers, pipes, binary,@@ -67,18 +68,11 @@ 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+ hs-source-dirs: benchmarks+ default-language: Haskell2010+ build-depends: base >=4.6 && <4.9,+ b-tree,+ pipes,+ criterion >=0.8 && <2.0 default-language: Haskell2010
+ benchmarks/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
+ src/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.
+ src/BTree/BinaryFile.hs view
@@ -0,0 +1,88 @@+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+{-# INLINE writeWithHeader #-}++annotate :: Monad m => String -> ExceptT String m a -> ExceptT String m a+annotate ann = fmapLT ((ann++": ")++)++runGetT :: Monad m => B.Get a -> LBS.ByteString -> ExceptT String m a+runGetT _get bs = do+ case B.runGetOrFail _get bs of+ Left (_, _, e) -> throwE e+ Right (_, _, a) -> return 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)+ -> ExceptT 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) $+ throwE "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
+ src/BTree/BinaryList.hs view
@@ -0,0 +1,90 @@+{-# 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 qualified Data.Binary.Builder as BB+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 BB.empty)+ where+ go :: Int -> Producer a m r -> BB.Builder+ -> Producer LBS.ByteString m (Header, (BinaryList a, r))+ go !n prod accum = do+ result <- lift $ next prod+ case result of+ Left r -> do+ let hdr = Header (fromIntegral n)+ yield $ BB.toLazyByteString accum+ return (hdr, (BinList fname, r))+ Right (a, prod')+ | n `mod` 100 == 0 -> do+ yield $ BB.toLazyByteString accum+ go (n+1) prod' (B.execPut $ B.put a)+ | otherwise ->+ go (n+1) prod' (accum `BB.append` B.execPut (B.put a))+{-# INLINE toBinaryList #-}++withHeader :: MonadIO m+ => BinaryList a -> (Header -> Handle -> m b) -> ExceptT String m b+withHeader (BinList fname) action = readWithHeader fname action++length :: MonadIO m => BinaryList a -> ExceptT 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 -> ExceptT 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'+{-# INLINE stream #-}
+ src/BTree/BuildUnordered.hs view
@@ -0,0 +1,113 @@+{-# 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+ -> ExceptT String m ()+fromUnorderedToFile scratch maxChunk order dest producer = {-# SCC fromUnorderedToFile #-} do+ bList <- lift (execStateT (fillLists producer) []) >>= {-# SCC goMerge #-} goMerge+ size <- BL.length bList+ stream <- {-# SCC stream #-} BL.stream bList+ lift $ {-# SCC buildTree #-} 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 = {-# SCC fillLists #-} do+ fname <- liftIO $ tempFilePath scratch "chunk.list"+ (leaves, rest) <- lift $ takeChunk maxChunk prod+ (bList, ()) <- lift $ BL.toBinaryList fname $ each $ S.toAscList leaves+ modify (bList:)+ case rest of+ Left r -> return r+ Right nextProd -> fillLists nextProd++ goMerge :: [BL.BinaryList (BLeaf k e)] -> ExceptT 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''+{-# INLINE fromUnorderedToFile #-}++-- | 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+{-# INLINE splitChunks #-}++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] -> ExceptT 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+{-# INLINE mergeLists #-}++-- | 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'+{-# INLINE takeChunk #-}
+ src/BTree/Builder.hs view
@@ -0,0 +1,237 @@+{-# 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 = {-# SCC putBS #-} 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'+{-# INLINE putBS #-}++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)+{-# INLINE next' #-}++-- | 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 = {-# SCC buildNodes #-}+ 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+{-# INLINE buildNodes #-}++-- | 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+{-# INLINE buildTree #-}++dropUpstream :: Monad m => Proxy X () () b m r -> Proxy X () b' b m r+dropUpstream = {-# SCC dropUpstream #-} go+ where+ go producer = do+ n <- lift $ next producer+ case n of+ Left r -> return r+ Right (a, producer') -> respond a >> go producer'+{-# INLINE dropUpstream #-}++-- | 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)+{-# INLINE fromOrderedToFile #-}++-- | 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)
+ src/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 = runExceptT $ do+ d <- fmapLT show $ tryIO $ mmapFileByteString fname Nothing+ ExceptT $ 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
+ src/BTree/Merge.hs view
@@ -0,0 +1,63 @@+{-# 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+{-# INLINE mergeLeaves #-}++-- | 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+{-# INLINE mergeTrees #-}++-- | 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)
+ src/BTree/Types.hs view
@@ -0,0 +1,105 @@+{-# 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+ {-# INLINE get #-}+ put (OnDisk off) = put off+ {-# INLINE put #-}++-- | 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+ {-# INLINE compare #-}++instance (Binary k, Binary e) => Binary (BLeaf k e) where+ get = BLeaf <$> get <*> get+ {-# INLINE get #-}+ put (BLeaf k e) = put k >> put e+ {-# INLINE put #-}++-- | '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)+ {-# INLINE get #-}++ put (Node e0 es) = putWord8 0 >> put e0 >> put es+ put (Leaf (BLeaf k0 e)) = putWord8 1 >> put k0 >> put e+ {-# INLINE put #-}++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
+ src/BTree/Walk.hs view
@@ -0,0 +1,53 @@+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+{-# INLINE 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+{-# INLINE walkLeaves #-}++-- | 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+{-# INLINE walkNodes #-}++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+{-# INLINE walkNodesWithOffset #-}