btree 0.2 → 0.3
raw patch · 10 files changed
+1689/−1313 lines, 10 filesdep +MonadRandomdep −compact-mutabledep −prim-arraydep ~basedep ~primitive
Dependencies added: MonadRandom
Dependencies removed: compact-mutable, prim-array
Dependency ranges changed: base, primitive
Files
- bench/Main.hs +32/−30
- btree.cabal +10/−14
- src/ArrayList.hs +322/−0
- src/BTree.hs +1/−1
- src/BTree/Array.hs +0/−21
- src/BTree/Compact.hs +0/−590
- src/BTree/Contractible.hs +0/−516
- src/BTree/Linear.hs +15/−14
- src/BTree/Store.hs +1011/−0
- test/Spec.hs +298/−127
bench/Main.hs view
@@ -11,24 +11,27 @@ ) where import qualified BTree.Linear as BTL-import qualified BTree.Compact as BTC+import qualified BTree.Store as BTS 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+import Foreign.Ptr (Ptr)+import Data.Int -- this specialization does not seem to work.+-- relying on specialize pragmas is the worst.+-- {-# SPECIALIZE BTS.modifyWithPtr :: BTS.BTree Int Int -> Int -> (Either () (Ptr Int -> Int -> IO ())) -> (Ptr Int -> Int -> IO ((),BTS.Decision)) -> IO ((), BTS.BTree Int Int) #-} -- {-# 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 multiplier = 20 let total = 200000 * multiplier- range = 10000000 * multiplier+ range = 1000000 * multiplier lookups = 100000 * multiplier putStrLn $ concat [ "This benchmark will insert "@@ -40,33 +43,34 @@ , "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+ -- 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+ BTS.with_ $ \b0 -> do buildStart <- getTime Monotonic- (b,ctx) <- offHeapBTree token total range+ b1 <- offHeapBTree b0 total range buildEnd <- getTime Monotonic performGC start <- getTime Monotonic- x <- lookupManyOffHeap lookups b+ x <- lookupManyOffHeap lookups b1 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))+ return b1 lookupMany :: Int -> BTL.BTree RealWorld Int Int -> BTL.Context RealWorld -> IO Int lookupMany total b ctx = go 0 0@@ -77,12 +81,12 @@ go (n + 1) (s + fromMaybe 0 m) else return s -lookupManyOffHeap :: Int -> BTC.BTree Int Int RealWorld c -> IO Int+lookupManyOffHeap :: Int -> BTS.BTree Int Int -> IO Int lookupManyOffHeap total b = go 0 0 where go !n !s = if n < total then do- m <- BTC.lookup b n+ m <- BTS.lookup b n go (n + 1) (s + fromMaybe 0 m) else return s @@ -100,19 +104,17 @@ go 0 b0 offHeapBTree ::- Token c + BTS.BTree Int Int -> 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+ -> IO (BTS.BTree Int Int)+offHeapBTree b0 total range = do let go !n !b = if n < total then do let x = mod (hashWithSalt mySalt n) range- b' <- BTC.insert ctx b x x+ b' <- BTS.insert b x x go (n + 1) b'- else return (b,ctx)+ else return b go 0 b0
btree.cabal view
@@ -1,8 +1,8 @@ name: btree-version: 0.2-synopsis: B-Tree on the compact heap+version: 0.3+synopsis: B-Tree on Unmanaged Heap -- description:-homepage: https://github.com/andrewthad/b-plus-tree#readme+homepage: https://github.com/andrewthad/btree license: BSD3 license-file: LICENSE author: Andrew Martin@@ -15,19 +15,16 @@ library hs-source-dirs: src+ ghc-options: -O2 exposed-modules: BTree+ BTree.Store BTree.Linear- -- BTree.Generic- BTree.Array- BTree.Compact- BTree.Contractible+ ArrayList build-depends:- base >= 4.10 && < 4.11+ base >= 4.9 && < 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+ , primitive >= 0.6.1 && < 0.7 default-language: Haskell2010 test-suite test@@ -37,7 +34,6 @@ build-depends: base , btree- , prim-array , tasty , tasty-smallcheck , tasty-hunit@@ -45,20 +41,20 @@ , containers , transformers , primitive- , compact-mutable , hashable+ , MonadRandom -- ghc-options: -threaded -rtsopts -with-rtsopts=-N default-language: Haskell2010 benchmark bench type: exitcode-stdio-1.0+ ghc-options: -O2 build-depends: base , btree , clock , hashable , ghc-prim- , compact-mutable default-language: Haskell2010 hs-source-dirs: bench main-is: Main.hs
+ src/ArrayList.hs view
@@ -0,0 +1,322 @@+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE UnboxedTuples #-}++module ArrayList+ ( ArrayList+ , size+ , with+ , new+ , free+ , pushR+ , pushArrayR+ , popL+ , dropWhileL+ , dropWhileScanL+ , dropScanL+ , dropL+ , dump+ , dumpMap+ , dumpList+ , clear+ , showDebug+ ) where++import Foreign.Ptr+import Foreign.Storable+import Foreign.Marshal.Array+import Data.Primitive.PrimArray+import GHC.Prim (RealWorld)+import GHC.Int (Int(..))+import GHC.Ptr (Ptr(..))+import GHC.IO (IO(..))+import GHC.Prim ((*#),copyAddrToByteArray#)+import Data.Primitive hiding (sizeOf)+import Data.Primitive.Types+import Data.Bits (unsafeShiftR)+import BTree.Store (Initialize(..),Deinitialize(..))+import Control.Monad (when)+import Data.Primitive.Ptr (copyPtrToMutablePrimArray)+import qualified Data.Primitive as PM+import qualified Foreign.Marshal.Alloc as FMA+import qualified Foreign.Storable as FS++-- Special value for Ptr: If the pointer is null, then+-- it is understood that the ArrayList is of length 0.+-- For such an ArrayList, it is understood that+-- all of the fields should be 0.+data ArrayList a = ArrayList+ {-# UNPACK #-} !Int -- start index (in elements)+ {-# UNPACK #-} !Int -- used length (in elements)+ {-# UNPACK #-} !Int -- buffer length (in elements)+ {-# UNPACK #-} !(Ptr a) -- all the data++instance Storable (ArrayList a) where+ sizeOf _ = wordSz * 4+ alignment _ = wordSz+ {-# INLINE peek #-}+ peek ptr = ArrayList+ <$> peek (castPtr ptr)+ <*> peek (plusPtr ptr wordSz)+ <*> peek (plusPtr ptr (wordSz + wordSz))+ <*> peek (plusPtr ptr (wordSz + wordSz + wordSz))+ {-# INLINE poke #-}+ poke ptr (ArrayList a b c d) = do+ poke (castPtr ptr) a+ poke (plusPtr ptr wordSz) b+ poke (plusPtr ptr (wordSz + wordSz)) c+ poke (plusPtr ptr (wordSz + wordSz + wordSz)) d++instance Storable a => Initialize (ArrayList a) where+ {-# INLINE initialize #-}+ initialize (Ptr addr#) = setAddr (Addr addr#) 4 (0 :: Int)++wordSz :: Int+wordSz = PM.sizeOf (undefined :: Int)+ +initialSize :: Int+initialSize = 4++size :: ArrayList a -> Int+size (ArrayList _ len _ _) = len++with :: Storable a => (ArrayList a -> IO (ArrayList a,b)) -> IO b+with f = do+ initial <- new+ (final,a) <- f initial+ free final+ return a++new :: forall a. Storable a => IO (ArrayList a)+new = return (ArrayList 0 0 0 nullPtr)++{-# INLINABLE pushR #-}+pushR :: forall a. Storable a => ArrayList a -> a -> IO (ArrayList a)+pushR (ArrayList start len bufLen ptr) a = if start + len < bufLen+ then do+ poke (advancePtr ptr (start + len)) a+ return (ArrayList start (len + 1) bufLen ptr)+ else if+ | len == 0 -> do+ when (bufLen /= 0) (fail "ArrayList.pushR: invariant violated")+ when (start /= 0) (fail "ArrayList.pushR: invariant violated")+ when (ptr /= nullPtr) (fail "ArrayList.pushR: invariant violated")+ ptr <- FMA.mallocBytes (FS.sizeOf (undefined :: a) * initialSize)+ poke ptr a+ return (ArrayList 0 1 initialSize ptr)+ | len < half bufLen -> do+ moveArray ptr (advancePtr ptr start) len+ poke (advancePtr ptr len) a+ return (ArrayList 0 (len + 1) bufLen ptr)+ | otherwise -> do+ newPtr <- FMA.mallocBytes (FS.sizeOf (undefined :: a) * bufLen * 2)+ moveArray newPtr (advancePtr ptr start) len+ FMA.free ptr+ poke (advancePtr newPtr len) a+ return (ArrayList 0 (len + 1) (bufLen * 2) newPtr)++{-# INLINE pushArrayR #-}+pushArrayR :: forall a. (Storable a, Prim a) => ArrayList a -> PrimArray a -> IO (ArrayList a)+pushArrayR (ArrayList start len bufLen ptr) as =+ -- I think this should actually be less than or equal to+ if start + len + asLen < bufLen+ then do+ copyPrimArrayToPtr (advancePtr ptr (start + len)) as 0 asLen+ return (ArrayList start (len + asLen) bufLen ptr)+ else if+ -- this might give poor guarentees concerning worst+ -- case behaviors, but whatever for now.+ | len == 0 -> do+ when (bufLen /= 0) (fail "ArrayList.pushArrayR: invariant violated")+ when (start /= 0) (fail "ArrayList.pushArrayR: invariant violated")+ when (ptr /= nullPtr) (fail "ArrayList.pushArrayR: invariant violated")+ let newBufLen = twiceUntilExceeds initialSize asLen+ ptr <- FMA.mallocBytes (FS.sizeOf (undefined :: a) * newBufLen)+ copyPrimArrayToPtr ptr as 0 asLen+ return (ArrayList 0 asLen newBufLen ptr)+ | len < half bufLen && asLen < half bufLen -> do+ moveArray ptr (advancePtr ptr start) len+ copyPrimArrayToPtr (advancePtr ptr len) as 0 asLen+ return (ArrayList 0 (len + asLen) bufLen ptr)+ | otherwise -> do+ let newBufLen = twiceUntilExceeds (2 * bufLen) (len + asLen)+ newPtr <- FMA.mallocBytes (FS.sizeOf (undefined :: a) * newBufLen)+ moveArray newPtr (advancePtr ptr start) len+ FMA.free ptr+ copyPrimArrayToPtr (advancePtr newPtr len) as 0 asLen+ return (ArrayList 0 (len + asLen) newBufLen newPtr)+ where+ asLen = sizeofPrimArray as++twiceUntilExceeds :: Int -> Int -> Int+twiceUntilExceeds !i !limit = go i where + go !n = if n > limit+ then n+ else go (n * 2)+ ++popL :: forall a. Storable a => ArrayList a -> IO (ArrayList a, Maybe a)+popL xs@(ArrayList start len bufLen ptr)+ | len < 1 = return (xs, Nothing)+ | otherwise = do+ a <- peek (advancePtr ptr start)+ newArrList <- minimizeMemory (ArrayList (start + 1) (len - 1) bufLen ptr)+ return (newArrList, Just a)++{-# INLINE dropWhileL #-}+dropWhileL :: forall a. Storable a+ => ArrayList a+ -> (a -> IO Bool) -- ^ predicate+ -> IO (ArrayList a,Int)+dropWhileL (ArrayList start len bufLen ptr) p = do+ let go :: Int -> IO Int+ go !i = if i < len+ then do+ a <- peek (advancePtr ptr (start + i))+ b <- p a+ if b+ then go (i + 1)+ else return i+ else return i+ dropped <- go 0+ newArrList <- minimizeMemory $ ArrayList (start + dropped) (len - dropped) bufLen ptr+ return (newArrList,dropped)++{-# INLINE dropWhileScanL #-}+dropWhileScanL :: forall a b. Storable a+ => ArrayList a+ -> b+ -> (b -> a -> IO (Bool,b))+ -> IO (ArrayList a,Int,b)+dropWhileScanL (ArrayList start len bufLen ptr) b0 p = do+ let go :: Int -> b -> IO (Int,b)+ go !i !b = if i < len+ then do+ !a <- peek (advancePtr ptr (start + i))+ (!shouldContinue,!b') <- p b a+ if shouldContinue+ then go (i + 1) b'+ else return (i,b')+ else return (i,b)+ (dropped,b') <- go 0 b0+ newArrList <- minimizeMemory $ ArrayList (start + dropped) (len - dropped) bufLen ptr+ return (newArrList,dropped,b')++{-# INLINE dropScanL #-}+dropScanL :: forall a b. Storable a+ => ArrayList a+ -> Int+ -> b+ -> (b -> a -> IO b)+ -> IO (ArrayList a, b)+dropScanL (ArrayList start len bufLen ptr) n b0 p = do+ let !m = min n len+ let go :: Int -> b -> IO b+ go !i !b = if i < m+ then do+ a <- peek (advancePtr ptr (start + i))+ b' <- p b a+ go (i + 1) b'+ else return b+ b' <- go 0 b0+ newArrList <- minimizeMemory $ ArrayList (start + m) (len - m) bufLen ptr+ return (newArrList,b')++{-# INLINE dropL #-}+dropL :: forall a. Storable a => ArrayList a -> Int -> IO (ArrayList a)+dropL (ArrayList start len bufLen ptr) n = do+ let m = min n len+ minimizeMemory $ ArrayList (start + m) (len - m) bufLen ptr++{-# INLINE minimizeMemory #-}+minimizeMemory :: forall a. Storable a => ArrayList a -> IO (ArrayList a)+minimizeMemory xs@(ArrayList start len bufLen ptr)+ -- We do not drop below a certain size, since then we would+ -- end up doing frequent reallocations. Although, once the size+ -- reaches zero, we deallocate entirely since this can save a lot+ -- of memory when we have many empty ArrayLists.+ | len == 0 = do+ FMA.free ptr+ return (ArrayList 0 0 0 nullPtr)+ | bufLen <= initialSize = return xs+ | len < eighth bufLen = do+ newPtr <- FMA.mallocBytes (FS.sizeOf (undefined :: a) * div bufLen 2)+ moveArray newPtr (advancePtr ptr start) len+ FMA.free ptr+ return (ArrayList 0 len (div bufLen 2) newPtr)+ | otherwise = return (ArrayList start len bufLen ptr)+ ++{-# INLINE half #-}+half :: Int -> Int+half x = unsafeShiftR x 1++{-# INLINE quarter #-}+quarter :: Int -> Int+quarter x = unsafeShiftR x 2++{-# INLINE eighth #-}+eighth :: Int -> Int+eighth x = unsafeShiftR x 3++{-# INLINE sixteenth #-}+sixteenth :: Int -> Int+sixteenth x = unsafeShiftR x 4++-- | This should not be used in production code.+dumpList :: (Prim a, Storable a) => ArrayList a -> IO (ArrayList a, [a])+dumpList xs@(ArrayList _ len _ _) = do+ marr <- newPrimArray len+ newXs <- dump xs marr 0+ arr <- unsafeFreezePrimArray marr+ return (newXs,primArrayToListN len arr)++primArrayToListN :: forall a. Prim a => Int -> PrimArray a -> [a]+primArrayToListN len arr = go 0+ where+ go :: Int -> [a]+ go !ix = if ix < len+ then indexPrimArray arr ix : go (ix + 1)+ else []+ +-- | Deletes all elements from the linked list, copying them+-- into the buffer specified by the pointer. Returns an+-- empty linked list.+dump :: (Prim a, Storable a)+ => ArrayList a -> MutablePrimArray RealWorld a -> Int -> IO (ArrayList a)+dump xs@(ArrayList start len _ ptr) marr ix = do+ copyPtrToMutablePrimArray marr ix (advancePtr ptr start) len+ clear xs++-- | Dump the elements into a 'MutablePrimArray', mapping over them+-- first. This is a fairly niche function.+dumpMap :: (Storable a, Prim b)+ => ArrayList a -> (a -> b) -> MutablePrimArray RealWorld b -> Int -> IO (ArrayList a)+dumpMap xs@(ArrayList start len _ ptr) f marr ix = do+ let go :: Int -> IO ()+ go !i = if i < len+ then do+ a <- peekElemOff ptr (start + i)+ writePrimArray marr (ix + i) (f a)+ else return ()+ go 0+ clear xs++-- | Does not affect the contents of the ArrayList+showDebug :: forall a. (Prim a, Storable a, Show a) => ArrayList a -> IO String+showDebug (ArrayList start len _ ptr) = do+ marr <- newPrimArray len+ copyPtrToMutablePrimArray marr 0 (plusPtr ptr (start * PM.sizeOf (undefined :: a))) len+ arr <- unsafeFreezePrimArray marr+ return (show (primArrayToListN len arr :: [a]))++clear :: Storable a => ArrayList a -> IO (ArrayList a)+clear xs@(ArrayList _ len _ _) = dropL xs len++-- | Final consumer of the ArrayList.+free :: ArrayList a -> IO ()+free (ArrayList _ _ _ ptr) = FMA.free ptr+
src/BTree.hs view
@@ -18,7 +18,7 @@ ) where import Prelude hiding (lookup)-import Data.Primitive hiding (fromList)+import Data.Primitive (Prim) import Control.Monad.ST import Data.Primitive.MutVar
− src/BTree/Array.hs
@@ -1,21 +0,0 @@-{-# 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
@@ -1,590 +0,0 @@-{-# 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- , Decision(..)- , new- , debugMap- , insert- , modifyWithM- , lookup- , toAscList- , foldrWithKey- ) 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--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 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 (BNode k v) where- unsafeContractedUnliftedPtrCount# _ = 4#- unsafeContractedByteCount# _ = sizeOf# (undefined :: Int) *# 2#- readContractedArray# ba aa ix s1 =- let ixByte = ix *# 2#- ixPtr = ix *# 4#- in case readIntArray# ba (ixByte +# 0#) s1 of- (# s2, sz #) -> case readIntArray# ba (ixByte +# 1#) s2 of- (# s3, toggle #) -> case readMutableByteArrayArray# aa (ixPtr +# 0#) s3 of- (# 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 *# 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#) 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 BNode.-data FlattenedContents k v s c = FlattenedContents- {-# UNPACK #-} !Int- {-# UNPACK #-} !(MutablePrimArray s v)- {-# UNPACK #-} !(ContractedMutableArray (BNode k v) s c)--data Contents k v s c- = ContentsValues {-# UNPACK #-} !(MutablePrimArray s v)- | ContentsNodes {-# UNPACK #-} !(ContractedMutableArray (BNode 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 (BNode 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 (BNode k v) s c-demandFlattenedContentsNodes (FlattenedContents _ _ nodes) = nodes--data Insert k v s c- = Ok- !v- {-# UNPACK #-} !Int -- new size of left child- | Split- {-# NOUNPACK #-} !(BNode k v s c)- !k- !v- {-# UNPACK #-} !Int- -- ^ The new node that will go to the right,- -- the key propagated to the parent,- -- the inserted value, updated sizing info for the left child--{-# INLINE mkBTree #-}-mkBTree :: PrimMonad m- => Token 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 (BNode k v (PrimState m) c)-mkBTree token garbage a b c = do- let !garbageValues = coercePrimArray b- !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)- => Token c- -> Int -- ^ degree, must be at least 3- -> m (BTree k v (PrimState m) c)-new !token !degree = do- if degree < 3- then error "Btree.new: max nodes per child cannot be less than 3"- else return ()- !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- node <- mkBTree token garbageNodes 0 keys (ContentsValues values)- return (BTree degree node)---- {-# 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 (BTree _ theNode) k = go theNode- where- 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- 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 node--{-# INLINE insert #-}-insert :: (Ord k, Prim k, Prim v, PrimMonad m)- => Token c- -> BTree k v (PrimState m) c- -> k- -> v- -> m (BTree k v (PrimState m) c)-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 :: 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)- => 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 !token (BTree !degree !root) !k !newValue alter = do- !ins <- go root- case ins of- 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 { _bnodeSize = newLeftSize }- !newRoot@(BNode _ _ (FlattenedContents _ _ cmptRootChildren)) <- mkBTree token newRootChildren 1 newRootKeys (ContentsNodes newRootChildren)- writeContractedArray cmptRootChildren 0 leftNode- writeContractedArray cmptRootChildren 1 rightNode- return (v,BTree degree newRoot)- where- 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- if ix < 0- then do- let !gtIx = decodeGtIndex ix- !v = newValue- if sz < degree - 1- then do- -- We have enough space- unsafeInsertPrimArray sz gtIx k keys- unsafeInsertPrimArray sz gtIx v values- 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- 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- copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize- copyMutablePrimArray rightValues 0 leftValues leftSize rightSize- unsafeInsertPrimArray leftSize gtIx k leftKeys- unsafeInsertPrimArray leftSize gtIx v leftValues- else do- -- 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- 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' sz)- ContentsNodes nodes -> do- !(!gtIx,!isEq) <- findIndexGte keys k sz- -- case e of- -- Right _ -> error "write Right case"- -- Left gtIx -> do- let !nodeIx = if isEq then gtIx + 1 else gtIx- !node <- readContractedArray nodes nodeIx- !ins <- go node- case ins of- 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--- * 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)- => BTree k v (PrimState m) c- -> m String-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- pairStrs <- showPairs sz keys values- return (map (\s -> (level,s)) pairStrs)- ContentsNodes nodes -> do- pairs <- pairForM sz keys nodes- $ \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- 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)- 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)- => 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)--foldrWithKey :: forall m k v b c. (PrimMonad m, Ord k, Prim k, Prim v)- => (k -> v -> b -> m b)- -> b- -> BTree k v (PrimState m) c- -> m b-foldrWithKey f b0 (BTree _ root) = flip go b0 root- where- 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--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
@@ -1,516 +0,0 @@-{-# 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- , Decision(..)- , new- , modifyWithM- , lookup- , foldrWithKey- ) 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--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 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 (BNode 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, sz #) -> case readIntArray# ba (ixByte +# 1#) s2 of- (# s3, toggle #) -> case readMutableByteArrayArray# aa (ixPtr +# 0#) s3 of- (# 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 *# 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#) 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 BNode.-data FlattenedContents (k :: *) (v :: * -> Heap -> *) (s :: *) (c :: Heap) = FlattenedContents- {-# UNPACK #-} !Int- {-# UNPACK #-} !(ContractedMutableArray 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 (BNode 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 (BNode 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 (BNode k v) s c-demandFlattenedContentsNodes (FlattenedContents _ _ nodes) = nodes--data Insert k v s c- = Ok- !(v s c)- {-# UNPACK #-} !Int -- new size of left child- | Split- {-# NOUNPACK #-} !(BNode k v s c)- !k- !(v 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 for the left child--{-# INLINE mkBTree #-}-mkBTree :: PrimMonad m- => Token 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 (BNode k v (PrimState m) c)-mkBTree token garbage a b c = do- let !garbageValues = coerceContactedArray garbage- !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)- => Token c- -> Int -- ^ degree, must be at least 3- -> m (BTree k v (PrimState m) c)-new !token !degree = do- if degree < 3- then error "Btree.new: max nodes per child cannot be less than 3"- else return ()- !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- node <- mkBTree token garbageNodes 0 keys (ContentsValues values)- return (BTree degree node)---- {-# 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 (BTree _ theNode) k = go theNode- where- 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- 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--data Decision a = Keep | Replace !a---- 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 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 !token (BTree !degree !root) !k !newValue alter = do- !ins <- go root- case ins of- 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 { _bnodeSize = newLeftSize }- !newRoot@(BNode _ _ (FlattenedContents _ _ cmptRootChildren)) <- mkBTree token newRootChildren 1 newRootKeys (ContentsNodes newRootChildren)- writeContractedArray cmptRootChildren 0 leftNode- writeContractedArray cmptRootChildren 1 rightNode- return (v,BTree degree newRoot)- where- 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- 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- unsafeInsertPrimArray sz gtIx k keys- unsafeInsertContractedArray sz gtIx v values- 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- rightKeys' <- newPrimArray (degree - 1)- rightValues' <- newContractedArray token (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- copyMutablePrimArray rightKeys 0 leftKeys leftSize rightSize- copyContractedMutableArray rightValues 0 leftValues leftSize rightSize- unsafeInsertPrimArray leftSize gtIx k leftKeys- unsafeInsertContractedArray leftSize gtIx v leftValues- else do- -- 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- !propagated <- readPrimArray rightKeys 0- 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' sz)- ContentsNodes nodes -> do- !(!gtIx,!isEq) <- findIndexGte keys k sz- -- case e of- -- Right _ -> error "write Right case"- -- Left gtIx -> do- let !nodeIx = if isEq then gtIx + 1 else gtIx- !node <- readContractedArray nodes nodeIx- !ins <- go node- case ins of- 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--- * 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--foldrWithKey :: forall m k v b c. (PrimMonad m, Ord k, Prim k, Contractible v)- => (k -> v (PrimState m) c -> b -> m b)- -> b- -> BTree k v (PrimState m) c- -> m b-foldrWithKey f b0 (BTree _ root) = flip go b0 root- where- 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--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
@@ -18,10 +18,11 @@ ) where import Prelude hiding (lookup)-import Data.Primitive hiding (fromList) import Data.Primitive.MutVar import Control.Monad import Data.Foldable (foldlM)+import Data.Primitive (MutableArray,Prim)+import qualified Data.Primitive as P import Data.Primitive.PrimArray import Control.Monad.Primitive@@ -69,7 +70,7 @@ return (Just v) ContentsNodes nodes -> do ix <- findIndexBetween keys k sz- go =<< readArray nodes ix+ go =<< P.readArray nodes ix data Insert s k v = Ok !v@@ -138,7 +139,7 @@ go :: Int -> b -> m b go !ix !b1 = if ix >= 0 then do- a <- readArray arr ix+ a <- P.readArray arr ix b2 <- f a b1 go (ix - 1) b2 else return b1@@ -178,9 +179,9 @@ newRootSz <- newMutVar 1 newRootKeys <- newPrimArray (degree - 1) writePrimArray newRootKeys 0 newRootKey- newRootChildren <- newArray degree uninitializedNode- writeArray newRootChildren 0 leftNode- writeArray newRootChildren 1 rightNode+ newRootChildren <- P.newArray degree uninitializedNode+ P.writeArray newRootChildren 0 leftNode+ P.writeArray newRootChildren 1 rightNode let newRoot = BTree newRootSz newRootKeys (ContentsNodes newRootChildren) return (v,newRoot) where@@ -244,7 +245,7 @@ -- case e of -- Right _ -> error "write Right case" -- Left gtIx -> do- node <- readArray nodes (if isEq then gtIx + 1 else gtIx)+ node <- P.readArray nodes (if isEq then gtIx + 1 else gtIx) ins <- go node case ins of Ok v -> return (Ok v)@@ -260,14 +261,14 @@ leftNodes = nodes middleKey <- readPrimArray keys middleIx rightKeys :: MutablePrimArray (PrimState m) k <- newPrimArray (degree - 1)- rightNodes <- newArray degree uninitializedNode+ rightNodes <- P.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+ P.copyMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize unsafeInsertPrimArray leftSize gtIx propagated leftKeys unsafeInsertArray (leftSize + 1) (gtIx + 1) rightNode leftNodes writeMutVar szRef (leftSize + 1)@@ -277,7 +278,7 @@ -- 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+ P.copyMutableArray rightNodes 0 leftNodes (leftSize + 1) rightSize unsafeInsertPrimArray (rightSize - 1) (gtIx - leftSize - 1) propagated rightKeys unsafeInsertArray rightSize (gtIx - leftSize) rightNode rightNodes writeMutVar szRef leftSize@@ -347,8 +348,8 @@ -> 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+ P.copyMutableArray marr (i + 1) marr i (sz - i)+ P.writeArray marr i x -- Inserts a value at the designated index, -- shifting everything after it to the right.@@ -407,7 +408,7 @@ 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+ BTree lastSzRef lastKeys lastContents <- P.readArray nodes sz lastSz <- readMutVar lastSzRef lastStrs <- go (level + 1) lastSz lastKeys lastContents -- return (nextStrs ++ [(level,show k)])@@ -427,7 +428,7 @@ go ix = if ix < sz then do a <- readPrimArray marr1 ix- c <- readArray marr2 ix+ c <- P.readArray marr2 ix b <- f a c bs <- go (ix + 1) return (b : bs)
+ src/BTree/Store.hs view
@@ -0,0 +1,1011 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE MultiWayIf #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}++-- {-# OPTIONS_GHC -Wall -Werror -O2 #-}++module BTree.Store+ ( BTree+ , Initialize(..)+ , Deinitialize(..)+ , Decision(..)+ , new+ , free+ , with+ , with_+ , lookup+ , insert+ , modifyWithM_+ , modifyWithM+ , modifyWithPtr+ , foldrWithKey+ , toAscList+ -- * Weird Operations+ , index+ , indexNode+ -- * Force inlining+ , inlineModifyWithPtr+ , inlineModifyWithM+ ) where++import Prelude hiding (lookup)+import Foreign.Storable+import Foreign.Ptr+import Foreign.Marshal.Alloc hiding (free)+import Foreign.Marshal.Array+import Data.Bits+import Data.Word+import Data.Int+import GHC.Ptr (Ptr(..))+import GHC.Magic (inline)+import qualified Data.Primitive.Addr as PA+import qualified Foreign.Marshal.Alloc as FMA++data BTree k v = BTree + !Int -- height+ !(Ptr (Node k v)) -- root node++data Node k v++class Storable a => Initialize a where+ initialize :: Ptr a -> IO ()+ -- ^ Initialize the memory at a pointer. An implementation+ -- of this function may do nothing, or if the data contains+ -- more pointers, @initialize@ may allocate additional memory.+ initializeElemOff :: Ptr a -> Int -> IO ()+ -- ^ Can be overridden for efficiency+ initializeElemOff ptr ix = do+ initialize (plusPtr ptr (ix * sizeOf (undefined :: a)) :: Ptr a)+ initializeElems :: Ptr a -> Int -> IO ()+ -- ^ Initialize a pointer representing an array with+ -- a given number of elements. This has a default implementation+ -- but may be overriden for efficency.+ initializeElems ptr n = go 0+ where+ go !i = if i < n+ then do+ initialize (plusPtr ptr (i * sizeOf (undefined :: a)) :: Ptr a)+ go (i + 1)+ else return ()++class Storable a => Deinitialize a where+ deinitialize :: Ptr a -> IO ()+ deinitializeElemOff :: Ptr a -> Int -> IO ()+ -- ^ Can be overridden for efficiency+ deinitializeElemOff ptr ix =+ deinitialize (plusPtr ptr (ix * sizeOf (undefined :: a)) :: Ptr a)+ -- ^ Free any memory in the data structure pointed to.+ deinitializeElems :: Ptr a -> Int -> IO ()+ -- ^ Free any memory pointed to by elements of the array.+ -- This has a default implementation+ -- but may be overriden for efficency.+ deinitializeElems ptr n = go 0+ where+ go !i = if i < n+ then do+ deinitialize (plusPtr ptr (i * sizeOf (undefined :: a)) :: Ptr a)+ go (i + 1)+ else return ()++instance Storable (BTree k v) where+ sizeOf _ = 2 * sizeOf (undefined :: Int)+ alignment _ = alignment (undefined :: Int)+ peek ptr = do+ height <- peekElemOff (castPtr ptr :: Ptr Int) 0+ root <- peekElemOff (castPtr ptr :: Ptr (Ptr (Node k v))) 1+ return (BTree height root)+ poke ptr (BTree height root) = do+ pokeElemOff (castPtr ptr :: Ptr Int) 0 height+ pokeElemOff (castPtr ptr :: Ptr (Ptr (Node k v))) 1 root++-- this instance relies on Int and Ptr being the same+-- size. this seems to be true for everything that+-- GHC targets.+--+-- This instance bypasses the check on the size of the keys+-- and values. This is not good.+instance Initialize (BTree k v) where+ initialize ptr = do+ pokeElemOff (castPtr ptr :: Ptr Int) 0 (0 :: Int)+ pokeElemOff (castPtr ptr :: Ptr (Ptr (Node k v))) 1 =<< newNode 0++instance (Storable k, Deinitialize v) => Deinitialize (BTree k v) where+ deinitialize ptr = do+ bt <- peek ptr+ free bt++newtype Uninitialized a = Uninitialized a+ deriving (Storable)++instance Storable a => Initialize (Uninitialized a) where+ initialize _ = return ()+ initializeElemOff _ _ = return ()+ initializeElems _ _ = return ()++instance Storable a => Deinitialize (Uninitialized a) where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Initialize Word8 where+ initialize _ = return ()+ initializeElemOff _ _ = return ()+ initializeElems _ _ = return ()++instance Deinitialize Word8 where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Initialize Word16 where+ initialize _ = return ()+ initializeElemOff _ _ = return ()+ initializeElems _ _ = return ()++instance Deinitialize Word16 where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Initialize Word64 where+ initialize _ = return ()+ initializeElemOff _ _ = return ()+ initializeElems _ _ = return ()++instance Deinitialize Word64 where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()+++instance Initialize Word where+ initialize ptr = poke ptr (0 :: Word)+ initializeElemOff ptr off = pokeElemOff ptr off (0 :: Word)+ initializeElems ptr elemLen = PA.setAddr (ptrToAddr ptr) elemLen (0 :: Word)++instance Deinitialize Word where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Initialize Int where+ initialize ptr = poke ptr (0 :: Int)+ initializeElemOff ptr off = pokeElemOff ptr off (0 :: Int)+ initializeElems ptr elemLen = PA.setAddr (ptrToAddr ptr) elemLen (0 :: Int)++instance Initialize Int64 where+ initialize ptr = poke ptr (0 :: Int64)+ initializeElemOff ptr off = pokeElemOff ptr off (0 :: Int64)+ initializeElems ptr elemLen = PA.setAddr (ptrToAddr ptr) elemLen (0 :: Int64)++ptrToAddr :: Ptr a -> PA.Addr+ptrToAddr (Ptr x) = PA.Addr x++instance Initialize Word32 where+ initialize _ = return ()+ initializeElemOff _ _ = return ()+ initializeElems _ _ = return ()++instance Deinitialize Word32 where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Deinitialize Int where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Deinitialize Int64 where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()++instance Initialize Char where+ initialize _ = return ()+ initializeElemOff _ _ = return ()+ initializeElems _ _ = return ()++instance Deinitialize Char where+ deinitialize _ = return ()+ deinitializeElemOff _ _ = return ()+ deinitializeElems _ _ = return ()+++newtype Arr a = Arr { getArr :: Ptr a }+data KeysValues k v = KeysValues !(Arr k) !(Arr v)+data KeysNodes k v = KeysNodes !(Arr k) !(Arr (Ptr (Node k v)))++new :: forall k v. (Storable k, Storable v) => IO (BTree k v)+new = do+ -- we only calculate these degrees so that we can do one+ -- upfront check instead of check every time we call insert,+ -- which would be weird. This also helps us see the failure+ -- sooner.+ let childDegree = calcChildDegree (undefined :: Ptr (Node k v))+ branchDegree = calcBranchDegree (undefined :: Ptr (Node k v))+ if childDegree < minimumDegree+ then fail $ "Btree.new: child degree cannot be less than " ++ show minimumDegree+ else return ()+ if branchDegree < minimumDegree+ then fail $ "Btree.new: branch degree cannot be less than " ++ show minimumDegree+ else return ()+ ptr <- newNode 0+ return (BTree 0 ptr)++minimumDegree :: Int+minimumDegree = 6++-- | Release all memory allocated by the b-tree. Do not attempt+-- to use the b-tree after calling this.+free :: forall k v. (Storable k, Deinitialize v) => BTree k v -> IO ()+free (BTree height root) = go height root+ where+ branchDegree :: Int+ !branchDegree = calcBranchDegree root+ childDegree :: Int+ !childDegree = calcChildDegree root+ go :: Int -> Ptr (Node k v) -> IO ()+ go n ptrNode = if n > 0+ then do+ sz <- readNodeSize ptrNode+ let KeysNodes _ nodes = readNodeKeysNodes branchDegree ptrNode+ arrMapM_ (go (n - 1)) (sz + 1) nodes+ FMA.free ptrNode+ else do+ sz <- readNodeSize ptrNode+ let KeysValues _ values = readNodeKeysValues childDegree ptrNode+ deinitializeElems (getArr values) sz+ FMA.free ptrNode++with :: (Storable k, Initialize v, Deinitialize v) => (BTree k v -> IO (a, BTree k v)) -> IO a+with f = do+ initial <- new+ (a,final) <- f initial+ free final+ return a++with_ :: (Storable k, Initialize v, Deinitialize v) => (BTree k v -> IO (BTree k v)) -> IO ()+with_ f = do+ initial <- new+ final <- f initial+ free final++newNode :: + Int -- ^ initial size, if you pick something greater than 0,+ -- you need to write to those indices after calling this.+ -> IO (Ptr (Node k v))+newNode n = do+ -- We would really like to ensure that this is aligned to a+ -- 4k boundary, but malloc does not guarentee this. I think+ -- that posix_memalign should work, but whatever.+ ptr <- mallocBytes maxSize+ poke ptr n+ return (castPtr ptr)+ +readArr :: Storable a => Arr a -> Int -> IO a+readArr (Arr ptr) ix = peekElemOff ptr ix++writeArr :: Storable a => Arr a -> Int -> a -> IO ()+writeArr (Arr ptr) ix a = pokeElemOff ptr ix a++readNodeSize :: Ptr (Node k v) -> IO Int+readNodeSize ptr = peek (castPtr ptr)++writeNodeSize :: Ptr (Node k v) -> Int -> IO ()+writeNodeSize ptr sz = poke (castPtr ptr) sz++readNodeKeys :: forall k v. Storable k => Ptr (Node k v) -> Arr k+readNodeKeys ptr1 =+ let ptr2 = plusPtr ptr1 (sizeOf (undefined :: Int))+ ptr3 = alignPtr ptr2 (alignment (undefined :: k))+ in Arr ptr3++readNodeKeysValues :: forall k v. Storable k => Int -> Ptr (Node k v) -> KeysValues k v+readNodeKeysValues degree ptr1 = + let keys = readNodeKeys ptr1+ ptr2 = plusPtr (getArr keys) (sizeOf (undefined :: k) * (degree - 1))+ ptr3 = alignPtr ptr2 (alignment (undefined :: k))+ in KeysValues keys (Arr ptr3)++readNodeKeysNodes :: forall k v. Storable k => Int -> Ptr (Node k v) -> KeysNodes k v+readNodeKeysNodes degree ptr1 = + let keys = readNodeKeys ptr1+ ptr2 = plusPtr (getArr keys) (sizeOf (undefined :: k) * (degree - 1))+ ptr3 = alignPtr ptr2 (alignment (undefined :: (Ptr (Node k v))))+ in KeysNodes keys (Arr ptr3)++maxSize :: Int+maxSize = 4096 - 2 * sizeOf (undefined :: Int)+-- maxSize = 200++-- not actually sure if this is really correct.+{-# INLINE calcBranchDegree #-}+calcBranchDegree :: forall k v. (Storable k, Storable v) => Ptr (Node k v) -> Int+calcBranchDegree _ = calcBranchDegreeInt (sizeOf (undefined :: k)) (alignment (undefined :: k))++{-# INLINE calcBranchDegreeInt #-}+calcBranchDegreeInt :: Int -> Int -> Int+calcBranchDegreeInt keySz keyAlignment = + let space = maxSize - max (sizeOf (undefined :: Int)) keyAlignment - sizeOf (undefined :: Ptr a)+ allowedNodes = quot space (sizeOf (undefined :: Ptr a) + keySz)+ in allowedNodes++-- not actually sure if this is really correct. need to think about this math+-- a little more. Or I guess I could write something that does a brute force+-- consideration of all the possible sizes and alignment. That would convince me.+{-# INLINE calcChildDegree #-}+calcChildDegree :: forall k v. (Storable k, Storable v) => Ptr (Node k v) -> Int+calcChildDegree _ = calcChildDegreeInt+ (sizeOf (undefined :: k))+ (alignment (undefined :: k))+ (sizeOf (undefined :: v))++{-# INLINE calcChildDegreeInt #-}+calcChildDegreeInt :: Int -> Int -> Int -> Int+calcChildDegreeInt keySz keyAlignment valueSz = + let space = maxSize - max (sizeOf (undefined :: Int)) keyAlignment - valueSz+ allowedValues = quot space (valueSz + keySz)+ in allowedValues + 1 -- add one because of the meaning we assign to degree++{-# INLINABLE lookup #-}+-- {-# SPECIALIZE lookup :: BTree Int Int -> Int -> IO (Maybe Int) #-}+-- {-# SPECIALIZE lookup :: BTree Int64 Int -> Int64 -> IO (Maybe Int) #-}+-- {-# SPECIALIZE lookup :: BTree Word32 Int -> Word32 -> IO (Maybe Int) #-}+-- {-# SPECIALIZE lookup :: BTree Word16 Int -> Word16 -> IO (Maybe Int) #-}+lookup :: forall k v. (Ord k, Storable k, Storable v)+ => BTree k v -> k -> IO (Maybe v)+lookup (BTree height rootNode) k = go height rootNode+ where+ branchDegree :: Int+ !branchDegree = calcBranchDegree rootNode+ childDegree :: Int+ childDegree = calcChildDegree rootNode+ go :: Int -> Ptr (Node k v) -> IO (Maybe v)+ go !n !ptrNode = if n > 0+ then do+ !sz <- readNodeSize ptrNode+ let !(KeysNodes keys nodes) = readNodeKeysNodes branchDegree ptrNode+ !ix <- findIndexOfGtElem keys k sz+ !node <- readArr nodes ix+ go (n - 1) node+ else do+ !sz <- readNodeSize ptrNode+ let !(KeysValues keys values) = readNodeKeysValues childDegree ptrNode+ !ix <- findIndex keys k sz+ if ix < 0+ then return Nothing+ else do+ !v <- readArr values ix+ return (Just v)++index :: forall k v. (Storable k, Storable v) => BTree k v -> (Int -> Int) -> Int -> IO v+index (BTree height rootNode) f = go height rootNode+ where+ branchDegree :: Int+ !branchDegree = calcBranchDegree rootNode+ go :: Int -> Ptr (Node k v) -> Int -> IO v+ go !n !ptrNode !k = if n > 0+ then do+ !sz <- readNodeSize ptrNode+ let !ix = mod k sz+ let !(KeysNodes keys nodes) = readNodeKeysNodes branchDegree ptrNode+ !node <- readArr nodes ix+ go (n - 1) node (f k)+ else do+ !sz <- readNodeSize ptrNode+ let !(KeysValues keys !values) = readNodeKeysValues (calcChildDegree rootNode) ptrNode+ readArr values (mod k sz)++-- This function is only provided so that I can randomly choose+-- a leaf of the B-Tree and garbage collect old things.+indexNode :: forall k v. (Storable k, Storable v) => BTree k v -> (Int -> Int) -> Int -> IO (Ptr v, Int)+indexNode (BTree height rootNode) f = go height rootNode+ where+ branchDegree :: Int+ !branchDegree = calcBranchDegree rootNode+ go :: Int -> Ptr (Node k v) -> Int -> IO (Ptr v, Int)+ go !n !ptrNode !k = if n > 0+ then do+ !sz <- readNodeSize ptrNode+ let !ix = mod k sz+ let !(KeysNodes keys nodes) = readNodeKeysNodes branchDegree ptrNode+ !node <- readArr nodes ix+ go (n - 1) node (f k)+ else do+ !sz <- readNodeSize ptrNode+ let !(KeysValues keys !values) = readNodeKeysValues (calcChildDegree rootNode) ptrNode+ return (getArr values, sz)++data Insert k v r+ = Ok !r+ | Split !(Ptr (Node k v)) !k !r+ -- The new node that will go to the right,+ -- the key propagated to the parent,+ -- the inserted value+ | TooSmall !r+ | TotallyEmpty !(Ptr (Node k v)) !r+ -- The node has zero keys left. Its sole child+ -- is provided.++{-# INLINE insert #-}+insert :: (Ord k, Storable k, Initialize v)+ => BTree k v+ -> k+ -> v+ -> IO (BTree k v)+insert !m !k !v = do+ !(!(),!node) <- modifyWithPtr m k+ (Right (\ptr ix -> pokeElemOff ptr ix v))+ (\ptr ix -> pokeElemOff ptr ix v >> return ((),Keep))+ return node++-- delete :: (Ord k, Storable k, Regioned v)+-- => BTree k v+-- -> k+-- -> IO (BTree k v)+-- delete !m !k = do+-- !(!(),!node) <- modifyWithPtr m k+-- (Left ())+-- (\_ _ -> return ((),Delete))+-- return node++data Decision = Keep | Delete++-- data Position = Next | Prev++{-# INLINE modifyWithM_ #-}+modifyWithM_ :: forall k v. (Ord k, Storable k, Initialize v)+ => BTree k v + -> k+ -> (v -> IO v) -- ^ value modification, happens for newly inserted elements and for previously existing elements+ -> IO (BTree k v)+modifyWithM_ bt k alter = do+ (_, bt') <- modifyWithPtr bt k+ (Right (\ptr ix -> peekElemOff ptr ix >>= alter >>= pokeElemOff ptr ix))+ (\ptr ix -> peekElemOff ptr ix >>= alter >>= pokeElemOff ptr ix >>= \_ -> return ((),Keep))+ return bt'++{-# INLINE modifyWithM #-}+modifyWithM :: forall k v a. (Ord k, Storable k, Initialize v)+ => BTree k v + -> k+ -> (v -> IO (a, v)) -- ^ value modification, happens for newly inserted elements and for previously existing elements+ -> IO (a, BTree k v)+modifyWithM bt k alter = do+ (a, bt') <- modifyWithPtr bt k+ (Right (\ptr ix -> do+ (a,v') <- alter =<< peekElemOff ptr ix+ pokeElemOff ptr ix v'+ return a+ ))+ (\ptr ix -> do+ (a,v') <- alter =<< peekElemOff ptr ix+ pokeElemOff ptr ix v'+ return (a,Keep)+ )+ return (a,bt')++{-# INLINE inlineModifyWithM #-}+inlineModifyWithM :: forall k v a. (Ord k, Storable k, Initialize v)+ => BTree k v + -> k+ -> (v -> IO (a, v)) -- ^ value modification, happens for newly inserted elements and for previously existing elements+ -> IO (a, BTree k v)+inlineModifyWithM bt k alter = do+ (a, bt') <- inlineModifyWithPtr bt k+ (Right (\ptr ix -> do+ (a,v') <- alter =<< peekElemOff ptr ix+ pokeElemOff ptr ix v'+ return a+ ))+ (\ptr ix -> do+ (a,v') <- alter =<< peekElemOff ptr ix+ pokeElemOff ptr ix v'+ return (a,Keep)+ )+ return (a,bt')++{-# NOINLINE modifyWithPtr #-}+modifyWithPtr :: forall k v r. (Ord k, Storable k, Initialize v)+ => BTree k v + -> k+ -> (Either r (Ptr v -> Int -> IO r)) -- ^ modifications to newly inserted value+ -> (Ptr v -> Int -> IO (r,Decision)) -- ^ modification to value if key is found+ -> IO (r, BTree k v)+modifyWithPtr a b c d = inlineModifyWithPtr a b c d++{-# INLINE inlineModifyWithPtr #-}+inlineModifyWithPtr :: forall k v r. (Ord k, Storable k, Initialize v)+ => BTree k v + -> k+ -> (Either r (Ptr v -> Int -> IO r)) -- ^ modifications to newly inserted value+ -> (Ptr v -> Int -> IO (r,Decision)) -- ^ modification to value if key is found+ -> IO (r, BTree k v)+inlineModifyWithPtr (BTree !height !root) !k !mpostInitializeElemOff alterElemOff = do+ !ins <- go height root+ case ins of+ Ok !r -> return (r, BTree height root)+ TotallyEmpty child r -> do+ FMA.free root+ return (r, BTree (height - 1) child)+ -- if the root is too small, we do not care. The root+ -- can have any number of keys greater than 1.+ TooSmall !r -> return (r, BTree 0 root)+ Split !rightNode !newRootKey !v -> do+ newRoot <- newNode 1+ let KeysNodes keys nodes = readNodeKeysNodes branchDegree newRoot+ leftNode = root+ writeArr keys 0 newRootKey+ writeArr nodes 0 leftNode+ writeArr nodes 1 rightNode+ return (v,BTree (height + 1) newRoot)+ where+ childDegree :: Int+ !childDegree = calcChildDegree root+ branchDegree :: Int+ !branchDegree = calcBranchDegree root+ go :: Int -> Ptr (Node k v) -> IO (Insert k v r)+ go n ptrNode = if n > 0+ then do+ sz <- readNodeSize ptrNode+ let KeysNodes keys nodes = readNodeKeysNodes branchDegree ptrNode+ !gtIx <- findIndexOfGtElem keys k sz+ !node <- readArr nodes gtIx+ !ins <- go (n - 1) node+ case ins of+ Ok !r -> return (Ok r)+ TotallyEmpty _ _ -> fail "TotallyEmpty: handle this in go"+ TooSmall !r -> do+ if n == 1+ then + if | gtIx >= sz -> do+ if (gtIx /= sz) then fail "bad logic found: gtIx must be sz" else return ()+ childSz <- readNodeSize node+ let KeysValues childKeys childValues = readNodeKeysValues childDegree node+ prevPtrNode <- readArr nodes (gtIx - 1)+ prevSz <- readNodeSize prevPtrNode+ let KeysValues prevKeys prevValues = readNodeKeysValues childDegree prevPtrNode+ if childSz + prevSz < childDegree+ then do+ mergeIntoLeft prevKeys prevValues prevSz childKeys childValues childSz+ writeNodeSize prevPtrNode (childSz + prevSz)+ FMA.free node+ if sz < 2+ then do+ -- whatever code handles this one level up needs+ -- to remember to call free on the now-obsolete+ -- branch node. + return (TotallyEmpty prevPtrNode r)+ else do+ -- putStrLn $ "size of nodes: " ++ show sz+ _ <- fail "merging arrays"+ removeArr sz (sz - 1) keys -- first key+ removeArr (sz + 1) sz nodes -- right child of first key+ writeNodeSize ptrNode (sz - 1)+ continue+ else do+ -- putStrLn $ "child size: " ++ show childSz+ -- putStrLn $ "next size: " ++ show nextSz+ (newPrevSz,newChildSz) <- balanceArrays prevKeys prevValues prevSz childKeys childValues childSz+ writeNodeSize prevPtrNode newPrevSz+ writeNodeSize node newChildSz+ readArr childKeys 0 >>= writeArr keys (sz - 1)+ continue+ | gtIx > 0 -> fail "write me now"+ -- childSz <- readNodeSize node+ -- let KeysValues childKeys childValues = readNodeKeysValues childDegree node+ -- nextPtrNode <- readArr nodes (gtIx + 1)+ -- nextSz <- readNodeSize nextPtrNode+ -- let KeysValues nextKeys nextValues = readNodeKeysValues childDegree nextPtrNode+ -- prevPtrNode <- readArr nodes (gtIx - 1)+ -- prevSz <- readNodeSize prevPtrNode+ -- let KeysValues prevKeys prevValues = readNodeKeysValues childDegree prevPtrNode+ -- if nextSz > prevSz+ -- then runNext + -- else runPrev+ | otherwise -> do -- gtIx must be 0+ if (gtIx /= 0) then fail "bad logic found" else return ()+ childSz <- readNodeSize node+ let KeysValues childKeys childValues = readNodeKeysValues childDegree node+ nextPtrNode <- readArr nodes 1+ nextSz <- readNodeSize nextPtrNode+ let KeysValues nextKeys nextValues = readNodeKeysValues childDegree nextPtrNode+ if childSz + nextSz < childDegree+ then do+ mergeIntoLeft childKeys childValues childSz nextKeys nextValues nextSz+ writeNodeSize node (childSz + nextSz)+ FMA.free nextPtrNode+ -- _ <- fail "after call free"+ if sz < 2+ then do+ -- whatever code handles this one level up needs+ -- to remember to call free on the now-obsolete+ -- branch node. + return (TotallyEmpty node r)+ else do+ -- putStrLn $ "size of nodes: " ++ show sz+ _ <- fail "merging arrays"+ removeArr sz 0 keys -- first key+ removeArr (sz + 1) 1 nodes -- right child of first key+ writeNodeSize ptrNode (sz - 1)+ continue+ else do+ -- putStrLn $ "child size: " ++ show childSz+ -- putStrLn $ "next size: " ++ show nextSz+ _ <- fail "balancing arrays"+ (newChildSz,newNextSz) <- balanceArrays childKeys childValues childSz nextKeys nextValues nextSz+ writeNodeSize nextPtrNode newNextSz+ writeNodeSize node newChildSz+ readArr nextKeys 0 >>= writeArr keys 0+ continue+ else fail "write code for branch handling a branch merge"+ where+ continue :: IO (Insert k v r)+ continue = do+ newSz <- readNodeSize ptrNode+ let minimumBranchSz = half branchDegree - 1+ if newSz < minimumBranchSz+ then return (TooSmall r)+ else return (Ok r)+ -- runNext :: Position -> Int -> Ptr (Node k v) -> Int -> IO (Insert k v r)+ -- runNext _pos _keyIx _neighborPtrNode _neighborSz = fail "write runNext"+ -- childSz <- readNodeSize node+ -- let KeysValues childKeys childValues = readNodeKeysValues childDegree node+ -- let KeysValues neighborKeys neighborValues = readNodeKeysValues childDegree neighborPtrNode+ -- let preservedPtr = case pos of+ -- Next -> node+ -- Prev -> neighborPtrNode+ -- let destroyedPtr = case pos of+ -- Next -> neighborPtrNode+ -- Prev -> node+ -- let destroyedPtrIx = case pos of+ -- Next -> neighborIx - 1+ -- Prev -> neighborIx+ -- if childSz + nextSz < childDegree+ -- then do+ -- case pos of+ -- Next -> mergeIntoLeft childKeys childValues childSz neighborKeys neighborValues neighborSz+ -- Prev -> mergeIntoLeft neighborKeys neighborValues neighborSz childKeys childValues childSz+ -- writeNodeSize preservedPtr (childSz + neighborSz)+ -- FMA.free destroyedPtr+ -- -- _ <- fail "after call free"+ -- if sz < 2+ -- then return (TotallyEmpty preservedPtr r)+ -- else do+ -- -- putStrLn $ "size of nodes: " ++ show sz+ -- _ <- fail "merging arrays"+ -- removeArr sz 0 keys -- first key+ -- removeArr (sz + 1) 1 nodes -- right child of first key+ -- writeNodeSize ptrNode (sz - 1)+ -- continue+ -- else do+ -- -- putStrLn $ "child size: " ++ show childSz+ -- -- putStrLn $ "next size: " ++ show nextSz+ -- _ <- fail "balancing arrays"+ -- (newChildSz,newNextSz) <- balanceArrays childKeys childValues childSz nextKeys nextValues nextSz+ -- writeNodeSize nextPtrNode newNextSz+ -- writeNodeSize node newChildSz+ -- readArr nextKeys 0 >>= writeArr keys 0+ -- continue+ Split !rightNode !propagated !v -> if sz < branchDegree - 1+ then do+ insertArr sz gtIx propagated keys+ insertArr (sz + 1) (gtIx + 1) rightNode nodes+ writeNodeSize ptrNode (sz + 1)+ return (Ok v)+ else do+ let !middleIx = half sz+ !leftKeys = keys+ !leftNodes = nodes+ !middleKey <- readArr keys middleIx+ let !leftSize = middleIx+ !rightSize = sz - leftSize+ (!actualLeftSz,!actualRightSz) = if middleIx >= gtIx+ then (leftSize + 1, rightSize - 1)+ else (leftSize, rightSize)+ newNodePtr <- newNode actualRightSz+ writeNodeSize ptrNode actualLeftSz+ let KeysNodes rightKeys rightNodes = readNodeKeysNodes branchDegree newNodePtr+ if middleIx >= gtIx+ then do+ copyArr rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyArr rightNodes 0 leftNodes (leftSize + 1) rightSize+ insertArr leftSize gtIx propagated leftKeys+ insertArr (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.+ copyArr rightKeys 0 leftKeys (leftSize + 1) (rightSize - 1)+ copyArr rightNodes 0 leftNodes (leftSize + 1) rightSize+ insertArr (rightSize - 1) (gtIx - leftSize - 1) propagated rightKeys+ insertArr rightSize (gtIx - leftSize) rightNode rightNodes+ return (Split newNodePtr middleKey v)+ else do+ sz <- readNodeSize ptrNode+ let !(KeysValues !keys !values) = readNodeKeysValues childDegree ptrNode+ !ix <- findIndex keys k sz+ if ix < 0+ then case mpostInitializeElemOff of+ Left r -> return (Ok r)+ Right postInitializeElemOff -> do+ let !gtIx = decodeGtIndex ix+ if sz < childDegree - 1+ then do+ -- We have enough space+ insertArr sz gtIx k keys+ r <- insertInitArr sz gtIx values $ \thePtr theIx -> do+ initializeElemOff thePtr theIx+ postInitializeElemOff thePtr theIx+ writeNodeSize ptrNode (sz + 1)+ return (Ok r)+ else do+ -- We do not have enough space. The node must be split.+ let !leftSize = half sz+ !rightSize = sz - leftSize+ !leftKeys = keys+ !leftValues = values+ let (newLeftSz,actualRightSz) = if gtIx < leftSize+ then (leftSize + 1, rightSize)+ else (leftSize,rightSize + 1)+ newNodePtr <- newNode actualRightSz+ writeNodeSize ptrNode newLeftSz+ let KeysValues rightKeys rightValues = readNodeKeysValues childDegree newNodePtr+ r <- if gtIx < leftSize+ then do+ copyArr rightKeys 0 leftKeys leftSize rightSize+ copyArr rightValues 0 leftValues leftSize rightSize+ insertArr leftSize gtIx k leftKeys+ insertInitArr leftSize gtIx leftValues $ \thePtr theIx -> do+ initializeElemOff thePtr theIx+ postInitializeElemOff thePtr theIx+ else do+ -- 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.+ copyArr rightKeys 0 leftKeys leftSize rightSize+ copyArr rightValues 0 leftValues leftSize rightSize+ insertArr rightSize (gtIx - leftSize) k rightKeys+ insertInitArr rightSize (gtIx - leftSize) rightValues $ \thePtr theIx -> do+ initializeElemOff thePtr theIx+ postInitializeElemOff thePtr theIx+ !propagated <- readArr rightKeys 0+ return (Split newNodePtr propagated r)+ else do+ -- The key was already present in this leaf node+ !(r,dec) <- alterElemOff (getArr values) ix+ case dec of+ Keep -> return (Ok r)+ Delete -> fail "write the delete code for b tree" -- do+ -- let newSize = sz - 1+ -- minimumChildSz = half childDegree+ -- writeNodeSize ptrNode newSize+ -- removeArr sz ix keys+ -- removeArrDeinit sz ix values+ -- if newSize < minimumChildSz+ -- then return (TooSmall r)+ -- else return (Ok r)++-- this is used when one of the arrays is too small. The+-- caller of this function must ensure in advance that+-- the arrays will end up being appropriately sized+-- after the balancing.+{-# INLINE balanceArrays #-}+balanceArrays :: (Storable k, Storable v) => Arr k -> Arr v -> Int -> Arr k -> Arr v -> Int -> IO (Int,Int)+balanceArrays arrA valA szA arrB valB szB = do+ let newSzA = half (szA + szB)+ newSzB = szA + szB - newSzA+ deltaA = newSzA - szA+ deltaB = negate deltaA+ if deltaA > 0 + then do+ copyArr arrA szA arrB 0 deltaA+ copyArr arrB 0 arrB deltaA (szB - deltaA)+ copyArr valA szA valB 0 deltaA+ copyArr valB 0 valB deltaA (szB - deltaA)+ else do+ copyArr arrB deltaB arrB 0 szB+ copyArr arrB 0 arrA (szA - deltaB) deltaB+ copyArr valB deltaB valB 0 szB+ copyArr valB 0 valA (szA - deltaB) deltaB+ return (newSzA,newSzB)++-- After this operation, all of the values are in the first+-- provided array. The second one should be considered unusable+-- and it should be freed from memory soon.+{-# INLINE mergeIntoLeft #-}+mergeIntoLeft :: (Storable k, Storable v)+ => Arr k -> Arr v -> Int -> Arr k -> Arr v -> Int -> IO ()+mergeIntoLeft arrA valA szA arrB valB szB = do+ copyArr arrA szA arrB 0 szB+ copyArr valA szA valB 0 szB++{-# INLINE copyArr #-}+copyArr :: forall a. Storable a+ => Arr a -- ^ dest+ -> Int -- ^ dest offset+ -> Arr a -- ^ source+ -> Int -- ^ source offset+ -> Int -- ^ length+ -> IO ()+copyArr (Arr dest) doff (Arr src) soff len = moveArray+ (advancePtr dest doff)+ (advancePtr src soff)+ len++{-# INLINE insertArr #-}+insertArr :: Storable a+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> a -- ^ Value+ -> Arr a -- ^ Array to modify+ -> IO ()+insertArr !sz !i !x !arr = do+ copyArr arr (i + 1) arr i (sz - i)+ writeArr arr i x++-- {-# INLINE removeArrDeinit #-}+-- removeArrDeinit :: Deinitialize a+-- => Int -- ^ Size of the original array+-- -> Int -- ^ Index+-- -> Arr a -- ^ Array to modify+-- -> IO ()+-- removeArrDeinit !sz !i !arr = do+-- deinitializeElemOff (getArr arr) i+-- copyArr arr i arr (i + 1) (sz - i - 1)++{-# INLINE removeArr #-}+removeArr :: Storable a+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> Arr a -- ^ Array to modify+ -> IO ()+removeArr !sz !i !arr = do+ copyArr arr i arr (i + 1) (sz - i - 1)++{-# INLINE insertInitArr #-}+insertInitArr :: forall a r. Storable a+ => Int -- ^ Size of the original array+ -> Int -- ^ Index+ -> Arr a -- ^ Array to modify+ -> (Ptr a -> Int -> IO r)+ -> IO r+insertInitArr !sz !i !arr@(Arr ptr0) f = do+ copyArr arr (i + 1) arr i (sz - i)+ f ptr0 i++-- | 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.+{-# INLINE findIndexOfGtElem #-}+findIndexOfGtElem :: (Ord a, Storable a) => Arr a -> a -> Int -> IO Int+findIndexOfGtElem v needle sz = go 0 (sz - 1)+ where+ go :: Int -> Int -> IO Int+ go !lo !hi = if lo <= hi+ then do+ let !mid = lo + half (hi - lo)+ !val <- readArr v mid+ if | val == needle -> return (mid + 1)+ | val < needle -> go (mid + 1) hi+ | otherwise -> go lo (mid - 1)+ else return lo++-- 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)+{-# INLINE findIndex #-}+findIndex :: (Ord a, Storable a)+ => Arr a+ -> a+ -> Int+ -> IO Int -- (Either Int Int)+findIndex !marr !needle !sz = go 0+ where+ {-# INLINE go #-}+ go :: Int -> IO Int+ go !i = if i < sz+ then do+ !a <- readArr marr i+ case inline (compare a needle) of+ LT -> go (i + 1)+ EQ -> return i+ GT -> return (encodeGtIndex i)+ else return (encodeGtIndex i)++foldrWithKey :: forall k v b. (Ord k, Storable k, Storable v)+ => (k -> v -> b -> IO b)+ -> b+ -> BTree k v+ -> IO b+foldrWithKey f b0 (BTree height root) = go height root b0+ where+ branchDegree :: Int+ !branchDegree = calcBranchDegree root+ childDegree :: Int+ childDegree = calcChildDegree root+ go :: Int -> Ptr (Node k v) -> b -> IO b+ go !n !ptrNode b = do+ sz <- readNodeSize ptrNode+ if n > 0+ then do+ let KeysNodes _ nodes = readNodeKeysNodes branchDegree ptrNode+ foldrArray (sz + 1) (go (n - 1)) b nodes+ else do+ let KeysValues keys values = readNodeKeysValues childDegree ptrNode+ foldrPrimArrayPairs sz f b keys values++foldrPrimArrayPairs :: forall k v b. (Ord k, Storable k, Storable v)+ => Int -- ^ length of arrays+ -> (k -> v -> b -> IO b)+ -> b+ -> Arr k+ -> Arr v+ -> IO b+foldrPrimArrayPairs len f b0 ks vs = go (len - 1) b0+ where+ go :: Int -> b -> IO b+ go !ix !b1 = if ix >= 0+ then do+ k <- readArr ks ix+ v <- readArr vs ix+ b2 <- f k v b1+ go (ix - 1) b2+ else return b1++foldrArray :: forall a b. Storable a+ => Int -- ^ length of array+ -> (a -> b -> IO b)+ -> b+ -> Arr a+ -> IO b+foldrArray len f b0 arr = go (len - 1) b0+ where+ go :: Int -> b -> IO b+ go !ix !b1 = if ix >= 0+ then do+ a <- readArr arr ix+ b2 <- f a b1+ go (ix - 1) b2+ else return b1++arrMapM_ :: (Storable a) => (a -> IO b) -> Int -> Arr a -> IO ()+arrMapM_ f len arr = go 0+ where+ go :: Int -> IO ()+ go i = if i < len+ then do+ _ <- f =<< readArr arr i+ go (i + 1)+ else return ()+ ++{-# INLINE encodeGtIndex #-}+encodeGtIndex :: Int -> Int+encodeGtIndex i = negate i - 1++{-# INLINE decodeGtIndex #-}+decodeGtIndex :: Int -> Int+decodeGtIndex x = negate x - 1++{-# INLINE half #-}+half :: Int -> Int+half x = unsafeShiftR x 1++-- | This is provided for convenience but is not something+-- typically useful in production code.+toAscList :: forall k v. (Ord k, Storable k, Storable v)+ => BTree k v+ -> IO [(k,v)]+toAscList = foldrWithKey f []+ where+ f :: k -> v -> [(k,v)] -> IO [(k,v)]+ f k v xs = return ((k,v) : xs)
test/Spec.hs view
@@ -1,3 +1,8 @@+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE BangPatterns #-} {-# LANGUAGE ScopedTypeVariables #-}@@ -17,43 +22,58 @@ import Control.Monad.ST import Debug.Trace import Control.Monad.Trans.Except+import Control.Monad.Trans.Maybe 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 Foreign.Storable+import GHC.TypeLits+import Foreign.Ptr+import Control.Monad.Random.Strict hiding (fromList)+import Data.Bifunctor+import GHC.Exts (fromList) 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 BTree.Store as BTS+import qualified ArrayList as AL 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 ()+ BTS.with_ $ \bt0 -> do+ bt1 <- BTS.modifyWithM_ bt0 (4 :: Int) $ \bti0 -> do+ bti1 <- BTS.insert bti0 'x' (7 :: Int)+ bti2 <- BTS.insert bti1 'z' (7 :: Int)+ bti3 <- BTS.insert bti2 'y' (7 :: Int)+ return bti3+ bt2 <- BTS.modifyWithM_ bt1 (2 :: Int) $ \bti0 -> do+ bti1 <- BTS.insert bti0 'a' (7 :: Int)+ bti2 <- BTS.insert bti1 'b' (7 :: Int)+ bti3 <- BTS.insert bti2 'c' (7 :: Int)+ return bti3+ mint <- runMaybeT $ do+ bti <- MaybeT (BTS.lookup bt2 4)+ MaybeT (BTS.lookup bti 'x')+ print mint+ return bt2+ -- BTS.toAscList bt2 >>= print + -- BTS.with_ $ \bt0 -> do+ -- bt1 <- BTS.insert bt0 (4 :: Int) 'x'+ -- bt2 <- BTS.insert bt1 3 'z'+ -- BTS.toAscList bt2 >>= print + -- return bt2 defaultMain tests- basicBenchmarks+ -- basicBenchmarks putStrLn "Finished test suite" tests :: TestTree@@ -63,48 +83,139 @@ 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)+ (forall x. (Hashable x, Show x, Ord x, Eq x, BTS.Initialize x, BTS.Deinitialize x, Bounded x, Integral x) => [x] -> 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 3 "small maps with 256 bit keys and values, all permutations, no splitting"+ (over (series :: Series IO [Padded 256]) f) , testPropDepth 4 "small maps of degree 3, all permutations, one split"- (over (series :: Series IO [Positive Int]) (f 3))+ (over (series :: Series IO [Padded 256]) f) , testPropDepth 7 "small maps of degree 3, all permutations"- (over (series :: Series IO [Positive Int]) (f 3))+ (over (series :: Series IO [Padded 256]) f) , testPropDepth 7 "small maps of degree 4, all permutations"- (over (series :: Series IO [Positive Int]) (f 4))+ (over (series :: Series IO [Padded 256]) f) , testPropDepth 10 "medium maps of degree 3, few permutations"- (over doubletonSeriesA (f 3))+ (over (doubletonSeriesA (Proxy :: Proxy 256)) f) , testPropDepth 10 "medium maps of degree 4, few permutations"- (over doubletonSeriesA (f 4))+ (over (doubletonSeriesA (Proxy :: Proxy 256)) f) , testPropDepth 10 "medium maps of degree 3, repeat keys likely, few permutations"- (over doubletonSeriesB (f 3))+ (over (doubletonSeriesB (Proxy :: Proxy 256)) f) , testPropDepth 10 "medium maps of degree 4, repeat keys likely, few permutations"- (over doubletonSeriesB (f 4))+ (over (doubletonSeriesB (Proxy :: Proxy 256)) f) , testPropDepth 150 "large maps of degree 3, repeat keys certain, one permutation"- (over singletonSeriesB (f 3))+ (over (singletonSeriesB (Proxy :: Proxy 256)) f) , testPropDepth 150 "large maps of degree 6, one permutation"- (over singletonSeriesA (f 6))+ (over (singletonSeriesA (Proxy :: Proxy 128)) f) , testPropDepth 150 "large maps of degree 7, repeat keys certain, one permutation"- (over singletonSeriesB (f 7))+ (over (singletonSeriesB (Proxy :: Proxy 128)) f)+ , testPropDepth 200 "large maps" (over word32Series f)+ -- , testPropDepth 1050 "large maps with Word16" (over word16SeriesSingles f) ] +arraylistTests :: [TestTree]+arraylistTests =+ [ testPropDepth 10 "arraylist inserts followed by dump (short)" (over word16Series arrayListInsertions)+ , testPropDepth 150 "arraylist inserts followed by dump (long)" (over word32Series arrayListInsertions)+ , testPropDepth 150 "arraylist inserts followed by repeated pop (long)" (over word32Series pushPop)+ , testPropDepth 50 "arraylist dropWhile" (over word32Series arrayListDropWhile)+ , testPropDepth 50 "insert array" (over word32Series arrayListInsertArray)+ , testPropDepth 100 "insert big array" (over word32Series arrayListInsertBigArray)+ , testPropDepth 100 "insert big arrays" (over word32Series arrayListInsertArrays)+ -- , testPropDepth 150 "arraylist push, pop, twice (long)" (over word32Series pushPopTwice)+ ]+ 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))+ [ testGroup "unmanaged heap" (smallcheckTests orderingStorable)+ , testGroup "unmanaged heap nested" (smallcheckTests orderingNested)+ -- the diverse ones take too long to run+ -- , testGroup "unmanaged heap nested diverse" (smallcheckTests orderingNestedDiverse)+ -- deletion does not work yet+ -- , testGroup "unmanaged heap deletions" (smallcheckTests deletionStorable)+ , testGroup "arraylist" arraylistTests ] +arrayListInsertions :: (Eq a, Show a, Prim a, Storable a) => [a] -> Either String String+arrayListInsertions xs = unsafePerformIO $ AL.with $ \a0 -> do+ a1 <- foldlM AL.pushR a0 xs+ (a2,ys) <- AL.dumpList a1+ return $ (,) a2 $ if xs == ys+ then Right "good"+ else Left ("expected " ++ show xs ++ " but got " ++ show ys)++pushPop :: forall a. (Eq a, Show a, Prim a, Storable a) => [a] -> Either String String+pushPop xs = unsafePerformIO $ AL.with $ \a0 -> do+ a1 <- foldlM AL.pushR a0 xs+ let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])+ go al = do+ (al',m) <- AL.popL al+ case m of+ Nothing -> return (al',[])+ Just a -> fmap (second (a:)) (go al')+ (a2,ys) <- go a1+ return $ (,) a2 $ if xs == ys+ then Right "good"+ else Left $ "expected " ++ show xs ++ " but got " ++ show ys++arrayListDropWhile :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a) => [a] -> Either String String+arrayListDropWhile xs = unsafePerformIO $ AL.with $ \a0 ->+ case deterministicShuffle xs of+ [] -> return (a0, Right "good")+ x : _ -> do+ a1 <- foldlM AL.pushR a0 xs+ (a2,_) <- AL.dropWhileL a1 (\y -> return (y /= x))+ (a3,ys) <- AL.dumpList a2+ let expected = L.dropWhile (/= x) xs+ return $ (,) a3 $ if expected == ys+ then Right "good"+ else Left ("expected " ++ show expected ++ " but got " ++ show ys ++ " using pivot of " ++ show x)+ +arrayListInsertArray :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a)+ => [a] -> Either String String+arrayListInsertArray xs = unsafePerformIO $ AL.with $ \a0 -> do+ a1 <- foldlM AL.pushArrayR a0 (map P.singletonPrimArray xs)+ let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])+ go al = do+ (al',m) <- AL.popL al+ case m of+ Nothing -> return (al',[])+ Just a -> fmap (second (a:)) (go al')+ (a2,ys) <- go a1+ return $ (,) a2 $ if xs == ys+ then Right "good"+ else Left $ "expected " ++ show xs ++ " but got " ++ show ys+ +arrayListInsertBigArray :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a)+ => [a] -> Either String String+arrayListInsertBigArray xs = unsafePerformIO $ AL.with $ \a0 -> do+ a1 <- AL.pushArrayR a0 (fromList xs)+ let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])+ go al = do+ (al',m) <- AL.popL al+ case m of+ Nothing -> return (al',[])+ Just a -> fmap (second (a:)) (go al')+ (a2,ys) <- go a1+ return $ (,) a2 $ if xs == ys+ then Right "good"+ else Left $ "expected " ++ show xs ++ " but got " ++ show ys++arrayListInsertArrays :: forall a. (Hashable a, Eq a, Show a, Prim a, Storable a)+ => [a] -> Either String String+arrayListInsertArrays xs = unsafePerformIO $ AL.with $ \a0 -> do+ a1 <- AL.pushArrayR a0 (fromList xs)+ a2 <- AL.pushArrayR a1 (fromList xs)+ let go :: AL.ArrayList a -> IO (AL.ArrayList a, [a])+ go al = do+ (al',m) <- AL.popL al+ case m of+ Nothing -> return (al',[])+ Just a -> fmap (second (a:)) (go al')+ (a3,zs) <- go a2+ return $ (,) a3 $ if zs == (xs ++ xs)+ then Right "good"+ else Left $ "expected " ++ show (xs ++ xs) ++ " but got " ++ show zs+ unitTests :: TestTree unitTests = testGroup "Unit tests" [ testCase "put followed by get (tests lookup,insert,toAscList)" $ do@@ -143,9 +254,12 @@ 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- _ <- BTC.new token 5 :: IO (BTC.BTree Word Word RealWorld _)- return ()+ , testCase "ArrayList dropWhileScanL on empty" $ do+ xs <- AL.new+ (xs',n,r) <- AL.dropWhileScanL xs (55 :: Word32) (\b a -> return (True,b + a))+ n @?= 0+ r @?= 55+ AL.free xs' ] testPropDepth :: Testable IO a => Int -> String -> a -> TestTree@@ -178,31 +292,30 @@ else Left ("looked up " ++ show x ++ " but found wrong value " ++ show y) return (r1 >> r2) -lookupAfterInsertCompact :: (Show n, Ord n, Prim n)+lookupAfterInsertUnmanaged :: (Show n, Ord n, BTS.Initialize n, BTS.Deinitialize n) => Int -- ^ degree of b-tree -> [Positive n] -- ^ values to insert -> Either Reason Reason-lookupAfterInsertCompact degree xs' =+lookupAfterInsertUnmanaged 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- m0 <- BTC.new c degree- m1 <- foldlM (\ !m !x -> BTC.insert c m x x) m0 xs+ in fmap (const "good") $ unsafePerformIO $ BTS.with $ \m0 -> do+ m1 <- foldlM (\ !m !x -> BTS.insert m x x) m0 xs r1 <- foldlM (\e x -> case e of Right () -> do- BTC.lookup m1 x >>= \case+ BTS.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+ r2 <- runExceptT $ forM_ xs $ \x -> lift (BTS.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)+ return (r1 >> r2, m1) ordering :: (Show n, Ord n, Prim n)@@ -221,55 +334,109 @@ 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+-- 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+-- 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)++orderingStorable :: (Hashable x, Show x, Eq x, Ord x, Storable x, BTS.Initialize x, BTS.Deinitialize x)+ => [x] -- ^ 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- 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)+orderingStorable xs = + let expected = map (\x -> (x,x)) $ S.toAscList $ S.fromList xs+ result = unsafePerformIO $ BTS.with $ \m0 -> do+ m1 <- foldlM (\ !m !x -> BTS.insert m x x) m0 xs+ actual <- BTS.toAscList m1+ let e = if actual == expected+ then Right "good"+ else Left (notice (show expected) (show actual) "layout not available")+ return (e,m1)+ in result +-- this does all insertions followed by all deletions+-- deletionStorable :: KnownNat n+-- => [Padded n] -- ^ values to insert+-- -> Either Reason Reason+-- deletionStorable xs = +-- let expected = map (\x -> (x,x)) $ S.toAscList $ S.fromList xs+-- result = unsafePerformIO $ BTS.with $ \m0 -> do+-- m1 <- foldlM (\ !m !x -> BTS.insert m x x) m0 xs+-- m2 <- foldlM (\ !m !x -> BTS.delete m x) m1 (deterministicShuffle xs)+-- actual <- BTS.toAscList m2+-- let e = if actual == []+-- then Right "good"+-- else Left (notice "empty list" (show actual) "layout not available")+-- return (e,m2)+-- in result++ -- 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+orderingNested :: (Bounded x, Integral x, Hashable x, Show x, Eq x, Ord x, Storable x, BTS.Initialize x, BTS.Deinitialize x)+ => [x] -- ^ values to insert -> Either Reason Reason-orderingNested degree xs' = - let xs = map getPositive xs'- e = runST $ withToken $ \c -> do- m0 <- BTT.new c degree+orderingNested xs = + let e = unsafePerformIO $ BTS.with $ \m0 -> do m1 <- foldlM (\ !mtop !x -> do let subValues = take 10 (iterate (fromIntegral . hashWithSalt 13 . (+ div maxBound 3)) x)- foldM ( \ !m !y -> do- (_,t) <- BTT.modifyWithM c m x (BTC.new c degree) $ \mbottom -> do- fmap BTT.Replace (BTC.insert c mbottom y y)- return t+ foldM + ( \ !m !y -> BTS.modifyWithM_ m x $ \mbottom ->+ BTS.insert mbottom y y ) mtop subValues ) m0 xs- runExceptT $ forM_ xs $ \x -> do- m <- lift $ BTT.lookup m1 x + e <- runExceptT $ forM_ xs $ \x -> do+ m <- lift $ BTS.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+ n <- lift $ BTS.lookup b x case n of Nothing -> ExceptT (return (Left ("could not find " ++ show x ++ " in bottom b-tree"))) Just k -> return ()+ return (e,m1) in fmap (const "good") e +orderingNestedDiverse :: (Bounded x, Integral x, Hashable x, Show x, Eq x, Ord x, Storable x, BTS.Initialize x, BTS.Deinitialize x)+ => [x] -- ^ values to insert+ -> Either Reason Reason+orderingNestedDiverse xs = + let e = unsafePerformIO $ BTS.with $ \m0 -> do+ let topSub = 600 :: Word32+ subValues = enumFromTo 0 topSub+ m1 <- foldlM+ (\ !mtop !x -> do+ foldM + ( \ !m !y -> BTS.modifyWithM_ m x $ \mbottom ->+ BTS.insert mbottom y y+ ) mtop subValues+ ) m0 xs+ e <- runExceptT $ forM_ xs $ \x -> do+ m <- lift $ BTS.lookup m1 x + case m of+ Nothing -> ExceptT (return (Left ("could not find " ++ show x ++ " in top b-tree")))+ Just b -> do+ n <- lift $ BTS.lookup b topSub+ case n of+ Nothing -> ExceptT (return (Left ("could not find " ++ show x ++ " in bottom b-tree")))+ Just k -> return ()+ return (e,m1)+ in fmap (const "good") e+ notice :: String -> String -> String -> String notice expected actual layout = concat [ "expected: " , expected- , ", actual: "+ , ",\n actual: " , actual , ", layout:\n" , layout@@ -281,57 +448,61 @@ (\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)+doubletonSeriesA :: Proxy n -> Series m [Padded n]+doubletonSeriesA _ = (fmap.fmap) Padded (scanSeries (\n -> [n + 9787, n + 29059]) 0) -doubletonSeriesB :: Series m [Positive Word8]-doubletonSeriesB = (fmap.fmap) Positive (scanSeries (\n -> [n + 89, n + 71]) 0)+doubletonSeriesB :: Proxy n -> Series m [Padded n]+doubletonSeriesB _ = (fmap.fmap) Padded (scanSeries (\n -> [n + 89, n + 71]) 0) -singletonSeriesA :: Series m [Positive Word16]-singletonSeriesA = (fmap.fmap) Positive (scanSeries (\n -> [n + 26399]) 0)+singletonSeriesA :: Proxy n -> Series m [Padded n]+singletonSeriesA _ = (fmap.fmap) Padded (scanSeries (\n -> [n + 26399]) 0) -singletonSeriesB :: Series m [Positive Word8]-singletonSeriesB = (fmap.fmap) Positive (scanSeries (\n -> [n + 73]) 0)+singletonSeriesB :: Proxy n -> Series m [Padded n]+singletonSeriesB _ = (fmap.fmap) Padded (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- 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 c m y y- go (ix + 1) m'- else return ()- go 0 m0- getSizeOfCompact c+word16Series :: Series m [Word16]+word16Series = (scanSeries (\n -> [n + 89, n + 71]) 0) -sizeAfterRepeatedInserts :: Int -> IO Word -sizeAfterRepeatedInserts total = withToken $ \c -> do- m0 <- BTC.new c 8- let go !ix !m = if ix < total- then do- -- same key every time- m' <- BTC.insert c m (99 :: Int) (ix :: Int)- go (ix + 1) m'- else return ()- go 0 m0- getSizeOfCompact c+word32Series :: Series m [Word32]+word32Series = (scanSeries (\n -> [n + 73]) 0) -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)- +word16SeriesSingles :: Series m [Word16]+word16SeriesSingles = (scanSeries (\n -> [n + 73]) 0)++word32SeriesAlt :: Series m [Word32]+word32SeriesAlt = (scanSeries (\n -> [n + 73, n + 89]) 0)++newtype Padded (n :: Nat) = Padded Word+ deriving (Eq,Ord,Bounded,Hashable,Integral,Real,Num,Enum)++instance KnownNat n => Storable (Padded n) where+ sizeOf _ = fromInteger (natVal (Proxy :: Proxy n))+ alignment _ = fromInteger (natVal (Proxy :: Proxy n))+ peek ptr = fmap Padded (peek (castPtr ptr))+ poke ptr (Padded w) = poke (castPtr ptr) w++instance KnownNat n => BTS.Initialize (Padded n) where+ initialize _ = return ()++instance KnownNat n => BTS.Deinitialize (Padded n) where+ deinitialize _ = return ()++instance Show (Padded n) where+ show (Padded w) = show w++instance Monad m => Serial m (Padded n) where+ series = fmap (\(Positive n) -> Padded (intToWord n)) series++intToWord :: Int -> Word+intToWord = fromIntegral++deterministicShuffle :: Hashable a => [a] -> [a]+deterministicShuffle xs = evalRand (shuffle xs) (mkStdGen (hash xs))++shuffle :: [a] -> Rand StdGen [a]+shuffle [] = return []+shuffle xs = do+ randomPosition <- getRandomR (0, length xs - 1)+ let (left, (a:right)) = splitAt randomPosition xs+ fmap (a:) (shuffle (left ++ right))