diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,32 @@ +Copyright (c) 2017, Henri Verroken, Steven Keuchel + +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 Henri Verroken or Steven Keuchel 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/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,18 @@ +haskey-btree +============ + +[![Travis](https://travis-ci.org/haskell-haskey/haskey-btree.svg?branch=master)](https://travis-ci.org/haskell-haskey/haskey-btree) +[![Coverage](https://coveralls.io/repos/github/haskell-haskey/haskey-btree/badge.svg?branch=master)](https://coveralls.io/github/haskell-haskey/haskey-btree?branch=master) +[![Hackage](https://img.shields.io/hackage/v/haskey-btree.svg?maxAge=2592000)](https://hackage.haskell.org/package/haskey-btree) +[![Stackage Nightly](http://stackage.org/package/haskey-btree/badge/nightly)](http://stackage.org/nightly/package/haskey-btree) +[![Stackage LTS](http://stackage.org/package/haskey-btree/badge/lts)](http://stackage.org/lts/package/haskey-btree) + +B+-tree implementation in Haskell. + +This package provides two B+-tree implementations. The first one is a pure +B+-tree of a specific order, while the second one is an impure one backed +by a page allocator. + +This project is part of the [haskey](https://github.com/haskell-haskey/haskey) +project. The [haskey](https://github.com/haskell-haskey/haskey) repository +contains more information on how to use this library. 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/haskey-btree.cabal b/haskey-btree.cabal new file mode 100644 --- /dev/null +++ b/haskey-btree.cabal @@ -0,0 +1,141 @@ +name: haskey-btree +version: 0.1.0.0 +synopsis: B+-tree implementation in Haskell. +description: + This package provides two B+-tree implementations. The first one is a pure + B+-tree of a specific order, while the second one is an impure one backed + by a page allocator. + . + For more information on how to use this package, visit + +homepage: https://github.com/haskell-haskey/haskey-btree +license: BSD3 +license-file: LICENSE +author: Henri Verroken, Steven Keuchel +maintainer: steven.keuchel@gmail.com +copyright: Copyright (c) 2017, Henri Verroken, Steven Keuchel +category: Database +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + default-language: Haskell2010 + ghc-options: -Wall + exposed-modules: + Data.BTree.Alloc + Data.BTree.Alloc.Class + Data.BTree.Alloc.Debug + Data.BTree.Impure + Data.BTree.Impure.Delete + Data.BTree.Impure.Fold + Data.BTree.Impure.Insert + Data.BTree.Impure.Lookup + Data.BTree.Impure.NonEmpty + Data.BTree.Impure.Overflow + Data.BTree.Impure.Setup + Data.BTree.Impure.Structures + Data.BTree.Primitives + Data.BTree.Primitives.Exception + Data.BTree.Primitives.Height + Data.BTree.Primitives.Ids + Data.BTree.Primitives.Index + Data.BTree.Primitives.Key + Data.BTree.Primitives.Leaf + Data.BTree.Primitives.Value + Data.BTree.Pure + Data.BTree.Pure.Setup + + other-modules: + Data.BTree.Utils.List + Data.BTree.Utils.Map + Data.BTree.Utils.Vector + + other-extensions: + DataKinds + DeriveFoldable + DeriveFunctor + DeriveTraversable + GADTs + KindSignatures + MultiWayIf + ScopedTypeVariables + StandaloneDeriving + + build-depends: + base >=4.7 && <5, + binary >=0.6 && <0.9 || >0.9 && <1, + bytestring >=0.10 && <1, + containers >=0.5 && <1, + focus >=0.1.2 && <0.2, + hashable >=1.2 && <1.3, + list-t >=0.2 && <2, + mtl >=2.1 && <3, + semigroups >=0.12 && <1, + stm >=2.1 && <3, + transformers >=0.3 && <1, + vector >=0.10 && <1 + + +test-suite haskey-btree-properties + main-is: Properties.hs + type: exitcode-stdio-1.0 + other-modules: + Properties.Impure.Fold + Properties.Impure.Insert + Properties.Impure.Structures + Properties.Primitives.Height + Properties.Primitives.Ids + Properties.Primitives.Index + Properties.Primitives.Leaf + Properties.Pure + Properties.Utils + + build-depends: + base >=4.7 && <5, + binary >=0.6 && <0.9 || >0.9 && <1, + bytestring >=0.10 && <1, + containers >=0.5 && <1, + data-ordlist >=0.4 && <1, + mtl >=2.1 && <3, + transformers >=0.3 && <1, + vector >=0.10 && <1, + + HUnit >=1.3 && <2, + QuickCheck >=2 && <3, + test-framework >=0.8 && <1, + test-framework-hunit >=0.3 && <1, + test-framework-quickcheck2 >=0.3 && <1, + haskey-btree + + default-language: Haskell2010 + ghc-options: -Wall + hs-source-dirs: tests + +test-suite haskey-btree-integration + main-is: Integration.hs + type: exitcode-stdio-1.0 + other-modules: + Integration.WriteOpenRead.Debug + Integration.WriteOpenRead.Transactions + + build-depends: + base >=4.7 && <5, + binary >=0.6 && <0.9 || >0.9 && <1, + containers >=0.5 && <1, + mtl >=2.1 && <3, + transformers >=0.3 && <1, + + QuickCheck >=2 && <3, + test-framework >=0.8 && <1, + test-framework-quickcheck2 >=0.3 && <1, + haskey-btree + + default-language: Haskell2010 + ghc-options: -Wall + hs-source-dirs: tests + +source-repository head + type: git + location: https://github.com/haskell-haskey/haskey-btree diff --git a/src/Data/BTree/Alloc.hs b/src/Data/BTree/Alloc.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Alloc.hs @@ -0,0 +1,6 @@ +-- | Page allocators that manage all physical pages. +module Data.BTree.Alloc ( + module Data.BTree.Alloc.Class +) where + +import Data.BTree.Alloc.Class diff --git a/src/Data/BTree/Alloc/Class.hs b/src/Data/BTree/Alloc/Class.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Alloc/Class.hs @@ -0,0 +1,67 @@ +-- | A page allocator manages all physical pages. +module Data.BTree.Alloc.Class ( + -- * Classes + AllocReaderM(..) +, AllocM(..) +) where + +import Prelude hiding (max, min, pred) + +import Control.Applicative (Applicative) + +import Data.Word (Word64) + +import Data.BTree.Impure.Structures +import Data.BTree.Primitives + +-------------------------------------------------------------------------------- + +-- | A page allocator that can read physical pages. +class (Applicative m, Monad m) => AllocReaderM m where + -- | Read a page and return the actual node. + readNode :: (Key key, Value val) + => Height height + -> NodeId height key val + -> m (Node height key val) + + -- | Read an overflow page. + readOverflow :: (Value val) + => OverflowId + -> m val + +-- | A page allocator that can write physical pages. +class AllocReaderM m => AllocM m where + -- | A function that calculates the hypothetical size of a node, if it were + -- to be written to a page (regardless of the maximum page size). + nodePageSize :: (Key key, Value val) + => m (Height height -> Node height key val -> PageSize) + + -- | The maximum page size the allocator can handle. + maxPageSize :: m PageSize + + -- | The maximum key size + maxKeySize :: m Word64 + + -- | The maximum value size + maxValueSize :: m Word64 + + -- | Allocate a new page for a node, and write the node to the page. + allocNode :: (Key key, Value val) + => Height height + -> Node height key val + -> m (NodeId height key val) + + -- | Free the page belonging to the node. + freeNode :: Height height + -> NodeId height key val + -> m () + + -- | Allocate a new overflow page, and write the value to the page. + allocOverflow :: (Value val) + => val + -> m OverflowId + + -- | Free an overflow page. + freeOverflow :: OverflowId -> m () + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Alloc/Debug.hs b/src/Data/BTree/Alloc/Debug.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Alloc/Debug.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +-- | An in memory allocator for debugging and testing purposes. +module Data.BTree.Alloc.Debug where + +import Control.Applicative (Applicative, (<$>)) +import Control.Monad.Identity +import Control.Monad.IO.Class +import Control.Monad.State + +import Data.Binary.Put (runPut) +import Data.Map (Map, (!)) +import Data.Word (Word32) +import qualified Data.ByteString.Lazy as BL +import qualified Data.Map as M + +import Unsafe.Coerce (unsafeCoerce) + +import Data.BTree.Alloc.Class +import Data.BTree.Impure +import Data.BTree.Impure.Structures +import Data.BTree.Primitives + +data SomeNode = forall h k v. SomeNode (Height h) (Node h k v) + +getSomeNode :: SomeNode -> Node h k v +getSomeNode (SomeNode _ n) = unsafeCoerce n + +data SomeVal = forall v. SomeVal v + +getSomeVal :: SomeVal -> v +getSomeVal (SomeVal v) = unsafeCoerce v + +data Pages = Pages { + pagesNodes :: Map PageId SomeNode + , pagesOverflow :: Map Word32 SomeVal + } + +emptyPages :: Pages +emptyPages = Pages { + pagesNodes = M.empty + , pagesOverflow = M.empty + } + +newtype DebugT m a = DebugT { runDebugT :: StateT Pages m a } + deriving (Functor, Applicative, Monad, MonadIO, MonadState Pages) + +runDebug :: Pages -> DebugT Identity a -> (a, Pages) +runDebug pages = runIdentity . flip runStateT pages . runDebugT + +evalDebug :: Pages -> DebugT Identity a -> a +evalDebug pages = fst . runDebug pages + +instance (Functor m, Monad m) => AllocReaderM (DebugT m) where + readNode _ nid = do + n <- gets (\pgs -> pagesNodes pgs ! nodeIdToPageId nid) + return $ getSomeNode n + + readOverflow (_, c) = do + v <- gets (\pgs -> pagesOverflow pgs ! c) + return $ getSomeVal v + +instance (Functor m, Monad m) => AllocM (DebugT m) where + nodePageSize = return $ \h -> case viewHeight h of + UZero -> fromIntegral . BL.length . runPut . putLeafNode + USucc _ -> fromIntegral . BL.length . runPut . putIndexNode + + maxPageSize = return 256 + maxKeySize = return 20 + maxValueSize = return 20 + + allocNode h n = do + pid <- fromIntegral <$> gets (M.size . pagesNodes) + let n' = SomeNode h n + modify $ \pgs -> pgs { pagesNodes = M.insert pid n' (pagesNodes pgs) } + return $ pageIdToNodeId pid + + freeNode _ _ = return () + + allocOverflow v = do + let v' = SomeVal v + c <- fromIntegral <$> gets (M.size . pagesOverflow) + modify $ \pgs -> pgs { pagesOverflow = M.insert c v' (pagesOverflow pgs) } + return (0, c) + + freeOverflow _ = return () diff --git a/src/Data/BTree/Impure.hs b/src/Data/BTree/Impure.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure.hs @@ -0,0 +1,61 @@ +-- | An impure B+-tree implementation. +-- +-- This module contains the implementation of a B+-tree that is backed by a +-- page allocator (see "Data.BTree.Alloc"). +module Data.BTree.Impure ( + -- * Structures + Tree(..) +, Node(..) + + -- * Construction +, empty +, fromList +, fromMap + + -- * Manipulation +, insertTree +, insertTreeMany +, deleteTree + + -- * Lookup +, lookupTree +, lookupMinTree + + -- * Folds +, foldr +, foldrM +, foldrWithKey +, foldrWithKeyM +, foldMap +, toList +) where + +import Prelude hiding (foldr, foldMap) + +import Data.Map (Map) +import qualified Data.Map as M + +import Data.BTree.Alloc.Class +import Data.BTree.Impure.Delete (deleteTree) +import Data.BTree.Impure.Structures (Tree(..), Node(..)) +import Data.BTree.Impure.Fold (foldr, foldrM, foldrWithKey, foldrWithKeyM, foldMap, toList) +import Data.BTree.Impure.Insert (insertTree, insertTreeMany) +import Data.BTree.Impure.Lookup (lookupTree, lookupMinTree) + +import Data.BTree.Primitives + +-- | Create an empty tree. +empty :: Tree k v +empty = Tree zeroHeight Nothing + +-- | Create a tree from a list. +fromList :: (AllocM m, Key k, Value v) + => [(k, v)] + -> m (Tree k v) +fromList = fromMap . M.fromList + +-- | Create a tree from a map. +fromMap :: (AllocM m, Key k, Value v) + => Map k v + -> m (Tree k v) +fromMap kvs = insertTreeMany kvs empty diff --git a/src/Data/BTree/Impure/Delete.hs b/src/Data/BTree/Impure/Delete.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Delete.hs @@ -0,0 +1,132 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | Algorithms related to deletion from an impure B+-tree. +module Data.BTree.Impure.Delete where + +import Data.Monoid +import Data.Traversable (traverse) +import qualified Data.Map as M + +import Data.BTree.Alloc.Class +import Data.BTree.Impure.Insert +import Data.BTree.Impure.Setup +import Data.BTree.Impure.Structures +import Data.BTree.Primitives.Exception +import Data.BTree.Primitives + +-------------------------------------------------------------------------------- + +-- | Check whether a node needs to be merged. +nodeNeedsMerge :: Node height key val -> Bool +nodeNeedsMerge (Idx children) = + indexNumKeys children < minIdxKeys +nodeNeedsMerge (Leaf items) = + M.size items < minLeafItems + +-- | Merge two nodes. +mergeNodes :: (AllocM m, Key key, Value val) + => Height height + -> Node height key val + -> key + -> Node height key val + -> m (Index key (Node height key val)) +mergeNodes _ (Leaf leftItems) _middleKey (Leaf rightItems) = + splitLeaf (leftItems <> rightItems) +mergeNodes h (Idx leftIdx) middleKey (Idx rightIdx) = + splitIndex h (mergeIndex leftIdx middleKey rightIdx) + +-------------------------------------------------------------------------------- + +deleteRec :: forall height key val m. (AllocM m, Key key, Value val) + => key + -> Height height + -> NodeId height key val + -> m (Node height key val) +deleteRec key = fetchAndGo + where + fetchAndGo :: forall hgt. Height hgt + -> NodeId hgt key val + -> m (Node hgt key val) + fetchAndGo hgt nid = do + node <- readNode hgt nid + freeNode hgt nid + recurse hgt node + + recurse :: forall hgt. Height hgt + -> Node hgt key val + -> m (Node hgt key val) + recurse hgt (Idx children) = do + let (ctx, childId) = valView key children + subHeight = decrHeight hgt + newChild <- fetchAndGo subHeight childId + let childNeedsMerge = nodeNeedsMerge newChild + if | childNeedsMerge, Just (rKey, rChildId, rCtx) <- rightView ctx -> do + rChild <- readNode subHeight rChildId + freeNode subHeight rChildId + newChildren <- mergeNodes subHeight newChild rKey rChild + newChildrenIds <- traverse (allocNode subHeight) newChildren + return (Idx (putIdx rCtx newChildrenIds)) + | childNeedsMerge, Just (lCtx, lChildId, lKey) <- leftView ctx -> do + lChild <- readNode subHeight lChildId + freeNode subHeight lChildId + newChildren <- mergeNodes subHeight lChild lKey newChild + newChildrenIds <- traverse (allocNode subHeight) newChildren + return (Idx (putIdx lCtx newChildrenIds)) + -- No left or right sibling? This is a constraint violation. Also + -- this couldn't be the root because it would've been shrunk + -- before. + | childNeedsMerge -> throw $ TreeAlgorithmError "deleteRec" + "constraint violation, found an index node with a single child" + | otherwise -> do + newChildId <- allocNode subHeight newChild + return (Idx (putVal ctx newChildId)) + recurse _hgt (Leaf items) = + case M.lookup key items of + Nothing -> return $ Leaf items + Just (RawValue _) -> return $ Leaf (M.delete key items) + Just (OverflowValue oid) -> do + freeOverflow oid + return $ Leaf (M.delete key items) + +-------------------------------------------------------------------------------- + +-- | Delete a node from the tree. +deleteTree :: (AllocM m, Key key, Value val) + => key + -> Tree key val + -> m (Tree key val) +deleteTree k tree + | Tree + { treeRootId = Nothing + } <- tree + = return tree + | Tree + { treeHeight = height + , treeRootId = Just rootId + } <- tree + = do + newRootNode <- deleteRec k height rootId + case newRootNode of + Idx index + | Just childNodeId <- fromSingletonIndex index -> + return $! Tree + { treeHeight = decrHeight height + , treeRootId = Just childNodeId + } + Leaf items + | M.null items -> + return $! Tree + { treeHeight = zeroHeight + , treeRootId = Nothing + } + _ -> do + newRootNodeId <- allocNode height newRootNode + return $! Tree + { treeHeight = height + , treeRootId = Just newRootNodeId + } + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Impure/Fold.hs b/src/Data/BTree/Impure/Fold.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Fold.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE GADTs #-} +-- | Algorithms related to folding over an impure B+-tree. +module Data.BTree.Impure.Fold where + +import Prelude hiding (foldr, foldl) + +import Data.Map (Map) +import Data.Monoid (Monoid, (<>), mempty) +import qualified Data.Map as M +import qualified Data.Foldable as F + +import Data.BTree.Alloc.Class +import Data.BTree.Impure.Overflow +import Data.BTree.Impure.Structures +import Data.BTree.Primitives + +-------------------------------------------------------------------------------- + +-- | Perform a right-associative fold over the tree. +foldr :: (AllocReaderM m, Key k, Value a) + => (a -> b -> b) -> b -> Tree k a -> m b +foldr f = foldrM (\a b -> return (f a b)) + +-- | Perform a right-associative fold over the tree key-value pairs. +foldrWithKey :: (AllocReaderM m, Key k, Value a) + => (k -> a -> b -> b) -> b -> Tree k a -> m b +foldrWithKey f = foldrWithKeyM (\k a b -> return (f k a b)) + +-- | Perform a monadic right-associative fold over the tree. +foldrM :: (AllocReaderM m, Key k, Value a) + => (a -> b -> m b) -> b -> Tree k a -> m b +foldrM f = foldrWithKeyM (const f) + +-- | Perform a monadic right-assiciative fold over the tree key-value pairs. +foldrWithKeyM :: (AllocReaderM m, Key k, Value a) + => (k -> a -> b -> m b) -> b -> Tree k a -> m b +foldrWithKeyM _ x (Tree _ Nothing) = return x +foldrWithKeyM f x (Tree h (Just nid)) = foldrIdWithKeyM f x h nid + +foldrIdWithKeyM :: (AllocReaderM m, Key k, Value a) + => (k -> a -> b -> m b) -> b -> Height h -> NodeId h k a -> m b +foldrIdWithKeyM f x h nid = readNode h nid >>= foldrNodeWithKeyM f x h + +foldrNodeWithKeyM :: (AllocReaderM m, Key k, Value a) + => (k -> a -> b -> m b) -> b -> Height h -> Node h k a -> m b +foldrNodeWithKeyM f x _ (Leaf items) = + fromLeafItems items >>= foldrLeafItemsWithKeyM f x +foldrNodeWithKeyM f x h (Idx idx) = + F.foldrM (\nid x' -> foldrIdWithKeyM f x' (decrHeight h) nid) x idx + +foldrLeafItemsWithKeyM :: (AllocReaderM m, Key k, Value a) + => (k -> a -> b -> m b) -> b -> Map k a -> m b +foldrLeafItemsWithKeyM f x items = M.foldlWithKey f' return items x + where f' m k a z = f k a z >>= m + +-------------------------------------------------------------------------------- + +-- | Map each value of the tree to a monoid, and combine the results. +foldMap :: (AllocReaderM m, Key k, Value a, Monoid c) + => (a -> c) -> Tree k a -> m c +foldMap f = foldr ((<>) . f) mempty + +-- | Convert an impure B+-tree to a list of key-value pairs. +toList :: (AllocReaderM m, Key k, Value a) + => Tree k a -> m [(k, a)] +toList = foldrWithKey (\k v xs -> (k, v):xs) [] + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Impure/Insert.hs b/src/Data/BTree/Impure/Insert.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Insert.hs @@ -0,0 +1,186 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | Algorithms related to inserting key-value pairs in an impure B+-tree. +module Data.BTree.Impure.Insert where + +import Data.Map (Map) +import Data.Traversable (traverse) +import qualified Data.Map as M + +import Data.BTree.Alloc.Class +import Data.BTree.Impure.Overflow +import Data.BTree.Impure.Structures +import Data.BTree.Primitives.Exception +import Data.BTree.Primitives + +-------------------------------------------------------------------------------- + +-- | Split an index node. +-- +-- This function is partial. It fails when the original index cannot be split, +-- because it does not contain enough elements (underflow). +splitIndex :: (AllocM m, Key key, Value val) => + Height ('S height) -> + Index key (NodeId height key val) -> + m (Index key (Node ('S height) key val)) +splitIndex h index = do + m <- maxPageSize + nodePageSize' <- nodePageSize + let binPred n = nodePageSize' h n <= m + case extendIndexPred binPred Idx index of + Just extIndex -> return extIndex + Nothing -> throw $ TreeAlgorithmError "splitIndex" + "splitting failed, underflow" + +-- | Split a leaf node. +-- +-- This function is partial. It fails when the original leaf cannot be split, +-- because it does not contain enough elements (underflow). +splitLeaf :: (AllocM m, Key key, Value val) => + LeafItems key val -> + m (Index key (Node 'Z key val)) +splitLeaf items = do + m <- maxPageSize + nodePageSize' <- nodePageSize + let binPred n = nodePageSize' zeroHeight n <= m + case splitLeafManyPred binPred Leaf items of + Just v -> return v + Nothing -> throw $ TreeAlgorithmError "splitLeaf" + "splitting failed, underflow" + +-------------------------------------------------------------------------------- + +insertRec :: forall m height key val. (AllocM m, Key key, Value val) + => key + -> val + -> Height height + -> NodeId height key val + -> m (Index key (NodeId height key val)) +insertRec k v = fetch + where + fetch :: forall hgt. + Height hgt + -> NodeId hgt key val + -> m (Index key (NodeId hgt key val)) + fetch hgt nid = do + node <- readNode hgt nid + freeNode hgt nid + case node of + Idx children -> do + let (ctx,childId) = valView k children + newChildIdx <- fetch (decrHeight hgt) childId + newChildren <- splitIndex hgt (putIdx ctx newChildIdx) + traverse (allocNode hgt) newChildren + Leaf items -> do + v' <- toLeafValue v + traverse (allocNode hgt) =<< splitLeaf (M.insert k v' items) + +insertRecMany :: forall m height key val. (AllocM m, Key key, Value val) + => Height height + -> Map key val + -> NodeId height key val + -> m (Index key (NodeId height key val)) +insertRecMany h kvs nid + | M.null kvs = return (singletonIndex nid) + | otherwise = do + n <- readNode h nid + freeNode h nid + case n of + Idx idx -> do + let dist = distribute kvs idx + newIndex <- dist `bindIndexM` uncurry (insertRecMany (decrHeight h)) + newChildren <- splitIndex h newIndex + traverse (allocNode h) newChildren + Leaf items -> do + kvs' <- toLeafItems kvs + traverse (allocNode h) =<< splitLeaf (M.union kvs' items) + +-------------------------------------------------------------------------------- + +-- | Insert a key-value pair in an impure B+-tree. +-- +-- You are responsible to make sure the key is smaller than 'maxKeySize', +-- otherwise a 'KeyTooLargeError' can (but not always will) be thrown. +insertTree :: (AllocM m, Key key, Value val) + => key + -> val + -> Tree key val + -> m (Tree key val) +insertTree key val tree + | Tree + { treeHeight = height + , treeRootId = Just rootId + } <- tree + = do + newRootIdx <- insertRec key val height rootId + case fromSingletonIndex newRootIdx of + Just newRootId -> + return $! Tree + { treeHeight = height + , treeRootId = Just newRootId + } + Nothing -> do + -- Root got split, so allocate a new root node. + let newHeight = incrHeight height + newRootId <- allocNode newHeight Idx + { idxChildren = newRootIdx } + return $! Tree + { treeHeight = newHeight + , treeRootId = Just newRootId + } + | Tree + { treeRootId = Nothing + } <- tree + = do -- Allocate new root node + leafItems' <- toLeafItems $ M.singleton key val + newRootId <- allocNode zeroHeight Leaf + { leafItems = leafItems' + } + return $! Tree + { treeHeight = zeroHeight + , treeRootId = Just newRootId + } + +-- | Bulk insert a bunch of key-value pairs in an impure B+-tree. +-- +-- You are responsible to make sure all keys is smaller than 'maxKeySize', +-- otherwise a 'KeyTooLargeError' can (but not always will) be thrown. +insertTreeMany :: (AllocM m, Key key, Value val) + => Map key val + -> Tree key val + -> m (Tree key val) +insertTreeMany kvs tree + | Tree + { treeHeight = height + , treeRootId = Just rootId + } <- tree + = do + newRootIdx <- insertRecMany height kvs rootId + fixUp height newRootIdx + | Tree { treeRootId = Nothing } <- tree + = do + kvs' <- toLeafItems kvs + idx <- traverse (allocNode zeroHeight) =<< splitLeaf kvs' + fixUp zeroHeight $! idx + +-- | Fix up the root node of a tree. +-- +-- Fix up the root node of a tree, where all other nodes are valid, but the +-- root node may contain more items than allowed. Do this by repeatedly +-- splitting up the root node. +fixUp :: (AllocM m, Key key, Value val) + => Height height + -> Index key (NodeId height key val) + -> m (Tree key val) +fixUp h idx = case fromSingletonIndex idx of + Just newRootNid -> + return $! Tree { treeHeight = h + , treeRootId = Just newRootNid } + Nothing -> do + let newHeight = incrHeight h + children <- splitIndex newHeight idx + childrenNids <- traverse (allocNode newHeight) children + fixUp newHeight $! childrenNids + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Impure/Lookup.hs b/src/Data/BTree/Impure/Lookup.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Lookup.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | Algorithms related to looking up key-value pairs in an impure B+-tree. +module Data.BTree.Impure.Lookup where + +import qualified Data.Map as M + +import Control.Applicative ((<$>)) + +import Data.BTree.Alloc.Class +import Data.BTree.Impure.Overflow +import Data.BTree.Impure.Structures +import Data.BTree.Primitives + +-------------------------------------------------------------------------------- + +lookupRec :: forall m height key val. (AllocReaderM m, Key key, Value val) + => key + -> Height height + -> NodeId height key val + -> m (Maybe val) +lookupRec k = fetchAndGo + where + fetchAndGo :: forall hgt. + Height hgt -> + NodeId hgt key val -> + m (Maybe val) + fetchAndGo hgt nid = + readNode hgt nid >>= go hgt + + go :: forall hgt. + Height hgt -> + Node hgt key val -> + m (Maybe val) + go hgt (Idx children) = do + let (_ctx,childId) = valView k children + fetchAndGo (decrHeight hgt) childId + go _hgt (Leaf items) = + case M.lookup k items of Nothing -> return Nothing + Just v -> Just <$> fromLeafValue v + +-- | Lookup a value in an impure B+-tree. +lookupTree :: forall m key val. (AllocReaderM m, Key key, Value val) + => key + -> Tree key val + -> m (Maybe val) +lookupTree k tree + | Tree + { treeHeight = height + , treeRootId = Just rootId + } <- tree + = lookupRec k height rootId + | Tree + { treeRootId = Nothing + } <- tree + = return Nothing + +-------------------------------------------------------------------------------- + +-- | The minimal key of the map, returns 'Nothing' if the map is empty. +lookupMinTree :: (AllocReaderM m, Key key, Value val) + => Tree key val + -> m (Maybe (key, val)) +lookupMinTree tree + | Tree { treeRootId = Nothing } <- tree = return Nothing + | Tree { treeHeight = height + , treeRootId = Just rootId } <- tree + = lookupMinRec height rootId + where + lookupMinRec :: (AllocReaderM m, Key key, Value val) + => Height height + -> NodeId height key val + -> m (Maybe (key, val)) + lookupMinRec h nid = readNode h nid >>= \case + Idx children -> let (_, childId) = valViewMin children in + lookupMinRec (decrHeight h) childId + Leaf items -> case lookupMin items of + Nothing -> return Nothing + Just (k, v) -> do + v' <- fromLeafValue v + return $ Just (k, v') + + lookupMin m | M.null m = Nothing + | otherwise = Just $! M.findMin m + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Impure/NonEmpty.hs b/src/Data/BTree/Impure/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/NonEmpty.hs @@ -0,0 +1,90 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE StandaloneDeriving #-} +-- | Non empty wrapper around the impure 'Tree'. +module Data.BTree.Impure.NonEmpty ( + -- * Structures + NonEmptyTree(..) +, Node(..) + + -- * Conversions +, fromTree +, toTree +, nonEmptyToList + + -- * Construction +, fromNonEmptyList + + -- * Inserting +, insertNonEmptyTree +, insertNonEmptyTreeMany +) where + +import Control.Applicative ((<$>), (<*>)) + +import Data.Binary +import Data.List.NonEmpty (NonEmpty((:|))) +import Data.Map (Map) +import Data.Maybe (fromJust) +import Data.Typeable (Typeable) +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M + +import Data.BTree.Alloc.Class +import Data.BTree.Impure (Tree(..), Node(..), insertTree, insertTreeMany, empty, toList) +import Data.BTree.Primitives + +-- | A non-empty variant of 'Tree'. +data NonEmptyTree key val where + NonEmptyTree :: { -- | A term-level witness for the type-level height index. + treeHeight :: Height height + , -- | An empty tree is represented by 'Nothing'. Otherwise it's + -- 'Just' a 'NodeId' pointer the root. + treeRootId :: NodeId height key val + } -> NonEmptyTree key val + deriving (Typeable) + +deriving instance (Show key, Show val) => Show (NonEmptyTree key val) + +instance (Value k, Value v) => Value (NonEmptyTree k v) where + +instance Binary (NonEmptyTree key val) where + put (NonEmptyTree h root) = put h >> put root + get = NonEmptyTree <$> get <*> get + +-- | Convert a 'Tree' into a 'NonEmptyTree'. +fromTree :: Tree key val -> Maybe (NonEmptyTree key val) +fromTree (Tree h n) = case n of + Nothing -> Nothing + Just root -> Just $ NonEmptyTree h root + +-- | Convert a 'NonEmptyTree' into a 'Tree'. +toTree :: NonEmptyTree key val -> Tree key val +toTree (NonEmptyTree h n) = Tree h (Just n) + +-- | Create a 'NonEmptyTree' from a 'NonEmpty' list. +fromNonEmptyList :: (AllocM m, Key k, Value v) + => NonEmpty (k, v) + -> m (NonEmptyTree k v) +fromNonEmptyList (x :| xs) = fromJust . fromTree <$> insertTreeMany (M.fromList (x:xs)) empty + +-- | Insert an item into a 'NonEmptyTree' +insertNonEmptyTree :: (AllocM m, Key k, Value v) + => k + -> v + -> NonEmptyTree k v + -> m (NonEmptyTree k v) +insertNonEmptyTree k v tree = fromJust . fromTree <$> insertTree k v (toTree tree) + +-- | Bulk insert a bunch of key-value pairs into a 'NonEmptyTree'. +insertNonEmptyTreeMany :: (AllocM m, Key k, Value v) + => Map k v + -> NonEmptyTree k v + -> m (NonEmptyTree k v) +insertNonEmptyTreeMany kvs tree = fromJust . fromTree <$> insertTreeMany kvs (toTree tree) + +-- | Convert a non-empty tree to a list of key-value pairs. +nonEmptyToList :: (AllocReaderM m, Key k, Value v) + => NonEmptyTree k v + -> m (NonEmpty (k, v)) +nonEmptyToList tree = NE.fromList <$> toList (toTree tree) diff --git a/src/Data/BTree/Impure/Overflow.hs b/src/Data/BTree/Impure/Overflow.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Overflow.hs @@ -0,0 +1,38 @@ +-- | Functions related to overflow pages. +module Data.BTree.Impure.Overflow where + +import Prelude hiding (max, mapM) + +import Control.Applicative ((<$>)) + +import Data.Binary (encode) +import Data.Map (Map) +import Data.Traversable (mapM) +import qualified Data.ByteString.Lazy as BL + +import Data.BTree.Alloc.Class +import Data.BTree.Impure.Structures +import Data.BTree.Primitives + +toLeafValue :: (AllocM m, Value v) + => v + -> m (LeafValue v) +toLeafValue v = do + max <- maxValueSize + if BL.length (encode v) <= fromIntegral max + then return $ RawValue v + else OverflowValue <$> allocOverflow v + +fromLeafValue :: (AllocReaderM m, Value v) + => LeafValue v + -> m v +fromLeafValue (RawValue v) = return v +fromLeafValue (OverflowValue oid) = readOverflow oid + + +toLeafItems :: (AllocM m, Value v) => Map k v -> m (LeafItems k v) +toLeafItems = mapM toLeafValue + + +fromLeafItems :: (AllocReaderM m, Value v) => LeafItems k v -> m (Map k v) +fromLeafItems = mapM fromLeafValue diff --git a/src/Data/BTree/Impure/Setup.hs b/src/Data/BTree/Impure/Setup.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Setup.hs @@ -0,0 +1,11 @@ +-- | Setup of an impure B+-tree +module Data.BTree.Impure.Setup where + +minFanout :: Int +minFanout = 2 + +minLeafItems :: Int +minLeafItems = minFanout + +minIdxKeys :: Int +minIdxKeys = minFanout - 1 diff --git a/src/Data/BTree/Impure/Structures.hs b/src/Data/BTree/Impure/Structures.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Impure/Structures.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +-- | Basic structures of an impure B+-tree. +module Data.BTree.Impure.Structures ( + -- * Structures + Tree(..) +, Node(..) +, LeafItems +, LeafValue(..) + + -- * Binary encoding +, putLeafNode +, getLeafNode +, putIndexNode +, getIndexNode + + -- * Casting +, castNode +, castNode' +, castValue +) where + +import Control.Applicative ((<$>), (<*>)) +import Control.Monad (replicateM) + +import Data.Binary (Binary(..), Put, Get) +import Data.Bits ((.|.), shiftL, shiftR) +import Data.Map (Map) +import Data.Proxy (Proxy(..)) +import Data.Typeable (Typeable, typeRep, cast) +import Data.Word (Word8, Word32) +import qualified Data.Map as M + +import Numeric (showHex) + +import Unsafe.Coerce + +import Data.BTree.Primitives + +-------------------------------------------------------------------------------- + +-- | A B+-tree. +-- +-- This is a simple wrapper around a root 'Node'. The type-level height is +-- existentially quantified, but a term-level witness is stores. +data Tree key val where + Tree :: { -- | A term-level witness for the type-level height index. + treeHeight :: Height height + , -- | An empty tree is represented by 'Nothing'. Otherwise it's + -- 'Just' a 'NodeId' pointer the root. + treeRootId :: Maybe (NodeId height key val) + } -> Tree key val + deriving (Typeable) + +data LeafValue v = RawValue v | OverflowValue OverflowId + deriving (Eq, Show) + +instance Binary v => Binary (LeafValue v) where + put (RawValue v) = put (0x00 :: Word8) >> put v + put (OverflowValue v) = put (0x01 :: Word8) >> put v + + get = (get :: Get Word8) >>= \case + 0x00 -> RawValue <$> get + 0x01 -> OverflowValue <$> get + t -> fail $ "unknown leaf value: " ++ showHex t "" + +type LeafItems k v = Map k (LeafValue v) + +-- | A node in a B+-tree. +-- +-- Nodes are parameterized over the key and value types and are additionally +-- indexed by their height. All paths from the root to the leaves have the same +-- length. The height is the number of edges from the root to the leaves, +-- i.e. leaves are at height zero and index nodes increase the height. +-- +-- Sub-trees are represented by a 'NodeId' that are used to resolve the actual +-- storage location of the sub-tree node. +data Node height key val where + Idx :: { idxChildren :: Index key (NodeId height key val) + } -> Node ('S height) key val + Leaf :: { leafItems :: LeafItems key val + } -> Node 'Z key val + deriving (Typeable) + +instance (Eq key, Eq val) => Eq (Node height key val) where + Leaf x == Leaf y = x == y + Idx x == Idx y = x == y + +deriving instance (Show key, Show val) => Show (Node height key val) +deriving instance (Show key, Show val) => Show (Tree key val) + +instance (Value k, Value v) => Value (Tree k v) where + +-------------------------------------------------------------------------------- + +instance Binary (Tree key val) where + put (Tree height rootId) = put height >> put rootId + get = Tree <$> get <*> get + +-- | Encode a 'Leaf' 'Node'. +putLeafNode :: (Binary key, Binary val) => Node 'Z key val -> Put +putLeafNode (Leaf items) = do + encodeSize $ fromIntegral (M.size items) + mapM_ put $ M.toList items + where + encodeSize :: Word32 -> Put + encodeSize s = put msb1 >> put msb2 >> put msb3 + where + msb1 = fromIntegral $ s `shiftR` 16 :: Word8 + msb2 = fromIntegral $ s `shiftR` 8 :: Word8 + msb3 = fromIntegral s :: Word8 + +-- | Decode a 'Leaf' 'Node'. +getLeafNode :: (Ord key, Binary key, Binary val) => Height 'Z -> Get (Node 'Z key val) +getLeafNode _ = do + v <- decodeSize <$> get + l <- replicateM (fromIntegral v) get + return $ Leaf (M.fromList l) + where + decodeSize :: (Word8, Word8, Word8) -> Word32 + decodeSize (msb1, msb2, msb3) = msb1' .|. msb2' .|. msb3' + where + msb1' = (fromIntegral msb1 :: Word32) `shiftL` 16 + msb2' = (fromIntegral msb2 :: Word32) `shiftL` 8 + msb3' = fromIntegral msb3 :: Word32 + +-- | Encode an 'Idx' 'Node'. +putIndexNode :: (Binary key, Binary val) => Node ('S n) key val -> Put +putIndexNode (Idx idx) = put idx + +-- | Decode an 'Idx' 'Node'. +getIndexNode :: (Binary key, Binary val) => Height ('S n) -> Get (Node ('S n) key val) +getIndexNode _ = Idx <$> get + +-------------------------------------------------------------------------------- + +-- | Cast a node to a different type. +-- +-- Essentially this is just a drop-in replacement for 'Data.Typeable.cast'. +castNode :: forall n key1 val1 height1 key2 val2 height2. + (Typeable key1, Typeable val1, Typeable key2, Typeable val2) + => Height height1 -- ^ Term-level witness for the source height. + -> Height height2 -- ^ Term-level witness for the target height. + -> n height1 key1 val1 -- ^ Node to cast. + -> Maybe (n height2 key2 val2) +castNode height1 height2 n + | typeRep (Proxy :: Proxy key1) == typeRep (Proxy :: Proxy key2) + , typeRep (Proxy :: Proxy val1) == typeRep (Proxy :: Proxy val2) + , fromHeight height1 == fromHeight height2 + = Just (unsafeCoerce n) + | otherwise + = Nothing + +-- | Cast a node to one of the available types. +castNode' :: forall n h k v. + (Typeable k, Typeable v) + => Height h -- ^ Term-level witness for the source height + -> n h k v -- ^ Node to cast. + -> Either (n 'Z k v) (n ('S h) k v) +castNode' h n + | Just v <- castNode h zeroHeight n = Left v + | otherwise = Right (unsafeCoerce n) + +-------------------------------------------------------------------------------- + +-- | Cast a value to a different type. +-- +-- Essentially this is just a drop-in replacement for +-- 'Data.Typeable.cast'. +castValue :: (Typeable v1, Typeable v2) => v1 -> Maybe v2 +castValue = cast diff --git a/src/Data/BTree/Primitives.hs b/src/Data/BTree/Primitives.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives.hs @@ -0,0 +1,17 @@ +-- | Primitive data structures and algorithms needed for both the pure +-- ("Data.BTree.Pure") and impure ("Data.BTree.Impure") B+-tree implementation. +module Data.BTree.Primitives ( + module Data.BTree.Primitives.Height +, module Data.BTree.Primitives.Ids +, module Data.BTree.Primitives.Index +, module Data.BTree.Primitives.Key +, module Data.BTree.Primitives.Leaf +, module Data.BTree.Primitives.Value +) where + +import Data.BTree.Primitives.Height +import Data.BTree.Primitives.Ids +import Data.BTree.Primitives.Index +import Data.BTree.Primitives.Key +import Data.BTree.Primitives.Leaf +import Data.BTree.Primitives.Value diff --git a/src/Data/BTree/Primitives/Exception.hs b/src/Data/BTree/Primitives/Exception.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Exception.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE DeriveDataTypeable #-} +-- | A collection of exceptions that can be raised in the pure algorithms in +-- "Data.BTree.Primitives", "Data.BTree.Impure" and "Data.BTree.Pure". +module Data.BTree.Primitives.Exception ( + -- * Re-exports + throw + + -- * Custom exceptions +, TreeAlgorithmError(..) +, KeyTooLargeError(..) +) where + +import Control.Exception (Exception, throw) + +import Data.Typeable (Typeable) + +-- | An exception raised when the pure modification algorithms are called using +-- invalid state. +-- +-- This exception is only raised when a the library contains a bug. +-- +-- The first argument is a function name indicating the location of the error. +-- The second argument is the description of the error. +data TreeAlgorithmError = TreeAlgorithmError String String deriving (Typeable) + +instance Show TreeAlgorithmError where + show (TreeAlgorithmError loc msg) = + loc ++ ": " ++ msg ++ + " (TreeAlgorithmError, this indicates a " ++ + "bug in the haskey-btree library, please report)" + +instance Exception TreeAlgorithmError where + +-- | An exception thrown when the keys inserted in the database are larger than +-- 'Data.BTree.Alloc.Class.maxKeySize'. +-- +-- Note that this exception can be thrown long after the key violating the +-- maximum key size was inserted. It is only detected when the tree +-- modification algorithms try to split the node containing that key. +-- +-- Increase the page size to fix this problem. +data KeyTooLargeError = KeyTooLargeError deriving (Show, Typeable) + +instance Exception KeyTooLargeError where diff --git a/src/Data/BTree/Primitives/Height.hs b/src/Data/BTree/Primitives/Height.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Height.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE GADTs #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE KindSignatures #-} + +module Data.BTree.Primitives.Height where + +import Data.Binary (Binary) +import Data.Word (Word8) +import Unsafe.Coerce + +-------------------------------------------------------------------------------- + +data Nat = Z | S Nat + +newtype Height (h :: Nat) = Height { fromHeight :: Word8 } + deriving (Binary, Eq, Ord) + +instance Show (Height h) where + showsPrec p = showsPrec p . fromHeight + +zeroHeight :: Height 'Z +zeroHeight = Height 0 +{-# INLINE zeroHeight #-} + +incrHeight :: Height h -> Height ('S h) +incrHeight = Height . (+1) . fromHeight +{-# INLINE incrHeight #-} + +decrHeight :: Height ('S h) -> Height h +decrHeight = Height . (+(-1)) . fromHeight +{-# INLINE decrHeight #-} + +-------------------------------------------------------------------------------- + +data UHeight (height :: Nat) :: * where + UZero :: UHeight 'Z + USucc :: Height height -> UHeight ('S height) + +viewHeight :: Height height -> UHeight height +viewHeight (Height 0) = unsafeCoerce UZero +viewHeight (Height n) = unsafeCoerce (USucc (Height (n-1))) + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Primitives/Ids.hs b/src/Data/BTree/Primitives/Ids.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Ids.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE KindSignatures #-} + +module Data.BTree.Primitives.Ids where + +import Data.BTree.Primitives.Height +import Data.BTree.Primitives.Key +import Data.BTree.Primitives.Value + +import Data.Binary (Binary) +import Data.Hashable (Hashable) +import Data.Typeable (Typeable) +import Data.Word +import Numeric (showHex) + +-------------------------------------------------------------------------------- + +-- | Reference to a stored page. +newtype PageId = PageId { fromPageId :: Word64 } + deriving (Eq, Ord, Binary, Num, Value, Key, Typeable) + +-- | Reference to a stored overflow page. +-- +-- An overflow id is the combination of the transaction id that +-- generated it, and a counter. +type OverflowId = (TxId, Word32) + +-- | Type used to indicate the size of storage pools. +newtype PageCount = PageCount { fromPageCount :: Word64 } + deriving (Eq, Ord, Binary, Num, Enum, Typeable) + +-- | Type used to indicate the size of a single physical page in bytes. +newtype PageSize = PageSize { fromPageSize :: Word32 } + deriving (Eq, Ord, Show, Binary, Num, Enum, Real, Integral, Typeable) + +-- | Reference to a stored 'Node'. +-- +-- 'NodeId' has phantom type arguments for the parameters of 'Node' to be able +-- to enforce consistency. In a setting with a single storage pool this 'Id' +-- will essentially be a 'PageId' with just the extra typing. In a multi +-- storage pool setting 'NodeId's will additionally have to be resolved to +-- 'PageId's by the node allocator. +newtype NodeId (height :: Nat) key val = NodeId { fromNodeId :: Word64 } + deriving (Eq, Ord, Binary, Num) + +-- | Convert a 'NodeId' to a 'PageId' +nodeIdToPageId :: NodeId height key val -> PageId +nodeIdToPageId = PageId . fromNodeId + +-- | Convert a 'PageId' to a 'NodeId' +pageIdToNodeId :: PageId -> NodeId height key val +pageIdToNodeId = NodeId . fromPageId + +-- | Transaction ids that are used as revision numbers. +newtype TxId = TxId { fromTxId :: Word64 } + deriving (Eq, Ord, Binary, Num, Hashable, Value, Key, Typeable) + +instance Show PageId where + showsPrec _ (PageId n) = showString "0x" . showHex n +instance Show PageCount where + showsPrec _ (PageCount n) = showString "0x" . showHex n +instance Show (NodeId height key val) where + showsPrec _ (NodeId n) = showString "0x" . showHex n +instance Show TxId where + showsPrec _ (TxId n) = showString "0x" . showHex n + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Primitives/Index.hs b/src/Data/BTree/Primitives/Index.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Index.hs @@ -0,0 +1,300 @@ +{-# LANGUAGE DeriveFoldable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveTraversable #-} + +module Data.BTree.Primitives.Index where + +import Control.Applicative ((<$>)) +import Control.Monad.Identity (runIdentity) + +import Data.Binary (Binary(..), Put) +import Data.Bits ((.|.), shiftL, shiftR) +import Data.Foldable (Foldable) +import Data.Monoid +import Data.Traversable (Traversable) +import Data.Vector (Vector) +import Data.Word (Word8, Word32) +import qualified Data.Map as M +import qualified Data.Vector as V + +import Data.BTree.Primitives.Exception +import Data.BTree.Utils.List (safeLast) +import Data.BTree.Utils.Vector (isStrictlyIncreasing, vecUncons, vecUnsnoc) + +-------------------------------------------------------------------------------- + +-- | The 'Index' encodes the internal structure of an index node. +-- +-- The index is abstracted over the type 'node' of sub-trees. The keys and +-- nodes are stored in separate 'Vector's and the keys are sorted in strictly +-- increasing order. There should always be one more sub-tree than there are +-- keys. Hence structurally the smallest 'Index' has one sub-tree and no keys, +-- but a valid B+-tree index node will have at least two sub-trees and one key. +data Index key node = Index !(Vector key) !(Vector node) + deriving (Eq, Functor, Foldable, Show, Traversable) + +instance (Binary k, Binary n) => Binary (Index k n) where + put (Index keys nodes) = do + encodeSize $ fromIntegral (V.length keys) + V.mapM_ put keys + V.mapM_ put nodes + where + encodeSize :: Word32 -> Put + encodeSize s = put msb1 >> put msb2 >> put msb3 + where + msb1 = fromIntegral $ s `shiftR` 16 :: Word8 + msb2 = fromIntegral $ s `shiftR` 8 :: Word8 + msb3 = fromIntegral s :: Word8 + + get = do + numKeys <- decodeSize <$> get + keys <- V.replicateM (fromIntegral numKeys) get + values <- V.replicateM (fromIntegral numKeys + 1) get + return $ Index keys values + where + decodeSize :: (Word8, Word8, Word8) -> Word32 + decodeSize (msb1, msb2, msb3) = msb1' .|. msb2' .|. msb3' + where + msb1' = (fromIntegral msb1 :: Word32) `shiftL` 16 + msb2' = (fromIntegral msb2 :: Word32) `shiftL` 8 + msb3' = fromIntegral msb3 :: Word32 + +-- | Return the number of keys in this 'Index'. +indexNumKeys :: Index key val -> Int +indexNumKeys (Index keys _vals) = V.length keys + +-- | Return the number of values stored in this 'Index'. +indexNumVals :: Index key val -> Int +indexNumVals (Index _keys vals) = V.length vals + +-- | Validate the key/node count invariant of an 'Index'. +validIndex :: Ord key => Index key node -> Bool +validIndex (Index keys nodes) = + V.length keys + 1 == V.length nodes && + isStrictlyIncreasing keys + +-- | Validate the size of an 'Index'. +validIndexSize :: Ord key => Int -> Int -> Index key node -> Bool +validIndexSize minIdxKeys maxIdxKeys idx@(Index keys _) = + validIndex idx && V.length keys >= minIdxKeys && V.length keys <= maxIdxKeys + +-- | Split an index node. +-- +-- This function splits an index node into two new nodes at the given key +-- position @numLeftKeys@ and returns the resulting indices and the key +-- separating them. Eventually this should take the binary size of serialized +-- keys and sub-tree pointers into account. See also 'splitLeaf' in +-- "Data.BTree.Primitives.Leaf". +splitIndexAt :: Int -> Index key val -> (Index key val, key, Index key val) +splitIndexAt numLeftKeys (Index keys vals) + | (leftKeys, middleKeyAndRightKeys) <- V.splitAt numLeftKeys keys + , (leftVals, rightVals) <- V.splitAt (numLeftKeys+1) vals + = case vecUncons middleKeyAndRightKeys of + Just (middleKey,rightKeys) -> + (Index leftKeys leftVals, middleKey, Index rightKeys rightVals) + Nothing -> throw $ + TreeAlgorithmError "splitIndex" "cannot split an empty index" + +-- | Split an index many times. +-- +-- This function splits an 'Index' node into smaller pieces. Each resulting +-- sub-'Index' has between @maxIdxKeys/2@ and @maxIdxKeys@ inclusive values and +-- is additionally applied to the function @f@. +-- +-- This is the dual of a monadic bind and is also known as the `extended` +-- function of extendable functors. See "Data.Functor.Extend" in the +-- "semigroupoids" package. +-- +-- prop> bindIndex (extendedIndex n id idx) id == idx +extendedIndex :: Int -> (Index k b -> a) -> Index k b -> Index k a +extendedIndex maxIdxKeys f = go + where + maxIdxVals = maxIdxKeys + 1 + + go index + | numVals <= maxIdxVals + = singletonIndex (f index) + | numVals <= 2*maxIdxVals + = case splitIndexAt (div numVals 2 - 1) index of + (leftIndex, middleKey, rightIndex) -> + indexFromList [middleKey] [f leftIndex, f rightIndex] + | otherwise + = case splitIndexAt maxIdxKeys index of + (leftIndex, middleKey, rightIndex) -> + mergeIndex (singletonIndex (f leftIndex)) + middleKey (go rightIndex) + where + numVals = indexNumVals index + +extendIndexPred :: (a -> Bool) -> + (Index k b -> a) -> Index k b -> Maybe (Index k a) +extendIndexPred p f = go + where + go index + | let indexEnc = f index + , p indexEnc + = Just (singletonIndex indexEnc) + | indexNumKeys index <= 2 + = -- Cannot split node with only 2 keys, increase page size + throw KeyTooLargeError + | otherwise + = do + let numKeys = indexNumKeys index + (leftEnc, (middleKey, right)) <- safeLast $ + takeWhile (p . fst) + [ (leftEnc, (middleKey, right)) + | i <- [1..numKeys-2] -- left and right must contain at least one key + , let (left,middleKey,right) = splitIndexAt i index + leftEnc = f left + ] + rightEnc <- go right + return $! mergeIndex (singletonIndex leftEnc) middleKey rightEnc + +-- | Merge two indices. +-- +-- Merge two indices 'leftIndex', 'rightIndex' given a discriminating key +-- 'middleKey', i.e. such that '∀ (k,v) ∈ leftIndex. k < middleKey' and +-- '∀ (k,v) ∈ rightIndex. middleKey <= k'. +-- +-- 'mergeIndex' is a partial inverse of splitIndex, i.e. +-- prop> splitIndex is == (left,mid,right) => mergeIndex left mid right == is +mergeIndex :: Index key val -> key -> Index key val -> Index key val +mergeIndex (Index leftKeys leftVals) middleKey (Index rightKeys rightVals) = + Index + (leftKeys <> V.singleton middleKey <> rightKeys) + (leftVals <> rightVals) + +-- | Create an index from key-value lists. +-- +-- The internal invariants of the 'Index' data structure apply. That means +-- there is one more value than there are keys and keys are ordered in strictly +-- increasing order. +indexFromList :: [key] -> [val] -> Index key val +indexFromList ks vs = Index (V.fromList ks) (V.fromList vs) + +-- | Create an index with a single value. +singletonIndex :: val -> Index key val +singletonIndex = Index V.empty . V.singleton + +-- | Test if the index consists of a single value. +-- +-- Returns the element if the index is a singleton. Otherwise fails. +-- +-- prop> fromSingletonIndex (singletonIndex val) == Just val +fromSingletonIndex :: Index key val -> Maybe val +fromSingletonIndex (Index _keys vals) = + if V.length vals == 1 then Just $! V.unsafeHead vals else Nothing + +-------------------------------------------------------------------------------- + +-- | Bind an index +-- +-- prop> bindIndex idx singletonIndex == idx +bindIndex :: Index k a -> (a -> Index k b) -> Index k b +bindIndex idx f = runIdentity $ bindIndexM idx (return . f) + +bindIndexM :: (Functor m, Monad m) + => Index k a + -> (a -> m (Index k b)) + -> m (Index k b) +bindIndexM (Index ks vs) f = case vecUncons vs of + Just (v, vtail) -> do + i <- f v + V.foldM' g i (V.zip ks vtail) + where + g acc (k , w) = mergeIndex acc k <$> f w + Nothing -> + throw $ TreeAlgorithmError "bindIndexM" "cannot bind an empty Index" + +-------------------------------------------------------------------------------- + +-- | Representation of one-hole contexts of 'Index'. +-- +-- Just one val removes. All keys are present. +-- +-- prop> V.length leftVals == V.length lefyKeys +-- prop> V.length rightVals == V.length rightKeys +data IndexCtx key val = IndexCtx + { indexCtxLeftKeys :: !(Vector key) + , indexCtxRightKeys :: !(Vector key) + , indexCtxLeftVals :: !(Vector val) + , indexCtxRightVals :: !(Vector val) + } + deriving (Functor, Foldable, Show, Traversable) + +putVal :: IndexCtx key val -> val -> Index key val +putVal ctx val = + Index + (indexCtxLeftKeys ctx <> indexCtxRightKeys ctx) + (indexCtxLeftVals ctx <> V.singleton val <> indexCtxRightVals ctx) + +putIdx :: IndexCtx key val -> Index key val -> Index key val +putIdx ctx (Index keys vals) = + Index + (indexCtxLeftKeys ctx <> keys <> indexCtxRightKeys ctx) + (indexCtxLeftVals ctx <> vals <> indexCtxRightVals ctx) + +valView :: Ord key => key -> Index key val -> (IndexCtx key val, val) +valView key (Index keys vals) + | (leftKeys,rightKeys) <- V.span (<=key) keys + , n <- V.length leftKeys + , (leftVals,valAndRightVals) <- V.splitAt n vals + , Just (val,rightVals) <- vecUncons valAndRightVals + = ( IndexCtx + { indexCtxLeftKeys = leftKeys + , indexCtxRightKeys = rightKeys + , indexCtxLeftVals = leftVals + , indexCtxRightVals = rightVals + }, + val + ) + | otherwise + = throw $ TreeAlgorithmError "valView" "cannot split an empty index" + +valViewMin :: Index key val -> (IndexCtx key val, val) +valViewMin (Index keys vals) + | Just (val, rightVals) <- vecUncons vals + = ( IndexCtx + { indexCtxLeftKeys = V.empty + , indexCtxRightKeys = keys + , indexCtxLeftVals = V.empty + , indexCtxRightVals = rightVals + }, + val + ) + | otherwise + = throw $ TreeAlgorithmError "valViewMin" "cannot split an empty index" + +-- | Distribute a map of key-value pairs over an index. +distribute :: Ord k => M.Map k v -> Index k node -> Index k (M.Map k v, node) +distribute kvs (Index keys nodes) + | a <- V.imap rangeTail (Nothing `V.cons` V.map Just keys) + , b <- V.map (uncurry rangeHead) (V.zip (V.map Just keys `V.snoc` Nothing) a) + = Index keys b + where + rangeTail idx Nothing = (kvs, nodes V.! idx) + rangeTail idx (Just key) = (takeWhile' (>= key) kvs, nodes V.! idx) + rangeHead Nothing (tail', node) = (tail', node) + rangeHead (Just key) (tail', node) = (takeWhile' (< key) tail', node) + + takeWhile' :: (k -> Bool) -> M.Map k v -> M.Map k v + takeWhile' p = fst . M.partitionWithKey (\k _ -> p k) + +leftView :: IndexCtx key val -> Maybe (IndexCtx key val, val, key) +leftView ctx = do + (leftVals, leftVal) <- vecUnsnoc (indexCtxLeftVals ctx) + (leftKeys, leftKey) <- vecUnsnoc (indexCtxLeftKeys ctx) + return (ctx { indexCtxLeftKeys = leftKeys + , indexCtxLeftVals = leftVals + }, leftVal, leftKey) + +rightView :: IndexCtx key val -> Maybe (key, val, IndexCtx key val) +rightView ctx = do + (rightVal, rightVals) <- vecUncons (indexCtxRightVals ctx) + (rightKey, rightKeys) <- vecUncons (indexCtxRightKeys ctx) + return (rightKey, rightVal, + ctx { indexCtxRightKeys = rightKeys + , indexCtxRightVals = rightVals + }) + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Primitives/Key.hs b/src/Data/BTree/Primitives/Key.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Key.hs @@ -0,0 +1,56 @@ + +module Data.BTree.Primitives.Key where + +import Data.ByteString (ByteString) +import Data.Int +import Data.Word +import qualified Data.ByteString as BS +import qualified Data.ByteString.Unsafe as BS + +import Data.BTree.Primitives.Exception +import Data.BTree.Primitives.Value + +-------------------------------------------------------------------------------- + +class (Ord k, Value k) => Key k where + -- | Given two keys 'a', 'b' such that 'a < b' compute two new keys 'a2', + -- 'b2' such that 'a <= a2 < b2 <= b'. Obviously this always holds for 'a2 + -- == a' and 'b2 = b' but for 'ByteString's we can potentially find smaller + -- 'a2' and 'b2'. If 'a' equals 'b', the behaviour is undefined. + narrow :: k -> k -> (k,k) + narrow = (,) + +instance Key () +instance Key Bool +instance Key Double +instance Key Float +instance Key Int8 +instance Key Int16 +instance Key Int32 +instance Key Int64 +instance Key Integer +instance Key Word8 +instance Key Word16 +instance Key Word32 +instance Key Word64 + +instance Key ByteString where + narrow a b = + case (compare n na, compare n nb) of + -- So the n+1th byte is the first distinguishing byte. + (LT,LT) -> (BS.unsafeTake (n+1) a, BS.unsafeTake (n+1) b) + -- In this case 'a' is a prefix of 'b'. Can't do anything for a, but we + -- can shorten 'b'. + (EQ,LT) -> (a, BS.unsafeTake (n+1) b) + -- Inputs violate the invariant a throw $ TreeAlgorithmError "narrow (Binary)" $ concat + ["Key ByteString: can't narrow ", show a, " and ", show b] + where + na = BS.length a + nb = BS.length b + -- Length of the longest Common prefix + n = length (takeWhile id (BS.zipWith (==) a b)) + +instance (Key a, Key b) => Key (a, b) where + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Primitives/Leaf.hs b/src/Data/BTree/Primitives/Leaf.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Leaf.hs @@ -0,0 +1,78 @@ +module Data.BTree.Primitives.Leaf where + +import Control.Applicative ((<$>)) + +import Data.Map (Map) +import qualified Data.Map as M + +import Data.BTree.Primitives.Exception +import Data.BTree.Primitives.Index +import Data.BTree.Primitives.Key +import Data.BTree.Utils.List (safeLast) +import Data.BTree.Utils.Map (mapInits, mapSplitAt) + +-------------------------------------------------------------------------------- + +-- | Split a leaf many times until the predicate is satisfied. +-- +-- This function ensures that the for each returned leaf, the predicate is +-- satisfied, or returns 'Nothing' when it can't be satisfied. +splitLeafManyPred :: (Key key) + => (a -> Bool) + -> (Map key val -> a) + -> Map key val + -> Maybe (Index key a) +splitLeafManyPred p f = go + where + go items + | indexEnc <- f items + , p indexEnc + = Just (singletonIndex indexEnc) + | otherwise + = do + left <- lstForWhich (p . f) inits' + let right = items `M.difference` left + mergeIndex (singletonIndex (f left)) + (fst $ M.findMin right) + <$> go right + where + inits' = tail (mapInits items) + + lstForWhich :: (a -> Bool) -> [a] -> Maybe a + lstForWhich g xs = safeLast $ takeWhile g xs + +-- | Split a leaf many times. +-- +-- This function ensures that the for each returned leaf, the amount of +-- items <= maxLeafItems (and >= minLeafItems, except when the original +-- leaf had less than minLeafItems items. +splitLeafMany :: Key key => Int -> (Map key val -> a) -> Map key val -> Index key a +splitLeafMany maxLeafItems f items + | M.size items <= maxLeafItems + = singletonIndex (f items) + | M.size items <= 2*maxLeafItems + , numLeft <- div (M.size items) 2 + , (leftLeaf, rightLeaf) <- mapSplitAt numLeft items + , Just ((key,_), _) <- M.minViewWithKey rightLeaf + = indexFromList [key] [f leftLeaf, f rightLeaf] + | (keys, maps) <- split' items ([], []) + = indexFromList keys (map f maps) + where + split' :: Key key => Map key val -> ([key], [Map key val]) -> ([key], [Map key val]) + split' m (keys, leafs) + | M.size m > 2*maxLeafItems + , (leaf, rem') <- mapSplitAt maxLeafItems m + , (key, _) <- M.findMin rem' + = split' rem' (key:keys, leaf:leafs) + | M.size m > maxLeafItems + , numLeft <- div (M.size m) 2 + , (left, right) <- mapSplitAt numLeft m + , (key, _) <- M.findMin right + = split' M.empty (key:keys, right:(left:leafs)) + | M.null m + = (reverse keys, reverse leafs) + | otherwise + = throw $ TreeAlgorithmError "splitLeafMany" + "constraint violation, got a Map with <= maxLeafItems" + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Primitives/Value.hs b/src/Data/BTree/Primitives/Value.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Primitives/Value.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module Data.BTree.Primitives.Value where + +import Control.Applicative ((<$>),(<*>)) +import Data.Binary (Binary) +import Data.ByteString (ByteString) +import Data.Int +import Data.Proxy (Proxy (..)) +import Data.Typeable +import Data.Word + +-------------------------------------------------------------------------------- + +class (Binary v, Show v, Typeable v) => Value v where + -- | 'Just' with the size in bytes if 'v' is a fixed sized value, 'Nothing' + -- if 'v' is variable sized. + fixedSize :: Proxy v -> Maybe Int + fixedSize _ = Nothing + +instance Value () where fixedSize _ = Just 0 +instance Value Bool where fixedSize _ = Just 1 +instance Value Char where fixedSize _ = Just 4 +instance Value Double where fixedSize _ = Just 8 +instance Value Float where fixedSize _ = Just 4 +instance Value Int8 where fixedSize _ = Just 1 +instance Value Int16 where fixedSize _ = Just 2 +instance Value Int32 where fixedSize _ = Just 4 +instance Value Int64 where fixedSize _ = Just 8 +instance Value Word8 where fixedSize _ = Just 1 +instance Value Word16 where fixedSize _ = Just 2 +instance Value Word32 where fixedSize _ = Just 4 +instance Value Word64 where fixedSize _ = Just 8 + +instance Value ByteString +instance Value Integer + +instance (Value k1, Value k2) => Value (k1,k2) where + fixedSize _ = + (+) <$> fixedSize (Proxy :: Proxy k1) + <*> fixedSize (Proxy :: Proxy k2) + +instance Value v => Value [v] where + fixedSize = const Nothing +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Pure.hs b/src/Data/BTree/Pure.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Pure.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | A pure in-memory B+-tree implementation. +module Data.BTree.Pure ( + -- * Structures + module Data.BTree.Pure.Setup +, Tree(..) +, Node(..) + + -- * Manipulations +, empty +, singleton +, fromList +, insert +, insertMany +, delete + + -- * Lookup +, lookup +, findWithDefault +, member +, notMember + + -- * Properties +, null +, size + + -- * Folds +, foldrWithKey +, toList +) where + +import Prelude hiding (lookup, null) + +import Data.BTree.Primitives.Exception +import Data.BTree.Primitives.Height +import Data.BTree.Primitives.Index +import Data.BTree.Primitives.Key +import Data.BTree.Primitives.Leaf +import Data.BTree.Pure.Setup + +import Data.Map (Map) +import Data.Maybe (isJust, isNothing, fromMaybe) +import Data.Monoid +import qualified Data.Foldable as F +import qualified Data.List as L +import qualified Data.Map as M + +-------------------------------------------------------------------------------- + +-- | A pure B+-tree. +-- +-- This is a simple wrapper around a root 'Node'. An empty tree is represented +-- by 'Nothing'. Otherwise it's 'Just' the root. The height is existentially +-- quantified. +data Tree key val where + Tree :: !TreeSetup + -> Maybe (Node height key val) + -> Tree key val + + +-- | A node in a B+-tree. +-- +-- Nodes are parameterized over the key and value types and are additionally +-- indexed by their height. All paths from the root to the leaves have the same +-- length. The height is the number of edges from the root to the leaves, +-- i.e. leaves are at height zero and index nodes increase the height. +data Node (height :: Nat) key val where + Idx :: { idxChildren :: Index key (Node height key val) + } -> Node ('S height) key val + Leaf :: { leafItems :: Map key val + } -> Node 'Z key val + +deriving instance (Show key, Show val) => Show (Node height key val) +deriving instance (Show key, Show val) => Show (Tree key val) + +-- | The empty tree. +empty :: TreeSetup -> Tree key val +empty setup = Tree setup Nothing + +-- | Construct a tree containg one element. +singleton :: Key k => TreeSetup -> k -> v -> Tree k v +singleton s k v = insert k v (empty s) + +-- | /O(n*log n)/. Construct a B-tree from a list of key\/value pairs. +-- +-- If the list contains duplicate keys, the last pair for a duplicate key is +-- kept. +fromList :: Key k => TreeSetup -> [(k,v)] -> Tree k v +fromList s = L.foldl' (flip $ uncurry insert) (empty s) + +-------------------------------------------------------------------------------- + +-- | Insert a key-value pair into a tree. +-- +-- When inserting a new entry, the leaf it is inserted to and the index nodes +-- on the path to the leaf potentially need to be split. Instead of returning +-- the outcome, 1 node or 2 nodes (with a discriminating key), we return an +-- 'Index' of these nodes. +-- +-- If the key already existed in the tree, it is overwritten. +insert :: Key k => k -> v -> Tree k v -> Tree k v +insert k d (Tree setup (Just rootNode)) + | newRootIdx <- insertRec setup k d rootNode + = case fromSingletonIndex newRootIdx of + Just newRootNode -> + -- The result from the recursive insert is a single node. Use + -- this as a new root. + Tree setup (Just newRootNode) + Nothing -> + -- The insert resulted in a index with multiple nodes, i.e. + -- the splitting propagated to the root. Create a new 'Idx' + -- node with the index. This increments the height. + Tree setup (Just (Idx newRootIdx)) +insert k d (Tree setup Nothing) + = -- Inserting into an empty tree creates a new singleton tree. + Tree setup (Just (Leaf (M.singleton k d))) + +insertRec :: Key key + => TreeSetup + -> key + -> val + -> Node height key val + -> Index key (Node height key val) +insertRec setup key val (Idx children) + | -- Punch a hole into the index at the sub-tree we recurse into. + (ctx, child) <- valView key children + , newChildIdx <- insertRec setup key val child + = -- Fill the hole with the resulting 'Index' from the recursive call + -- and then check if the split needs to be propagated. + splitIndex setup (putIdx ctx newChildIdx) +insertRec setup key val (Leaf items) + = splitLeaf setup (M.insert key val items) + +-- | Insert a bunch of key-value pairs simultaneously. +insertMany :: Key k => Map k v -> Tree k v -> Tree k v +insertMany kvs (Tree setup (Just rootNode)) + = fixUp setup $ insertRecMany setup kvs rootNode +insertMany kvs (Tree setup Nothing) + = fixUp setup $ splitLeaf setup kvs + +insertRecMany :: Key key + => TreeSetup + -> Map key val + -> Node height key val + -> Index key (Node height key val) +insertRecMany setup kvs (Idx idx) + | M.null kvs + = singletonIndex (Idx idx) + | dist <- distribute kvs idx + = splitIndex setup (dist `bindIndex` uncurry (insertRecMany setup)) +insertRecMany setup kvs (Leaf items) + = splitLeaf setup (M.union kvs items) + +-- | Fix up the root node of a tree. +-- +-- Fix up the root node of a tree, where all other nodes are valid, but the +-- root node may contain more items than allowed. Do this by repeatedly +-- splitting up the root node. +fixUp :: Key key + => TreeSetup + -> Index key (Node height key val) + -> Tree key val +fixUp setup idx = case fromSingletonIndex idx of + Just newRootNode -> Tree setup (Just newRootNode) + Nothing -> fixUp setup (splitIndex setup idx) + +-------------------------------------------------------------------------------- + +-- | /O(n)/. Fold key\/value pairs in the B-tree. +foldrWithKey :: forall k v w. (k -> v -> w -> w) -> w -> Tree k v -> w +foldrWithKey f z0 (Tree _ mbRoot) = maybe z0 (go z0) mbRoot + where + go :: w -> Node h k v -> w + go z1 (Leaf items) = M.foldrWithKey f z1 items + go z1 (Idx index) = F.foldr (flip go) z1 index + +-- | /O(n)/. Convert the B-tree to a sorted list of key\/value pairs. +toList :: Tree k v -> [(k,v)] +toList = foldrWithKey (\k v kvs -> (k,v):kvs) [] + +-------------------------------------------------------------------------------- + +-- | Delete a key-value pair from the tree. +delete :: Key k => k -> Tree k v -> Tree k v +delete _key (Tree setup Nothing) = Tree setup Nothing +delete key (Tree setup (Just rootNode)) = case deleteRec setup key rootNode of + Idx index + | Just childNode <- fromSingletonIndex index -> Tree setup (Just childNode) + Leaf items + | M.null items -> Tree setup Nothing + newRootNode -> Tree setup (Just newRootNode) + +deleteRec :: Key k + => TreeSetup + -> k + -> Node n k v + -> Node n k v +deleteRec setup key (Idx children) + | childNeedsMerge, Just (rKey, rChild, rCtx) <- rightView ctx + = Idx (putIdx rCtx (mergeNodes setup newChild rKey rChild)) + | childNeedsMerge, Just (lCtx, lChild, lKey) <- leftView ctx + = Idx (putIdx lCtx (mergeNodes setup lChild lKey newChild)) + -- No left or right sibling? This is a constraint violation. Also + -- this couldn't be the root because it would've been shrunk + -- before. + | childNeedsMerge + = throw $ TreeAlgorithmError "deleteRec" + "constraint violation, found an index node with a single child" + | otherwise = Idx (putVal ctx newChild) + where + (ctx, child) = valView key children + newChild = deleteRec setup key child + childNeedsMerge = nodeNeedsMerge setup newChild +deleteRec _ key (Leaf items) + = Leaf (M.delete key items) + +nodeNeedsMerge :: TreeSetup -> Node height key value -> Bool +nodeNeedsMerge setup = \case + Idx children -> indexNumKeys children < minIdxKeys setup + Leaf items -> M.size items < minLeafItems setup + +mergeNodes :: Key key + => TreeSetup + -> Node height key val + -> key + -> Node height key val + -> Index key (Node height key val) +mergeNodes setup (Leaf leftItems) _middleKey (Leaf rightItems) = + splitLeaf setup (leftItems <> rightItems) +mergeNodes setup (Idx leftIdx) middleKey (Idx rightIdx) = + splitIndex setup (mergeIndex leftIdx middleKey rightIdx) + +-------------------------------------------------------------------------------- + +lookupRec :: Key key + => key + -> Node height key val + -> Maybe val +lookupRec key (Idx children) + | (_, childNode) <- valView key children + = lookupRec key childNode +lookupRec key (Leaf items) + = M.lookup key items + +-- | Lookup a value in the tree. +lookup :: Key k => k -> Tree k v -> Maybe v +lookup _ (Tree _ Nothing) = Nothing +lookup k (Tree _ (Just n)) = lookupRec k n + +-- | Lookup a value in the tree, or return a default. +findWithDefault :: Key k => v -> k -> Tree k v -> v +findWithDefault v k = fromMaybe v . lookup k + +-- | Check whether a key is present in the tree. +member :: Key k => k -> Tree k v -> Bool +member k = isJust . lookup k + +-- | Check whether a key is not present in the tree. +notMember :: Key k => k -> Tree k v -> Bool +notMember k = isNothing . lookup k + +-------------------------------------------------------------------------------- + +-- | Check whether the tree is empty. +null :: Tree k v -> Bool +null (Tree _ n) = isNothing n + +sizeNode :: Node n k v -> Int +sizeNode (Leaf items) = M.size items +sizeNode (Idx nodes) = F.sum (fmap sizeNode nodes) + +-- | The size of a tree. +size :: Tree k v -> Int +size (Tree _ Nothing) = 0 +size (Tree _ (Just n)) = sizeNode n + +-------------------------------------------------------------------------------- + +-- | Make a tree node foldable over its value. +instance F.Foldable (Tree key) where + foldMap _ (Tree _ Nothing) = mempty + foldMap f (Tree _ (Just n)) = F.foldMap f n + +instance F.Foldable (Node height key) where + foldMap f (Idx idx) = + F.foldMap (F.foldMap f) idx + + foldMap f (Leaf items) = F.foldMap f items + +-------------------------------------------------------------------------------- + +splitIndex :: TreeSetup + -> Index key (Node height key val) + -> Index key (Node ('S height) key val) +splitIndex setup = extendedIndex (maxIdxKeys setup) Idx + +splitLeaf :: Key key + => TreeSetup + -> Map key val + -> Index key (Node 'Z key val) +splitLeaf setup = splitLeafMany (maxLeafItems setup) Leaf + +-------------------------------------------------------------------------------- diff --git a/src/Data/BTree/Pure/Setup.hs b/src/Data/BTree/Pure/Setup.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Pure/Setup.hs @@ -0,0 +1,52 @@ +-- | This module contains structures relating the the setup of a pure B+-tree. +module Data.BTree.Pure.Setup ( + -- * Setup + TreeSetup(..) + + -- * Predefined setups +, twoThreeSetup +, setupWithMinimumDegreeOf +) where + +-- | Setup of a pure B+-tree. +data TreeSetup = TreeSetup { + minFanout :: Int + , maxFanout :: Int + , minIdxKeys :: Int + , maxIdxKeys :: Int + , minLeafItems :: Int + , maxLeafItems :: Int + } deriving (Show) + +-- | Setup of a 2-3 tree. +twoThreeSetup :: TreeSetup +twoThreeSetup = TreeSetup { + minFanout = minFanout' + , maxFanout = maxFanout' + , minIdxKeys = minFanout' - 1 + , maxIdxKeys = maxFanout' - 1 + , minLeafItems = minFanout' + , maxLeafItems = 2*minFanout' - 1 + } + where + minFanout' = 2 + maxFanout' = 2*minFanout' - 1 + +-- | Setup of a B+-tree with a certain minimum degree, as defined in CLRS. +-- +-- To get for example a 2-3-4 tree, use +-- +-- >>> setupWithMinimumDegreeOf 2 +-- +setupWithMinimumDegreeOf :: Int -> TreeSetup +setupWithMinimumDegreeOf deg = TreeSetup { + minFanout = minFanout' + , maxFanout = maxFanout' + , minIdxKeys = minFanout' - 1 + , maxIdxKeys = maxFanout' - 1 + , minLeafItems = minFanout' - 1 + , maxLeafItems = maxFanout' - 1 + } + where + minFanout' = deg + maxFanout' = 2*deg diff --git a/src/Data/BTree/Utils/List.hs b/src/Data/BTree/Utils/List.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Utils/List.hs @@ -0,0 +1,6 @@ +module Data.BTree.Utils.List where + +import Data.Maybe (listToMaybe) + +safeLast :: [a] -> Maybe a +safeLast = listToMaybe . reverse diff --git a/src/Data/BTree/Utils/Map.hs b/src/Data/BTree/Utils/Map.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Utils/Map.hs @@ -0,0 +1,13 @@ +module Data.BTree.Utils.Map where + +import Data.List (inits) +import Data.Map (Map) +import qualified Data.Map as M + +mapSplitAt :: Eq k => Int -> Map k v -> (Map k v, Map k v) +mapSplitAt i m + | (l,r) <- splitAt i (M.toList m) + = (M.fromAscList l, M.fromAscList r) + +mapInits :: Ord k => Map k v -> [Map k v] +mapInits = map M.fromList . inits . M.toList diff --git a/src/Data/BTree/Utils/Vector.hs b/src/Data/BTree/Utils/Vector.hs new file mode 100644 --- /dev/null +++ b/src/Data/BTree/Utils/Vector.hs @@ -0,0 +1,24 @@ +module Data.BTree.Utils.Vector where + +import Data.List (inits) +import Data.Vector (Vector) +import qualified Data.Vector as V + +vecUncons :: Vector a -> Maybe (a, Vector a) +vecUncons v + | V.null v = Nothing + | otherwise = Just (V.unsafeHead v, V.unsafeTail v) + +vecUnsnoc :: Vector a -> Maybe (Vector a, a) +vecUnsnoc v + | V.null v = Nothing + | otherwise = Just (V.unsafeInit v, V.unsafeLast v) + +vecInits :: Vector a -> Vector (Vector a) +vecInits = V.map V.fromList . V.fromList . inits . V.toList + +isStrictlyIncreasing :: Ord key => Vector key -> Bool +isStrictlyIncreasing ks = case vecUncons ks of + Just (h,t) -> + snd $ V.foldl' (\(lb,res) next -> (next, res && lb < next)) (h, True) t + Nothing -> True diff --git a/tests/Integration.hs b/tests/Integration.hs new file mode 100644 --- /dev/null +++ b/tests/Integration.hs @@ -0,0 +1,15 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE RankNTypes #-} +module Main (main) where + +import Test.Framework (Test, defaultMain) + +import qualified Integration.WriteOpenRead.Debug + +tests :: [Test] +tests = + [ Integration.WriteOpenRead.Debug.tests + ] + +main :: IO () +main = defaultMain tests diff --git a/tests/Integration/WriteOpenRead/Debug.hs b/tests/Integration/WriteOpenRead/Debug.hs new file mode 100644 --- /dev/null +++ b/tests/Integration/WriteOpenRead/Debug.hs @@ -0,0 +1,113 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +module Integration.WriteOpenRead.Debug where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck + +import Control.Applicative ((<$>)) +import Control.Monad +import Control.Monad.Identity +import Control.Monad.Trans.State + +import Data.Binary (Binary) +import Data.Foldable (foldlM) +import Data.Map (Map) +import Data.Typeable (Typeable) +import Data.Word (Word8) +import qualified Data.Map as M + +import GHC.Generics (Generic) + +import Data.BTree.Alloc.Class +import Data.BTree.Alloc.Debug +import Data.BTree.Impure +import Data.BTree.Primitives +import qualified Data.BTree.Impure as Tree + +import Integration.WriteOpenRead.Transactions + +tests :: Test +tests = testGroup "WriteOpenRead.Concurrent" + [ testProperty "debug allocator" prop_debug_allocator + ] + +data AllocatorState k v = AllocatorState { + allocatorStatePages :: Pages + , allocatorStateTree :: Tree k v } + +prop_debug_allocator :: Property +prop_debug_allocator = forAll genTestSequence $ \(TestSequence txs) -> + let s = AllocatorState emptyPages Tree.empty + m = runIdentity $ evalStateT (runSeq txs) s + in + m `seq` True + where + runSeq = foldlM writeReadTest M.empty + + writeReadTest :: Map Integer TestValue + -> TestTransaction Integer TestValue + -> StateT (AllocatorState Integer TestValue) + Identity + (Map Integer TestValue) + writeReadTest m tx = do + openAndWrite tx + read' <- openAndRead + let expected = testTransactionResult m tx + if read' == M.toList expected + then return expected + else error $ "error:" + ++ "\n after: " ++ show tx + ++ "\n expectd: " ++ show (M.toList expected) + ++ "\n got: " ++ show read' + + openAndRead = do + pages <- gets allocatorStatePages + tree <- gets allocatorStateTree + return $ evalDebug pages (readAll tree) + + openAndWrite tx = do + pages <- gets allocatorStatePages + tree <- gets allocatorStateTree + + let (tree', pages') = runDebug pages (doTx tree tx) + put AllocatorState { + allocatorStatePages = pages' + , allocatorStateTree = tree' } + +readAll :: (AllocM m, Key k, Value v) + => Tree k v + -> m [(k, v)] +readAll = Tree.toList + +doTx :: (AllocM m, Key k, Value v) + => Tree k v + -> TestTransaction k v + -> m (Tree k v) +doTx tree (TestTransaction actions) = + foldl (>=>) return (map writeAction actions) tree + where + writeAction (Insert k v) = insertTree k v + writeAction (Replace k v) = insertTree k v + writeAction (Delete k) = deleteTree k + +-------------------------------------------------------------------------------- + +-- | Value used for testing. +-- +-- This value will overflow 20% of the time. +newtype TestValue = TestValue (Either Integer [Word8]) + deriving (Eq, Generic, Show, Typeable) + +instance Binary TestValue where +instance Value TestValue where + +instance Arbitrary TestValue where + arbitrary = + TestValue <$> frequency [(80, Left <$> small), (20, Right <$> big)] + where + small = arbitrary + big = arbitrary + +-------------------------------------------------------------------------------- diff --git a/tests/Integration/WriteOpenRead/Transactions.hs b/tests/Integration/WriteOpenRead/Transactions.hs new file mode 100644 --- /dev/null +++ b/tests/Integration/WriteOpenRead/Transactions.hs @@ -0,0 +1,105 @@ +{-# LANGUAGE RecordWildCards #-} +module Integration.WriteOpenRead.Transactions where + +import Test.QuickCheck + +import Control.Applicative ((<$>), (<*>), pure) +import Control.Monad.State + +import Data.List (inits) +import Data.Map (Map) +import qualified Data.Map as M + +-------------------------------------------------------------------------------- + +newtype TestSequence k v = TestSequence [TestTransaction k v] + deriving (Show) + +data TransactionSetup = TransactionSetup { sequenceInsertFrequency :: !Int + , sequenceReplaceFrequency :: !Int + , sequenceDeleteFrequency :: !Int } + deriving (Show) + +deleteHeavySetup :: TransactionSetup +deleteHeavySetup = TransactionSetup { sequenceInsertFrequency = 35 + , sequenceReplaceFrequency = 20 + , sequenceDeleteFrequency = 45 } + +insertHeavySetup :: TransactionSetup +insertHeavySetup = TransactionSetup { sequenceInsertFrequency = 12 + , sequenceReplaceFrequency = 4 + , sequenceDeleteFrequency = 4 } + +genTransactionSetup :: Gen TransactionSetup +genTransactionSetup = + frequency [(45, return deleteHeavySetup), + (45, return insertHeavySetup)] + +newtype TestTransaction k v = TestTransaction [TestAction k v] + deriving (Show) + +testTransactionResult :: Ord k => Map k v -> TestTransaction k v -> Map k v +testTransactionResult m (TestTransaction actions) + = foldl (flip doAction) m actions + +data TestAction k v = Insert k v + | Replace k v + | Delete k + deriving (Show) + +doAction :: Ord k => TestAction k v -> Map k v -> Map k v +doAction action m + | Insert k v <- action = M.insert k v m + | Replace k v <- action = M.insert k v m + | Delete k <- action = M.delete k m + +genTestTransaction :: (Ord k, Arbitrary k, Arbitrary v) => Map k v -> TransactionSetup -> Gen (TestTransaction k v, Map k v) +genTestTransaction db TransactionSetup{..} = sized $ \n -> do + k <- choose (0, n) + (m, actions) <- execStateT (replicateM k next) (db, []) + tx <- TestTransaction <$> pure (reverse actions) + return (tx, m) + where + genAction :: (Ord k, Arbitrary k, Arbitrary v) + => Map k v + -> Gen (TestAction k v) + genAction m + | M.null m = genInsert + | otherwise = frequency [(sequenceInsertFrequency, genInsert ), + (sequenceReplaceFrequency, genReplace m), + (sequenceDeleteFrequency, genDelete m )] + + genInsert :: (Arbitrary k, Arbitrary v) => Gen (TestAction k v) + genInsert = Insert <$> arbitrary <*> arbitrary + genReplace m = Replace <$> elements (M.keys m) <*> arbitrary + genDelete m = Delete <$> elements (M.keys m) + + next :: (Ord k, Arbitrary k, Arbitrary v) + => StateT (Map k v, [TestAction k v]) Gen () + next = do + (m, actions) <- get + action <- lift $ genAction m + put (doAction action m, action:actions) + +shrinkTestTransaction :: (Ord k, Arbitrary k, Arbitrary v) + => TestTransaction k v + -> [TestTransaction k v] +shrinkTestTransaction (TestTransaction actions) = map TestTransaction (init (inits actions)) + +genTestSequence :: (Ord k, Arbitrary k, Arbitrary v) => Gen (TestSequence k v) +genTestSequence = sized $ \n -> do + k <- choose (0, n) + (_, txs) <- execStateT (replicateM k next) (M.empty, []) + return $ TestSequence (reverse txs) + where + next :: (Ord k, Arbitrary k, Arbitrary v) + => StateT (Map k v, [TestTransaction k v]) Gen () + next = do + (m, txs) <- get + (tx, m') <- lift $ genTransactionSetup >>= genTestTransaction m + put (m', tx:txs) + +shrinkTestSequence :: (Ord k, Arbitrary k, Arbitrary v) + => TestSequence k v + -> [TestSequence k v] +shrinkTestSequence (TestSequence txs) = map TestSequence (shrinkList shrinkTestTransaction txs) diff --git a/tests/Properties.hs b/tests/Properties.hs new file mode 100644 --- /dev/null +++ b/tests/Properties.hs @@ -0,0 +1,31 @@ +module Main (main) where + +import Test.Framework (Test, defaultMain) + +import qualified Properties.Impure.Fold +import qualified Properties.Impure.Insert +import qualified Properties.Impure.Structures +import qualified Properties.Primitives.Height +import qualified Properties.Primitives.Ids +import qualified Properties.Primitives.Index +import qualified Properties.Primitives.Leaf +import qualified Properties.Pure + +-------------------------------------------------------------------------------- + +tests :: [Test] +tests = + [ Properties.Impure.Fold.tests + , Properties.Impure.Insert.tests + , Properties.Impure.Structures.tests + , Properties.Primitives.Height.tests + , Properties.Primitives.Ids.tests + , Properties.Primitives.Index.tests + , Properties.Primitives.Leaf.tests + , Properties.Pure.tests + ] + +main :: IO () +main = defaultMain tests + +-------------------------------------------------------------------------------- diff --git a/tests/Properties/Impure/Fold.hs b/tests/Properties/Impure/Fold.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Impure/Fold.hs @@ -0,0 +1,25 @@ +module Properties.Impure.Fold where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) + +import Control.Monad ((>=>)) + +import Data.Int +import qualified Data.Map as M + +import Data.BTree.Alloc.Debug +import Data.BTree.Impure.Insert +import qualified Data.BTree.Impure as Tree + +tests :: Test +tests = testGroup "Impure.Fold" + [ testProperty "foldable toList fromList" prop_foldable_toList_fromList + ] + +prop_foldable_toList_fromList :: [(Int64, Integer)] -> Bool +prop_foldable_toList_fromList kvs + | (v, _) <- runDebug emptyPages $ + foldl (>=>) return (map (uncurry insertTree) kvs) Tree.empty + >>= Tree.toList + = v == M.toList (M.fromList kvs) diff --git a/tests/Properties/Impure/Insert.hs b/tests/Properties/Impure/Insert.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Impure/Insert.hs @@ -0,0 +1,44 @@ +module Properties.Impure.Insert where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) + +import Control.Monad ((>=>)) + +import Data.Int +import Data.Word (Word8) +import qualified Data.Map as M + +import Data.BTree.Alloc.Debug +import Data.BTree.Impure.Insert +import qualified Data.BTree.Impure as Tree + +tests :: Test +tests = testGroup "Impure.Insert" + [ testProperty "insertTreeMany" prop_insertTreeMany + , testProperty "insertOverflows" prop_insertOverflows + ] + +prop_insertTreeMany :: [(Int64, Integer)] -> [(Int64, Integer)] -> Bool +prop_insertTreeMany xs ys = ty1 == ty2 + where + tx = insertAll xs Tree.empty + + ty1 = evalDebug emptyPages $ + tx + >>= insertAll ys + >>= Tree.toList + + ty2 = evalDebug emptyPages $ + tx + >>= insertTreeMany (M.fromList ys) + >>= Tree.toList + + insertAll kvs = foldl (>=>) return (map (uncurry insertTree) kvs) + +prop_insertOverflows :: M.Map Int64 [Word8] -> Bool +prop_insertOverflows kvs + | v <- evalDebug emptyPages $ + insertTreeMany kvs Tree.empty + >>= Tree.toList + = v == M.toList kvs diff --git a/tests/Properties/Impure/Structures.hs b/tests/Properties/Impure/Structures.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Impure/Structures.hs @@ -0,0 +1,78 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +module Properties.Impure.Structures where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck + +import Control.Applicative ((<$>), (<*>)) + +import Data.Binary.Get (runGet) +import Data.Binary.Put (runPut) +import Data.Int +import Data.Typeable +import qualified Data.Binary as B + +import Data.BTree.Impure.Structures +import Data.BTree.Primitives + +import Properties.Primitives.Height (genNonZeroHeight) +import Properties.Primitives.Index () -- Arbitrary instance of Index +import Properties.Primitives.Ids () -- Arbitrary instance of NodeId + +tests :: Test +tests = testGroup "Impure.Structures" + [ testProperty "binary leafValue" prop_binary_leafValue + , testProperty "binary leafNode" prop_binary_leafNode + , testProperty "binary indexNode" prop_binary_indexNode + , testProperty "binary tree" prop_binary_tree + ] + +instance Arbitrary v => Arbitrary (LeafValue v) where + arbitrary = oneof [RawValue <$> arbitrary, OverflowValue <$> arbitrary] + +instance (Key k, Arbitrary k, Arbitrary v) => Arbitrary (Node 'Z k v) where + arbitrary = Leaf <$> arbitrary + +instance (Key k, Arbitrary k) => Arbitrary (Node ('S height) k v) where + arbitrary = Idx <$> arbitrary + +instance Arbitrary (Tree k v) where + arbitrary = Tree <$> arbitrary <*> arbitrary + +prop_binary_leafValue :: LeafValue Int64 -> Bool +prop_binary_leafValue xs = B.decode (B.encode xs) == xs + +prop_binary_leafNode :: Property +prop_binary_leafNode = forAll genLeafNode $ \leaf -> + runGet (getLeafNode zeroHeight) (runPut (putLeafNode leaf)) == leaf + +genLeafNode :: Gen (Node 'Z Int64 Bool) +genLeafNode = Leaf <$> arbitrary + +prop_binary_indexNode :: Property +prop_binary_indexNode = forAll genIndexNode $ \(h, idx) -> + runGet (getIndexNode h) (runPut (putIndexNode idx)) == idx + +genIndexNode :: Gen (Height ('S h), Node ('S h) Int64 Bool) +genIndexNode = do + h <- genNonZeroHeight + n <- Idx <$> arbitrary + return (h, n) + +prop_binary_tree :: Tree Int64 Bool -> Bool +prop_binary_tree t = B.decode (B.encode t) `treeEqShape` t + +-------------------------------------------------------------------------------- + +-- | Compare the shape of a 'Tree' structure +treeEqShape :: (Typeable key, Typeable val) + => Tree key val + -> Tree key val + -> Bool +Tree hx Nothing `treeEqShape` Tree hy Nothing = fromHeight hx == fromHeight hy +Tree hx (Just rx) `treeEqShape` Tree hy (Just ry) = + maybe False (== ry) $ castNode hx hy rx +Tree _ _ `treeEqShape` Tree _ _ = False diff --git a/tests/Properties/Primitives/Height.hs b/tests/Properties/Primitives/Height.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Primitives/Height.hs @@ -0,0 +1,28 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE StandaloneDeriving #-} +module Properties.Primitives.Height (tests, genNonZeroHeight) where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck + +import Data.BTree.Primitives.Height + +import Properties.Utils (testBinary) + +deriving instance Arbitrary (Height h) + +genNonZeroHeight :: Gen (Height h) +genNonZeroHeight = suchThat arbitrary $ \h -> case viewHeight h of + UZero -> False + USucc _ -> True + +tests :: Test +tests = testGroup "Primitives.Height" + [ testProperty "binary" prop_binary + ] + +prop_binary :: Height h -> Bool +prop_binary = testBinary diff --git a/tests/Properties/Primitives/Ids.hs b/tests/Properties/Primitives/Ids.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Primitives/Ids.hs @@ -0,0 +1,43 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE StandaloneDeriving #-} +module Properties.Primitives.Ids (tests) where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck + +import Control.Applicative ((<$>)) + +import Data.Int + +import Data.BTree.Primitives.Ids + +import Properties.Utils (testBinary) + +instance Arbitrary PageSize where + arbitrary = PageSize . fromIntegral <$> elements pows + where + -- minimum page size is 128 (fits at + -- least two keys in an index node) + pows = ((2 :: Int) ^) <$> ([7..12] :: [Int]) + +deriving instance Arbitrary (NodeId height key val) +deriving instance Arbitrary PageId +deriving instance Arbitrary TxId + +prop_binary_nodeId :: NodeId h Int64 Bool -> Bool +prop_binary_nodeId = testBinary + +prop_binary_pageId :: PageId -> Bool +prop_binary_pageId = testBinary + +prop_binary_txId :: TxId -> Bool +prop_binary_txId = testBinary + +tests :: Test +tests = testGroup "Primitives.Ids" + [ testProperty "binary nodeId" prop_binary_nodeId + , testProperty "binary pageId" prop_binary_pageId + , testProperty "binary txId" prop_binary_txId + ] diff --git a/tests/Properties/Primitives/Index.hs b/tests/Properties/Primitives/Index.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Primitives/Index.hs @@ -0,0 +1,122 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE RecordWildCards #-} +module Properties.Primitives.Index (tests) where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck + +import Control.Applicative ((<$>)) + +import Data.Int +import Data.List (nub) +import Data.List.Ordered (isSortedBy) +import Data.Maybe (isNothing) +import Data.Monoid ((<>)) +import qualified Data.Binary as B +import qualified Data.ByteString.Lazy as BL +import qualified Data.Foldable as F +import qualified Data.Map as M +import qualified Data.Vector as V + +import Data.BTree.Primitives.Ids +import Data.BTree.Primitives.Index +import Data.BTree.Primitives.Key +import Data.BTree.Pure.Setup + +import Properties.Primitives.Ids () -- Arbitrary instance of PageSize + +instance (Key k, Arbitrary k, Arbitrary v) => Arbitrary (Index k v) where + arbitrary = do + keys <- V.fromList . nub <$> orderedList + vals <- V.fromList <$> vector (V.length keys + 1) + return (Index keys vals) + shrink (Index keys vals) = + [ Index newKeys newVals + | k <- [0..V.length keys - 1] + , let (preKeys,sufKeys) = V.splitAt k keys + newKeys = preKeys <> V.drop 1 sufKeys + (preVals,sufVals) = V.splitAt k vals + newVals = preVals <> V.drop 1 sufVals + ] + +tests :: Test +tests = testGroup "Primitives.Index" + [ testProperty "binary" prop_binary + , testProperty "validIndex arbitrary" prop_validIndex_arbitrary + , testProperty "validIndex singletonIndex" prop_validIndex_singletonIndex + , testProperty "mergeIndex splitIndexAt" prop_mergeIndex_splitIndexAt + , testProperty "fromSingletonIndex singletonIndex" + prop_fromSingletonIndex_singletonIndex + , testProperty "distribute" prop_distribute + , testProperty "extendedIndex" prop_extendedIndex + , testProperty "extendIndexPred" prop_extendIndexPred + , testProperty "bindIndex_extendedIndex" prop_bindIndex_extendedIndex + ] + +prop_binary :: Index Int64 Bool -> Bool +prop_binary x = x == B.decode (B.encode x) + +prop_validIndex_arbitrary :: Index Int64 Bool -> Bool +prop_validIndex_arbitrary = validIndex + +prop_validIndex_singletonIndex :: Int64 -> Bool +prop_validIndex_singletonIndex i = + validIndex (singletonIndex i :: Index Int64 Int64) + +prop_mergeIndex_splitIndexAt :: Property +prop_mergeIndex_splitIndexAt = + forAll (arbitrary `suchThat` (isNothing . fromSingletonIndex)) $ \ix -> + and [ mergeIndex left middle right == (ix :: Index Int64 Bool) + | k <- [0..indexNumKeys ix - 1] + , let (left, middle, right) = splitIndexAt k ix + ] + +prop_fromSingletonIndex_singletonIndex :: Int64 -> Bool +prop_fromSingletonIndex_singletonIndex i = + fromSingletonIndex (singletonIndex i) == Just i + +prop_distribute :: M.Map Int64 Int -> Index Int64 Int -> Bool +prop_distribute kvs idx + | idx'@(Index keys vs) <- distribute kvs idx + , x <- V.all pred1 $ V.zip keys (V.init $ V.map fst vs) + , y <- V.all pred2 $ V.zip keys (V.tail $ V.map fst vs) + , z <- M.unions (V.toList $ V.map fst vs) == kvs + , u <- validIndex idx' + = x && y && z && u + where + pred1 (key, sub) = M.null sub || fst (M.findMax sub) < key + pred2 (key, sub) = M.null sub || fst (M.findMin sub) >= key + +prop_extendedIndex :: Index Int64 Int -> Bool +prop_extendedIndex idx + | Index keys idxs <- extendedIndex maxIdxKeys id idx + , numKeyIdxsOK <- V.length idxs == 1 + V.length keys + , validIdxs <- V.all validIndex idxs + , keysMaxOK <- V.all (\(key, Index keys' _) -> V.last keys' < key) $ V.zip keys idxs + , keysMinOK <- V.all (\(key, Index keys' _) -> V.head keys' > key) $ V.zip keys (V.tail idxs) + , keysOrderOK <- isSortedBy (<) (V.toList keys) + , joinedNodesOK <- concatMap F.toList (V.toList idxs) == F.toList idx + = numKeyIdxsOK && validIdxs && keysMaxOK && keysMinOK && keysOrderOK && joinedNodesOK + where + TreeSetup{..} = twoThreeSetup + +prop_extendIndexPred :: PageSize -> Index Int64 Int -> Bool +prop_extendIndexPred (PageSize pageSize) idx + | indexNumVals idx <= 2 + = True + | Just (Index keys idxs) <- extendIndexPred pred' id idx + , numKeyIdxsOK <- V.length idxs == 1 + V.length keys + , validIdxs <- V.all validIndex idxs + , keysMaxOK <- V.all (\(key, Index keys' _) -> V.last keys' < key) $ V.zip keys idxs + , keysOrderOK <- isSortedBy (<) (V.toList keys) + , joinedNodesOK <- concatMap F.toList (V.toList idxs) == F.toList idx + = numKeyIdxsOK && validIdxs && keysMaxOK && keysOrderOK && joinedNodesOK + | otherwise + = False + where + pred' m' = BL.length (B.encode m') <= fromIntegral pageSize + +prop_bindIndex_extendedIndex :: Int -> Index Int64 Int -> Bool +prop_bindIndex_extendedIndex n idx = + bindIndex (extendedIndex (abs n + 1) id idx) id == idx diff --git a/tests/Properties/Primitives/Leaf.hs b/tests/Properties/Primitives/Leaf.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Primitives/Leaf.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE RecordWildCards #-} +module Properties.Primitives.Leaf (tests) where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) + +import Data.BTree.Primitives.Ids +import Data.BTree.Primitives.Index +import Data.BTree.Primitives.Leaf +import Data.BTree.Pure.Setup + +import Data.Int +import Data.List.Ordered (isSortedBy) +import qualified Data.Binary as B +import qualified Data.ByteString.Lazy as BL +import qualified Data.Map as M +import qualified Data.Vector as V + +import Properties.Primitives.Ids () -- Arbitrary instance of PageSize + +tests :: Test +tests = testGroup "Primitives.Leaf" + [ testProperty "splitLeafManyPred" prop_splitLeafManyPred + , testProperty "splitLeafMany" prop_splitLeafMany + ] + +prop_splitLeafManyPred :: PageSize -> M.Map Int64 Int -> Bool +prop_splitLeafManyPred (PageSize pageSize) m + | M.null m + = True + | Just (Index vkeys vitems) <- splitLeafManyPred pred' id m + , (keys, maps) <- (V.toList vkeys, V.toList vitems) + , numKeyMapsOK <- length maps == 1 + length keys + , predMapsOK <- all pred' maps && all ((>= 1) . M.size) maps + , keysMaxOK <- all (\(key, m') -> fst (M.findMax m') < key) $ zip keys maps + , keysMinOK <- all (\(key, m') -> fst (M.findMin m') >= key) $ zip keys (tail maps) + , keysOrderOK <- isSortedBy (<) keys + , joinedMapsOK <- M.unions maps == m + = numKeyMapsOK && predMapsOK && keysMaxOK && keysMinOK && keysOrderOK && joinedMapsOK + | otherwise + = False + where + pred' m' = BL.length (B.encode m') <= fromIntegral pageSize + +prop_splitLeafMany :: M.Map Int64 Int -> Bool +prop_splitLeafMany m + | M.size m <= maxLeafItems = True + | Index vkeys vitems <- splitLeafMany maxLeafItems id m + , (keys, maps) <- (V.toList vkeys, V.toList vitems) + , numKeyMapsOK <- length maps == 1 + length keys + , sizeMapsOK <- all (\m' -> M.size m' >= minLeafItems && M.size m' <= maxLeafItems) maps + , keysMaxOK <- all (\(key, m') -> fst (M.findMax m') < key) $ zip keys maps + , keysMinOK <- all (\(key, m') -> fst (M.findMin m') >= key) $ zip keys (tail maps) + , keysOrderOK <- isSortedBy (<) keys + , joinedMapsOK <- M.unions maps == m + = numKeyMapsOK && sizeMapsOK && keysMaxOK && keysMinOK && keysOrderOK && joinedMapsOK + where + TreeSetup{..} = twoThreeSetup + diff --git a/tests/Properties/Pure.hs b/tests/Properties/Pure.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Pure.hs @@ -0,0 +1,94 @@ +{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE LambdaCase #-} +module Properties.Pure where + +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck + +import Data.BTree.Primitives.Index +import Data.BTree.Primitives.Key +import Data.BTree.Pure +import qualified Data.BTree.Pure as Tree + +import Control.Applicative ((<$>)) + +import Data.Function (on) +import Data.Int +import Data.List (nubBy) +import Data.Monoid (Sum(..)) +import qualified Data.Foldable as F +import qualified Data.Map as M + +tests :: Test +tests = testGroup "Pure" + [ testProperty "foldable" prop_foldable + , testProperty "validTree fromList" prop_validTree_fromList + , testProperty "foldableToList fromList" prop_foldableToList_fromList + , testProperty "toList fromList" prop_toList_fromList + , testProperty "insertMany" prop_insertMany + , testProperty "insert insertMany" prop_insert_insertMany + , testProperty "lookup insert" prop_lookup_insert + ] + +instance (Key k, Arbitrary k, Arbitrary v) => Arbitrary (Tree.Tree k v) where + arbitrary = Tree.fromList testSetup <$> arbitrary + shrink = map (Tree.fromList testSetup) . shrink . Tree.toList + +prop_foldable :: [(Int64, Int)] -> Bool +prop_foldable xs = F.foldMap snd xs' == F.foldMap id (Tree.fromList testSetup xs') + where xs' = nubByFstEq . map(\x -> (fst x, Sum $ snd x)) $ xs + +prop_validTree_fromList :: [(Int64, Int)] -> Bool +prop_validTree_fromList xs = validTree (Tree.fromList testSetup xs) + +prop_foldableToList_fromList :: [(Int64, Int)] -> Bool +prop_foldableToList_fromList xs = + F.toList (Tree.fromList testSetup xs) == + F.toList (M.fromList xs) + +prop_toList_fromList :: [(Int64, Int)] -> Bool +prop_toList_fromList xs = + Tree.toList (Tree.fromList testSetup xs) == + M.toList (M.fromList xs) + +prop_insertMany :: [(Int64, Int)] -> [(Int64, Int)] -> Bool +prop_insertMany xs ys + | isValid <- validTree txy + , equiv <- Tree.toList txy == M.toList mxy + = isValid && equiv + where + (mx, my) = (M.fromList xs, M.fromList ys) + mxy = M.union mx my + ty = Tree.fromList testSetup ys + txy = Tree.insertMany mx ty + +prop_insert_insertMany :: M.Map Int64 Int -> Tree.Tree Int64 Int -> Bool +prop_insert_insertMany kvs t = + Tree.toList (Tree.insertMany kvs t) == + Tree.toList (foldl (flip $ uncurry Tree.insert) t (M.toList kvs)) + +prop_lookup_insert :: Int64 -> Int -> Tree.Tree Int64 Int -> Bool +prop_lookup_insert k v t = Tree.lookup k (Tree.insert k v t) == Just v + +nubByFstEq :: Eq a => [(a, b)] -> [(a, b)] +nubByFstEq = nubBy ((==) `on` fst) + +-- | Check whether a given tree is valid. +validTree :: Ord key => Tree key val -> Bool +validTree (Tree _ Nothing) = True +validTree (Tree setup (Just (Leaf items))) = M.size items <= maxLeafItems setup +validTree (Tree setup (Just (Idx idx))) = + validIndexSize 1 (maxIdxKeys setup) idx && F.all (validNode setup) idx + +-- | Check whether a (non-root) node is valid. +validNode :: Ord key => TreeSetup -> Node height key val -> Bool +validNode setup = \case + Leaf items -> M.size items >= minLeafItems setup && + M.size items <= maxLeafItems setup + Idx idx -> validIndexSize (minIdxKeys setup) (maxIdxKeys setup) idx && + F.all (validNode setup) idx + +testSetup :: TreeSetup +testSetup = twoThreeSetup diff --git a/tests/Properties/Utils.hs b/tests/Properties/Utils.hs new file mode 100644 --- /dev/null +++ b/tests/Properties/Utils.hs @@ -0,0 +1,6 @@ +module Properties.Utils where + +import qualified Data.Binary as B + +testBinary :: (Eq a, B.Binary a) => a -> Bool +testBinary x = B.decode (B.encode x) == x