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