btree 0.1.0.0 → 0.2
raw patch · 4 files changed
+317/−421 lines, 4 files
Files
- btree.cabal +1/−1
- src/BTree/Compact.hs +151/−194
- src/BTree/Contractible.hs +149/−206
- test/Spec.hs +16/−20
btree.cabal view
@@ -1,5 +1,5 @@ name: btree-version: 0.1.0.0+version: 0.2 synopsis: B-Tree on the compact heap -- description: homepage: https://github.com/andrewthad/b-plus-tree#readme
src/BTree/Compact.hs view
@@ -10,19 +10,16 @@ {-# OPTIONS_GHC -O2 -Wall -Werror -fno-warn-unused-imports #-} - module BTree.Compact ( BTree- , Context(..)- , Sizing(..) , Decision(..) , new- , newContext , debugMap , insert , modifyWithM , lookup , toAscList+ , foldrWithKey ) where import Prelude hiding (lookup)@@ -45,86 +42,79 @@ 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))- }+data BTree k v s (c :: Heap) = BTree+ {-# UNPACK #-} !Int -- degree+ {-# UNPACK #-} !(BNode k v 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)+data BNode k v s (c :: Heap) = BNode+ { _bnodeSize :: {-# UNPACK #-} !Int -- size, number of keys present in node+ , _bnodeKeys :: {-# UNPACK #-} !(MutablePrimArray s k)+ , _bnodeContents :: {-# 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#+instance Contractible (BNode k v) where+ unsafeContractedUnliftedPtrCount# _ = 4# unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 2# readContractedArray# ba aa ix s1 = let ixByte = ix *# 2#- ixPtr = ix *# 5#+ ixPtr = ix *# 4# in case readIntArray# ba (ixByte +# 0#) s1 of- (# s2, szIx #) -> case readIntArray# ba (ixByte +# 1#) s2 of+ (# s2, sz #) -> 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 =+ (# s4, keys #) -> case readMutableByteArrayArray# aa (ixPtr +# 1#) s4 of+ (# s5, values #) -> case readMutableByteArrayArray# aa (ixPtr +# 2#) s5 of+ (# s6, nodesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 3#) s6 of+ (# s7, nodesPtrs #) ->+ (# s7, (BNode (I# sz) (MutablePrimArray keys) (FlattenedContents (I# toggle) (MutablePrimArray values) (ContractedMutableArray nodesBytes nodesPtrs))) #)+ writeContractedArray# ba aa ix (BNode (I# sz) (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+ ixPtr = ix *# 4#+ in case writeIntArray# ba (ixByte +# 0#) sz 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)-+ s3 -> case writeMutableByteArrayArray# aa (ixPtr +# 0#) keys s3 of+ s4 -> case writeMutableByteArrayArray# aa (ixPtr +# 1#) values s4 of+ s5 -> case writeMutableByteArrayArray# aa (ixPtr +# 2#) nodesBytes s5 of+ s6 -> writeMutableArrayArrayArray# aa (ixPtr +# 3#) nodesPtrs s6 +instance Contractible (BTree k v) where+ unsafeContractedUnliftedPtrCount# _ = 4#+ unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 3#+ readContractedArray# ba aa ix s1 =+ let ixByte = ix *# 3#+ ixPtr = ix *# 4#+ in case readIntArray# ba (ixByte +# 0#) s1 of+ (# s2, sz #) -> case readIntArray# ba (ixByte +# 1#) s2 of+ (# s3, toggle #) -> case readIntArray# ba (ixByte +# 2#) s3 of+ (# s4, degree #) -> case readMutableByteArrayArray# aa (ixPtr +# 0#) s4 of+ (# s5, keys #) -> case readMutableByteArrayArray# aa (ixPtr +# 1#) s5 of+ (# s6, values #) -> case readMutableByteArrayArray# aa (ixPtr +# 2#) s6 of+ (# s7, nodesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 3#) s7 of+ (# s8, nodesPtrs #) ->+ (# s8, BTree (I# degree) (BNode (I# sz) (MutablePrimArray keys) (FlattenedContents (I# toggle) (MutablePrimArray values) (ContractedMutableArray nodesBytes nodesPtrs))) #)+ writeContractedArray# ba aa ix (BTree (I# degree) (BNode (I# sz) (MutablePrimArray keys) (FlattenedContents (I# toggle) (MutablePrimArray values) (ContractedMutableArray nodesBytes nodesPtrs)))) s1 =+ let ixByte = ix *# 3#+ ixPtr = ix *# 4#+ in case writeIntArray# ba (ixByte +# 0#) sz s1 of+ s2 -> case writeIntArray# ba (ixByte +# 1#) toggle s2 of+ s3 -> case writeIntArray# ba (ixByte +# 2#) degree s3 of+ s4 -> case writeMutableByteArrayArray# aa (ixPtr +# 0#) keys s4 of+ s5 -> case writeMutableByteArrayArray# aa (ixPtr +# 1#) values s5 of+ s6 -> case writeMutableByteArrayArray# aa (ixPtr +# 2#) nodesBytes s6 of+ s7 -> writeMutableArrayArrayArray# aa (ixPtr +# 3#) nodesPtrs s7+ -- We manually flatten this sum type so that it can be unpacked--- into BTree.+-- into BNode. data FlattenedContents k v s c = FlattenedContents {-# UNPACK #-} !Int {-# UNPACK #-} !(MutablePrimArray s v)- {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)+ {-# UNPACK #-} !(ContractedMutableArray (BNode k v) s c) data Contents k v s c = ContentsValues {-# UNPACK #-} !(MutablePrimArray s v)- | ContentsNodes {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)+ | ContentsNodes {-# UNPACK #-} !(ContractedMutableArray (BNode k v) s c) {-# INLINE flattenContentsToContents #-} flattenContentsToContents :: @@ -140,7 +130,7 @@ {-# INLINE contentsToFlattenContents #-} contentsToFlattenContents :: MutablePrimArray s v -- ^ garbage value- -> ContractedMutableArray (BTree k v) s c -- ^ garbage value+ -> ContractedMutableArray (BNode k v) s c -- ^ garbage value -> Contents k v s c -> FlattenedContents k v s c contentsToFlattenContents !garbageValues !garbageNodes !c = case c of@@ -150,63 +140,63 @@ -- | 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 k v s c -> ContractedMutableArray (BNode k v) s c demandFlattenedContentsNodes (FlattenedContents _ _ nodes) = nodes data Insert k v s c- = Ok !v+ = Ok+ !v+ {-# UNPACK #-} !Int -- new size of left child | Split- {-# NOUNPACK #-} !(BTree k v s c)+ {-# NOUNPACK #-} !(BNode k v s c) !k !v- {-# UNPACK #-} !(Sizing s c)+ {-# UNPACK #-} !Int -- ^ The new node that will go to the right, -- the key propagated to the parent,- -- the inserted value, updated sizing info.+ -- the inserted value, updated sizing info for the left child {-# INLINE mkBTree #-} mkBTree :: PrimMonad m => Token c- -> ContractedMutableArray (BTree k v) (PrimState m) c -- ^ garbage value- -> Sizing (PrimState m) c+ -> ContractedMutableArray (BNode k v) (PrimState m) c -- ^ garbage value+ -> Int -- Sizing (PrimState m) c -> MutablePrimArray (PrimState m) k -- ^ keys -> Contents k v (PrimState m) c- -> m (BTree k v (PrimState m) c)+ -> m (BNode k v (PrimState m) c) mkBTree token garbage a b c = do let !garbageValues = coercePrimArray b- !bt = BTree a b (contentsToFlattenContents garbageValues garbage c)+ !bt = BNode 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+ => Token c+ -> Int -- ^ degree, must be at least 3 -> m (BTree k v (PrimState m) c)-new (Context !degree !token !szRef) = do+new !token !degree = 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)+ node <- mkBTree token garbageNodes 0 keys (ContentsValues values)+ return (BTree degree node) --- {-# SPECIALIZE lookup :: BTree RealWorld Int Int c -> Int -> IO (Maybe Int) #-}+-- {-# SPECIALIZE lookup :: BNode 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+lookup (BTree _ theNode) k = go 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+ go :: BNode k v (PrimState m) c -> m (Maybe v)+ go (BNode sz keys c@(FlattenedContents _tog _ _)) = do case flattenContentsToContents c of ContentsValues values -> do ix <- findIndex keys k sz@@ -218,57 +208,48 @@ 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-+ go node {-# INLINE insert #-} insert :: (Ord k, Prim k, Prim v, PrimMonad m)- => Context (PrimState m) c+ => Token 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))+insert !token !m !k !v = do+ !(!_,!node) <- modifyWithM token 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) #-}+{-# SPECIALIZE modifyWithM :: Token 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+ => Token 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.+modifyWithM !token (BTree !degree !root) !k !newValue alter = do !ins <- go root case ins of- Ok v -> return (v,root)- Split !rightNode newRootKey v sizing -> do- writeNodeSize sizing 1+ Ok !v !newNodeSz -> return (v,BTree degree (root { _bnodeSize = newNodeSz }))+ Split !rightNode !newRootKey !v !newLeftSize -> do 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)+ let !leftNode = root { _bnodeSize = newLeftSize }+ !newRoot@(BNode _ _ (FlattenedContents _ _ cmptRootChildren)) <- mkBTree token newRootChildren 1 newRootKeys (ContentsNodes newRootChildren) writeContractedArray cmptRootChildren 0 leftNode writeContractedArray cmptRootChildren 1 rightNode- !newSizing <- nextSizing token sizing- writeMutVar sizingRef newSizing- return (v,newRoot)+ return (v,BTree degree 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+ go :: BNode k v (PrimState m) c -> m (Insert k v (PrimState m) c)+ go (BNode !sz !keys !c) = do case flattenContentsToContents c of ContentsValues !values -> do !ix <- findIndex keys k sz@@ -279,29 +260,28 @@ 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)+ return (Ok v (sz + 1)) 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)+ rightKeys' <- newPrimArray (degree - 1)+ rightValues' <- newPrimArray (degree - 1)+ let (newLeftSz,actualRightSz) = if gtIx < leftSize+ then (leftSize + 1, rightSize)+ else (leftSize,rightSize + 1)+ !newTree@(BNode _ rightKeys (FlattenedContents _ rightValues _)) <- mkBTree token (demandFlattenedContentsNodes c) actualRightSz rightKeys' (ContentsValues rightValues') 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@@ -311,65 +291,67 @@ 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)+ return (Split newTree propagated v newLeftSz) 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')+ return (Ok v' sz) 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)+ let !nodeIx = if isEq then gtIx + 1 else gtIx+ !node <- readContractedArray nodes nodeIx !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)+ Ok !v !newNodeSz -> do+ when (newNodeSz /= _bnodeSize node) $ do+ writeContractedArray nodes nodeIx (node { _bnodeSize = newNodeSz })+ return (Ok v sz)+ Split !rightNode !propagated !v !newNodeSz -> do+ when (newNodeSz /= _bnodeSize node) $ do+ writeContractedArray nodes nodeIx (node { _bnodeSize = newNodeSz })+ 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 (sz + 1))+ else do+ let !middleIx = div sz 2+ !leftKeys = keys+ !leftNodes = nodes+ !middleKey <- readPrimArray keys middleIx+ !rightKeysOnHeap <- newPrimArray (degree - 1)+ !rightNodes' <- newContractedArray token degree -- uninitializedNode+ let !leftSize = middleIx+ !rightSize = sz - leftSize+ (!actualLeftSz,!actualRightSz) = if middleIx >= gtIx+ then (leftSize + 1, rightSize - 1)+ else (leftSize, rightSize)+ -- _ <- error ("die: " ++ show actualRightSz ++ " " ++ show sz ++ " " ++ show actualLeftSz)+ !x@(BNode _ rightKeys (FlattenedContents _ _ rightNodes)) <- mkBTree token rightNodes' actualRightSz rightKeysOnHeap (ContentsNodes rightNodes')+ 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+ 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+ return (Split x middleKey v actualLeftSz) -- Preconditions: -- * marr is sorted low to high@@ -454,12 +436,9 @@ 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+ => 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+debugMap (BTree _ (BNode !rootSz !rootKeys !rootContents)) = do 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@@ -467,13 +446,11 @@ return (map (\s -> (level,s)) pairStrs) ContentsNodes nodes -> do pairs <- pairForM sz keys nodes- $ \k (BTree theNextSizing nextKeys nextContents) -> do- nextSz <- readNodeSize theNextSizing+ $ \k (BNode nextSz nextKeys nextContents) -> do 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+ BNode lastSz lastKeys lastContents <- readContractedArray nodes sz lastStrs <- go (level + 1) lastSz lastKeys lastContents -- return (nextStrs ++ [(level,show k)]) return ([(level, "start")] ++ concat pairs ++ lastStrs)@@ -518,42 +495,22 @@ -- | 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+ => 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+foldrWithKey f b0 (BTree _ 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+ go :: BNode k v (PrimState m) c -> b -> m b+ go (BNode sz keys c) !b = do case flattenContentsToContents c of ContentsValues values -> foldrPrimArrayPairs sz f b keys values ContentsNodes nodes -> foldrArray (sz + 1) go b nodes
src/BTree/Contractible.hs view
@@ -13,13 +13,11 @@ module BTree.Contractible ( BTree- , Context , Decision(..) , new- , newContext , modifyWithM , lookup- , toAscList+ , foldrWithKey ) where import Prelude hiding (lookup)@@ -33,7 +31,6 @@ 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@@ -43,87 +40,83 @@ 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))--- }+data BTree k (v :: * -> Heap -> *) s (c :: Heap) = BTree+ {-# UNPACK #-} !Int -- degree+ {-# UNPACK #-} !(BNode k v 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)+data BNode k (v :: * -> Heap -> *) s (c :: Heap) = BNode+ { _bnodeSize :: {-# UNPACK #-} !Int -- size, number of keys present in node+ , _bnodeKeys :: {-# UNPACK #-} !(MutablePrimArray s k)+ , _bnodeContents :: {-# 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#+instance Contractible (BNode k v) where+ unsafeContractedUnliftedPtrCount# _ = 5# unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 2# readContractedArray# ba aa ix s1 = let ixByte = ix *# 2#- ixPtr = ix *# 6#+ ixPtr = ix *# 5# in case readIntArray# ba (ixByte +# 0#) s1 of- (# s2, szIx #) -> case readIntArray# ba (ixByte +# 1#) s2 of+ (# s2, sz #) -> 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 =+ (# s4, keys #) -> case readMutableByteArrayArray# aa (ixPtr +# 1#) s4 of+ (# s5, valuesBytes #) -> case readMutableByteArrayArray# aa (ixPtr +# 2#) s5 of+ (# s6, nodesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 3#) s6 of+ (# s7, nodesPtrs #) -> case readMutableArrayArrayArray# aa (ixPtr +# 4#) s7 of+ (# s8, valuesPtrs #) ->+ (# s8, (BNode (I# sz) (MutablePrimArray keys) (FlattenedContents (I# toggle) (ContractedMutableArray valuesBytes valuesPtrs) (ContractedMutableArray nodesBytes nodesPtrs))) #)+ writeContractedArray# ba aa ix (BNode (I# sz) (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+ ixPtr = ix *# 5#+ in case writeIntArray# ba (ixByte +# 0#) sz 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)-+ s3 -> case writeMutableByteArrayArray# aa (ixPtr +# 0#) keys s3 of+ s4 -> case writeMutableByteArrayArray# aa (ixPtr +# 1#) valuesBytes s4 of+ s5 -> case writeMutableByteArrayArray# aa (ixPtr +# 2#) nodesBytes s5 of+ s6 -> case writeMutableArrayArrayArray# aa (ixPtr +# 3#) nodesPtrs s6 of+ s7 -> writeMutableArrayArrayArray# aa (ixPtr +# 4#) valuesPtrs s7 +instance Contractible (BTree k v) where+ unsafeContractedUnliftedPtrCount# _ = 5#+ unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 3#+ readContractedArray# ba aa ix s1 =+ let ixByte = ix *# 3#+ ixPtr = ix *# 5#+ in case readIntArray# ba (ixByte +# 0#) s1 of+ (# s2, sz #) -> case readIntArray# ba (ixByte +# 1#) s2 of+ (# s3, toggle #) -> case readIntArray# ba (ixByte +# 2#) s3 of+ (# s4, degree #) -> case readMutableByteArrayArray# aa (ixPtr +# 0#) s4 of+ (# s5, keys #) -> case readMutableByteArrayArray# aa (ixPtr +# 1#) s5 of+ (# s6, valuesBytes #) -> case readMutableByteArrayArray# aa (ixPtr +# 2#) s6 of+ (# s7, nodesBytes #) -> case readMutableArrayArrayArray# aa (ixPtr +# 3#) s7 of+ (# s8, nodesPtrs #) -> case readMutableArrayArrayArray# aa (ixPtr +# 4#) s8 of+ (# s9, valuesPtrs #) ->+ (# s9, BTree (I# degree) (BNode (I# sz) (MutablePrimArray keys) (FlattenedContents (I# toggle) (ContractedMutableArray valuesBytes valuesPtrs) (ContractedMutableArray nodesBytes nodesPtrs))) #)+ writeContractedArray# ba aa ix (BTree (I# degree) (BNode (I# sz) (MutablePrimArray keys) (FlattenedContents (I# toggle) (ContractedMutableArray valuesBytes valuesPtrs) (ContractedMutableArray nodesBytes nodesPtrs)))) s1 =+ let ixByte = ix *# 3#+ ixPtr = ix *# 5#+ in case writeIntArray# ba (ixByte +# 0#) sz s1 of+ s2 -> case writeIntArray# ba (ixByte +# 1#) toggle s2 of+ s3 -> case writeIntArray# ba (ixByte +# 2#) degree s3 of+ s4 -> case writeMutableByteArrayArray# aa (ixPtr +# 0#) keys s4 of+ s5 -> case writeMutableByteArrayArray# aa (ixPtr +# 1#) valuesBytes s5 of+ s6 -> case writeMutableByteArrayArray# aa (ixPtr +# 2#) nodesBytes s6 of+ s7 -> case writeMutableArrayArrayArray# aa (ixPtr +# 3#) nodesPtrs s7 of+ s8 -> writeMutableArrayArrayArray# aa (ixPtr +# 4#) valuesPtrs s8+ -- We manually flatten this sum type so that it can be unpacked--- into BTree.+-- into BNode. data FlattenedContents (k :: *) (v :: * -> Heap -> *) (s :: *) (c :: Heap) = FlattenedContents {-# UNPACK #-} !Int {-# UNPACK #-} !(ContractedMutableArray v s c)- {-# UNPACK #-} !(ContractedMutableArray (BTree k v) s c)+ {-# UNPACK #-} !(ContractedMutableArray (BNode 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)+ | ContentsNodes {-# UNPACK #-} !(ContractedMutableArray (BNode k v) s c) {-# INLINE flattenContentsToContents #-} flattenContentsToContents :: @@ -139,7 +132,7 @@ {-# INLINE contentsToFlattenContents #-} contentsToFlattenContents :: ContractedMutableArray v s c -- ^ garbage value- -> ContractedMutableArray (BTree k v) s c -- ^ garbage value+ -> ContractedMutableArray (BNode k v) s c -- ^ garbage value -> Contents k v s c -> FlattenedContents k v s c contentsToFlattenContents !garbageValues !garbageNodes !c = case c of@@ -149,63 +142,63 @@ -- | 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 k v s c -> ContractedMutableArray (BNode k v) s c demandFlattenedContentsNodes (FlattenedContents _ _ nodes) = nodes -data Insert k (v :: * -> Heap -> *) s c- = Ok !(v s c)+data Insert k v s c+ = Ok+ !(v s c)+ {-# UNPACK #-} !Int -- new size of left child | Split- {-# NOUNPACK #-} !(BTree k v s c)+ {-# NOUNPACK #-} !(BNode k v s c) !k !(v s c)- {-# UNPACK #-} !(Sizing s c)+ {-# UNPACK #-} !Int -- ^ The new node that will go to the right, -- the key propagated to the parent,- -- the inserted value, updated sizing info.+ -- the inserted value, updated sizing info for the left child {-# INLINE mkBTree #-} mkBTree :: PrimMonad m => Token c- -> ContractedMutableArray (BTree k v) (PrimState m) c -- ^ garbage value- -> Sizing (PrimState m) c+ -> ContractedMutableArray (BNode k v) (PrimState m) c -- ^ garbage value+ -> Int -- Sizing (PrimState m) c -> MutablePrimArray (PrimState m) k -- ^ keys -> Contents k v (PrimState m) c- -> m (BTree k v (PrimState m) c)+ -> m (BNode k v (PrimState m) c) mkBTree token garbage a b c = do let !garbageValues = coerceContactedArray garbage- !bt = BTree a b (contentsToFlattenContents garbageValues garbage c)+ !bt = BNode 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+ => Token c+ -> Int -- ^ degree, must be at least 3 -> m (BTree k v (PrimState m) c)-new (Context !degree !token !szRef) = do+new !token !degree = 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)+ node <- mkBTree token garbageNodes 0 keys (ContentsValues values)+ return (BTree degree node) --- {-# SPECIALIZE lookup :: BTree RealWorld Int Int c -> Int -> IO (Maybe Int) #-}+-- {-# SPECIALIZE lookup :: BNode 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+lookup (BTree _ 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+ go :: BNode k v (PrimState m) c -> m (Maybe (v (PrimState m) c))+ go (BNode sz keys c@(FlattenedContents _tog _ _)) = do case flattenContentsToContents c of ContentsValues values -> do ix <- findIndex keys k sz@@ -219,55 +212,35 @@ !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+data Decision a = Keep | Replace !a --- | 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.+-- When we turn on this specialize pragma, it gets way faster+-- for the particular case.+-- {-# SPECIALIZE modifyWithM :: Token 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 :: * -> Heap -> *) (c :: Heap). (Ord k, Prim k, Contractible v, PrimMonad m)- => Context (PrimState m) c+modifyWithM :: forall m k v c. (Ord k, Prim k, Contractible v, PrimMonad m)+ => Token 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.+modifyWithM !token (BTree !degree !root) !k !newValue alter = do !ins <- go root case ins of- Ok v -> return (v,root)- Split !rightNode newRootKey v sizing -> do- writeNodeSize sizing 1+ Ok !v !newNodeSz -> return (v,BTree degree (root { _bnodeSize = newNodeSz }))+ Split !rightNode !newRootKey !v !newLeftSize -> do 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)+ let !leftNode = root { _bnodeSize = newLeftSize }+ !newRoot@(BNode _ _ (FlattenedContents _ _ cmptRootChildren)) <- mkBTree token newRootChildren 1 newRootKeys (ContentsNodes newRootChildren) writeContractedArray cmptRootChildren 0 leftNode writeContractedArray cmptRootChildren 1 rightNode- !newSizing <- nextSizing token sizing- writeMutVar sizingRef newSizing- return (v,newRoot)+ return (v,BTree degree 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+ go :: BNode k v (PrimState m) c -> m (Insert k v (PrimState m) c)+ go (BNode !sz !keys !c) = do case flattenContentsToContents c of ContentsValues !values -> do !ix <- findIndex keys k sz@@ -280,30 +253,28 @@ 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)+ return (Ok v (sz + 1)) 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')+ let (newLeftSz,actualRightSz) = if gtIx < leftSize+ then (leftSize + 1, rightSize)+ else (leftSize,rightSize + 1)+ !newTree@(BNode _ rightKeys (FlattenedContents _ rightValues _)) <- mkBTree token (demandFlattenedContentsNodes c) actualRightSz 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@@ -313,64 +284,67 @@ 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)+ return (Split newTree propagated v newLeftSz) 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')+ return (Ok v' sz) 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)+ let !nodeIx = if isEq then gtIx + 1 else gtIx+ !node <- readContractedArray nodes nodeIx !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)+ Ok !v !newNodeSz -> do+ when (newNodeSz /= _bnodeSize node) $ do+ writeContractedArray nodes nodeIx (node { _bnodeSize = newNodeSz })+ return (Ok v sz)+ Split !rightNode !propagated !v !newNodeSz -> do+ when (newNodeSz /= _bnodeSize node) $ do+ writeContractedArray nodes nodeIx (node { _bnodeSize = newNodeSz })+ 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 (sz + 1))+ else do+ let !middleIx = div sz 2+ !leftKeys = keys+ !leftNodes = nodes+ !middleKey <- readPrimArray keys middleIx+ !rightKeysOnHeap <- newPrimArray (degree - 1)+ !rightNodes' <- newContractedArray token degree -- uninitializedNode+ let !leftSize = middleIx+ !rightSize = sz - leftSize+ (!actualLeftSz,!actualRightSz) = if middleIx >= gtIx+ then (leftSize + 1, rightSize - 1)+ else (leftSize, rightSize)+ -- _ <- error ("die: " ++ show actualRightSz ++ " " ++ show sz ++ " " ++ show actualLeftSz)+ !x@(BNode _ rightKeys (FlattenedContents _ _ rightNodes)) <- mkBTree token rightNodes' actualRightSz rightKeysOnHeap (ContentsNodes rightNodes')+ 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+ 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+ return (Split x middleKey v actualLeftSz) -- Preconditions: -- * marr is sorted low to high@@ -454,46 +428,15 @@ 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+foldrWithKey f b0 (BTree _ 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+ go :: BNode k v (PrimState m) c -> b -> m b+ go (BNode sz keys c) !b = do case flattenContentsToContents c of ContentsValues values -> foldrPrimArrayPairs sz f b keys values ContentsNodes nodes -> foldrArray (sz + 1) go b nodes
test/Spec.hs view
@@ -68,6 +68,8 @@ smallcheckTests f = [ testPropDepth 3 "small maps of degree 3, all permutations, no splitting" (over (series :: Series IO [Positive Int]) (f 3))+ , testPropDepth 4 "small maps of degree 3, all permutations, one split"+ (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"@@ -142,8 +144,7 @@ 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 _)+ _ <- BTC.new token 5 :: IO (BTC.BTree Word Word RealWorld _) return () ] @@ -185,9 +186,8 @@ 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+ m0 <- BTC.new c degree+ m1 <- foldlM (\ !m !x -> BTC.insert c m x x) m0 xs r1 <- foldlM (\e x -> case e of Right () -> do BTC.lookup m1 x >>= \case@@ -229,30 +229,28 @@ 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+ m0 <- BTC.new c degree+ m1 <- foldlM (\ !m !x -> BTC.insert c m x x) m0 xs+ (,) <$> BTC.toAscList m1 <*> BTC.debugMap 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)+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+ m0 <- BTT.new c degree 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)+ (_,t) <- BTT.modifyWithM c m x (BTC.new c degree) $ \mbottom -> do+ fmap BTT.Replace (BTC.insert c mbottom y y) return t ) mtop subValues ) m0 xs@@ -297,13 +295,12 @@ 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+ m0 <- BTC.new c degree 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+ m' <- BTC.insert c m y y go (ix + 1) m' else return () go 0 m0@@ -311,12 +308,11 @@ sizeAfterRepeatedInserts :: Int -> IO Word sizeAfterRepeatedInserts total = withToken $ \c -> do- ctx <- BTC.newContext 8 c- m0 <- BTC.new ctx+ m0 <- BTC.new c 8 let go !ix !m = if ix < total then do -- same key every time- m' <- BTC.insert ctx m (99 :: Int) (ix :: Int)+ m' <- BTC.insert c m (99 :: Int) (ix :: Int) go (ix + 1) m' else return () go 0 m0