btree (empty) → 0.1.0.0
raw patch · 11 files changed
+2321/−0 lines, 11 filesdep +basedep +btreedep +clocksetup-changed
Dependencies added: base, btree, clock, compact-mutable, containers, ghc-prim, hashable, prim-array, primitive, smallcheck, tasty, tasty-hunit, tasty-smallcheck, transformers
Files
- LICENSE +30/−0
- README.md +1/−0
- Setup.hs +2/−0
- bench/Main.hs +124/−0
- btree.cabal +68/−0
- src/BTree.hs +92/−0
- src/BTree/Array.hs +21/−0
- src/BTree/Compact.hs +633/−0
- src/BTree/Contractible.hs +573/−0
- src/BTree/Linear.hs +436/−0
- test/Spec.hs +341/−0
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright Andrew Martin (c) 2017++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 Andrew Martin 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.
+ README.md view
@@ -0,0 +1,1 @@+# b-plus-tree
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ bench/Main.hs view
@@ -0,0 +1,124 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE PartialTypeSignatures #-}+{-# LANGUAGE TypeFamilies #-}++{-# OPTIONS_GHC -O2 #-}++module Main+ ( main+ ) where++import qualified BTree.Linear as BTL+import qualified BTree.Compact as BTC+import Control.Monad+import Data.Primitive.Compact (Token,withToken)+import GHC.Prim+import System.Mem (performGC)+import Data.Hashable+import Data.Maybe+import System.Clock++-- this specialization does not seem to work.+-- {-# SPECIALIZE BTC.modifyWithM :: BTC.Context RealWorld c -> BTC.BTree RealWorld Int Int c -> Int -> (Maybe Int -> IO Int) -> IO (Int, BTC.BTree RealWorld Int Int c) #-}++main :: IO ()+main = do+ putStrLn "Starting benchmark suite"+ let multiplier = 5+ let total = 200000 * multiplier+ range = 10000000 * multiplier+ lookups = 100000 * multiplier+ putStrLn $ concat+ [ "This benchmark will insert "+ , show total+ , " numbers into a b-tree. The range of these "+ , " numbers is from 0 to "+ , show (range - 1)+ , ". Then, we try looking up the numbers from "+ , "0 to "+ , show (lookups - 1)+ ]+ replicateM_ 1 $ do+ buildStart <- getTime Monotonic+ (b,ctx) <- onHeapBTree total range+ buildEnd <- getTime Monotonic+ performGC+ start <- getTime Monotonic+ x <- lookupMany lookups b ctx+ end <- getTime Monotonic+ putStrLn ("Accumulated sum (not a benchmark): " ++ show x)+ putStrLn "On-heap tree, Amount of time taken to build: "+ putStrLn (showTimeSpec (diffTimeSpec buildEnd buildStart))+ putStrLn "On-heap tree, Amount of time taken for lookups: "+ putStrLn (showTimeSpec (diffTimeSpec end start))+ performGC+ withToken $ \token -> do+ buildStart <- getTime Monotonic+ (b,ctx) <- offHeapBTree token total range+ buildEnd <- getTime Monotonic+ performGC+ start <- getTime Monotonic+ x <- lookupManyOffHeap lookups b+ end <- getTime Monotonic+ putStrLn ("Accumulated sum (not a benchmark): " ++ show x)+ putStrLn "Off-heap tree, Amount of time taken to build: "+ putStrLn (showTimeSpec (diffTimeSpec buildEnd buildStart))+ putStrLn "Off-heap tree, Amount of time taken for lookups: "+ putStrLn (showTimeSpec (diffTimeSpec end start))+ +lookupMany :: Int -> BTL.BTree RealWorld Int Int -> BTL.Context RealWorld -> IO Int+lookupMany total b ctx = go 0 0+ where+ go !n !s = if n < total+ then do+ m <- BTL.lookup ctx b n+ go (n + 1) (s + fromMaybe 0 m) + else return s++lookupManyOffHeap :: Int -> BTC.BTree Int Int RealWorld c -> IO Int+lookupManyOffHeap total b = go 0 0+ where+ go !n !s = if n < total+ then do+ m <- BTC.lookup b n+ go (n + 1) (s + fromMaybe 0 m) + else return s+ +onHeapBTree :: Int -> Int+ -> IO (BTL.BTree RealWorld Int Int, BTL.Context RealWorld)+onHeapBTree total range = do+ let ctx = BTL.Context 100+ b0 <- BTL.new ctx+ let go !n !b = if n < total+ then do+ let x = mod (hashWithSalt mySalt n) range+ b' <- BTL.insert ctx b x x+ go (n + 1) b'+ else return (b,ctx)+ go 0 b0++offHeapBTree ::+ Token c + -> Int+ -> Int+ -> IO (BTC.BTree Int Int RealWorld c, BTC.Context RealWorld c)+offHeapBTree token total range = do+ ctx <- BTC.newContext 100 token+ b0 <- BTC.new ctx+ let go !n !b = if n < total+ then do+ let x = mod (hashWithSalt mySalt n) range+ b' <- BTC.insert ctx b x x+ go (n + 1) b'+ else return (b,ctx)+ go 0 b0+++mySalt :: Int+mySalt = 2343++showTimeSpec :: TimeSpec -> String+showTimeSpec (TimeSpec s ns) = + show (fromIntegral s + (fromIntegral ns / 1000000000) :: Double)
+ btree.cabal view
@@ -0,0 +1,68 @@+name: btree+version: 0.1.0.0+synopsis: B-Tree on the compact heap+-- description:+homepage: https://github.com/andrewthad/b-plus-tree#readme+license: BSD3+license-file: LICENSE+author: Andrew Martin+maintainer: andrew.thaddeus@gmail.com+copyright: 2017 Andrew Martin+category: Web+build-type: Simple+extra-source-files: README.md+cabal-version: >=1.10++library+ hs-source-dirs: src+ exposed-modules:+ BTree+ BTree.Linear+ -- BTree.Generic+ BTree.Array+ BTree.Compact+ BTree.Contractible+ build-depends:+ base >= 4.10 && < 4.11+ , ghc-prim >= 0.5 && < 0.6+ , primitive >= 0.6.2 && < 0.7+ , prim-array >= 0.2 && < 0.3+ , compact-mutable >= 0.1 && < 0.2+ default-language: Haskell2010++test-suite test+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Spec.hs+ build-depends:+ base+ , btree+ , prim-array+ , tasty+ , tasty-smallcheck+ , tasty-hunit+ , smallcheck+ , containers+ , transformers+ , primitive+ , compact-mutable+ , hashable+ -- ghc-options: -threaded -rtsopts -with-rtsopts=-N+ default-language: Haskell2010++benchmark bench+ type: exitcode-stdio-1.0+ build-depends:+ base+ , btree+ , clock+ , hashable+ , ghc-prim+ , compact-mutable+ default-language: Haskell2010+ hs-source-dirs: bench+ main-is: Main.hs++source-repository head+ type: git+ location: https://github.com/andrewthad/b-plus-tree
+ src/BTree.hs view
@@ -0,0 +1,92 @@+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE BangPatterns #-}++{-# OPTIONS_GHC -Wall -Werror #-}++module BTree+ ( BTree+ , Context(..)+ , lookup+ , insert+ , modifyWithM+ , new+ , foldrWithKey+ , toAscList+ , fromList+ , debugMap+ ) where++import Prelude hiding (lookup)+import Data.Primitive hiding (fromList)+import Control.Monad.ST+import Data.Primitive.MutVar++import qualified BTree.Linear as BTL++data BTree s k v = BTree+ !(MutVar s (BTL.BTree s k v)) -- The actual B Tree+ !(Context s) -- Context++newtype Context s = Context (BTL.Context s)++new :: (Prim k, Prim v)+ => Context s -- ^ Max number of children per node+ -> ST s (BTree s k v)+new outerCtx@(Context ctx) = do+ rootRef <- newMutVar =<< BTL.new ctx+ return (BTree rootRef outerCtx)++lookup :: forall s k v. (Ord k, Prim k, Prim v)+ => BTree s k v -> k -> ST s (Maybe v)+lookup (BTree nodeRef (Context ctx)) k = do + root <- readMutVar nodeRef+ BTL.lookup ctx root k++insert :: (Ord k, Prim k, Prim v)+ => BTree s k v+ -> k+ -> v+ -> ST s ()+insert m k v = modifyWithM m k (\_ -> return v) >> return ()++-- | This is provided for completeness but is not something+-- typically useful in producetion code.+toAscList :: forall s k v. (Ord k, Prim k, Prim v)+ => BTree s k v+ -> ST s [(k,v)]+toAscList (BTree rootRef (Context ctx))+ = BTL.toAscList ctx =<< readMutVar rootRef++fromList :: (Ord k, Prim k, Prim v)+ => Context s -> [(k,v)] -> ST s (BTree s k v)+fromList ctxOuter@(Context ctx) xs = do+ root <- BTL.fromList ctx xs+ rootRef <- newMutVar root+ return (BTree rootRef ctxOuter)++foldrWithKey :: forall s k v b. (Ord k, Prim k, Prim v)+ => (k -> v -> b -> ST s b)+ -> b+ -> BTree s k v+ -> ST s b+foldrWithKey f b0 (BTree rootRef (Context ctx)) =+ readMutVar rootRef >>= BTL.foldrWithKey f b0 ctx++modifyWithM :: forall s k v. (Ord k, Prim k, Prim v)+ => BTree s k v+ -> k+ -> (Maybe v -> ST s v)+ -> ST s v+modifyWithM (BTree rootRef (Context ctx)) k alter = do+ root <- readMutVar rootRef+ (v,newRoot) <- BTL.modifyWithM ctx root k alter+ writeMutVar rootRef newRoot+ return v++debugMap :: forall s k v. (Prim k, Prim v, Show k, Show v)+ => BTree s k v+ -> ST s String+debugMap (BTree rootRef (Context ctx))+ = BTL.debugMap ctx =<< readMutVar rootRef+
+ src/BTree/Array.hs view
@@ -0,0 +1,21 @@+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE MagicHash #-}++module BTree.Array + (+ ) where++import Data.Kind (Type)+import Data.Primitive.PrimArray+import Data.Primitive.Array+import Data.Primitive.Compact+import Data.Primitive.Types+import Control.Monad.Primitive+import Data.Proxy+import Data.Primitive.Compact+import GHC.Prim+import GHC.Types+
+ src/BTree/Compact.hs view
@@ -0,0 +1,633 @@+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE UnboxedSums #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE DataKinds #-}++{-# OPTIONS_GHC -O2 -Wall -Werror -fno-warn-unused-imports #-}+++module BTree.Compact+ ( BTree+ , Context(..)+ , Sizing(..)+ , Decision(..)+ , new+ , newContext+ , debugMap+ , insert+ , modifyWithM+ , lookup+ , toAscList+ ) where++import Prelude hiding (lookup)+import Data.Primitive hiding (fromList)+import Control.Monad+import Data.Foldable (foldlM)+import Data.Primitive.Compact+import Data.Word+import Control.Monad.ST+import Control.Monad.Primitive+import GHC.Prim+import Data.Bits (unsafeShiftR)++import Data.Primitive.PrimRef+import Data.Primitive.PrimArray+import Data.Primitive.MutVar+import GHC.Ptr (Ptr(..))+import GHC.Int (Int(..))+import Numeric (showHex)++import qualified Data.List as L++-- One easy improvement I would like to make is to change+-- the way that sizing is being handled. Now that all of+-- the BTrees get serialized to bytearrays (and arrayarrays),+-- we should just be able to stick the size directly+-- into the BTree without doing the weird indirection trick.+-- The only tricky thing is that we will have to update the+-- size of a node on our way back up after an insertion.+-- This will required modifying the Insert data type.++data Context s c = Context+ { _contextDegree :: {-# UNPACK #-} !Int+ , _contextToken :: {-# UNPACK #-} !(Token c)+ , _contextSizing :: {-# UNPACK #-} !(MutVar s (Sizing s c))+ }++-- Use mkBTree instead. Using this for pattern matching is ok. +data BTree k v s (c :: Heap)+ = BTree+ {-# UNPACK #-} !(Sizing s c) -- block and index for current size+ {-# UNPACK #-} !(MutablePrimArray s k)+ {-# UNPACK #-} !(FlattenedContents k v s c)++-- In defining this instance, we make the assumption that an+-- Addr and an Int have the same size.+instance Contractible (BTree k v) where+ unsafeContractedUnliftedPtrCount# _ = 5#+ unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 2#+ readContractedArray# ba aa ix s1 =+ let ixByte = ix *# 2#+ ixPtr = ix *# 5#+ in case readIntArray# ba (ixByte +# 0#) s1 of+ (# s2, szIx #) -> case readIntArray# ba (ixByte +# 1#) s2 of+ (# s3, toggle #) -> case readMutableByteArrayArray# aa (ixPtr +# 0#) s3 of+ (# s4, szBlock #) -> case readMutableByteArrayArray# aa (ixPtr +# 1#) s4 of+ (# s5, keys #) -> case readMutableByteArrayArray# aa (ixPtr +# 2#) s5 of+ (# s6, values #) -> case readMutableByteArrayArray# aa (ixPtr +# 3#) s6 of+ (# s7, nodesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 4#) s7 of+ (# s8, nodesPtrs #) ->+ (# s8, (BTree (Sizing (I# szIx) (MutablePrimArray szBlock)) (MutablePrimArray keys) (FlattenedContents (I# toggle) (MutablePrimArray values) (ContractedMutableArray nodesBytes nodesPtrs))) #)+ + writeContractedArray# ba aa ix (BTree (Sizing (I# szIx) (MutablePrimArray szBlock)) (MutablePrimArray keys) (FlattenedContents (I# toggle) (MutablePrimArray values) (ContractedMutableArray nodesBytes nodesPtrs))) s1 =+ let ixByte = ix *# 2#+ ixPtr = ix *# 5#+ in case writeIntArray# ba (ixByte +# 0#) szIx s1 of+ s2 -> case writeIntArray# ba (ixByte +# 1#) toggle s2 of+ s3 -> case writeMutableByteArrayArray# aa (ixPtr +# 0#) szBlock s3 of+ s4 -> case writeMutableByteArrayArray# aa (ixPtr +# 1#) keys s4 of+ s5 -> case writeMutableByteArrayArray# aa (ixPtr +# 2#) values s5 of+ s6 -> case writeMutableByteArrayArray# aa (ixPtr +# 3#) nodesBytes s6 of+ s7 -> writeMutableArrayArrayArray# aa (ixPtr +# 4#) nodesPtrs s7+ ++data Sizing s (c :: Heap) = Sizing+ {-# UNPACK #-} !Int+ -- The array index does not live in the compact region+ {-# UNPACK #-} !(MutablePrimArray s Word16)+ -- This array must live in the compact region that the+ -- token in the Context refers to.++packedSizesCount :: Int+packedSizesCount = 2040++newContext :: (PrimMonad m) => Int -> Token c -> m (Context (PrimState m) c)+newContext deg token = do+ !sizes0 <- compactAddGeneral token =<< newPrimArray packedSizesCount+ let !sizing0 = Sizing 0 sizes0+ ref <- newMutVar sizing0+ return (Context deg token ref) -- newCompactArray' newKeyArray newValueArray)+++-- We manually flatten this sum type so that it can be unpacked+-- into BTree.+data FlattenedContents k v s c = FlattenedContents+ {-# UNPACK #-} !Int+ {-# UNPACK #-} !(MutablePrimArray s v)+ {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)++data Contents k v s c+ = ContentsValues {-# UNPACK #-} !(MutablePrimArray s v)+ | ContentsNodes {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)++{-# INLINE flattenContentsToContents #-}+flattenContentsToContents :: + FlattenedContents k v s c+ -> Contents k v s c+flattenContentsToContents (FlattenedContents i values nodes) =+ if i == 0+ then ContentsValues values+ else ContentsNodes nodes++-- | This one is a little trickier. We have to provide garbage+-- to fill in the unused slot.+{-# INLINE contentsToFlattenContents #-}+contentsToFlattenContents :: + MutablePrimArray s v -- ^ garbage value+ -> ContractedMutableArray (BTree k v) s c -- ^ garbage value+ -> Contents k v s c+ -> FlattenedContents k v s c+contentsToFlattenContents !garbageValues !garbageNodes !c = case c of+ ContentsValues values -> FlattenedContents 0 values garbageNodes+ ContentsNodes nodes -> FlattenedContents 1 garbageValues nodes ++-- | Get the nodes out, even if they are garbage. This is used+-- to get a garbage value when needed.+{-# INLINE demandFlattenedContentsNodes #-}+demandFlattenedContentsNodes :: FlattenedContents k v s c -> ContractedMutableArray (BTree k v) s c+demandFlattenedContentsNodes (FlattenedContents _ _ nodes) = nodes++data Insert k v s c+ = Ok !v+ | Split+ {-# NOUNPACK #-} !(BTree k v s c)+ !k+ !v+ {-# UNPACK #-} !(Sizing s c)+ -- ^ The new node that will go to the right,+ -- the key propagated to the parent,+ -- the inserted value, updated sizing info.++{-# INLINE mkBTree #-}+mkBTree :: PrimMonad m+ => Token c+ -> ContractedMutableArray (BTree k v) (PrimState m) c -- ^ garbage value+ -> Sizing (PrimState m) c+ -> MutablePrimArray (PrimState m) k -- ^ keys+ -> Contents k v (PrimState m) c+ -> m (BTree k v (PrimState m) c)+mkBTree token garbage a b c = do+ let !garbageValues = coercePrimArray b+ !bt = BTree a b (contentsToFlattenContents garbageValues garbage c)+ compactAddGeneral token bt++coercePrimArray :: MutablePrimArray s a -> MutablePrimArray s b+coercePrimArray (MutablePrimArray a) = MutablePrimArray a++new :: (PrimMonad m, Prim k, Prim v)+ => Context (PrimState m) c+ -> m (BTree k v (PrimState m) c)+new (Context !degree !token !szRef) = do+ if degree < 3+ then error "Btree.new: max nodes per child cannot be less than 3"+ else return ()+ !sizing0 <- readMutVar szRef+ writeNodeSize sizing0 0+ writeMutVar szRef =<< nextSizing token sizing0+ !keys <- newPrimArray (degree - 1)+ !values <- newPrimArray (degree - 1)+ -- it kind of pains me that this is needed, but since+ -- we only do it once when calling @new@, it should+ -- not hurt performance at all.+ !garbageNodes <- newContractedArray token 0+ mkBTree token garbageNodes sizing0 keys (ContentsValues values)++-- {-# SPECIALIZE lookup :: BTree RealWorld Int Int c -> Int -> IO (Maybe Int) #-}+{-# INLINABLE lookup #-}+lookup :: forall m k v c. (PrimMonad m, Ord k, Prim k, Prim v)+ => BTree k v (PrimState m) c -> k -> m (Maybe v)+lookup theNode k = go 0 theNode+ where+ go :: Int -> BTree k v (PrimState m) c -> m (Maybe v)+ go !n (BTree sizing@(Sizing _szIx _) keys c@(FlattenedContents _tog _ _)) = do+ sz <- readNodeSize sizing+ case flattenContentsToContents c of+ ContentsValues values -> do+ ix <- findIndex keys k sz+ if ix < 0+ then return Nothing+ else do+ v <- readPrimArray values ix+ return (Just v)+ ContentsNodes nodes -> do+ ix <- findIndexOfGtElem keys k sz+ !node <- readContractedArray nodes ix+ go (n + 1) node++_addrToPtr :: Addr -> Ptr Word8+_addrToPtr (Addr a) = Ptr a+++{-# INLINE insert #-}+insert :: (Ord k, Prim k, Prim v, PrimMonad m)+ => Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> k+ -> v+ -> m (BTree k v (PrimState m) c)+insert !ctx !m !k !v = do+ !(!_,!node) <- modifyWithM ctx m k v (\_ -> return (Replace v))+ return node++data Decision a = Keep | Replace !a++-- When we turn on this specialize pragma, it gets way faster+-- for the particular case.+{-# SPECIALIZE modifyWithM :: Context RealWorld c -> BTree Int Int RealWorld c -> Int -> Int -> (Int -> IO (Decision Int)) -> IO (Int, BTree Int Int RealWorld c) #-}+{-# INLINABLE modifyWithM #-}+modifyWithM :: forall m k v c. (Ord k, Prim k, Prim v, PrimMonad m)+ => Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> k+ -> v -- ^ value to insert if key not found+ -> (v -> m (Decision v)) -- ^ modification to value if key is found+ -> m (v, BTree k v (PrimState m) c)+modifyWithM (Context !degree !token !sizingRef) !root !k !newValue alter = do+ -- I believe I have been enlightened.+ !ins <- go root+ case ins of+ Ok v -> return (v,root)+ Split !rightNode newRootKey v sizing -> do+ writeNodeSize sizing 1+ newRootKeys <- newPrimArray (degree - 1)+ writePrimArray newRootKeys 0 newRootKey+ !newRootChildren <- newContractedArray token degree+ let !leftNode = root+ !newRoot@(BTree _ _ (FlattenedContents _ _ cmptRootChildren)) <- mkBTree token newRootChildren sizing newRootKeys (ContentsNodes newRootChildren)+ writeContractedArray cmptRootChildren 0 leftNode+ writeContractedArray cmptRootChildren 1 rightNode+ !newSizing <- nextSizing token sizing+ writeMutVar sizingRef newSizing+ return (v,newRoot)+ where+ go :: BTree k v (PrimState m) c -> m (Insert k v (PrimState m) c)+ go (BTree !szRef !keys !c) = do+ !sz <- readNodeSize szRef+ case flattenContentsToContents c of+ ContentsValues !values -> do+ !ix <- findIndex keys k sz+ if ix < 0+ then do+ let !gtIx = decodeGtIndex ix+ !v = newValue+ if sz < degree - 1+ then do+ -- We have enough space+ writeNodeSize szRef (sz + 1)+ unsafeInsertPrimArray sz gtIx k keys+ unsafeInsertPrimArray sz gtIx v values+ return (Ok v)+ else do+ -- We do not have enough space. The node must be split.+ let !leftSize = div sz 2+ !rightSize = sz - leftSize+ !leftKeys = keys+ !leftValues = values+ rightSzRef <- readMutVar sizingRef+ rightKeys <- newPrimArray (degree - 1)+ rightValues <- newPrimArray (degree - 1)+ if gtIx < leftSize+ then do+ writeNodeSize rightSzRef rightSize+ copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize+ copyMutablePrimArray rightValues 0 leftValues leftSize rightSize+ unsafeInsertPrimArray leftSize gtIx k leftKeys+ unsafeInsertPrimArray leftSize gtIx v leftValues+ writeNodeSize szRef (leftSize + 1)+ else do+ writeNodeSize rightSzRef (rightSize + 1)+ -- Currently, we're copying from left to right and+ -- then doing another copy from right to right. We+ -- might be able to do better. We could do the same number+ -- of memcpys but copy fewer total elements and not+ -- have the slowdown caused by overlap.+ copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize+ copyMutablePrimArray rightValues 0 leftValues leftSize rightSize+ unsafeInsertPrimArray rightSize (gtIx - leftSize) k rightKeys+ unsafeInsertPrimArray rightSize (gtIx - leftSize) v rightValues+ writeNodeSize szRef leftSize+ !propagated <- readPrimArray rightKeys 0+ !newSizing <- nextSizing token rightSzRef+ !newTree <- mkBTree token (demandFlattenedContentsNodes c) rightSzRef rightKeys (ContentsValues rightValues)+ return (Split newTree propagated v newSizing)+ else do+ !v <- readPrimArray values ix+ !dec <- alter v+ !v' <- case dec of+ Keep -> return v+ Replace v' -> writePrimArray values ix v' >> return v'+ return (Ok v')+ ContentsNodes nodes -> do+ !(!gtIx,!isEq) <- findIndexGte keys k sz+ -- case e of+ -- Right _ -> error "write Right case"+ -- Left gtIx -> do+ !node <- readContractedArray nodes (if isEq then gtIx + 1 else gtIx)+ !ins <- go node+ case ins of+ Ok !v -> return (Ok v)+ Split !rightNode !propagated !v !sizing -> if sz < degree - 1+ then do+ unsafeInsertPrimArray sz gtIx propagated keys+ unsafeInsertContractedArray (sz + 1) (gtIx + 1) rightNode nodes+ writeNodeSize szRef (sz + 1)+ writeMutVar sizingRef sizing+ return (Ok v)+ else do+ let !middleIx = div sz 2+ !leftKeys = keys+ !leftNodes = nodes+ !rightSzRef = sizing+ !middleKey <- readPrimArray keys middleIx+ !rightKeysOnHeap <- newPrimArray (degree - 1)+ !rightNodes' <- newContractedArray token degree -- uninitializedNode+ !x@(BTree _ rightKeys (FlattenedContents _ _ rightNodes)) <- mkBTree token rightNodes' rightSzRef rightKeysOnHeap (ContentsNodes rightNodes')+ let !leftSize = middleIx+ !rightSize = sz - leftSize+ if middleIx >= gtIx+ then do+ copyMutablePrimArray rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyContractedMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize+ unsafeInsertPrimArray leftSize gtIx propagated leftKeys+ unsafeInsertContractedArray (leftSize + 1) (gtIx + 1) rightNode leftNodes+ writeNodeSize szRef (leftSize + 1)+ writeNodeSize rightSzRef (rightSize - 1)+ else do+ -- Currently, we're copying from left to right and+ -- then doing another copy from right to right. We can do better.+ -- There is a similar note further up.+ copyMutablePrimArray rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyContractedMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize+ unsafeInsertPrimArray (rightSize - 1) (gtIx - leftSize - 1) propagated rightKeys+ unsafeInsertContractedArray rightSize (gtIx - leftSize) rightNode rightNodes+ writeNodeSize szRef leftSize+ writeNodeSize rightSzRef rightSize+ !newSizing <- nextSizing token rightSzRef+ return (Split x middleKey v newSizing)++-- Preconditions:+-- * marr is sorted low to high+-- * sz is less than or equal to the true size of marr+-- The returned value is either+-- * in the inclusive range [0,sz - 1], indicates a match+-- * a negative number x, indicates that the first greater+-- element is found at index ((negate x) - 1)+findIndex :: forall m a. (PrimMonad m, Ord a, Prim a)+ => MutablePrimArray (PrimState m) a+ -> a+ -> Int+ -> m Int -- (Either Int Int)+findIndex !marr !needle !sz = go 0+ where+ go :: Int -> m Int+ go !i = if i < sz+ then do+ !a <- readPrimArray marr i+ case compare a needle of+ LT -> go (i + 1)+ EQ -> return i+ GT -> return (encodeGtIndex i)+ else return (encodeGtIndex i)++{-# INLINE encodeGtIndex #-}+encodeGtIndex :: Int -> Int+encodeGtIndex i = negate i - 1++{-# INLINE decodeGtIndex #-}+decodeGtIndex :: Int -> Int+decodeGtIndex x = negate x - 1++-- | The second value in the tuple is true when+-- the index match was exact.+findIndexGte :: forall m a. (Ord a, Prim a, PrimMonad m)+ => MutablePrimArray (PrimState m) a -> a -> Int -> m (Int,Bool)+findIndexGte !marr !needle !sz = go 0+ where+ go :: Int -> m (Int,Bool)+ go !i = if i < sz+ then do+ !a <- readPrimArray marr i+ case compare a needle of+ LT -> go (i + 1)+ EQ -> return (i,True)+ GT -> return (i,False)+ else return (i,False)++-- | This is a linear-cost search in an sorted array.+-- findIndexBetween :: forall m a. (PrimMonad m, Ord a, Prim a)+-- => MutablePrimArray (PrimState m) a -> a -> Int -> m Int+-- findIndexBetween !marr !needle !sz = go 0+-- where+-- go :: Int -> m Int+-- go !i = if i < sz+-- then do+-- a <- readPrimArray marr i+-- if a > needle+-- then return i+-- else go (i + 1)+-- else return i -- i should be equal to sz++-- Inserts a value at the designated index,+-- shifting everything after it to the right.+--+-- Example:+-- -----------------------------+-- | a | b | c | d | e | X | X |+-- -----------------------------+-- unsafeInsertPrimArray 5 3 'k' marr+--+unsafeInsertPrimArray ::+ (Prim a, PrimMonad m)+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> a -- ^ Value+ -> MutablePrimArray (PrimState m) a -- ^ Array to modify+ -> m ()+unsafeInsertPrimArray !sz !i !x !marr = do+ copyMutablePrimArray marr (i + 1) marr i (sz - i)+ writePrimArray marr i x++debugMap :: forall m k v c. (Prim k, Prim v, Show k, Show v, PrimMonad m)+ => Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> m String+-- debugMap (Context _ _ _ _ _) BTreeUnused = return "BTreeUnused, should not happen"+debugMap (Context _ _ _) (BTree !rootSizing !rootKeys !rootContents) = do+ !rootSz <- readNodeSize rootSizing+ let go :: Int -> Int -> MutablePrimArray (PrimState m) k -> FlattenedContents k v (PrimState m) c -> m [(Int,String)]+ go level sz keys c = case flattenContentsToContents c of+ ContentsValues values -> do+ pairStrs <- showPairs sz keys values+ return (map (\s -> (level,s)) pairStrs)+ ContentsNodes nodes -> do+ pairs <- pairForM sz keys nodes+ $ \k (BTree theNextSizing nextKeys nextContents) -> do+ nextSz <- readNodeSize theNextSizing+ nextStrs <- go (level + 1) nextSz nextKeys nextContents+ return (nextStrs ++ [(level,show k)]) -- ++ " (Size: " ++ show nextSz ++ ")")])+ -- I think this should always end up being in bounds+ BTree lastSizing lastKeys lastContents <- readContractedArray nodes sz+ lastSz <- readNodeSize lastSizing+ lastStrs <- go (level + 1) lastSz lastKeys lastContents+ -- return (nextStrs ++ [(level,show k)])+ return ([(level, "start")] ++ concat pairs ++ lastStrs)+ allStrs <- go 0 rootSz rootKeys rootContents+ return $ unlines $ map (\(level,str) -> replicate (level * 2) ' ' ++ str) ((0,"root size: " ++ show rootSz) : allStrs)++pairForM :: forall m a b c d. (Prim a, PrimMonad m, Contractible d)+ => Int + -> MutablePrimArray (PrimState m) a + -> ContractedMutableArray d (PrimState m) c+ -> (a -> d (PrimState m) c -> m b)+ -> m [b]+pairForM sz marr1 marr2 f = go 0+ where+ go :: Int -> m [b]+ go ix = if ix < sz+ then do+ a <- readPrimArray marr1 ix+ c <- readContractedArray marr2 ix+ b <- f a c+ bs <- go (ix + 1)+ return (b : bs)+ else return []++showPairs :: forall m k v. (PrimMonad m, Show k, Show v, Prim k, Prim v)+ => Int -- size+ -> MutablePrimArray (PrimState m) k+ -> MutablePrimArray (PrimState m) v+ -> m [String]+showPairs sz keys values = go 0+ where+ go :: Int -> m [String]+ go ix = if ix < sz+ then do+ k <- readPrimArray keys ix+ v <- readPrimArray values ix+ let str = show k ++ ": " ++ show v+ strs <- go (ix + 1)+ return (str : strs)+ else return []++-- | This is provided for completeness but is not something+-- typically useful in production code.+toAscList :: forall m k v c. (PrimMonad m, Ord k, Prim k, Prim v)+ => Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> m [(k,v)]+toAscList = foldrWithKey f []+ where+ f :: k -> v -> [(k,v)] -> m [(k,v)]+ f k v xs = return ((k,v) : xs)++readNodeSize :: PrimMonad m => Sizing (PrimState m) c -> m Int+readNodeSize (Sizing ix m) = do+ w16 <- readPrimArray m ix+ return (fromIntegral w16)++nextSizing :: PrimMonad m => Token c -> Sizing (PrimState m) c -> m (Sizing (PrimState m) c)+nextSizing !token (Sizing !ix !marr) =+ if ix < packedSizesCount - 1+ then return (Sizing (ix + 1) marr)+ else do+ marr' <- compactAddGeneral token =<< newPrimArray packedSizesCount+ return (Sizing 0 marr')++writeNodeSize :: PrimMonad m => Sizing (PrimState m) c -> Int -> m ()+writeNodeSize (Sizing ix m) sz = writePrimArray m ix (fromIntegral sz)++foldrWithKey :: forall m k v b c. (PrimMonad m, Ord k, Prim k, Prim v)+ => (k -> v -> b -> m b)+ -> b+ -> Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> m b+foldrWithKey f b0 (Context _ _ _) root = flip go b0 root+ where+ go :: BTree k v (PrimState m) c -> b -> m b+ -- go BTreeUnused !b = return b -- should not happen+ go (BTree sizing keys c) !b = do+ sz <- readNodeSize sizing+ case flattenContentsToContents c of+ ContentsValues values -> foldrPrimArrayPairs sz f b keys values+ ContentsNodes nodes -> foldrArray (sz + 1) go b nodes++foldrPrimArrayPairs :: forall m k v b. (PrimMonad m, Ord k, Prim k, Prim v)+ => Int -- ^ length of arrays+ -> (k -> v -> b -> m b)+ -> b+ -> MutablePrimArray (PrimState m) k+ -> MutablePrimArray (PrimState m) v+ -> m b+foldrPrimArrayPairs len f b0 ks vs = go (len - 1) b0+ where+ go :: Int -> b -> m b+ go !ix !b1 = if ix >= 0+ then do+ k <- readPrimArray ks ix+ v <- readPrimArray vs ix+ b2 <- f k v b1+ go (ix - 1) b2+ else return b1++foldrArray :: forall m a b (c :: Heap). (PrimMonad m, Contractible a)+ => Int -- ^ length of array+ -> (a (PrimState m) c -> b -> m b)+ -> b+ -> ContractedMutableArray a (PrimState m) c+ -> m b+foldrArray len f b0 arr = go (len - 1) b0+ where+ go :: Int -> b -> m b+ go !ix !b1 = if ix >= 0+ then do+ a <- readContractedArray arr ix+ b2 <- f a b1+ go (ix - 1) b2+ else return b1++-- | This lookup is O(log n). It provides the index of the+-- first element greater than the argument.+-- Precondition, the array provided is sorted low to high.+{-# INLINABLE findIndexOfGtElem #-}+findIndexOfGtElem :: forall m a. (Ord a, Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> a -> Int -> m Int+findIndexOfGtElem v needle sz = go 0 (sz - 1)+ where+ go :: Int -> Int -> m Int+ go !lo !hi = if lo <= hi+ then do+ let !mid = lo + half (hi - lo)+ !val <- readPrimArray v mid+ if | val == needle -> return (mid + 1)+ | val < needle -> go (mid + 1) hi+ | otherwise -> go lo (mid - 1)+ else return lo++-- -- | This lookup is O(log n). It provides the index of the+-- -- match, or it returns (-1) to indicate that there was+-- -- no match.+-- {-# INLINABLE lookupSorted #-}+-- lookupSorted :: forall m a. (Ord a, Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> a -> m Int+-- lookupSorted v needle = do+-- sz <- getSizeofMutablePrimArray v+-- go (-1) 0 (sz - 1)+-- where+-- go :: Int -> Int -> Int -> m Int+-- go !result !lo !hi = if lo <= hi+-- then do+-- let !mid = lo + half (hi - lo)+-- !val <- readPrimArray v mid+-- if | val == needle -> go mid lo (mid - 1)+-- | val < needle -> go result (mid + 1) hi+-- | otherwise -> go result lo (mid - 1)+-- else return result++{-# INLINE half #-}+half :: Int -> Int+half x = unsafeShiftR x 1
+ src/BTree/Contractible.hs view
@@ -0,0 +1,573 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE UnboxedSums #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE DataKinds #-}++{-# OPTIONS_GHC -O2 -Wall -Werror -fno-warn-unused-imports #-}++module BTree.Contractible+ ( BTree+ , Context+ , Decision(..)+ , new+ , newContext+ , modifyWithM+ , lookup+ , toAscList+ ) where++import Prelude hiding (lookup)+import Data.Primitive hiding (fromList)+import Control.Monad+import Data.Foldable (foldlM)+import Data.Primitive.Compact+import Data.Word+import Control.Monad.ST+import Control.Monad.Primitive+import GHC.Prim+import Data.Bits (unsafeShiftR)++import BTree.Compact (Context(..),Sizing(..),Decision(..))+import Data.Primitive.PrimRef+import Data.Primitive.PrimArray+import Data.Primitive.MutVar+import GHC.Ptr (Ptr(..))+import GHC.Int (Int(..))+import Numeric (showHex)++import qualified Data.List as L++-- One easy improvement I would like to make is to change+-- the way that sizing is being handled. Now that all of+-- the BTrees get serialized to bytearrays (and arrayarrays),+-- we should just be able to stick the size directly+-- into the BTree without doing the weird indirection trick.+-- The only tricky thing is that we will have to update the+-- size of a node on our way back up after an insertion.+-- This will required modifying the Insert data type.++-- data Context s c = Context+-- { _contextDegree :: {-# UNPACK #-} !Int+-- , _contextToken :: {-# UNPACK #-} !(Token c)+-- , _contextSizing :: {-# UNPACK #-} !(MutVar s (Sizing s c))+-- }++-- Use mkBTree instead. Using this for pattern matching is ok. +data BTree (k :: *) (v :: * -> Heap -> *) (s :: *) (c :: Heap)+ = BTree+ {-# UNPACK #-} !(Sizing s c) -- block and index for current size+ {-# UNPACK #-} !(MutablePrimArray s k)+ {-# UNPACK #-} !(FlattenedContents k v s c)++-- In defining this instance, we make the assumption that an+-- Addr and an Int have the same size.+instance Contractible (BTree k v) where+ unsafeContractedUnliftedPtrCount# _ = 6#+ unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 2#+ readContractedArray# ba aa ix s1 =+ let ixByte = ix *# 2#+ ixPtr = ix *# 6#+ in case readIntArray# ba (ixByte +# 0#) s1 of+ (# s2, szIx #) -> case readIntArray# ba (ixByte +# 1#) s2 of+ (# s3, toggle #) -> case readMutableByteArrayArray# aa (ixPtr +# 0#) s3 of+ (# s4, szBlock #) -> case readMutableByteArrayArray# aa (ixPtr +# 1#) s4 of+ (# s5, keys #) -> case readMutableByteArrayArray# aa (ixPtr +# 2#) s5 of+ (# s6, valuesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 3#) s6 of+ (# s7, valuesPtrs #) -> case readMutableByteArrayArray# aa (ixPtr +# 4#) s7 of+ (# s8, nodesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 5#) s8 of+ (# s9, nodesPtrs #) ->+ (# s9, (BTree (Sizing (I# szIx) (MutablePrimArray szBlock)) (MutablePrimArray keys) (FlattenedContents (I# toggle) (ContractedMutableArray valuesBytes valuesPtrs) (ContractedMutableArray nodesBytes nodesPtrs))) #)+ writeContractedArray# ba aa ix (BTree (Sizing (I# szIx) (MutablePrimArray szBlock)) (MutablePrimArray keys) (FlattenedContents (I# toggle) (ContractedMutableArray valuesBytes valuesPtrs) (ContractedMutableArray nodesBytes nodesPtrs))) s1 =+ let ixByte = ix *# 2#+ ixPtr = ix *# 6#+ in case writeIntArray# ba (ixByte +# 0#) szIx s1 of+ s2 -> case writeIntArray# ba (ixByte +# 1#) toggle s2 of+ s3 -> case writeMutableByteArrayArray# aa (ixPtr +# 0#) szBlock s3 of+ s4 -> case writeMutableByteArrayArray# aa (ixPtr +# 1#) keys s4 of+ s5 -> case writeMutableByteArrayArray# aa (ixPtr +# 2#) valuesBytes s5 of+ s6 -> case writeMutableArrayArrayArray# aa (ixPtr +# 3#) valuesPtrs s6 of+ s7 -> case writeMutableByteArrayArray# aa (ixPtr +# 4#) nodesBytes s7 of+ s8 -> writeMutableArrayArrayArray# aa (ixPtr +# 5#) nodesPtrs s8+ ++-- data Sizing s (c :: Heap) = Sizing+-- {-# UNPACK #-} !Int+-- -- The array index does not live in the compact region+-- {-# UNPACK #-} !(MutablePrimArray s Word16)+-- -- This array must live in the compact region that the+-- -- token in the Context refers to.++packedSizesCount :: Int+packedSizesCount = 2040++newContext :: (PrimMonad m) => Int -> Token c -> m (Context (PrimState m) c)+newContext deg token = do+ !sizes0 <- compactAddGeneral token =<< newPrimArray packedSizesCount+ let !sizing0 = Sizing 0 sizes0+ ref <- newMutVar sizing0+ return (Context deg token ref) -- newCompactArray' newKeyArray newValueArray)+++-- We manually flatten this sum type so that it can be unpacked+-- into BTree.+data FlattenedContents (k :: *) (v :: * -> Heap -> *) (s :: *) (c :: Heap) = FlattenedContents+ {-# UNPACK #-} !Int+ {-# UNPACK #-} !(ContractedMutableArray v s c)+ {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)++data Contents (k :: *) (v :: * -> Heap -> *) (s :: *) (c :: Heap)+ = ContentsValues {-# UNPACK #-} !(ContractedMutableArray v s c)+ | ContentsNodes {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)++{-# INLINE flattenContentsToContents #-}+flattenContentsToContents :: + FlattenedContents k v s c+ -> Contents k v s c+flattenContentsToContents (FlattenedContents i values nodes) =+ if i == 0+ then ContentsValues values+ else ContentsNodes nodes++-- | This one is a little trickier. We have to provide garbage+-- to fill in the unused slot.+{-# INLINE contentsToFlattenContents #-}+contentsToFlattenContents :: + ContractedMutableArray v s c -- ^ garbage value+ -> ContractedMutableArray (BTree k v) s c -- ^ garbage value+ -> Contents k v s c+ -> FlattenedContents k v s c+contentsToFlattenContents !garbageValues !garbageNodes !c = case c of+ ContentsValues values -> FlattenedContents 0 values garbageNodes+ ContentsNodes nodes -> FlattenedContents 1 garbageValues nodes ++-- | Get the nodes out, even if they are garbage. This is used+-- to get a garbage value when needed.+{-# INLINE demandFlattenedContentsNodes #-}+demandFlattenedContentsNodes :: FlattenedContents k v s c -> ContractedMutableArray (BTree k v) s c+demandFlattenedContentsNodes (FlattenedContents _ _ nodes) = nodes++data Insert k (v :: * -> Heap -> *) s c+ = Ok !(v s c)+ | Split+ {-# NOUNPACK #-} !(BTree k v s c)+ !k+ !(v s c)+ {-# UNPACK #-} !(Sizing s c)+ -- ^ The new node that will go to the right,+ -- the key propagated to the parent,+ -- the inserted value, updated sizing info.++{-# INLINE mkBTree #-}+mkBTree :: PrimMonad m+ => Token c+ -> ContractedMutableArray (BTree k v) (PrimState m) c -- ^ garbage value+ -> Sizing (PrimState m) c+ -> MutablePrimArray (PrimState m) k -- ^ keys+ -> Contents k v (PrimState m) c+ -> m (BTree k v (PrimState m) c)+mkBTree token garbage a b c = do+ let !garbageValues = coerceContactedArray garbage+ !bt = BTree a b (contentsToFlattenContents garbageValues garbage c)+ compactAddGeneral token bt++coerceContactedArray :: ContractedMutableArray a s c -> ContractedMutableArray b s c+coerceContactedArray (ContractedMutableArray a b) = ContractedMutableArray a b++new :: (PrimMonad m, Prim k, Contractible v)+ => Context (PrimState m) c+ -> m (BTree k v (PrimState m) c)+new (Context !degree !token !szRef) = do+ if degree < 3+ then error "Btree.new: max nodes per child cannot be less than 3"+ else return ()+ !sizing0 <- readMutVar szRef+ writeNodeSize sizing0 0+ writeMutVar szRef =<< nextSizing token sizing0+ !keys <- newPrimArray (degree - 1)+ !values <- newContractedArray token (degree - 1)+ -- it kind of pains me that this is needed, but since+ -- we only do it once when calling @new@, it should+ -- not hurt performance at all.+ !garbageNodes <- newContractedArray token 0+ mkBTree token garbageNodes sizing0 keys (ContentsValues values)++-- {-# SPECIALIZE lookup :: BTree RealWorld Int Int c -> Int -> IO (Maybe Int) #-}+{-# INLINABLE lookup #-}+lookup :: forall m k v c. (PrimMonad m, Ord k, Prim k, Contractible v)+ => BTree k v (PrimState m) c -> k -> m (Maybe (v (PrimState m) c))+lookup theNode k = go theNode+ where+ go :: BTree k v (PrimState m) c -> m (Maybe (v (PrimState m) c))+ go (BTree sizing@(Sizing _szIx _) keys c@(FlattenedContents _tog _ _)) = do+ sz <- readNodeSize sizing+ case flattenContentsToContents c of+ ContentsValues values -> do+ ix <- findIndex keys k sz+ if ix < 0+ then return Nothing+ else do+ v <- readContractedArray values ix+ return (Just v)+ ContentsNodes nodes -> do+ ix <- findIndexOfGtElem keys k sz+ !node <- readContractedArray nodes ix+ go node++_addrToPtr :: Addr -> Ptr Word8+_addrToPtr (Addr a) = Ptr a++-- -- the insert function will always cause memory leaks+-- with this data structure.+-- {-# INLINE insert #-}+-- insert :: (Ord k, Prim k, Contractible v, PrimMonad m)+-- => Context (PrimState m) c+-- -> BTree k v (PrimState m) c+-- -> k+-- -> v (PrimState m) c+-- -> m (BTree k v (PrimState m) c)+-- insert !ctx !m !k !v = do+-- !(!_,!node) <- modifyWithM ctx m k (return v) (\_ -> return (Replace v))+-- return node++-- | Important note: if the key is not found the default value will+-- be created and then immidiately have the modify function+-- applied to it as well. This is unusual, but it matches the+-- common use cases for this data structure.+{-# INLINABLE modifyWithM #-}+modifyWithM :: forall m k (v :: * -> Heap -> *) (c :: Heap). (Ord k, Prim k, Contractible v, PrimMonad m)+ => Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> k+ -> m (v (PrimState m) c) -- ^ value to insert if key not found+ -> (v (PrimState m) c -> m (Decision (v (PrimState m) c))) -- ^ modification to value if key is found+ -> m (v (PrimState m) c, BTree k v (PrimState m) c)+modifyWithM (Context !degree !token !sizingRef) !root !k !newValue alter = do+ -- I believe I have been enlightened.+ !ins <- go root+ case ins of+ Ok v -> return (v,root)+ Split !rightNode newRootKey v sizing -> do+ writeNodeSize sizing 1+ newRootKeys <- newPrimArray (degree - 1)+ writePrimArray newRootKeys 0 newRootKey+ !newRootChildren <- newContractedArray token degree+ let !leftNode = root+ !newRoot@(BTree _ _ (FlattenedContents _ _ cmptRootChildren)) <- mkBTree token newRootChildren sizing newRootKeys (ContentsNodes newRootChildren)+ writeContractedArray cmptRootChildren 0 leftNode+ writeContractedArray cmptRootChildren 1 rightNode+ !newSizing <- nextSizing token sizing+ writeMutVar sizingRef newSizing+ return (v,newRoot)+ where+ go :: BTree k v (PrimState m) c -> m (Insert k v (PrimState m) c)+ go (BTree !szRef !keys !c) = do+ !sz <- readNodeSize szRef+ case flattenContentsToContents c of+ ContentsValues !values -> do+ !ix <- findIndex keys k sz+ if ix < 0+ then do+ let !gtIx = decodeGtIndex ix+ v <- newValue >>= \v0 -> alter v0 >>= \case+ Keep -> return v0+ Replace v1 -> return v1+ if sz < degree - 1+ then do+ -- We have enough space+ writeNodeSize szRef (sz + 1)+ unsafeInsertPrimArray sz gtIx k keys+ unsafeInsertContractedArray sz gtIx v values+ return (Ok v)+ else do+ -- We do not have enough space. The node must be split.+ let !leftSize = div sz 2+ !rightSize = sz - leftSize+ !leftKeys = keys+ !leftValues = values+ rightSzRef <- readMutVar sizingRef+ rightKeys' <- newPrimArray (degree - 1)+ rightValues' <- newContractedArray token (degree - 1)+ !newTree@(BTree _ rightKeys (FlattenedContents _ rightValues _))<- mkBTree token (demandFlattenedContentsNodes c) rightSzRef rightKeys' (ContentsValues rightValues')+ if gtIx < leftSize+ then do+ writeNodeSize rightSzRef rightSize+ copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize+ copyContractedMutableArray rightValues 0 leftValues leftSize rightSize+ unsafeInsertPrimArray leftSize gtIx k leftKeys+ unsafeInsertContractedArray leftSize gtIx v leftValues+ writeNodeSize szRef (leftSize + 1)+ else do+ writeNodeSize rightSzRef (rightSize + 1)+ -- Currently, we're copying from left to right and+ -- then doing another copy from right to right. We+ -- might be able to do better. We could do the same number+ -- of memcpys but copy fewer total elements and not+ -- have the slowdown caused by overlap.+ copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize+ copyContractedMutableArray rightValues 0 leftValues leftSize rightSize+ unsafeInsertPrimArray rightSize (gtIx - leftSize) k rightKeys+ unsafeInsertContractedArray rightSize (gtIx - leftSize) v rightValues+ writeNodeSize szRef leftSize+ !propagated <- readPrimArray rightKeys 0+ !newSizing <- nextSizing token rightSzRef+ return (Split newTree propagated v newSizing)+ else do+ !v <- readContractedArray values ix+ !dec <- alter v+ !v' <- case dec of+ Keep -> return v+ Replace v' -> writeContractedArray values ix v' >> return v'+ return (Ok v')+ ContentsNodes nodes -> do+ !(!gtIx,!isEq) <- findIndexGte keys k sz+ -- case e of+ -- Right _ -> error "write Right case"+ -- Left gtIx -> do+ !node <- readContractedArray nodes (if isEq then gtIx + 1 else gtIx)+ !ins <- go node+ case ins of+ Ok !v -> return (Ok v)+ Split !rightNode !propagated !v !sizing -> if sz < degree - 1+ then do+ unsafeInsertPrimArray sz gtIx propagated keys+ unsafeInsertContractedArray (sz + 1) (gtIx + 1) rightNode nodes+ writeNodeSize szRef (sz + 1)+ writeMutVar sizingRef sizing+ return (Ok v)+ else do+ let !middleIx = div sz 2+ !leftKeys = keys+ !leftNodes = nodes+ !rightSzRef = sizing+ !middleKey <- readPrimArray keys middleIx+ !rightKeysOnHeap <- newPrimArray (degree - 1)+ !rightNodes' <- newContractedArray token degree -- uninitializedNode+ !x@(BTree _ rightKeys (FlattenedContents _ _ rightNodes)) <- mkBTree token rightNodes' rightSzRef rightKeysOnHeap (ContentsNodes rightNodes')+ let !leftSize = middleIx+ !rightSize = sz - leftSize+ if middleIx >= gtIx+ then do+ copyMutablePrimArray rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyContractedMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize+ unsafeInsertPrimArray leftSize gtIx propagated leftKeys+ unsafeInsertContractedArray (leftSize + 1) (gtIx + 1) rightNode leftNodes+ writeNodeSize szRef (leftSize + 1)+ writeNodeSize rightSzRef (rightSize - 1)+ else do+ -- Currently, we're copying from left to right and+ -- then doing another copy from right to right. We can do better.+ -- There is a similar note further up.+ copyMutablePrimArray rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyContractedMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize+ unsafeInsertPrimArray (rightSize - 1) (gtIx - leftSize - 1) propagated rightKeys+ unsafeInsertContractedArray rightSize (gtIx - leftSize) rightNode rightNodes+ writeNodeSize szRef leftSize+ writeNodeSize rightSzRef rightSize+ !newSizing <- nextSizing token rightSzRef+ return (Split x middleKey v newSizing)++-- Preconditions:+-- * marr is sorted low to high+-- * sz is less than or equal to the true size of marr+-- The returned value is either+-- * in the inclusive range [0,sz - 1], indicates a match+-- * a negative number x, indicates that the first greater+-- element is found at index ((negate x) - 1)+findIndex :: forall m a. (PrimMonad m, Ord a, Prim a)+ => MutablePrimArray (PrimState m) a+ -> a+ -> Int+ -> m Int -- (Either Int Int)+findIndex !marr !needle !sz = go 0+ where+ go :: Int -> m Int+ go !i = if i < sz+ then do+ !a <- readPrimArray marr i+ case compare a needle of+ LT -> go (i + 1)+ EQ -> return i+ GT -> return (encodeGtIndex i)+ else return (encodeGtIndex i)++{-# INLINE encodeGtIndex #-}+encodeGtIndex :: Int -> Int+encodeGtIndex i = negate i - 1++{-# INLINE decodeGtIndex #-}+decodeGtIndex :: Int -> Int+decodeGtIndex x = negate x - 1++-- | The second value in the tuple is true when+-- the index match was exact.+findIndexGte :: forall m a. (Ord a, Prim a, PrimMonad m)+ => MutablePrimArray (PrimState m) a -> a -> Int -> m (Int,Bool)+findIndexGte !marr !needle !sz = go 0+ where+ go :: Int -> m (Int,Bool)+ go !i = if i < sz+ then do+ !a <- readPrimArray marr i+ case compare a needle of+ LT -> go (i + 1)+ EQ -> return (i,True)+ GT -> return (i,False)+ else return (i,False)++-- | This is a linear-cost search in an sorted array.+-- findIndexBetween :: forall m a. (PrimMonad m, Ord a, Prim a)+-- => MutablePrimArray (PrimState m) a -> a -> Int -> m Int+-- findIndexBetween !marr !needle !sz = go 0+-- where+-- go :: Int -> m Int+-- go !i = if i < sz+-- then do+-- a <- readPrimArray marr i+-- if a > needle+-- then return i+-- else go (i + 1)+-- else return i -- i should be equal to sz++-- Inserts a value at the designated index,+-- shifting everything after it to the right.+--+-- Example:+-- -----------------------------+-- | a | b | c | d | e | X | X |+-- -----------------------------+-- unsafeInsertPrimArray 5 3 'k' marr+--+unsafeInsertPrimArray ::+ (Prim a, PrimMonad m)+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> a -- ^ Value+ -> MutablePrimArray (PrimState m) a -- ^ Array to modify+ -> m ()+unsafeInsertPrimArray !sz !i !x !marr = do+ copyMutablePrimArray marr (i + 1) marr i (sz - i)+ writePrimArray marr i x++-- | This is provided for completeness but is not something+-- typically useful in production code. This function is+-- particularly worthless in this setting because the values+-- are mutable.+toAscList :: forall m k v c. (PrimMonad m, Ord k, Prim k, Contractible v)+ => Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> m [(k,v (PrimState m) c)]+toAscList = foldrWithKey f []+ where+ f :: k -> v (PrimState m) c -> [(k,v (PrimState m) c)] -> m [(k,v (PrimState m) c)]+ f k v xs = return ((k,v) : xs)++readNodeSize :: PrimMonad m => Sizing (PrimState m) c -> m Int+readNodeSize (Sizing ix m) = do+ w16 <- readPrimArray m ix+ return (fromIntegral w16)++nextSizing :: PrimMonad m => Token c -> Sizing (PrimState m) c -> m (Sizing (PrimState m) c)+nextSizing !token (Sizing !ix !marr) =+ if ix < packedSizesCount - 1+ then return (Sizing (ix + 1) marr)+ else do+ marr' <- compactAddGeneral token =<< newPrimArray packedSizesCount+ return (Sizing 0 marr')++writeNodeSize :: PrimMonad m => Sizing (PrimState m) c -> Int -> m ()+writeNodeSize (Sizing ix m) sz = writePrimArray m ix (fromIntegral sz)++foldrWithKey :: forall m k v b c. (PrimMonad m, Ord k, Prim k, Contractible v)+ => (k -> v (PrimState m) c -> b -> m b)+ -> b+ -> Context (PrimState m) c+ -> BTree k v (PrimState m) c+ -> m b+foldrWithKey f b0 (Context _ _ _) root = flip go b0 root+ where+ go :: BTree k v (PrimState m) c -> b -> m b+ go (BTree sizing keys c) !b = do+ sz <- readNodeSize sizing+ case flattenContentsToContents c of+ ContentsValues values -> foldrPrimArrayPairs sz f b keys values+ ContentsNodes nodes -> foldrArray (sz + 1) go b nodes++foldrPrimArrayPairs :: forall m k v b c. (PrimMonad m, Ord k, Prim k, Contractible v)+ => Int -- ^ length of arrays+ -> (k -> v (PrimState m) c -> b -> m b)+ -> b+ -> MutablePrimArray (PrimState m) k+ -> ContractedMutableArray v (PrimState m) c+ -> m b+foldrPrimArrayPairs len f b0 ks vs = go (len - 1) b0+ where+ go :: Int -> b -> m b+ go !ix !b1 = if ix >= 0+ then do+ k <- readPrimArray ks ix+ v <- readContractedArray vs ix+ b2 <- f k v b1+ go (ix - 1) b2+ else return b1++foldrArray :: forall m a b (c :: Heap). (PrimMonad m, Contractible a)+ => Int -- ^ length of array+ -> (a (PrimState m) c -> b -> m b)+ -> b+ -> ContractedMutableArray a (PrimState m) c+ -> m b+foldrArray len f b0 arr = go (len - 1) b0+ where+ go :: Int -> b -> m b+ go !ix !b1 = if ix >= 0+ then do+ a <- readContractedArray arr ix+ b2 <- f a b1+ go (ix - 1) b2+ else return b1++-- | This lookup is O(log n). It provides the index of the+-- first element greater than the argument.+-- Precondition, the array provided is sorted low to high.+{-# INLINABLE findIndexOfGtElem #-}+findIndexOfGtElem :: forall m a. (Ord a, Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> a -> Int -> m Int+findIndexOfGtElem v needle sz = go 0 (sz - 1)+ where+ go :: Int -> Int -> m Int+ go !lo !hi = if lo <= hi+ then do+ let !mid = lo + half (hi - lo)+ !val <- readPrimArray v mid+ if | val == needle -> return (mid + 1)+ | val < needle -> go (mid + 1) hi+ | otherwise -> go lo (mid - 1)+ else return lo++-- -- | This lookup is O(log n). It provides the index of the+-- -- match, or it returns (-1) to indicate that there was+-- -- no match.+-- {-# INLINABLE lookupSorted #-}+-- lookupSorted :: forall m a. (Ord a, Prim a, PrimMonad m) => MutablePrimArray (PrimState m) a -> a -> m Int+-- lookupSorted v needle = do+-- sz <- getSizeofMutablePrimArray v+-- go (-1) 0 (sz - 1)+-- where+-- go :: Int -> Int -> Int -> m Int+-- go !result !lo !hi = if lo <= hi+-- then do+-- let !mid = lo + half (hi - lo)+-- !val <- readPrimArray v mid+-- if | val == needle -> go mid lo (mid - 1)+-- | val < needle -> go result (mid + 1) hi+-- | otherwise -> go result lo (mid - 1)+-- else return result++{-# INLINE half #-}+half :: Int -> Int+half x = unsafeShiftR x 1
+ src/BTree/Linear.hs view
@@ -0,0 +1,436 @@+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE BangPatterns #-}++{-# OPTIONS_GHC -Wall -Werror -fno-warn-unused-imports #-}++module BTree.Linear+ ( BTree+ , Context(..)+ , lookup+ , insert+ , modifyWithM+ , new+ , foldrWithKey+ , toAscList+ , fromList+ , debugMap+ ) where++import Prelude hiding (lookup)+import Data.Primitive hiding (fromList)+import Data.Primitive.MutVar+import Control.Monad+import Data.Foldable (foldlM)++import Data.Primitive.PrimArray+import Control.Monad.Primitive++data Context s = Context+ { contextDegree :: {-# UNPACK #-} !Int+ }++data BTree s k v = BTree+ !(MutVar s Int) -- current number of keys in this node+ !(MutablePrimArray s k)+ !(Contents s k v)++data Contents s k v+ = ContentsValues !(MutablePrimArray s v)+ | ContentsNodes !(MutableArray s (BTree s k v))++new :: (PrimMonad m, Prim k, Prim v)+ => Context (PrimState m) -- ^ Max number of children per node+ -> m (BTree (PrimState m) k v)+new (Context degree) = do+ if degree < 3+ then error "Btree.new: max nodes per child cannot be less than 3"+ else return ()+ szRef <- newMutVar 0+ keys <- newPrimArray (degree - 1)+ values <- newPrimArray (degree - 1)+ return (BTree szRef keys (ContentsValues values))++{-# INLINABLE lookup #-}+lookup :: forall m k v. (PrimMonad m, Ord k, Prim k, Prim v)+ => Context (PrimState m) -> BTree (PrimState m) k v -> k -> m (Maybe v)+lookup (Context _) theNode k = go theNode+ where+ go :: BTree (PrimState m) k v -> m (Maybe v)+ go (BTree szRef keys c) = do+ sz <- readMutVar szRef+ case c of+ ContentsValues values -> do+ e <- findIndex keys k sz+ case e of+ Left _ -> return Nothing+ Right ix -> do+ v <- readPrimArray values ix+ return (Just v)+ ContentsNodes nodes -> do+ ix <- findIndexBetween keys k sz+ go =<< readArray nodes ix++data Insert s k v+ = Ok !v+ | Split !(BTree s k v) !k !v+ -- ^ The new node that will go to the right,+ -- the key propagated to the parent,+ -- the inserted value.++uninitializedNode :: a+uninitializedNode = error "unitializedNode: this should not be forced, b+ tree implementation has a mistake."++{-# INLINE insert #-}+insert :: (PrimMonad m, Ord k, Prim k, Prim v)+ => Context (PrimState m)+ -> BTree (PrimState m) k v+ -> k+ -> v+ -> m (BTree (PrimState m) k v)+insert ctx m k v = do+ (_,node) <- modifyWithM ctx m k (\_ -> return v)+ return node++-- | This is provided for completeness but is not something+-- typically useful in producetion code.+toAscList :: forall m k v. (PrimMonad m, Ord k, Prim k, Prim v)+ => Context (PrimState m)+ -> BTree (PrimState m) k v+ -> m [(k,v)]+toAscList = foldrWithKey f []+ where+ f :: k -> v -> [(k,v)] -> m [(k,v)]+ f k v xs = return ((k,v) : xs)++fromList :: (PrimMonad m, Ord k, Prim k, Prim v)+ => Context (PrimState m) -> [(k,v)] -> m (BTree (PrimState m) k v)+fromList ctx xs = do+ root0 <- new ctx+ foldlM+ (\root (k,v) -> do+ insert ctx root k v+ ) root0 xs++foldrWithKey :: forall m k v b. (PrimMonad m, Ord k, Prim k, Prim v)+ => (k -> v -> b -> m b)+ -> b+ -> Context (PrimState m)+ -> BTree (PrimState m) k v+ -> m b+foldrWithKey f b0 (Context _) root = flip go b0 root+ where+ go :: BTree (PrimState m) k v -> b -> m b+ go (BTree szRef keys c) b = do+ sz <- readMutVar szRef+ case c of+ ContentsValues values -> foldrPrimArrayPairs sz f b keys values+ ContentsNodes nodes -> foldrArray (sz + 1) go b nodes++foldrArray :: forall m a b. (PrimMonad m)+ => Int -- ^ length of array+ -> (a -> b -> m b)+ -> b+ -> MutableArray (PrimState m) a+ -> m b+foldrArray len f b0 arr = go (len - 1) b0+ where+ go :: Int -> b -> m b+ go !ix !b1 = if ix >= 0+ then do+ a <- readArray arr ix+ b2 <- f a b1+ go (ix - 1) b2+ else return b1++foldrPrimArrayPairs :: forall m k v b. (PrimMonad m, Ord k, Prim k, Prim v)+ => Int -- ^ length of arrays+ -> (k -> v -> b -> m b)+ -> b+ -> MutablePrimArray (PrimState m) k+ -> MutablePrimArray (PrimState m) v+ -> m b+foldrPrimArrayPairs len f b0 ks vs = go (len - 1) b0+ where+ go :: Int -> b -> m b+ go !ix !b1 = if ix >= 0+ then do+ k <- readPrimArray ks ix+ v <- readPrimArray vs ix+ b2 <- f k v b1+ go (ix - 1) b2+ else return b1++{-# SPECIALIZE modifyWithM :: Context RealWorld -> BTree RealWorld Int Int -> Int -> (Maybe Int -> IO Int) -> IO (Int, BTree RealWorld Int Int) #-}+{-# INLINABLE modifyWithM #-}+modifyWithM :: forall m s k v. (PrimMonad m, Ord k, Prim k, Prim v)+ => Context s+ -> BTree (PrimState m) k v+ -> k+ -> (Maybe v -> m v)+ -> m (v, BTree (PrimState m) k v)+modifyWithM (Context degree) root k alter = do+ ins <- go root+ case ins of+ Ok v -> return (v,root)+ Split rightNode newRootKey v -> do+ let leftNode = root+ newRootSz <- newMutVar 1+ newRootKeys <- newPrimArray (degree - 1)+ writePrimArray newRootKeys 0 newRootKey+ newRootChildren <- newArray degree uninitializedNode+ writeArray newRootChildren 0 leftNode+ writeArray newRootChildren 1 rightNode+ let newRoot = BTree newRootSz newRootKeys (ContentsNodes newRootChildren)+ return (v,newRoot)+ where+ go :: BTree (PrimState m) k v -> m (Insert (PrimState m) k v)+ go (BTree szRef keys c) = do+ sz <- readMutVar szRef+ case c of+ ContentsValues values -> do+ e <- findIndex keys k sz+ case e of+ Left gtIx -> do+ v <- alter Nothing+ if sz < degree - 1+ then do+ -- We have enough space+ writeMutVar szRef (sz + 1)+ unsafeInsertPrimArray sz gtIx k keys+ unsafeInsertPrimArray sz gtIx v values+ return (Ok v)+ else do+ -- We do not have enough space. The node must be split.+ let leftSize = div sz 2+ rightSize = sz - leftSize+ leftKeys = keys+ leftValues = values+ if gtIx < leftSize+ then do+ rightKeys <- newPrimArray (degree - 1)+ rightValues <- newPrimArray (degree - 1)+ rightSzRef <- newMutVar rightSize+ copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize+ copyMutablePrimArray rightValues 0 leftValues leftSize rightSize+ unsafeInsertPrimArray leftSize gtIx k leftKeys+ unsafeInsertPrimArray leftSize gtIx v leftValues+ propagated <- readPrimArray rightKeys 0+ writeMutVar szRef (leftSize + 1)+ return (Split (BTree rightSzRef rightKeys (ContentsValues rightValues)) propagated v)+ else do+ rightKeys <- newPrimArray (degree - 1)+ rightValues <- newPrimArray (degree - 1)+ rightSzRef <- newMutVar (rightSize + 1)+ -- Currently, we're copying from left to right and+ -- then doing another copy from right to right. We+ -- might be able to do better. We could do the same number+ -- of memcpys but copy fewer total elements and not+ -- have the slowdown caused by overlap.+ copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize+ copyMutablePrimArray rightValues 0 leftValues leftSize rightSize+ unsafeInsertPrimArray rightSize (gtIx - leftSize) k rightKeys+ unsafeInsertPrimArray rightSize (gtIx - leftSize) v rightValues+ propagated <- readPrimArray rightKeys 0+ writeMutVar szRef leftSize+ return (Split (BTree rightSzRef rightKeys (ContentsValues rightValues)) propagated v)+ Right ix -> do+ v <- readPrimArray values ix+ v' <- alter (Just v)+ writePrimArray values ix v'+ return (Ok v')+ ContentsNodes nodes -> do+ (gtIx,isEq) <- findIndexGte keys k sz+ -- case e of+ -- Right _ -> error "write Right case"+ -- Left gtIx -> do+ node <- readArray nodes (if isEq then gtIx + 1 else gtIx)+ ins <- go node+ case ins of+ Ok v -> return (Ok v)+ Split rightNode propagated v -> if sz < degree - 1+ then do+ unsafeInsertPrimArray sz gtIx propagated keys+ unsafeInsertArray (sz + 1) (gtIx + 1) rightNode nodes+ writeMutVar szRef (sz + 1)+ return (Ok v)+ else do+ let middleIx = div sz 2+ leftKeys = keys+ leftNodes = nodes+ middleKey <- readPrimArray keys middleIx+ rightKeys :: MutablePrimArray (PrimState m) k <- newPrimArray (degree - 1)+ rightNodes <- newArray degree uninitializedNode+ rightSzRef <- newMutVar 0 -- this always gets replaced+ let leftSize = middleIx+ rightSize = sz - leftSize+ if middleIx >= gtIx+ then do+ copyMutablePrimArray rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize+ unsafeInsertPrimArray leftSize gtIx propagated leftKeys+ unsafeInsertArray (leftSize + 1) (gtIx + 1) rightNode leftNodes+ writeMutVar szRef (leftSize + 1)+ writeMutVar rightSzRef (rightSize - 1)+ else do+ -- Currently, we're copying from left to right and+ -- then doing another copy from right to right. We can do better.+ -- There is a similar note further up.+ copyMutablePrimArray rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize+ unsafeInsertPrimArray (rightSize - 1) (gtIx - leftSize - 1) propagated rightKeys+ unsafeInsertArray rightSize (gtIx - leftSize) rightNode rightNodes+ writeMutVar szRef leftSize+ writeMutVar rightSzRef rightSize+ return (Split (BTree rightSzRef rightKeys (ContentsNodes rightNodes)) middleKey v)+ +-- Preconditions:+-- * marr is sorted low to high+-- * sz is less than or equal to the true size of marr+-- The returned value is in the inclusive range [0,sz]+findIndexBetween :: forall m a. (PrimMonad m, Ord a, Prim a)+ => MutablePrimArray (PrimState m) a -> a -> Int -> m Int+findIndexBetween !marr !needle !sz = go 0+ where+ go :: Int -> m Int+ go !i = if i < sz+ then do+ a <- readPrimArray marr i+ if a > needle+ then return i+ else go (i + 1)+ else return i -- i should be equal to sz++-- Preconditions:+-- * marr is sorted low to high+-- * sz is less than or equal to the true size of marr+-- The returned value is either+-- * in the inclusive range [0,sz - 1]+-- * the value (-1), indicating that no match was found+findIndex :: forall m a. (PrimMonad m, Ord a, Prim a)+ => MutablePrimArray (PrimState m) a -> a -> Int -> m (Either Int Int)+findIndex !marr !needle !sz = go 0+ where+ go :: Int -> m (Either Int Int)+ go !i = if i < sz+ then do+ a <- readPrimArray marr i+ case compare a needle of+ LT -> go (i + 1)+ EQ -> return (Right i)+ GT -> return (Left i)+ else return (Left i)++-- | The second value in the tuple is true when+-- the index match was exact.+findIndexGte :: forall m a. (PrimMonad m, Ord a, Prim a)+ => MutablePrimArray (PrimState m) a -> a -> Int -> m (Int,Bool)+findIndexGte !marr !needle !sz = go 0+ where+ go :: Int -> m (Int,Bool)+ go !i = if i < sz+ then do+ a <- readPrimArray marr i+ case compare a needle of+ LT -> go (i + 1)+ EQ -> return (i,True)+ GT -> return (i,False)+ else return (i,False)++-- | Insert an element in the array, shifting the values right +-- of the index. The array size should be big enough for this+-- shift, this is not checked.+unsafeInsertArray :: (PrimMonad m)+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> a -- ^ Value+ -> MutableArray (PrimState m) a -- ^ Array to modify+ -> m ()+unsafeInsertArray sz i x marr = do+ copyMutableArray marr (i + 1) marr i (sz - i)+ writeArray marr i x++-- Inserts a value at the designated index,+-- shifting everything after it to the right.+--+-- Example:+-- -----------------------------+-- | a | b | c | d | e | X | X |+-- -----------------------------+-- unsafeInsertPrimArray 5 3 'k' marr+--+unsafeInsertPrimArray ::+ (PrimMonad m, Prim a)+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> a -- ^ Value+ -> MutablePrimArray (PrimState m) a -- ^ Array to modify+ -> m ()+unsafeInsertPrimArray sz i x marr = do+ copyMutablePrimArray marr (i + 1) marr i (sz - i)+ writePrimArray marr i x+++showPairs :: forall m k v. (PrimMonad m, Show k, Show v, Prim k, Prim v)+ => Int -- size+ -> MutablePrimArray (PrimState m) k+ -> MutablePrimArray (PrimState m) v+ -> m [String]+showPairs sz keys values = go 0+ where+ go :: Int -> m [String]+ go ix = if ix < sz+ then do+ k <- readPrimArray keys ix+ v <- readPrimArray values ix+ let str = show k ++ ": " ++ show v+ strs <- go (ix + 1)+ return (str : strs)+ else return []++-- | Show the internal structure of a Map, useful for debugging, not exported+debugMap :: forall m k v. (PrimMonad m, Prim k, Prim v, Show k, Show v)+ => Context (PrimState m)+ -> BTree (PrimState m) k v+ -> m String+debugMap (Context _) (BTree rootSzRef rootKeys rootContents) = do+ rootSz <- readMutVar rootSzRef+ let go :: Int -> Int -> MutablePrimArray (PrimState m) k -> Contents (PrimState m) k v -> m [(Int,String)]+ go level sz keys c = case c of+ ContentsValues values -> do+ pairStrs <- showPairs sz keys values+ return (map (\s -> (level,s)) pairStrs)+ ContentsNodes nodes -> do+ pairs <- pairForM sz keys nodes+ $ \k (BTree nextSzRef nextKeys nextContents) -> do+ nextSz <- readMutVar nextSzRef+ nextStrs <- go (level + 1) nextSz nextKeys nextContents+ return (nextStrs ++ [(level,show k)]) -- ++ " (Size: " ++ show nextSz ++ ")")])+ -- I think this should always end up being in bounds+ BTree lastSzRef lastKeys lastContents <- readArray nodes sz+ lastSz <- readMutVar lastSzRef+ lastStrs <- go (level + 1) lastSz lastKeys lastContents+ -- return (nextStrs ++ [(level,show k)])+ return ([(level, "start")] ++ concat pairs ++ lastStrs)+ allStrs <- go 0 rootSz rootKeys rootContents+ return $ unlines $ map (\(level,str) -> replicate (level * 2) ' ' ++ str) ((0,"root size: " ++ show rootSz) : allStrs)++pairForM :: forall m a b c. (PrimMonad m, Prim a)+ => Int + -> MutablePrimArray (PrimState m) a + -> MutableArray (PrimState m) c+ -> (a -> c -> m b)+ -> m [b]+pairForM sz marr1 marr2 f = go 0+ where+ go :: Int -> m [b]+ go ix = if ix < sz+ then do+ a <- readPrimArray marr1 ix+ c <- readArray marr2 ix+ b <- f a c+ bs <- go (ix + 1)+ return (b : bs)+ else return []++
+ test/Spec.hs view
@@ -0,0 +1,341 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE PartialTypeSignatures #-}+{-# LANGUAGE LambdaCase #-}++{-# OPTIONS_GHC -Wno-partial-type-signatures #-}++import Test.Tasty+import Test.Tasty.SmallCheck as SC+import Test.SmallCheck.Series+import Test.Tasty.HUnit+import Data.List+import Data.Ord+import Control.Monad+import Control.Monad.ST+import Debug.Trace+import Control.Monad.Trans.Except+import Control.Monad.Trans.Class+import Data.Word+import Data.Int+import Data.Proxy+import Data.Primitive.Types+import Data.Foldable+import Data.Primitive.Compact (withToken,getSizeOfCompact)+import System.IO.Unsafe+import Data.Hashable++import qualified Data.List as L+import qualified Data.List.NonEmpty as NE+import qualified BTree as B+import qualified BTree.Linear as BTL+import qualified BTree.Compact as BTC+import qualified BTree.Contractible as BTT+import qualified Data.Set as S+import qualified Data.Primitive.PrimArray as P++main :: IO ()+main = do+ putStrLn "Starting test suite"+ -- withToken $ \c -> do+ -- ctx <- BTC.newContext 3 c+ -- b0 <- BTC.new ctx :: IO (BTC.BTree Int Int RealWorld _)+ -- b1 <- BTC.insert ctx b0 (1 :: Int) (1 :: Int)+ -- b2 <- BTC.insert ctx b1 (2 :: Int) (2 :: Int)+ -- b3 <- BTC.insert ctx b2 (3 :: Int) (3 :: Int)+ -- b4 <- BTC.insert ctx b3 (4 :: Int) (4 :: Int)+ -- b5 <- BTC.insert ctx b4 (5 :: Int) (5 :: Int)+ -- b6 <- BTC.insert ctx b5 (6 :: Int) (6 :: Int)+ -- b7 <- BTC.insert ctx b6 (7 :: Int) (7 :: Int)+ -- print =<< BTC.lookup b7 3+ -- putStrLn =<< BTC.debugMap ctx b7+ -- return ()+ defaultMain tests+ basicBenchmarks+ putStrLn "Finished test suite"++tests :: TestTree+tests = testGroup "Tests" [unitTests,properties]++properties :: TestTree+properties = testGroup "Properties" [scProps]++smallcheckTests :: + (forall n. (Show n, Ord n, Prim n, Hashable n, Bounded n, Integral n) => Int -> [Positive n] -> Either Reason Reason)+ -> [TestTree]+smallcheckTests f = + [ testPropDepth 3 "small maps of degree 3, all permutations, no splitting"+ (over (series :: Series IO [Positive Int]) (f 3))+ , testPropDepth 7 "small maps of degree 3, all permutations"+ (over (series :: Series IO [Positive Int]) (f 3))+ , testPropDepth 7 "small maps of degree 4, all permutations"+ (over (series :: Series IO [Positive Int]) (f 4))+ , testPropDepth 10 "medium maps of degree 3, few permutations"+ (over doubletonSeriesA (f 3))+ , testPropDepth 10 "medium maps of degree 4, few permutations"+ (over doubletonSeriesA (f 4))+ , testPropDepth 10 "medium maps of degree 3, repeat keys likely, few permutations"+ (over doubletonSeriesB (f 3))+ , testPropDepth 10 "medium maps of degree 4, repeat keys likely, few permutations"+ (over doubletonSeriesB (f 4))+ , testPropDepth 150 "large maps of degree 3, repeat keys certain, one permutation"+ (over singletonSeriesB (f 3))+ , testPropDepth 150 "large maps of degree 6, one permutation"+ (over singletonSeriesA (f 6))+ , testPropDepth 150 "large maps of degree 7, repeat keys certain, one permutation"+ (over singletonSeriesB (f 7))+ ]++scProps :: TestTree+scProps = testGroup "smallcheck"+ [ testGroup "standard heap" (smallcheckTests ordering) + , testGroup "compact heap" (smallcheckTests orderingCompact)+ , testGroup "compact heap nested" (smallcheckTests orderingNested)+ , testPropDepth 7 "standard heap lookup"+ (over (series :: Series IO [Positive Int]) (lookupAfterInsert 3))+ , testPropDepth 500 "standard heap bigger lookup"+ (over singletonSeriesA (lookupAfterInsert 3))+ , testPropDepth 7 "compact heap lookup"+ (over (series :: Series IO [Positive Int]) (lookupAfterInsertCompact 3))+ , testPropDepth 500 "compact heap bigger lookup"+ (over singletonSeriesA (lookupAfterInsertCompact 10))+ ]++unitTests :: TestTree+unitTests = testGroup "Unit tests"+ [ testCase "put followed by get (tests lookup,insert,toAscList)" $ do+ let xs = [1,3,2,4,6,5 :: Word]+ xs' = map (\x -> (x,x)) xs+ e = runST $ runExceptT $ do+ b <- lift $ B.fromList (B.Context (BTL.Context 20)) xs'+ forM_ xs $ \k -> do+ mv <- lift $ B.lookup b k+ case mv of+ Nothing -> do+ flattened <- lift (B.toAscList b)+ ExceptT $ return $ Left $ concat+ [ "key "+ , show k+ , " was not found. Flattened tree: "+ , show flattened+ ]+ Just v -> if v == k+ then return ()+ else do+ flattened <- lift (B.toAscList b)+ ExceptT $ return $ Left $ concat+ [ "key "+ , show k + , " was found with non-matching value "+ , show v+ , ". Flattened tree: "+ , show flattened+ ]+ case e of+ Left err -> assertFailure err+ Right () -> return ()+ , testCase "insertions are sorted" $ do+ let xs = [1,3,2,4,6,5,19,11,7 :: Word]+ xs' = map (\x -> (x,x)) xs+ actual <- return (runST (B.fromList (B.Context (BTL.Context 4)) xs' >>= B.toAscList))+ actual @?= S.toAscList (S.fromList xs')+ , testCase "compact b-tree can be created" $ withToken $ \token -> do+ ctx <- BTC.newContext 5 token+ _ <- BTC.new ctx :: IO (BTC.BTree Word Word RealWorld _)+ return ()+ ]++testPropDepth :: Testable IO a => Int -> String -> a -> TestTree+testPropDepth n name = localOption (SmallCheckDepth n) . testProperty name++lookupAfterInsert :: (Show n, Ord n, Prim n)+ => Int -- ^ degree of b-tree+ -> [Positive n] -- ^ values to insert+ -> Either Reason Reason+lookupAfterInsert degree xs' =+ let xs = map getPositive xs'+ expected = map (\x -> (x,x)) $ S.toAscList $ S.fromList xs+ in fmap (const "good") $ runST $ do+ m <- B.new (B.Context (BTL.Context degree))+ forM_ xs $ \x -> do+ B.insert m x x+ r1 <- foldlM (\e x -> case e of+ Right () -> do+ B.lookup m x >>= \case+ Nothing -> return $ Left ("could not find " ++ show x ++ " after inserting it")+ Just y -> return $ if x == y+ then Right ()+ else Left ("looked up " ++ show x ++ " but found wrong value " ++ show y)+ Left err -> return (Left err)+ ) (Right ()) xs+ r2 <- runExceptT $ forM_ xs $ \x -> lift (B.lookup m x) >>= \case+ Nothing -> ExceptT $ return $ Left ("could not find " ++ show x ++ " after inserting it")+ Just y -> ExceptT $ return $ if x == y+ then Right ()+ else Left ("looked up " ++ show x ++ " but found wrong value " ++ show y)+ return (r1 >> r2)++lookupAfterInsertCompact :: (Show n, Ord n, Prim n)+ => Int -- ^ degree of b-tree+ -> [Positive n] -- ^ values to insert+ -> Either Reason Reason+lookupAfterInsertCompact degree xs' =+ let xs = map getPositive xs'+ expected = map (\x -> (x,x)) $ S.toAscList $ S.fromList xs+ in fmap (const "good") $ runST $ withToken $ \c -> do+ ctx <- BTC.newContext degree c+ m0 <- BTC.new ctx+ m1 <- foldlM (\ !m !x -> BTC.insert ctx m x x) m0 xs+ r1 <- foldlM (\e x -> case e of+ Right () -> do+ BTC.lookup m1 x >>= \case+ Nothing -> return $ Left ("could not find " ++ show x ++ " after inserting it")+ Just y -> return $ if x == y+ then Right ()+ else Left ("looked up " ++ show x ++ " but found wrong value " ++ show y)+ Left err -> return (Left err)+ ) (Right ()) xs+ r2 <- runExceptT $ forM_ xs $ \x -> lift (BTC.lookup m1 x) >>= \case+ Nothing -> ExceptT $ return $ Left ("could not find " ++ show x ++ " after inserting it")+ Just y -> ExceptT $ return $ if x == y+ then Right ()+ else Left ("looked up " ++ show x ++ " but found wrong value " ++ show y)+ return (r1 >> r2)+++ordering :: (Show n, Ord n, Prim n)+ => Int -- ^ degree of b-tree+ -> [Positive n] -- ^ values to insert+ -> Either Reason Reason+ordering degree xs' = + let xs = map getPositive xs'+ expected = map (\x -> (x,x)) $ S.toAscList $ S.fromList xs+ (actual,layout) = runST $ do+ m <- B.new (B.Context (BTL.Context degree))+ forM_ xs $ \x -> do+ B.insert m x x+ (,) <$> B.toAscList m <*> B.debugMap m+ in if actual == expected+ then Right "good"+ else Left (notice (show expected) (show actual) layout)++orderingCompact :: (Show n, Ord n, Prim n)+ => Int -- ^ degree of b-tree+ -> [Positive n] -- ^ values to insert+ -> Either Reason Reason+orderingCompact degree xs' = + let xs = map getPositive xs'+ expected = map (\x -> (x,x)) $ S.toAscList $ S.fromList xs+ (actual,layout) = runST $ withToken $ \c -> do+ ctx <- BTC.newContext degree c+ m0 <- BTC.new ctx+ m1 <- foldlM (\ !m !x -> BTC.insert ctx m x x) m0 xs+ (,) <$> BTC.toAscList ctx m1 <*> BTC.debugMap ctx m1+ in if actual == expected+ then Right "good"+ else Left (notice (show expected) (show actual) layout)++-- let us begin the most dangerous game.+orderingNested:: (Show n, Ord n, Prim n, Hashable n, Bounded n, Integral n)+ => Int -- ^ degree of b-tree+ -> [Positive n] -- ^ values to insert+ -> Either Reason Reason+orderingNested degree xs' = + let xs = map getPositive xs'+ e = runST $ withToken $ \c -> do+ ctx <- BTT.newContext degree c+ m0 <- BTT.new ctx+ m1 <- foldlM+ (\ !mtop !x -> do+ let subValues = take 10 (iterate (fromIntegral . hashWithSalt 13 . (+ div maxBound 3)) x)+ foldM ( \ !m !y -> do+ (_,t) <- BTT.modifyWithM ctx m x (BTC.new ctx) $ \mbottom -> do+ fmap BTT.Replace (BTC.insert ctx mbottom y y)+ return t+ ) mtop subValues+ ) m0 xs+ runExceptT $ forM_ xs $ \x -> do+ m <- lift $ BTT.lookup m1 x + case m of+ Nothing -> ExceptT (return (Left ("could not find " ++ show x ++ " in top b-tree")))+ Just b -> do+ n <- lift $ BTC.lookup b x+ case n of+ Nothing -> ExceptT (return (Left ("could not find " ++ show x ++ " in bottom b-tree")))+ Just k -> return ()+ in fmap (const "good") e++notice :: String -> String -> String -> String+notice expected actual layout = concat+ [ "expected: "+ , expected+ , ", actual: "+ , actual+ , ", layout:\n"+ , layout+ ]++scanSeries :: forall m a. (a -> [a]) -> a -> Series m [a]+scanSeries f x0 = generate $ \n ->+ map toList $ concat $ take n $ iterate+ (\ys -> ys >>= \xs@(x NE.:| _) -> f x >>= \z -> [z NE.:| (toList xs)])+ [x0 NE.:| []]++doubletonSeriesA :: Series m [Positive Word16]+doubletonSeriesA = (fmap.fmap) Positive (scanSeries (\n -> [n + 9787, n + 29059]) 0)++doubletonSeriesB :: Series m [Positive Word8]+doubletonSeriesB = (fmap.fmap) Positive (scanSeries (\n -> [n + 89, n + 71]) 0)++singletonSeriesA :: Series m [Positive Word16]+singletonSeriesA = (fmap.fmap) Positive (scanSeries (\n -> [n + 26399]) 0)++singletonSeriesB :: Series m [Positive Word8]+singletonSeriesB = (fmap.fmap) Positive (scanSeries (\n -> [n + 73]) 0)++sizeAfterInserts :: forall n. (Num n, Prim n, Ord n, Hashable n) => Proxy n -> n -> Int -> IO Word +sizeAfterInserts _ total degree = withToken $ \c -> do+ ctx <- BTC.newContext degree c+ m0 <- BTC.new ctx+ let go !ix !m = if ix < total+ then do+ let x = hashWithSalt 45237 (ix :: n)+ y = fromIntegral x :: n+ m' <- BTC.insert ctx m y y+ go (ix + 1) m'+ else return ()+ go 0 m0+ getSizeOfCompact c++sizeAfterRepeatedInserts :: Int -> IO Word +sizeAfterRepeatedInserts total = withToken $ \c -> do+ ctx <- BTC.newContext 8 c+ m0 <- BTC.new ctx+ let go !ix !m = if ix < total+ then do+ -- same key every time+ m' <- BTC.insert ctx m (99 :: Int) (ix :: Int)+ go (ix + 1) m'+ else return ()+ go 0 m0+ getSizeOfCompact c++basicBenchmarks :: IO ()+basicBenchmarks = do+ let degrees = [50,105]+ sizes = [10000,15000,30000]+ pairs = (,) <$> degrees <*> sizes+ forM_ pairs $ \(degree,size) -> do+ sz <- sizeAfterInserts (Proxy :: Proxy Int64) (fromIntegral size) degree+ putStrLn ("Bytes of " ++ show size ++ " distinct inserts (Int64) into b-tree of degree " ++ show degree ++ ": " ++ show sz)+ forM_ pairs $ \(degree,size) -> do+ sz <- sizeAfterInserts (Proxy :: Proxy Int32) (fromIntegral size) degree+ putStrLn ("Bytes of " ++ show size ++ " distinct inserts (Int32) into b-tree of degree " ++ show degree ++ ": " ++ show sz)+ putStrLn "Repeated Inserts"+ forM_ sizes $ \size -> do+ sz <- sizeAfterRepeatedInserts size+ putStrLn ("Bytes of " ++ show size ++ " repeated inserts into b-tree: " ++ show sz)+ +