diff --git a/b-tree.cabal b/b-tree.cabal --- a/b-tree.cabal +++ b/b-tree.cabal @@ -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 diff --git a/src/BTree.hs b/src/BTree.hs --- a/src/BTree.hs +++ b/src/BTree.hs @@ -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. diff --git a/src/BTree/BinaryFile.hs b/src/BTree/BinaryFile.hs --- a/src/BTree/BinaryFile.hs +++ b/src/BTree/BinaryFile.hs @@ -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 diff --git a/src/BTree/BinaryList.hs b/src/BTree/BinaryList.hs --- a/src/BTree/BinaryList.hs +++ b/src/BTree/BinaryList.hs @@ -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 diff --git a/src/BTree/BuildUnordered.hs b/src/BTree/BuildUnordered.hs --- a/src/BTree/BuildUnordered.hs +++ b/src/BTree/BuildUnordered.hs @@ -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 #-} diff --git a/src/BTree/Builder.hs b/src/BTree/Builder.hs --- a/src/BTree/Builder.hs +++ b/src/BTree/Builder.hs @@ -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@ diff --git a/src/BTree/Lookup.hs b/src/BTree/Lookup.hs --- a/src/BTree/Lookup.hs +++ b/src/BTree/Lookup.hs @@ -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 diff --git a/src/BTree/Merge.hs b/src/BTree/Merge.hs --- a/src/BTree/Merge.hs +++ b/src/BTree/Merge.hs @@ -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 diff --git a/src/BTree/Types.hs b/src/BTree/Types.hs --- a/src/BTree/Types.hs +++ b/src/BTree/Types.hs @@ -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 diff --git a/src/BTree/Walk.hs b/src/BTree/Walk.hs --- a/src/BTree/Walk.hs +++ b/src/BTree/Walk.hs @@ -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 #-} diff --git a/tests/QuickCheck.hs b/tests/QuickCheck.hs --- a/tests/QuickCheck.hs +++ b/tests/QuickCheck.hs @@ -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) + ] + ]