b-tree 0.1.2 → 0.1.3
raw patch · 11 files changed
+258/−148 lines, 11 filesdep +exceptionsdep +tastydep +tasty-quickcheckdep ~QuickCheckdep ~basedep ~binary
Dependencies added: exceptions, tasty, tasty-quickcheck
Dependency ranges changed: QuickCheck, base, binary, containers, directory, errors, lens, pipes, pipes-interleave, transformers, vector
Files
- b-tree.cabal +18/−15
- src/BTree.hs +3/−3
- src/BTree/BinaryFile.hs +36/−25
- src/BTree/BinaryList.hs +14/−5
- src/BTree/BuildUnordered.hs +25/−9
- src/BTree/Builder.hs +47/−44
- src/BTree/Lookup.hs +9/−1
- src/BTree/Merge.hs +5/−8
- src/BTree/Types.hs +50/−15
- src/BTree/Walk.hs +14/−12
- tests/QuickCheck.hs +37/−11
b-tree.cabal view
@@ -1,5 +1,5 @@ name: b-tree-version: 0.1.2+version: 0.1.3 synopsis: Immutable disk-based B* trees description: Immutable disk-based B* trees homepage: http://github.com/bgamari/b-tree@@ -17,14 +17,14 @@ location: git@github.com:bgamari/b-tree library- exposed-modules: BTree+ exposed-modules: BTree,+ BTree.BinaryList other-modules: BTree.Types, BTree.Merge, BTree.Builder, BTree.Lookup, BTree.Walk, BTree.BinaryFile,- BTree.BinaryList, BTree.BuildUnordered other-extensions: DeriveGeneric,@@ -36,19 +36,20 @@ GeneralizedNewtypeDeriving hs-source-dirs: src- build-depends: base >=4.6 && <4.9,+ build-depends: base >=4.6 && <4.11, mtl >=2.0 && <3.0,- pipes >=4.1 && <4.2,- pipes-interleave >= 0.1 && <1.0,+ pipes >=4.1 && <4.4,+ pipes-interleave >= 1.0 && <2.2, bytestring >=0.10 && <0.11,- binary >=0.7 && <0.8,- transformers >=0.3 && <0.5,- lens >=3.10 && <4.13,+ binary >=0.8.4 && <0.10,+ transformers >=0.3 && <0.6,+ lens >=3.10 && <4.16, containers >=0.5 && <0.6,- vector >=0.10 && <0.12,- errors >=2.0 && <2.1,+ vector >=0.10 && <0.13,+ errors >=2.0 && <2.3,+ exceptions >= 0.8 && < 0.9, filepath >=1.3 && <1.5,- directory >=1.2 && <1.3,+ directory >=1.2 && <1.4, mmap >=0.5 && <0.6 default-language: Haskell2010 ghc-options: -Wall@@ -58,19 +59,21 @@ main-is: QuickCheck.hs hs-source-dirs: tests default-language: Haskell2010- build-depends: base >=4.6 && <4.9,+ build-depends: base, containers, pipes, binary, b-tree,- QuickCheck+ QuickCheck >= 2.8.2,+ tasty,+ tasty-quickcheck benchmark btree-benchmark type: exitcode-stdio-1.0 main-is: Benchmark.hs hs-source-dirs: benchmarks default-language: Haskell2010- build-depends: base >=4.6 && <4.9,+ build-depends: base, b-tree, pipes, criterion >=0.8 && <2.0
src/BTree.hs view
@@ -1,3 +1,5 @@+-- | This package provides immutable B* trees targetting large data+-- sets requiring secondary storage. module BTree ( -- * Basic types BLeaf(..) , Size@@ -11,6 +13,7 @@ , open , fromByteString , lookup+ , size -- * Merging trees , mergeTrees , mergeLeaves@@ -26,6 +29,3 @@ 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
@@ -1,3 +1,5 @@+-- | This module provides helpers for emitting and reading binary files with+-- a trailing "header". module BTree.BinaryFile ( writeWithHeader , readWithHeader@@ -6,9 +8,11 @@ import Control.Monad (when) import Control.Error import Control.Monad.Trans.Class+import Control.Monad.Catch import Control.Applicative import Data.Word import System.IO+import Prelude import qualified Data.ByteString.Lazy as LBS import qualified Data.Binary as B@@ -16,9 +20,6 @@ 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@@ -36,21 +37,29 @@ 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+-- | Write the produced 'LBS.ByteString's to the file followed by the -- returned header-writeWithHeader :: (MonadIO m, B.Binary hdr)+writeWithHeader :: (MonadMask m, MonadIO m, B.Binary hdr) => FilePath -> Producer LBS.ByteString m (hdr, r) -> m r-writeWithHeader fname prod = do- h <- liftIO $ openFile fname WriteMode+writeWithHeader fname prod =+ bracket (liftIO $ openFile fname WriteMode) (liftIO . hClose)+ $ \hdl -> hWriteWithHeader hdl prod++-- | Write the produced 'LBS.ByteString's to the file followed by the+-- returned header+hWriteWithHeader :: (MonadIO m, B.Binary hdr)+ => Handle+ -> Producer LBS.ByteString m (hdr, r)+ -> m r+hWriteWithHeader h prod = do (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 #-} @@ -58,7 +67,7 @@ annotate ann = fmapLT ((ann++": ")++) runGetT :: Monad m => B.Get a -> LBS.ByteString -> ExceptT String m a-runGetT _get bs = do+runGetT _get bs = case B.runGetOrFail _get bs of Left (_, _, e) -> throwE e Right (_, _, a) -> return a@@ -66,23 +75,25 @@ -- | 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)+readWithHeader :: (MonadMask m, 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+ r <- lift $ bracket (liftIO $ openFile fname ReadMode) (liftIO . hClose) $ \h -> runExceptT $ do+ -- 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++ ExceptT $ return r
src/BTree/BinaryList.hs view
@@ -6,18 +6,20 @@ -- * Construction , toBinaryList -- * Fetching contents+ , open , stream -- * Other queries , length , filePath ) where -import Prelude hiding (length) import Control.Applicative import Control.Monad.Trans.Class+import Control.Monad.Catch import Control.Error import Data.Word import System.IO+import Prelude hiding (length) import qualified Data.ByteString.Lazy as LBS import qualified Data.Binary as B@@ -44,7 +46,7 @@ put (Header l) = B.putWord64le l -- | Encode the items of the given producer-toBinaryList :: forall m a r. (MonadIO m, B.Binary a)+toBinaryList :: forall m a r. (MonadMask m, 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)@@ -66,15 +68,22 @@ go (n+1) prod' (accum `BB.append` B.execPut (B.put a)) {-# INLINE toBinaryList #-} -withHeader :: MonadIO m+-- | Open a 'BinaryList' file.+--+-- TODO: Sanity checking at open time.+open :: FilePath -> BinaryList a+open = BinList++withHeader :: (MonadMask m, 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 :: (MonadMask m, 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)+stream :: forall m a. (B.Binary a, MonadMask m, MonadIO m) => BinaryList a -> ExceptT String m (Producer a m (Either String ())) stream bl = withHeader bl readContents where
src/BTree/BuildUnordered.hs view
@@ -3,9 +3,12 @@ {-# LANGUAGE FlexibleContexts #-} module BTree.BuildUnordered- ( fromUnorderedToFile ) where+ ( fromUnorderedToFile+ , fromUnorderedToList+ ) where import Control.Monad.Trans.State+import Control.Monad.Catch import Control.Error import Data.Traversable (forM) @@ -38,7 +41,8 @@ -- 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)+ (MonadMask m, 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@@ -46,13 +50,23 @@ -> 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+ bList <- fromUnorderedToList scratch maxChunk producer size <- BL.length bList stream <- {-# SCC stream #-} BL.stream bList lift $ {-# SCC buildTree #-} fromOrderedToFile order size dest stream liftIO $ removeFile $ BL.filePath bList+{-# INLINE fromUnorderedToFile #-}++fromUnorderedToList :: forall m a r.+ (MonadMask m, MonadIO m, B.Binary a, Ord a)+ => FilePath -- ^ Path to scratch directory+ -> Int -- ^ Maximum chunk size+ -> Producer a m r -- ^ 'Producer' of elements+ -> ExceptT String m (BL.BinaryList a)+fromUnorderedToList scratch maxChunk producer = do+ lift (execStateT (fillLists producer) []) >>= {-# SCC goMerge #-} goMerge where- fillLists :: Producer (BLeaf k e) m r -> StateT [BL.BinaryList (BLeaf k e)] m r+ fillLists :: Producer a m r -> StateT [BL.BinaryList a] m r fillLists prod = {-# SCC fillLists #-} do fname <- liftIO $ tempFilePath scratch "chunk.list" (leaves, rest) <- lift $ takeChunk maxChunk prod@@ -62,7 +76,7 @@ Left r -> return r Right nextProd -> fillLists nextProd - goMerge :: [BL.BinaryList (BLeaf k e)] -> ExceptT String m (BL.BinaryList (BLeaf k e))+ goMerge :: [BL.BinaryList a] -> ExceptT String m (BL.BinaryList a) goMerge [l] = return l goMerge ls = do ls'' <- forM (splitChunks maxChunkMerge ls) $ \ls'->do@@ -71,7 +85,7 @@ liftIO $ mapM_ (removeFile . BL.filePath) ls' return list goMerge ls''-{-# INLINE fromUnorderedToFile #-}+{-# INLINE fromUnorderedToList #-} -- | Split the list into chunks of bounded size and run each through a function splitChunks :: Int -> [a] -> [[a]]@@ -85,11 +99,13 @@ 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 :: (B.Binary a, Ord a, MonadMask m, 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)+ let prod = interleave (map throwLeft streams) (bList, ()) <- lift $ BL.toBinaryList dest prod return bList {-# INLINE mergeLists #-}
src/BTree/Builder.hs view
@@ -1,4 +1,5 @@ {-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE ScopedTypeVariables #-} module BTree.Builder ( buildNodes, putBS@@ -8,12 +9,14 @@ import Control.Monad.Trans.State.Strict import Control.Monad.IO.Class+import Control.Monad.Catch import Control.Monad import Data.Foldable as F import qualified Data.Sequence as Seq import Data.Sequence (Seq) +import Data.Word import Data.Ratio import Control.Lens import System.IO@@ -32,7 +35,7 @@ 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+putBS a0 = {-# SCC "putBS" #-} evalStateT (go a0) 0 where go a = do s <- get@@ -65,38 +68,38 @@ -- | 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)+optimalFill order size = go size where- go :: Int -> [[Int]]+ go :: Word64 -> [[Int]] go 0 = error "BTree.Builder.optimalFill: zero size" go n = let nNodes = ceiling (n % order')- order' = fromIntegral order :: Int+ order' = fromIntegral order :: Word64 nodes = let (nPerNode, leftover) = n `divMod` nNodes- in zipWith (+) (replicate nNodes nPerNode)- (replicate leftover 1 ++ repeat 0)+ in zipWith (+) (replicate (fromIntegral nNodes) (fromIntegral nPerNode))+ (replicate (fromIntegral leftover) 1 ++ repeat 0) rest = case nNodes of 1 -> [] _ -> go nNodes in nodes : rest +type BuildM k e m a = StateT [DepthState k e] (DiskProducer (BTree k OnDisk e) m) a+ -- | 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+buildNodes :: forall m k e r. Monad m => Order -> Size -> DiskProducer (BLeaf k e) m r- -> DiskProducer (BTree k OnDisk e) m (BTreeHeader k e)-buildNodes order size = {-# SCC buildNodes #-}+ -> DiskProducer (BTree k OnDisk e) m (Size, Maybe (OnDisk (BTree k OnDisk 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 :: Size -> DiskProducer (BLeaf k e) m r+ -> BuildM k e m (Size, Maybe (OnDisk (BTree k OnDisk e))) loop n producer = do _next <- lift $ lift $ next' producer case _next of@@ -109,16 +112,13 @@ OnDisk offset <- processNode k $ Leaf leaf loop (n-1) $ producer' (OnDisk offset) - isFilled :: Monad m- => StateT [DepthState k e] m Bool+ isFilled :: BuildM 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 :: BuildM k e m (OnDisk (BTree k OnDisk e)) emitNode = do (k0,node0):nodes <- zoom (singular _head) $ do nodes <- uses dNodes F.toList@@ -127,19 +127,18 @@ dMinFill %= tail return nodes - --when (null nodes)- -- $ error "BTree.Builder.buildNodes: Internal invariant broken: unexpected empty node"+ -- We used to check the invariants of (not $ null nodes), however this+ -- is wrong. nodes may be empty, for instance, when we are call from+ -- the [_] branch of flushAll.+ 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 :: k -> BTree k OnDisk e+ -> BuildM k e m (OnDisk (BTree k OnDisk e)) processNode startKey tree = do filled <- isFilled when filled $ void $ emitNode@@ -149,18 +148,16 @@ dNodeCount += 1 return offset - flushAll :: Monad m- => Size- -> StateT [DepthState k e]- (DiskProducer (BTree k OnDisk e) m)- (BTreeHeader k e)+ flushAll :: Size+ -> BuildM k e m (Size, Maybe (OnDisk (BTree k OnDisk e)))+ flushAll 0 = return (0, Nothing) 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+ return (realSize, Just root) d:_ -> do when (not $ Seq.null $ d^.dNodes) $ void $ emitNode zoom (singular _tail) $ flushAll realSize {-# INLINE buildNodes #-}@@ -171,12 +168,20 @@ => 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+buildTree order size producer+ | size < 0 = error "BTree.buildTree: Invalid tree size"+ | size == 0 = return zeroSizedHeader+ | otherwise = do+ (realSize, root) <- dropUpstream $ buildNodes order size (dropUpstream producer) >>~ putBS+ if realSize == 0+ then return zeroSizedHeader+ else return $ BTreeHeader magic 1 order realSize root+ where+ zeroSizedHeader = BTreeHeader magic 1 order 0 Nothing {-# INLINE buildTree #-} dropUpstream :: Monad m => Proxy X () () b m r -> Proxy X () b' b m r-dropUpstream = {-# SCC dropUpstream #-} go+dropUpstream = {-# SCC "dropUpstream" #-} go where go producer = do n <- lift $ next producer@@ -189,22 +194,20 @@ -- -- As the name suggests, this requires that the @Producer@ emits -- leaves in ascending key order.-fromOrderedToFile :: (MonadIO m, Binary e, Binary k)+fromOrderedToFile :: (MonadMask m, 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 ()+fromOrderedToFile order size fname producer =+ bracket (liftIO $ openFile fname WriteMode) (liftIO . hClose) $ \h -> do+ 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 where- invalidHeader = BTreeHeader 0 0 0 0 (OnDisk 0)+ invalidHeader = BTreeHeader 0 0 0 0 Nothing {-# INLINE fromOrderedToFile #-} -- | Build a B-tree into @ByteString@
src/BTree/Lookup.hs view
@@ -2,6 +2,7 @@ , open , fromByteString , lookup+ , size ) where import Prelude hiding (lookup)@@ -33,7 +34,10 @@ -- | 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)+lookup lt k =+ case lt ^. ltHeader . btRoot of+ Just root -> go $ fetch lt root+ Nothing -> Nothing where go (Leaf (BLeaf k' e)) | k' == k = Just e@@ -45,3 +49,7 @@ case takeWhile (\(k',_)->k' <= k) children of [] -> Nothing xs -> go $ fetch lt $ snd $ last xs++-- | How many keys are in a 'LookupTree'.+size :: LookupTree k e -> Size+size = _btSize . _ltHeader
src/BTree/Merge.hs view
@@ -1,19 +1,17 @@-{-# 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 Control.Monad.Catch import Data.Binary import Control.Lens import Pipes import Pipes.Interleave+import Prelude hiding (sum) import BTree.Types import BTree.Builder@@ -24,7 +22,7 @@ -- 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)+mergeLeaves :: (MonadMask m, 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@@ -33,17 +31,16 @@ mergeLeaves append destOrder destFile producers = do let size = sum $ map fst producers fromOrderedToFile destOrder size destFile $- mergeM (compare `on` key) doAppend (map snd producers)+ mergeM 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)+mergeTrees :: (MonadMask m, 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
src/BTree/Types.hs view
@@ -1,13 +1,24 @@-{-# LANGUAGE DeriveGeneric, FlexibleContexts, TemplateHaskell, UndecidableInstances, StandaloneDeriving #-}+{-# LANGUAGE DeriveGeneric #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE TemplateHaskell #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE StandaloneDeriving #-}+{-# LANGUAGE RecordWildCards #-} module BTree.Types where -import Data.Binary-import GHC.Generics-import Control.Monad (when) import Control.Applicative-import Control.Lens+import Data.Maybe (fromMaybe)+import GHC.Generics+import Control.Monad (when, replicateM) import Data.Int+import Prelude++import Data.Binary+import Data.Binary.Get+import Data.Binary.Put+import Control.Lens import qualified Data.ByteString as BS -- | An offset within the stream@@ -19,7 +30,7 @@ -- | 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.+-- | @'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@@ -33,7 +44,7 @@ -- | A tree leaf (e.g. key/value pair) data BLeaf k e = BLeaf !k !e- deriving (Generic)+ deriving (Generic, Functor) deriving instance (Show k, Show e) => Show (BLeaf k e) @@ -52,10 +63,10 @@ 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'+-- | @'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+-- 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)@@ -68,14 +79,22 @@ => Binary (BTree k f e) where get = do typ <- getWord8 case typ of- 0 -> Node <$> get <*> get+ 0 -> Node <$> get <*> getChildren 1 -> bleaf <$> get <*> get _ -> fail "BTree.Types/get: Unknown node type" where bleaf k v = Leaf (BLeaf k v)+ getChildren = do+ len <- getWord32be+ replicateM (fromIntegral len) $ (,) <$> get <*> get {-# INLINE get #-} - put (Node e0 es) = putWord8 0 >> put e0 >> put es- put (Leaf (BLeaf k0 e)) = putWord8 1 >> put k0 >> put e+ -- some versions of binary don't inline the Binary (,) instance, pitiful+ -- performance ensues+ put (Node e0 es) = do putWord8 0+ put e0+ putWord32be (fromIntegral $ length es)+ mapM_ (\(a,b) -> put a >> put b) es+ put (Leaf (BLeaf k0 e)) = putWord8 1 >> put k0 >> put e {-# INLINE put #-} magic :: Word64@@ -86,12 +105,28 @@ , _btVersion :: !Word64 , _btOrder :: !Order , _btSize :: !Size- , _btRoot :: !(OnDisk (BTree k OnDisk e))+ , _btRoot :: !(Maybe (OnDisk (BTree k OnDisk e)))+ -- ^ 'Nothing' represents an empty tree } deriving (Show, Eq, Generic) makeLenses ''BTreeHeader -instance Binary (BTreeHeader k e)+-- | It is critical that this encoding is of fixed size+instance Binary (BTreeHeader k e) where+ get = do+ _btMagic <- get+ _btVersion <- get+ _btOrder <- get+ _btSize <- get+ root <- get+ let _btRoot = if root == OnDisk 0 then Nothing else Just root+ return BTreeHeader {..}+ put (BTreeHeader {..}) = do+ put _btMagic+ put _btVersion+ put _btOrder+ put _btSize+ put $ fromMaybe (OnDisk 0) _btRoot validateHeader :: BTreeHeader k e -> Either String () validateHeader hdr = do
src/BTree/Walk.hs view
@@ -1,3 +1,5 @@+{-# LANGUAGE BangPatterns #-}+ module BTree.Walk ( walkLeaves , walkNodes , walkNodesWithOffset@@ -5,6 +7,7 @@ import BTree.Types import qualified Data.ByteString.Lazy as LBS+import qualified Data.ByteString as BS import Pipes import qualified Pipes.Prelude as PP import Data.Binary@@ -15,12 +18,10 @@ -- 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+filterLeaves = PP.mapFoldable getLeaf+ where+ getLeaf (Leaf leaf) = Just leaf+ getLeaf _ = Nothing {-# INLINE filterLeaves #-} -- | Iterate over the leaves of the given tree in ascending key order.@@ -41,13 +42,14 @@ 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+walkNodesWithOffset = go 0 . {-# SCC "buffer" #-}view ltData+ where go !offset bs =+ case runGetOrFail get (LBS.fromStrict bs) of Left (rest,_,err) -> return (rest, Just err)- Right (rest,o,a) -> do+ Right (_,o,a) -> do yield (offset, a)- if LBS.null rest- then return (rest, Nothing)+ let rest = BS.drop (fromIntegral o) bs+ if BS.null rest+ then return (LBS.fromStrict rest, Nothing) else go (offset+o) rest {-# INLINE walkNodesWithOffset #-}
tests/QuickCheck.hs view
@@ -1,25 +1,51 @@ import Test.QuickCheck+import Test.Tasty+import Test.Tasty.QuickCheck import qualified Data.Map as M-import Data.Foldable (foldl')+import Data.Foldable (foldl', foldMap) import Data.Binary (Binary) import Pipes hiding (discard)+import qualified Pipes.Prelude as PP import qualified BTree as BT -instance (Ord k, Arbitrary k, Arbitrary v) => Arbitrary (M.Map k v) where- arbitrary = do- ks <- arbitrary- vs <- arbitrary- return $ foldl' (\m (k,v)->M.insert k v m) M.empty (zip ks vs)- mapToBLeafs :: M.Map k v -> [BT.BLeaf k v] mapToBLeafs = map (\(k,v)->BT.BLeaf k v) . M.toAscList -test :: (Ord k, Eq v, Binary k, Binary v) => M.Map k v -> Property-test m | M.size m == 0 = discard-test m = ioProperty $ do+mapFromBLeafs :: Ord k => [BT.BLeaf k v] -> M.Map k v+mapFromBLeafs = foldMap (\(BT.BLeaf k v) -> M.singleton k v)++-- | Test lookup correctness of exact-sized input+testExact :: (Ord k, Eq v, Binary k, Binary v)+ => M.Map k v -> Property+testExact m = ioProperty $ do let len = fromIntegral $ M.size m bs <- BT.fromOrderedToByteString 10 len (each $ mapToBLeafs m) let Right bt = BT.fromByteString bs return $ all (\(k,v)->BT.lookup bt k == Just v) (M.assocs m) -main = quickCheck (test :: M.Map Int Int -> Property)+-- | Test lookup correctness of inexact-sized input+test :: (Ord k, Eq v, Binary k, Binary v)+ => BT.Size -> M.Map k v -> Property+test size m = ioProperty $ do+ bs <- BT.fromOrderedToByteString 10 size (each $ mapToBLeafs m)+ let Right bt = BT.fromByteString bs+ return $ all (\(k,v)->BT.lookup bt k == Just v) (take (fromIntegral size) $ M.toAscList m)++testWalk :: (Ord k, Eq v, Binary k, Binary v)+ => Positive Int -> M.Map k v -> Property+testWalk (Positive size) m = ioProperty $ do+ bs <- BT.fromOrderedToByteString 10 (fromIntegral size) (each $ mapToBLeafs m)+ let Right bt = BT.fromByteString bs+ xs = map (\(BT.BLeaf k v) -> (k, v)) $ PP.toList $ void $ BT.walkLeaves bt+ return $ xs == take size (M.toAscList m)++main :: IO ()+main = defaultMain $ testGroup "tests"+ [ testGroup "lookup"+ [ testProperty "inexact size" (test :: BT.Size -> M.Map Int Int -> Property)+ , testProperty "exact size" (testExact :: M.Map Int Int -> Property)+ ]+ , testGroup "walk"+ [ testProperty "walk" (testWalk :: Positive Int -> M.Map Int Int -> Property)+ ]+ ]