rrb-vector (empty) → 0.1.0.0
raw patch · 14 files changed
+1370/−0 lines, 14 filesdep +QuickCheckdep +basedep +deepseqsetup-changed
Dependencies added: QuickCheck, base, deepseq, gauge, hspec, indexed-traversable, primitive, rrb-vector
Files
- LICENSE +30/−0
- README.md +5/−0
- Setup.hs +2/−0
- bench/Main.hs +22/−0
- bench/Traverse.hs +21/−0
- rrb-vector.cabal +69/−0
- src/Data/RRBVector.hs +54/−0
- src/Data/RRBVector/Internal.hs +720/−0
- src/Data/RRBVector/Internal/Array.hs +207/−0
- src/Data/RRBVector/Internal/Buffer.hs +37/−0
- src/Data/RRBVector/Internal/Debug.hs +61/−0
- src/Data/RRBVector/Internal/Indexed.hs +26/−0
- src/Data/RRBVector/Internal/IntRef.hs +33/−0
- test/Spec.hs +83/−0
+ LICENSE view
@@ -0,0 +1,30 @@+Copyright konsumlamm (c) 2020++All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++ * Redistributions of source code must retain the above copyright+ notice, this list of conditions and the following disclaimer.++ * Redistributions in binary form must reproduce the above+ copyright notice, this list of conditions and the following+ disclaimer in the documentation and/or other materials provided+ with the distribution.++ * Neither the name of konsumlamm nor the names of other+ contributors may be used to endorse or promote products derived+ from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,5 @@+# rbb-vector++An implementation of a Relaxed Radix Balanced Vector (RRB-Vector).++For more information, see [`rrb-vector` on Hackage](https://hackage.haskell.org/package/rrb-vector).
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ bench/Main.hs view
@@ -0,0 +1,22 @@+import Data.Functor ((<&>))++import Gauge.Main++import qualified Data.RRBVector as RRB++main :: IO ()+main = defaultMain $ [10, 100, 1_000, 10_000, 100_000] <&> \n ->+ let v = RRB.fromList [1..n]+ idx = [n `div` 10, n `div` 2, n - n `div` 10]+ in bgroup (show n)+ [ bench "fromList" $ nf RRB.fromList [1..n]+ , bench "><" $ nf (\vec -> vec RRB.>< vec) v+ , bench "|>" $ nf (RRB.|> 42) v+ , bench "<|" $ nf (42 RRB.<|) v+ , bgroup "take" $ idx <&> \i -> bench (show i) $ nf (RRB.take i) v+ , bgroup "drop" $ idx <&> \i -> bench (show i) $ nf (RRB.drop i) v+ , bgroup "lookup" $ idx <&> \i -> bench (show i) $ nf (RRB.lookup i) v+ , bgroup "adjust" $ idx <&> \i -> bench (show i) $ nf (RRB.adjust i (+ 1)) v+ , bench "foldl" $ nf (foldl (+) 0) v+ , bench "foldr" $ nf (foldr (+) 0) v+ ]
+ bench/Traverse.hs view
@@ -0,0 +1,21 @@+{-# LANGUAGE BangPatterns #-}++import Data.Foldable (foldl', toList)+import Data.Functor++import Gauge.Main++import qualified Data.RRBVector as RRB++main :: IO ()+main = defaultMain $ [10, 100, 1_000] <&> \n ->+ let !v = RRB.fromList [1..n]+ in bgroup (show n)+ [ bench "f1 (where)" $ nf (\v -> f1 v v) v+ , bench "f2" $ nf (\v -> f2 v v) v+ ]+ where+ f1 xs ys = foldl' (\acc x -> acc RRB.>< RRB.fromList (replicate n x)) RRB.empty xs+ where+ n = length ys+ f2 xs ys = foldl' (\acc x -> acc RRB.>< RRB.fromList (replicate (length ys) x)) RRB.empty xs
+ rrb-vector.cabal view
@@ -0,0 +1,69 @@+name: rrb-vector+version: 0.1.0.0+synopsis: Efficient RRB-Vectors+description:+ An RRB-Vector is an efficient sequence data structure.+ It supports fast indexing, iteration, concatenation and splitting.+ .+ == Comparison with [Data.Sequence](https://hackage.haskell.org/package/containers/docs/Data-Sequence.html)+ .+ @Seq a@ is a container with a very similar API. RRB-Vectors are generally faster for indexing and iteration,+ while sequences are faster for access to the front/back (amortized \(O(1)\)).+homepage: https://github.com/konsumlamm/rrb-vector+bug-reports: https://github.com/konsumlamm/rrb-vector/issues+license: BSD3+license-file: LICENSE+author: konsumlamm+maintainer: konsumlamm@gmail.com+copyright: 2021 konsumlamm+category: Data Structures+build-type: Simple+extra-source-files: README.md+cabal-version: 2.0+tested-with: GHC == 8.4.4, GHC == 8.6.5, GHC == 8.8.4, GHC == 8.10.5, GHC == 9.0.1++source-repository head+ type: git+ location: https://github.com/konsumlamm/rbb-vector++library+ hs-source-dirs: src+ exposed-modules:+ Data.RRBVector+ Data.RRBVector.Internal.Debug+ other-modules:+ Data.RRBVector.Internal+ Data.RRBVector.Internal.Array+ Data.RRBVector.Internal.Buffer+ Data.RRBVector.Internal.Indexed+ Data.RRBVector.Internal.IntRef+ build-depends: base >= 4.11 && < 5, deepseq ^>= 1.4.3, indexed-traversable ^>= 0.1, primitive ^>= 0.7+ ghc-options: -O2 -Wall -Wno-name-shadowing+ default-language: Haskell2010++test-suite test+ hs-source-dirs: test+ main-is: Spec.hs+ type: exitcode-stdio-1.0+ ghc-options: -Wall -Wno-orphans -Wno-type-defaults+ build-depends: base >= 4.11 && < 5, rrb-vector, hspec, QuickCheck+ default-language: Haskell2010+ default-extensions: ExtendedDefaultRules, NumericUnderscores++benchmark rrb-bench+ hs-source-dirs: bench+ main-is: Main.hs+ type: exitcode-stdio-1.0+ default-language: Haskell2010+ ghc-options: -O2+ build-depends: base >= 4.11 && < 5, gauge, rrb-vector+ default-extensions: ExtendedDefaultRules, NumericUnderscores++benchmark traverse+ hs-source-dirs: bench+ main-is: Traverse.hs+ type: exitcode-stdio-1.0+ default-language: Haskell2010+ ghc-options: -O2+ build-depends: base >= 4.11 && < 5, gauge, rrb-vector, primitive+ default-extensions: ExtendedDefaultRules, NumericUnderscores
+ src/Data/RRBVector.hs view
@@ -0,0 +1,54 @@+{- |+The @'Vector' a@ type is an RRB-Vector of elements of type @a@.++This module should be imported qualified, to avoid name clashes with the "Prelude".++= Performance++The worst case running time complexities are given, with \(n\) referring to the number of elements in the vector+(or \(n_1\), \(n_2\), etc. for multiple vectors). Note that all logarithms are base 16,+so the constant factor for \(O(\log n)\) operations is quite small.++= Implementation++The implementation uses Relaxed-Radix-Balanced trees, as described by++* Nicolas Stucki, [\"Turning Relaxed Radix Balanced Vector from Theory into Practice for Scala Collections\"](https://github.com/nicolasstucki/scala-rrb-vector/blob/master/documents/Master%20Thesis%20-%20Nicolas%20Stucki%20-%20Turning%20Relaxed%20Radix%20Balanced%20Vector%20from%20Theory%20into%20Practice%20for%20Scala%20Collections.pdf), January 2015.++Currently, a branching factor of 16 is used. The tree is strict in its spine, but lazy in its elements.+-}++module Data.RRBVector+ ( Vector+ -- * Construction+ , empty, singleton, fromList+ -- ** Concatenation+ , (<|), (|>), (><)+ -- * Deconstruction+ , viewl, viewr+ -- * Indexing+ , lookup, index+ , (!?), (!)+ , update+ , adjust, adjust'+ , take, drop, splitAt+ , insertAt, deleteAt+ -- * With Index+ --+ -- | Reexported from [indexed-traversable](https://hackage.haskell.org/package/indexed-traversable).+ , module Data.Foldable.WithIndex+ , module Data.Functor.WithIndex+ , module Data.Traversable.WithIndex+ -- * Transformations+ , map, reverse+ -- * Zipping and unzipping+ , zip, zipWith, unzip+ ) where++import Prelude hiding (lookup, take, drop, splitAt, map, reverse, zip, zipWith, unzip)++import Data.Foldable.WithIndex+import Data.Functor.WithIndex+import Data.Traversable.WithIndex++import Data.RRBVector.Internal
+ src/Data/RRBVector/Internal.hs view
@@ -0,0 +1,720 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE CPP #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}++module Data.RRBVector.Internal+ ( Vector(..)+ , Tree(..)+ -- * Internal+ , blockShift, blockSize, treeSize, computeSizes, up+ -- * Construction+ , empty, singleton, fromList+ -- ** Concatenation+ , (<|), (|>), (><)+ -- * Deconstruction+ , viewl, viewr+ -- * Indexing+ , lookup, index+ , (!?), (!)+ , update+ , adjust, adjust'+ , take, drop, splitAt+ , insertAt, deleteAt+ -- * Transformations+ , map, reverse+ -- * Zipping and unzipping+ , zip, zipWith, unzip+ ) where++import Control.Applicative (Alternative, liftA2)+import qualified Control.Applicative+import Control.DeepSeq+import Control.Monad (when, MonadPlus)+import Control.Monad.ST (runST)+#if !(MIN_VERSION_base(4,13,0))+import Control.Monad.Fail (MonadFail(..))+#endif+import Control.Monad.Fix (MonadFix(..))+import Control.Monad.Zip (MonadZip(..))++import Data.Bits+import Data.Foldable (Foldable(..), for_)+import Data.Functor.Classes+import Data.Functor.Identity (Identity(..))+import Data.Maybe (fromMaybe)+import qualified Data.List as List+import qualified GHC.Exts as Exts+import GHC.Stack (HasCallStack)+import Text.Read+import Prelude hiding (lookup, map, take, drop, splitAt, head, last, reverse, zip, zipWith, unzip)++import Data.Functor.WithIndex+import Data.Foldable.WithIndex+import Data.Traversable.WithIndex++import Data.Primitive.PrimArray+import qualified Data.RRBVector.Internal.Array as A+import qualified Data.RRBVector.Internal.Buffer as Buffer+import Data.RRBVector.Internal.Indexed++infixr 5 ><+infixr 5 <|+infixl 5 |>++-- Invariant: Children of a Balanced node are always balanced.+-- A Leaf node is considered balanced.+-- Nodes are always non-empty.+data Tree a+ = Balanced !(A.Array (Tree a))+ | Unbalanced !(A.Array (Tree a)) !(PrimArray Int)+ | Leaf !(A.Array a)++-- | A vector.+--+-- The instances are based on those of @Seq@s, which are in turn based on those of lists.+data Vector a+ = Empty+ | Root+ !Int -- size+ !Int -- shift (blockShift * height)+ !(Tree a)++-- The number of bits used per level.+blockShift :: Int+blockShift = 4+{-# INLINE blockShift #-}++-- The maximum size of a block.+blockSize :: Int+blockSize = 1 `shiftL` blockShift++-- The mask used to extract the index into the array.+blockMask :: Int+blockMask = blockSize - 1++up :: Int -> Int+up sh = sh + blockShift+{-# INLINE up #-}++down :: Int -> Int+down sh = sh - blockShift+{-# INLINE down #-}++radixIndex :: Int -> Int -> Int+radixIndex i sh = i `shiftR` sh .&. blockMask+{-# INLINE radixIndex #-}++relaxedRadixIndex :: PrimArray Int -> Int -> Int -> (Int, Int)+relaxedRadixIndex sizes i sh =+ let guess = radixIndex i sh -- guess <= idx+ idx = loop guess+ subIdx = if idx == 0 then i else i - indexPrimArray sizes (idx - 1)+ in (idx, subIdx)+ where+ loop idx =+ let current = indexPrimArray sizes idx -- idx will always be in range for a well-formed tree+ in if i < current then idx else loop (idx + 1)+{-# INLINE relaxedRadixIndex #-}++treeToArray :: Tree a -> A.Array (Tree a)+treeToArray (Balanced arr) = arr+treeToArray (Unbalanced arr _) = arr+treeToArray (Leaf _) = error "treeToArray: leaf"++treeBalanced :: Tree a -> Bool+treeBalanced (Balanced _) = True+treeBalanced (Unbalanced _ _) = False+treeBalanced (Leaf _) = True++-- @treeSize sh@ is the size of a tree with shift @sh@.+treeSize :: Int -> Tree a -> Int+treeSize = go 0+ where+ go acc _ (Leaf arr) = acc + length arr+ go acc _ (Unbalanced _ sizes) = acc + indexPrimArray sizes (sizeofPrimArray sizes - 1)+ go acc sh (Balanced arr) =+ let i = length arr - 1+ in go (acc + i * (1 `shiftL` sh)) (down sh) (A.index arr i)+{-# INLINE treeSize #-}++-- @computeSizes sh@ turns an array into a tree node by computing the sizes of its subtrees.+-- @sh@ is the shift of the resulting tree.+computeSizes :: Int -> A.Array (Tree a) -> Tree a+computeSizes sh arr = runST $ do+ let len = length arr+ maxSize = 1 `shiftL` sh -- the maximum size of a subtree+ sizes <- newPrimArray len+ let loop acc isBalanced i+ | i < len =+ let subtree = A.index arr i+ size = treeSize (down sh) subtree+ acc' = acc + size+ isBalanced' = isBalanced && if i == len - 1 then treeBalanced subtree else size == maxSize+ in writePrimArray sizes i acc' *> loop acc' isBalanced' (i + 1)+ | otherwise = pure isBalanced+ isBalanced <- loop 0 True 0+ if isBalanced then+ pure $ Balanced arr+ else do+ sizes <- unsafeFreezePrimArray sizes -- safe because the mutable @sizes@ isn't used afterwards+ pure $ Unbalanced arr sizes++-- Integer log base 2.+log2 :: Int -> Int+log2 x = bitSizeMinus1 - countLeadingZeros x+ where+ bitSizeMinus1 = finiteBitSize (0 :: Int) - 1+{-# INLINE log2 #-}++instance Show1 Vector where+ liftShowsPrec sp sl p v = showsUnaryWith (liftShowsPrec sp sl) "fromList" p (toList v)++instance (Show a) => Show (Vector a) where+ showsPrec = showsPrec1++instance Read1 Vector where+ liftReadPrec rp rl = readData $ readUnaryWith (liftReadPrec rp rl) "fromList" fromList+ liftReadListPrec = liftReadListPrecDefault++instance (Read a) => Read (Vector a) where+ readPrec = readPrec1+ readListPrec = readListPrecDefault++instance Eq1 Vector where+ liftEq f v1 v2 = length v1 == length v2 && liftEq f (toList v1) (toList v2)++instance (Eq a) => Eq (Vector a) where+ (==) = eq1++instance Ord1 Vector where+ liftCompare f v1 v2 = liftCompare f (toList v1) (toList v2)++instance (Ord a) => Ord (Vector a) where+ compare = compare1++instance Semigroup (Vector a) where+ v1 <> v2 = v1 >< v2++instance Monoid (Vector a) where+ mempty = empty++instance Foldable Vector where+ foldr f acc = go+ where+ go Empty = acc+ go (Root _ _ tree) = foldrTree tree acc++ foldrTree (Balanced arr) acc' = foldr foldrTree acc' arr+ foldrTree (Unbalanced arr _) acc' = foldr foldrTree acc' arr+ foldrTree (Leaf arr) acc' = foldr f acc' arr+ {-# INLINE foldr #-}++ foldl f acc = go+ where+ go Empty = acc+ go (Root _ _ tree) = foldlTree acc tree++ foldlTree acc' (Balanced arr) = foldl foldlTree acc' arr+ foldlTree acc' (Unbalanced arr _) = foldl foldlTree acc' arr+ foldlTree acc' (Leaf arr) = foldl f acc' arr+ {-# INLINE foldl #-}++ foldr' f acc = go+ where+ go Empty = acc+ go (Root _ _ tree) = foldrTree' tree acc++ foldrTree' (Balanced arr) acc' = foldr' foldrTree' acc' arr+ foldrTree' (Unbalanced arr _) acc' = foldr' foldrTree' acc' arr+ foldrTree' (Leaf arr) acc' = foldr' f acc' arr+ {-# INLINE foldr' #-}++ foldl' f acc = go+ where+ go Empty = acc+ go (Root _ _ tree) = foldlTree' acc tree++ foldlTree' acc' (Balanced arr) = foldl' foldlTree' acc' arr+ foldlTree' acc' (Unbalanced arr _) = foldl' foldlTree' acc' arr+ foldlTree' acc' (Leaf arr) = foldl' f acc' arr+ {-# INLINE foldl' #-}++ null Empty = True+ null Root{} = False+ {-# INLINE null #-}++ length Empty = 0+ length (Root s _ _) = s+ {-# INLINE length #-}++instance FoldableWithIndex Int Vector where+ ifoldr f z0 v = foldr (\x g !i -> f i x (g (i + 1))) (const z0) v 0++ ifoldl f z0 v = foldl (\g x !i -> f i (g (i - 1)) x) (const z0) v (length v - 1)++instance Functor Vector where+ fmap = map+ x <$ v = fromList (replicate (length v) x)++instance FunctorWithIndex Int Vector where+ imap f v = runIdentity $ evalIndexed (traverse (Indexed . f') v) 0+ where+ f' x i = i `seq` WithIndex (i + 1) (Identity (f i x))++instance Traversable Vector where+ traverse _ Empty = pure Empty+ traverse f (Root size sh tree) = Root size sh <$> traverseTree tree+ where+ traverseTree (Balanced arr) = Balanced <$> A.traverse' traverseTree arr+ traverseTree (Unbalanced arr sizes) = Unbalanced <$> A.traverse' traverseTree arr <*> pure sizes+ traverseTree (Leaf arr) = Leaf <$> A.traverse f arr++instance TraversableWithIndex Int Vector where+ itraverse f v = evalIndexed (traverse (Indexed . f') v) 0+ where+ f' x i = i `seq` WithIndex (i + 1) (f i x)++instance Applicative Vector where+ pure = singleton+ fs <*> xs = foldl' (\acc f -> acc >< map f xs) empty fs+ liftA2 f xs ys = foldl' (\acc x -> acc >< map (f x) ys) empty xs+ xs *> ys = foldl' (\acc _ -> acc >< ys) empty xs+ xs <* ys = foldl' (\acc x -> acc >< fromList (replicate (length ys) x)) empty xs++instance Monad Vector where+ xs >>= f = foldl' (\acc x -> acc >< f x) empty xs++instance Alternative Vector where+ empty = empty+ (<|>) = (><)++instance MonadPlus Vector++instance MonadFail Vector where+ fail _ = empty++instance MonadFix Vector where+ mfix f = fromList $ fmap (\i -> let x = index i (f x) in x) [0..length (f err) - 1]+ where+ err = error "mfix for Data.RRBVector.Vector applied to strict function"++instance MonadZip Vector where+ mzipWith = zipWith+ mzip = zip+ munzip = unzip++instance Exts.IsList (Vector a) where+ type Item (Vector a) = a+ fromList = fromList+ toList = toList++instance (a ~ Char) => Exts.IsString (Vector a) where+ fromString = fromList++instance (NFData a) => NFData (Vector a) where+ rnf = rnf1++instance NFData1 Vector where+ liftRnf f = foldl' (\_ x -> f x) ()++-- | \(O(1)\). The empty vector.+--+-- > empty = fromList []+empty :: Vector a+empty = Empty++-- | \(O(1)\). A vector with a single element.+--+-- > singleton x = fromList [x]+singleton :: a -> Vector a+singleton x = Root 1 0 (Leaf $ A.singleton x)++-- | \(O(n)\). Create a new vector from a list.+fromList :: [a] -> Vector a+fromList [] = Empty+fromList [x] = singleton x+fromList ls = case nodes Leaf ls of+ [tree] -> Root (treeSize 0 tree) 0 tree -- tree is a single leaf+ ls' -> iterateNodes blockShift ls'+ where+ nodes f trees = runST $ do+ buffer <- Buffer.new blockSize+ let loop [] = do+ result <- Buffer.get buffer+ pure [f result]+ loop (t : ts) = do+ size <- Buffer.size buffer+ if size == blockSize then do+ result <- Buffer.get buffer+ Buffer.push buffer t+ rest <- loop ts+ pure (f result : rest)+ else do+ Buffer.push buffer t+ loop ts+ loop trees+ {-# INLINE nodes #-}++ iterateNodes sh trees = case nodes Balanced trees of+ [tree] -> Root (treeSize sh tree) sh tree+ trees' -> iterateNodes (up sh) trees'++-- | \(O(\log n)\). The element at the index or 'Nothing' if the index is out of range.+lookup :: Int -> Vector a -> Maybe a+lookup _ Empty = Nothing+lookup i (Root size sh tree)+ | i < 0 || i >= size = Nothing -- index out of range+ | otherwise = Just $ lookupTree i sh tree+ where+ lookupTree i sh (Balanced arr) = lookupTree i (down sh) (A.index arr (radixIndex i sh))+ lookupTree i sh (Unbalanced arr sizes) =+ let (idx, subIdx) = relaxedRadixIndex sizes i sh+ in lookupTree subIdx (down sh) (A.index arr idx)+ lookupTree i _ (Leaf arr) = A.index arr (i .&. blockMask)++-- | \(O(\log n)\). The element at the index. Calls 'error' if the index is out of range.+index :: HasCallStack => Int -> Vector a -> a+index i = fromMaybe (error "AMT.index: index out of range") . lookup i++-- | \(O(\log n)\). A flipped version of 'lookup'.+(!?) :: Vector a -> Int -> Maybe a+(!?) = flip lookup++-- | \(O(\log n)\). A flipped version of 'index'.+(!) :: HasCallStack => Vector a -> Int -> a+(!) = flip index++-- | \(O(\log n)\). Update the element at the index with a new element.+-- If the index is out of range, the original vector is returned.+update :: Int -> a -> Vector a -> Vector a+update _ _ Empty = Empty+update i x v@(Root size sh tree)+ | i < 0 || i >= size = v -- index out of range+ | otherwise = Root size sh (adjustTree i sh tree)+ where+ adjustTree i sh (Balanced arr) = Balanced (A.adjust' arr (radixIndex i sh) (adjustTree i (down sh)))+ adjustTree i sh (Unbalanced arr sizes) =+ let (idx, subIdx) = relaxedRadixIndex sizes i sh+ in Unbalanced (A.adjust' arr idx (adjustTree subIdx (down sh))) sizes+ adjustTree i _ (Leaf arr) = Leaf (A.update arr (i .&. blockMask) x)++-- | \(O(\log n)\). Adjust the element at the index by applying the function to it.+-- If the index is out of range, the original vector is returned.+adjust :: Int -> (a -> a) -> Vector a -> Vector a+adjust _ _ Empty = Empty+adjust i f v@(Root size sh tree)+ | i < 0 || i >= size = v -- index out of range+ | otherwise = Root size sh (adjustTree i sh tree)+ where+ adjustTree i sh (Balanced arr) = Balanced (A.adjust' arr (radixIndex i sh) (adjustTree i (down sh)))+ adjustTree i sh (Unbalanced arr sizes) =+ let (idx, subIdx) = relaxedRadixIndex sizes i sh+ in Unbalanced (A.adjust' arr idx (adjustTree subIdx (down sh))) sizes+ adjustTree i _ (Leaf arr) = Leaf (A.adjust arr (i .&. blockMask) f)++-- | \(O(\log n)\). Like 'adjust', but the result of the function is forced.+adjust' :: Int -> (a -> a) -> Vector a -> Vector a+adjust' _ _ Empty = Empty+adjust' i f v@(Root size sh tree)+ | i < 0 || i >= size = v -- index out of range+ | otherwise = Root size sh (adjustTree i sh tree)+ where+ adjustTree i sh (Balanced arr) = Balanced (A.adjust' arr (radixIndex i sh) (adjustTree i (down sh)))+ adjustTree i sh (Unbalanced arr sizes) =+ let (idx, subIdx) = relaxedRadixIndex sizes i sh+ in Unbalanced (A.adjust' arr idx (adjustTree subIdx (down sh))) sizes+ adjustTree i _ (Leaf arr) = Leaf (A.adjust' arr (i .&. blockMask) f)++-- | \(O(n)\). Apply the function to every element.+--+-- >>> map (+ 1) (fromList [1, 2, 3])+-- fromList [2,3,4]+map :: (a -> b) -> Vector a -> Vector b+map _ Empty = Empty+map f (Root size sh tree) = Root size sh (mapTree tree)+ where+ mapTree (Balanced arr) = Balanced (A.map' mapTree arr)+ mapTree (Unbalanced arr sizes) = Unbalanced (A.map' mapTree arr) sizes+ mapTree (Leaf arr) = Leaf (A.map f arr)++-- | \(O(n)\). Reverse the vector.+--+-- >>> reverse (fromList [1, 2, 3])+-- fromList [3,2,1]+reverse :: Vector a -> Vector a+reverse = fromList . foldl' (flip (:)) [] -- convert the vector to a reverse list and then rebuild++-- | \(O(\min(n_1, n_2))\). Take two vectors and return a vector of corresponding pairs.+-- If one input is longer, excess elements are discarded from the right end.+--+-- > zip = zipWith (,)+zip :: Vector a -> Vector b -> Vector (a, b)+zip v1 v2 = fromList $ List.zip (toList v1) (toList v2)++-- | \(O(\min(n_1, n_2))\). 'zipWith' generalizes 'zip' by zipping with the function.+zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c+zipWith f v1 v2 = fromList $ List.zipWith f (toList v1) (toList v2)++-- | \(O(n)\). Unzip a vector of pairs.+--+-- >>> unzip (fromList [(1, "a"), (2, "b"), (3, "c")])+-- (fromList [1,2,3],fromList ["a","b","c"])+unzip :: Vector (a, b) -> (Vector a, Vector b)+unzip v = (map fst v, map snd v)++-- | \(O(\log n)\). The first element and the vector without the first element, or 'Nothing' if the vector is empty.+--+-- >>> viewl (fromList [1, 2, 3])+-- Just (1,fromList [2,3])+viewl :: Vector a -> Maybe (a, Vector a)+viewl Empty = Nothing+viewl v@(Root _ _ tree) = let !tail = drop 1 v in Just (headTree tree, tail)+ where+ headTree (Balanced arr) = headTree (A.head arr)+ headTree (Unbalanced arr _) = headTree (A.head arr)+ headTree (Leaf arr) = A.head arr++-- | \(O(\log n)\). The vector without the last element and the last element, or 'Nothing' if the vector is empty.+--+-- >>> viewr (fromList [1, 2, 3])+-- Just (fromList [1,2],3)+viewr :: Vector a -> Maybe (Vector a, a)+viewr Empty = Nothing+viewr v@(Root size _ tree) = let !init = take (size - 1) v in Just (init, lastTree tree)+ where+ lastTree (Balanced arr) = lastTree (A.last arr)+ lastTree (Unbalanced arr _) = lastTree (A.last arr)+ lastTree (Leaf arr) = A.last arr++-- | \(O(\log n)\). Split the vector at the given index.+--+-- > splitAt n v = (take n v, drop n v)+splitAt :: Int -> Vector a -> (Vector a, Vector a)+splitAt n v =+ let !left = take n v+ !right = drop n v+ in (left, right)++-- | \(O(\log n)\). Insert an element at the given index.+insertAt :: Int -> a -> Vector a -> Vector a+insertAt i x v = let (left, right) = splitAt i v in (left |> x) >< right++-- | \(O(\log n)\). Delete the element at the given index.+deleteAt :: Int -> Vector a -> Vector a+deleteAt i v = let (left, right) = splitAt (i + 1) v in take i left >< right++-- concatenation++-- | \(O(\log \max(n_1, n_2))\). Concatenates two vectors.+--+-- >>> fromList [1, 2, 3] >< fromList [4, 5]+-- fromList [1,2,3,4,5]+(><) :: Vector a -> Vector a -> Vector a+Empty >< v = v+v >< Empty = v+Root size1 sh1 tree1 >< Root size2 sh2 tree2 =+ let maxShift = max sh1 sh2+ newTree = mergeTrees tree1 sh1 tree2 sh2+ in case singleTree newTree of+ Just newTree -> Root (size1 + size2) maxShift newTree+ Nothing -> Root (size1 + size2) (up maxShift) newTree+ where+ mergeTrees (Leaf arr1) _ (Leaf arr2) _ = Balanced $+ if length arr1 == blockSize then A.from2 (Leaf arr1) (Leaf arr2)+ else if length arr1 + length arr2 <= blockSize then A.singleton (Leaf (arr1 <> arr2))+ else+ let (left, right) = A.splitAt (arr1 <> arr2) blockSize+ in A.from2 (Leaf left) (Leaf right)+ mergeTrees tree1 sh1 tree2 sh2 = case compare sh1 sh2 of+ LT ->+ let right = treeToArray tree2+ (rightHead, rightTail) = viewl right+ merged = mergeTrees tree1 sh1 rightHead (down sh2)+ in mergeRebalance sh2 A.empty (treeToArray merged) rightTail+ GT ->+ let left = treeToArray tree1+ (leftInit, leftLast) = viewr left+ merged = mergeTrees leftLast (down sh1) tree2 sh2+ in mergeRebalance sh1 leftInit (treeToArray merged) A.empty+ EQ ->+ let left = treeToArray tree1+ right = treeToArray tree2+ (leftInit, leftLast) = viewr left+ (rightHead, rightTail) = viewl right+ merged = mergeTrees leftLast (down sh1) rightHead (down sh2)+ in mergeRebalance sh1 leftInit (treeToArray merged) rightTail+ where+ viewl arr = (A.head arr, A.drop arr 1)+ viewr arr = (A.take arr (length arr - 1), A.last arr)++ -- the type annotations are necessary to compile+ mergeRebalance :: forall a. Int -> A.Array (Tree a) -> A.Array (Tree a) -> A.Array (Tree a) -> Tree a+ mergeRebalance sh left center right+ | sh == blockShift = mergeRebalance' (\(Leaf arr) -> arr) Leaf+ | otherwise = mergeRebalance' treeToArray (computeSizes (down sh))+ where+ mergeRebalance' :: (Tree a -> A.Array t) -> (A.Array t -> Tree a) -> Tree a+ mergeRebalance' extract construct = runST $ do+ newRoot <- Buffer.new blockSize+ newSubtree <- Buffer.new blockSize+ newNode <- Buffer.new blockSize+ for_ (toList left ++ toList center ++ toList right) $ \subtree ->+ for_ (extract subtree) $ \x -> do+ lenNode <- Buffer.size newNode+ when (lenNode == blockSize) $ do+ pushTo construct newNode newSubtree+ lenSubtree <- Buffer.size newSubtree+ when (lenSubtree == blockSize) $ pushTo (computeSizes sh) newSubtree newRoot+ Buffer.push newNode x+ pushTo construct newNode newSubtree+ pushTo (computeSizes sh) newSubtree newRoot+ computeSizes (up sh) <$> Buffer.get newRoot+ {-# INLINE mergeRebalance' #-}++ pushTo f from to = do+ result <- Buffer.get from+ Buffer.push to (f result)+ {-# INLINE pushTo #-}++ singleTree (Balanced arr)+ | length arr == 1 = Just (A.head arr)+ singleTree (Unbalanced arr _)+ | length arr == 1 = Just (A.head arr)+ singleTree _ = Nothing++-- | \(O(\log n)\). Add an element to the left end of the vector.+--+-- >>> 1 <| fromList [2, 3, 4]+-- fromList [1,2,3,4]+(<|) :: a -> Vector a -> Vector a+x <| Empty = singleton x+x <| Root size sh tree+ | insertShift > sh = Root (size + 1) insertShift (computeSizes insertShift (A.from2 (newBranch x sh) tree))+ | otherwise = Root (size + 1) sh (consTree sh tree)+ where+ consTree sh (Balanced arr)+ | sh == insertShift = computeSizes sh (A.cons arr (newBranch x (down sh)))+ | otherwise = computeSizes sh (A.adjust' arr 0 (consTree (down sh)))+ consTree sh (Unbalanced arr _)+ | sh == insertShift = computeSizes sh (A.cons arr (newBranch x (down sh)))+ | otherwise = computeSizes sh (A.adjust' arr 0 (consTree (down sh)))+ consTree _ (Leaf arr) = Leaf $ A.cons arr x++ insertShift = computeShift size sh (up sh) tree++ -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)+ -- the index is computed for efficient calculation of the shift in a balanced subtree+ computeShift i sh min (Balanced _) =+ let newShift = (log2 i `div` blockShift) * blockShift+ in if newShift > sh then min else newShift+ computeShift _ sh min (Unbalanced arr sizes) =+ let i' = indexPrimArray sizes 0 -- the size of the first subtree+ newMin = if length arr < blockSize then sh else min+ in computeShift i' (down sh) newMin (A.head arr)+ computeShift _ _ min (Leaf arr) = if length arr < blockSize then 0 else min++-- | \(O(\log n)\). Add an element to the right end of the vector.+--+-- >>> fromList [1, 2, 3] |> 4+-- fromList [1,2,3,4]+(|>) :: Vector a -> a -> Vector a+Empty |> x = singleton x+Root size sh tree |> x+ | insertShift > sh = Root (size + 1) insertShift (computeSizes insertShift (A.from2 tree (newBranch x sh)))+ | otherwise = Root (size + 1) sh (snocTree sh tree)+ where+ snocTree sh (Balanced arr)+ | sh == insertShift = Balanced $ A.snoc arr (newBranch x (down sh)) -- the current subtree is fully balanced+ | otherwise = Balanced $ A.adjust' arr (length arr - 1) (snocTree (down sh))+ snocTree sh (Unbalanced arr sizes)+ | sh == insertShift = Unbalanced (A.snoc arr (newBranch x (down sh))) newSizesSnoc+ | otherwise = Unbalanced (A.adjust' arr (length arr - 1) (snocTree (down sh))) newSizesAdjust+ where+ -- snoc the last size + 1+ newSizesSnoc = runST $ do+ let lenSizes = sizeofPrimArray sizes+ newArr <- newPrimArray (lenSizes + 1)+ copyPrimArray newArr 0 sizes 0 lenSizes+ let lastSize = indexPrimArray sizes (lenSizes - 1)+ writePrimArray newArr lenSizes (lastSize + 1)+ unsafeFreezePrimArray newArr+ -- adjust the last size with (+ 1)+ newSizesAdjust = runST $ do+ let lenSizes = sizeofPrimArray sizes+ newArr <- newPrimArray lenSizes+ copyPrimArray newArr 0 sizes 0 lenSizes+ let lastSize = indexPrimArray sizes (lenSizes - 1)+ writePrimArray newArr (lenSizes - 1) (lastSize + 1)+ unsafeFreezePrimArray newArr+ snocTree _ (Leaf arr) = Leaf $ A.snoc arr x++ insertShift = computeShift size sh (up sh) tree++ -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)+ -- the index is computed for efficient calculation of the shift in a balanced subtree+ computeShift i sh min (Balanced _) =+ let newShift = (countTrailingZeros i `div` blockShift) * blockShift+ in if newShift > sh then min else newShift+ computeShift _ sh min (Unbalanced arr sizes) =+ let i' = indexPrimArray sizes (sizeofPrimArray sizes - 1) - indexPrimArray sizes (sizeofPrimArray sizes - 2) -- sizes has at least 2 elements, otherwise the node would be balanced+ newMin = if length arr < blockSize then sh else min+ in computeShift i' (down sh) newMin (A.last arr)+ computeShift _ _ min (Leaf arr) = if length arr < blockSize then 0 else min++-- create a new tree with shift @sh@+newBranch :: a -> Int -> Tree a+newBranch x = go+ where+ go 0 = Leaf $ A.singleton x+ go sh = Balanced $ A.singleton (go (down sh))+{-# INLINE newBranch #-}++-- splitting++-- | \(O(\log n)\). The first @i@ elements of the vector.+-- If @i@ is negative, the empty vector is returned. If the vector contains less than @i@ elements, the whole vector is returned.+take :: Int -> Vector a -> Vector a+take _ Empty = Empty+take n v@(Root size sh tree)+ | n <= 0 = empty+ | n >= size = v+ | otherwise = normalize $ Root n sh (takeTree (n - 1) sh tree)+ where+ -- the initial @i@ is @n - 1@ -- the index of the last element in the new tree+ takeTree i sh (Balanced arr) =+ let idx = radixIndex i sh+ newArr = A.take arr (idx + 1)+ in Balanced (A.adjust' newArr idx (takeTree i (down sh)))+ takeTree i sh (Unbalanced arr sizes) =+ let (idx, subIdx) = relaxedRadixIndex sizes i sh+ newArr = A.take arr (idx + 1)+ in computeSizes sh (A.adjust' newArr idx (takeTree subIdx (down sh)))+ takeTree i _ (Leaf arr) = Leaf (A.take arr ((i .&. blockMask) + 1))++-- | \(O(\log n)\). The vector without the first @i@ elements+-- If @i@ is negative, the whole vector is returned. If the vector contains less than @i@ elements, the empty vector is returned.+drop :: Int -> Vector a -> Vector a+drop _ Empty = Empty+drop n v@(Root size sh tree)+ | n <= 0 = v+ | n >= size = empty+ | otherwise = normalize $ Root (size - n) sh (dropTree n sh tree)+ where+ dropTree n sh (Balanced arr) =+ let idx = radixIndex n sh+ newArr = A.drop arr idx+ in computeSizes sh (A.adjust' newArr 0 (dropTree n (down sh)))+ dropTree n sh (Unbalanced arr sizes) =+ let (idx, subIdx) = relaxedRadixIndex sizes n sh+ newArr = A.drop arr idx+ in computeSizes sh (A.adjust' newArr 0 (dropTree subIdx (down sh)))+ dropTree n _ (Leaf arr) = Leaf (A.drop arr (n .&. blockMask))++normalize :: Vector a -> Vector a+normalize (Root size sh (Balanced arr))+ | length arr == 1 = normalize $ Root size (down sh) (A.head arr)+normalize (Root size sh (Unbalanced arr _))+ | length arr == 1 = normalize $ Root size (down sh) (A.head arr)+normalize v = v
+ src/Data/RRBVector/Internal/Array.hs view
@@ -0,0 +1,207 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE MagicHash #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE UnboxedTuples #-}++-- Warning: No bound checks are performed!++module Data.RRBVector.Internal.Array+ ( Array, MutableArray+ , empty, singleton, from2+ , index, head, last+ , update, adjust, adjust'+ , take, drop, splitAt+ , snoc, cons+ , map, map'+ , traverse, traverse'+ , new, read, write+ , freeze, thaw+ ) where++import Control.Applicative (Applicative(liftA2))+import Control.DeepSeq (NFData(..))+import Control.Monad (when)+import Control.Monad.ST+import Data.Foldable (Foldable(..))+import Data.Primitive.SmallArray+import Prelude hiding (take, drop, splitAt, head, last, map, traverse, read)++-- start length array+data Array a = Array !Int !Int !(SmallArray a)+data MutableArray s a = MutableArray !Int !Int !(SmallMutableArray s a)++instance Semigroup (Array a) where+ Array start1 len1 arr1 <> Array start2 len2 arr2 = Array 0 len' $ runSmallArray $ do+ sma <- newSmallArray len' uninitialized+ copySmallArray sma 0 arr1 start1 len1+ copySmallArray sma len1 arr2 start2 len2+ pure sma+ where+ !len' = len1 + len2++instance Foldable Array where+ foldr f = \z (Array start len arr) ->+ let !end = start + len+ go i+ | i == end = z+ | (# x #) <- indexSmallArray## arr i = f x (go (i + 1))+ in go start+ {-# INLINE foldr #-}++ foldl f = \z (Array start len arr) ->+ let go i+ | i < start = z+ | (# x #) <- indexSmallArray## arr i = f (go (i - 1)) x+ in go (start + len - 1)+ {-# INLINE foldl #-}++ foldr' f = \z (Array start len arr) ->+ let go i !acc+ | i < start = acc+ | (# x #) <- indexSmallArray## arr i = go (i - 1) (f x acc)+ in go (start + len - 1) z+ {-# INLINE foldr' #-}++ foldl' f = \z (Array start len arr) ->+ let !end = start + len+ go i !acc+ | i == end = acc+ | (# x #) <- indexSmallArray## arr i = go (i + 1) (f acc x)+ in go start z+ {-# INLINE foldl' #-}++ null arr = length arr == 0+ {-# INLINE null #-}++ length (Array _ len _) = len+ {-# INLINE length #-}++instance (NFData a) => NFData (Array a) where+ rnf = foldl' (\_ x -> rnf x) ()++uninitialized :: a+uninitialized = errorWithoutStackTrace "uninitialized"++empty :: Array a+empty = Array 0 0 $ runSmallArray (newSmallArray 0 uninitialized)++singleton :: a -> Array a+singleton x = Array 0 1 $ runSmallArray (newSmallArray 1 x)++from2 :: a -> a -> Array a+from2 x y = Array 0 2 $ runSmallArray $ do+ sma <- newSmallArray 2 x+ writeSmallArray sma 1 y+ pure sma++index :: Array a -> Int -> a+index (Array start _ arr) idx = indexSmallArray arr (start + idx)++update :: Array a -> Int -> a -> Array a+update (Array start len sa) idx x = Array 0 len $ runSmallArray $ do+ sma <- thawSmallArray sa start len+ writeSmallArray sma idx x+ pure sma++adjust :: Array a -> Int -> (a -> a) -> Array a+adjust (Array start len sa) idx f = Array 0 len $ runSmallArray $ do+ sma <- thawSmallArray sa start len+ x <- indexSmallArrayM sa (start + idx)+ writeSmallArray sma idx (f x)+ pure sma++adjust' :: Array a -> Int -> (a -> a) -> Array a+adjust' (Array start len sa) idx f = Array 0 len $ runSmallArray $ do+ sma <- thawSmallArray sa start len+ x <- indexSmallArrayM sa (start + idx)+ writeSmallArray sma idx $! f x+ pure sma++take :: Array a -> Int -> Array a+take (Array start _ arr) n = Array start n arr++drop :: Array a -> Int -> Array a+drop (Array start len arr) n = Array (start + n) (len - n) arr++splitAt :: Array a -> Int -> (Array a, Array a)+splitAt arr idx = (take arr idx, drop arr idx)++head :: Array a -> a+head arr = index arr 0++last :: Array a -> a+last arr = index arr (length arr - 1)++snoc :: Array a -> a -> Array a+snoc (Array _ len arr) x = Array 0 len' $ runSmallArray $ do+ sma <- newSmallArray len' x+ copySmallArray sma 0 arr 0 len+ pure sma+ where+ !len' = len + 1++cons :: Array a -> a -> Array a+cons (Array _ len arr) x = Array 0 len' $ runSmallArray $ do+ sma <- newSmallArray len' x+ copySmallArray sma 1 arr 0 len+ pure sma+ where+ !len' = len + 1++map :: (a -> b) -> Array a -> Array b+map f (Array start len arr) = Array 0 len $ runSmallArray $ do+ sma <- newSmallArray len uninitialized+ -- i is the index in arr, j is the index in sma+ let loop i j = when (j < len) $ do+ x <- indexSmallArrayM arr i+ writeSmallArray sma j (f x)+ loop (i + 1) (j + 1)+ loop start 0+ pure sma++map' :: (a -> b) -> Array a -> Array b+map' f (Array start len arr) = Array 0 len $ runSmallArray $ do+ sma <- newSmallArray len uninitialized+ -- i is the index in arr, j is the index in sma+ let loop i j = when (j < len) $ do+ x <- indexSmallArrayM arr i+ writeSmallArray sma j $! f x+ loop (i + 1) (j + 1)+ loop start 0+ pure sma++newtype STA a = STA (forall s. SmallMutableArray s a -> ST s (SmallArray a))++runSTA :: Int -> STA a -> Array a+runSTA len (STA m) = Array 0 len (runST $ newSmallArray len uninitialized >>= m)++traverse :: (Applicative f) => (a -> f b) -> Array a -> f (Array b)+traverse f (Array start len arr) =+ -- i is the index in arr, j is the index in sma+ let go i j+ | j == len = pure $ STA unsafeFreezeSmallArray+ | (# x #) <- indexSmallArray## arr i = liftA2 (\y (STA m) -> STA $ \sma -> writeSmallArray sma j y *> m sma) (f x) (go (i + 1) (j + 1))+ in runSTA len <$> go start 0++traverse' :: (Applicative f) => (a -> f b) -> Array a -> f (Array b)+traverse' f (Array start len arr) =+ -- i is the index in arr, j is the index in sma+ let go i j+ | j == len = pure $ STA unsafeFreezeSmallArray+ | (# x #) <- indexSmallArray## arr i = liftA2 (\ !y (STA m) -> STA $ \sma -> writeSmallArray sma j y *> m sma) (f x) (go (i + 1) (j + 1))+ in runSTA len <$> go start 0++new :: Int -> ST s (MutableArray s a)+new len = MutableArray 0 len <$> newSmallArray len uninitialized++read :: MutableArray s a -> Int -> ST s a+read (MutableArray start _ arr) idx = readSmallArray arr (start + idx)++write :: MutableArray s a -> Int -> a -> ST s ()+write (MutableArray start _ arr) idx = writeSmallArray arr (start + idx)++freeze :: MutableArray s a -> Int -> Int -> ST s (Array a)+freeze (MutableArray start _ arr) idx len = Array 0 len <$> freezeSmallArray arr (start + idx) len++thaw :: Array a -> Int -> Int -> ST s (MutableArray s a)+thaw (Array start _ arr) idx len = MutableArray 0 len <$> thawSmallArray arr (start + idx) len
+ src/Data/RRBVector/Internal/Buffer.hs view
@@ -0,0 +1,37 @@+module Data.RRBVector.Internal.Buffer+ ( Buffer+ , new+ , push+ , get+ , size+ ) where++import Control.Monad.ST++import Data.RRBVector.Internal.IntRef+import qualified Data.RRBVector.Internal.Array as A++data Buffer s a = Buffer !(A.MutableArray s a) !(IntRef s)++new :: Int -> ST s (Buffer s a)+new capacity = do+ buffer <- A.new capacity+ offset <- newIntRef 0+ pure (Buffer buffer offset)++push :: Buffer s a -> a -> ST s ()+push (Buffer buffer offset) x = do+ idx <- readIntRef offset+ A.write buffer idx x+ modifyIntRef offset (+ 1)++get :: Buffer s a -> ST s (A.Array a)+get (Buffer buffer offset) = do+ len <- readIntRef offset+ result <- A.freeze buffer 0 len+ writeIntRef offset 0+ pure result++size :: Buffer s a -> ST s Int+size (Buffer _ offset) = readIntRef offset+{-# INLInE size #-}
+ src/Data/RRBVector/Internal/Debug.hs view
@@ -0,0 +1,61 @@+{- |+This module contains some debug utilities. It should only be used for debugging/testing purposes.+-}++module Data.RRBVector.Internal.Debug+ ( showTree+ , fromListUnbalanced+ ) where++import Control.Monad.ST (runST)+import Data.Foldable (toList)+import Data.List (intercalate)+import Data.Primitive (primArrayToList)++import Data.RRBVector.Internal+import qualified Data.RRBVector.Internal.Buffer as Buffer++-- | \(O(n)\). Show the underlying tree of a vector.+showTree :: (Show a) => Vector a -> String+showTree Empty = "Empty"+showTree (Root size sh tree) = "Root {size = " ++ show size ++ ", shift = " ++ show sh ++ ", tree = " ++ debugShowTree tree ++ "}"+ where+ debugShowTree (Balanced arr) = "Balanced " ++ debugShowArray arr+ debugShowTree (Unbalanced arr sizes) = "Unbalanced " ++ debugShowArray arr ++ " (" ++ show (primArrayToList sizes) ++ ")"+ debugShowTree (Leaf arr) = "Leaf " ++ show (toList arr)++ debugShowArray arr = "[" ++ intercalate "," (fmap debugShowTree (toList arr)) ++ "]"++-- | \(O(n)\). Create a new unbalanced vector from a list.+--+-- Note that it is not possbible to create an invalid 'Vector' with this function.+fromListUnbalanced :: [a] -> Vector a+fromListUnbalanced [] = Empty+fromListUnbalanced [x] = singleton x+fromListUnbalanced ls = case nodes Leaf ls of+ [tree] -> Root (treeSize 0 tree) 0 tree -- tree is a single leaf+ ls' -> iterateNodes blockShift ls'+ where+ n = blockSize - 1++ nodes f trees = runST $ do+ buffer <- Buffer.new n+ let loop [] = do+ result <- Buffer.get buffer+ pure [f result]+ loop (t : ts) = do+ size <- Buffer.size buffer+ if size == n then do+ result <- Buffer.get buffer+ Buffer.push buffer t+ rest <- loop ts+ pure (f result : rest)+ else do+ Buffer.push buffer t+ loop ts+ loop trees+ {-# INLINE nodes #-}++ iterateNodes sh trees = case nodes (computeSizes sh) trees of+ [tree] -> Root (treeSize sh tree) sh tree+ trees' -> iterateNodes (up sh) trees'
+ src/Data/RRBVector/Internal/Indexed.hs view
@@ -0,0 +1,26 @@+module Data.RRBVector.Internal.Indexed where++-- TODO: use unboxed tuples?++data WithIndex a = WithIndex !Int a++-- | > Compose (State Int) f a+newtype Indexed f a = Indexed { runIndexed :: Int -> WithIndex (f a) }++instance Functor f => Functor (Indexed f) where+ fmap f (Indexed sf) = Indexed $ \s -> let WithIndex s' x = sf s in WithIndex s' (fmap f x)+ {-# INLINE fmap #-}++instance Applicative f => Applicative (Indexed f) where+ pure x = Indexed $ \s -> WithIndex s (pure x)+ {-# INLINE pure #-}++ Indexed sfa <*> Indexed sfb = Indexed $ \s ->+ let WithIndex s' f = sfa s+ WithIndex s'' x = sfb s'+ in WithIndex s'' (f <*> x)+ {-# INLINE (<*>) #-}++evalIndexed :: Indexed f a -> Int -> f a+evalIndexed (Indexed sf) x = let WithIndex _ y = sf x in y+{-# INLINE evalIndexed #-}
+ src/Data/RRBVector/Internal/IntRef.hs view
@@ -0,0 +1,33 @@+module Data.RRBVector.Internal.IntRef+ ( IntRef+ , newIntRef+ , readIntRef+ , writeIntRef+ , modifyIntRef+ ) where++import Control.Monad.ST+import Data.Primitive.PrimArray++newtype IntRef s = IntRef (MutablePrimArray s Int)++newIntRef :: Int -> ST s (IntRef s)+newIntRef i = do+ arr <- newPrimArray 1+ setPrimArray arr 0 1 i+ pure (IntRef arr)+{-# INLINE newIntRef #-}++readIntRef :: IntRef s -> ST s Int+readIntRef (IntRef arr) = readPrimArray arr 0+{-# INLINE readIntRef #-}++writeIntRef :: IntRef s -> Int -> ST s ()+writeIntRef (IntRef arr) = writePrimArray arr 0+{-# INLINE writeIntRef #-}++modifyIntRef :: IntRef s -> (Int -> Int) -> ST s ()+modifyIntRef (IntRef arr) f = do+ i <- readPrimArray arr 0+ writePrimArray arr 0 (f i)+{-# INLINE modifyIntRef #-}
+ test/Spec.hs view
@@ -0,0 +1,83 @@+import Data.Foldable (toList)+import Data.List (uncons)++import Test.Hspec+import Test.Hspec.QuickCheck+import Test.QuickCheck++import qualified Data.RRBVector as V+import Data.RRBVector.Internal.Debug (fromListUnbalanced)++default (Int)++instance (Arbitrary a) => Arbitrary (V.Vector a) where+ arbitrary = oneof [V.fromList <$> arbitrary, fromListUnbalanced <$> arbitrary]++lookupList :: Int -> [a] -> Maybe a+lookupList i ls+ | i < length ls = Just (ls !! i)+ | otherwise = Nothing++updateList :: Int -> a -> [a] -> [a]+updateList i x ls+ | i < length ls = let (left, _ : right) = splitAt i ls in left ++ (x : right)+ | otherwise = ls++adjustList :: Int -> (a -> a) -> [a] -> [a]+adjustList i f ls+ | i < length ls = let (left, x : right) = splitAt i ls in left ++ (f x : right)+ | otherwise = ls++unsnoc :: [a] -> Maybe ([a], a)+unsnoc [] = Nothing+unsnoc ls = Just (init ls, last ls)++main :: IO ()+main = hspec . modifyMaxSuccess maxN . modifyMaxSize maxN $ do+ prop "satisfies `fromList . toList == id`" $ \v -> V.fromList (toList v) === v+ prop "satisfies `toList . fromList == id`" $ \ls -> toList (V.fromList ls) === ls++ describe "lookup" $ do+ prop "gets the element at the index" $ \v (NonNegative i) -> V.lookup i v === lookupList i (toList v)+ prop "returns Nothing for negative indices" $ \v (Negative i) -> V.lookup i v === Nothing++ describe "update" $ do+ prop "updates the element at the index" $ \v (NonNegative i) x -> toList (V.update i x v) === updateList i x (toList v)+ prop "returns the vector for negative indices" $ \v (Negative i) x -> V.update i x v === v++ describe "adjust" $ do+ prop "adjusts the element at the index" $ \v (NonNegative i) -> toList (V.adjust i (+ 1) v) === adjustList i (+ 1) (toList v)+ prop "returns the vector for negative indices" $ \v (Negative i) -> V.adjust i (+ 1) v === v++ describe "><" $ do+ prop "concatenates two vectors" $ \v1 v2 -> toList (v1 V.>< v2) === toList v1 ++ toList v2+ prop "works for the empty vector" $ \v -> (V.empty V.>< v `shouldBe` v) .&&. (v V.>< V.empty `shouldBe` v)++ describe "|>" $ do+ prop "appends an element" $ \v x -> toList (v V.|> x) === toList v ++ [x]+ prop "works for the empty vector" $ \x -> V.empty V.|> x `shouldBe` V.singleton x++ describe "<|" $ do+ prop "prepends an element" $ \x v -> toList (x V.<| v) === x : toList v+ prop "works for the empty vector" $ \x -> x V.<| V.empty `shouldBe` V.singleton x++ describe "take" $ do+ prop "takes n elements" $ \v (Positive n) -> toList (V.take n v) === take n (toList v)+ prop "returns the empty vector for non-positive n" $ \v (NonPositive n) -> V.take n v === V.empty++ describe "drop" $ do+ prop "drops n elements" $ \v (Positive n) -> toList (V.drop n v) === drop n (toList v)+ prop "does nothing for non-positive n" $ \v (NonPositive n) -> V.drop n v === v++ describe "viewl" $ do+ prop "works like uncons" $ \v -> fmap (\(x, xs) -> (x, toList xs)) (V.viewl v) === uncons (toList v)+ prop "works for the empty vector" $ V.viewl V.empty `shouldBe` Nothing++ describe "viewr" $ do+ prop "works like unsnoc" $ \v -> fmap (\(xs, x) -> (toList xs, x)) (V.viewr v) === unsnoc (toList v)+ prop "works for the empty vector" $ V.viewr V.empty `shouldBe` Nothing++ describe "map" $ do+ prop "maps over the vector" $ \v -> toList (V.map (+ 1) v) === map (+ 1) (toList v)+ where+ maxN = const 10_000