diff --git a/vector/src/Data/Vector/Fusion/Bundle/Monadic.hs b/vector/src/Data/Vector/Fusion/Bundle/Monadic.hs index 99a3bdea..43e8d145 100644 --- a/vector/src/Data/Vector/Fusion/Bundle/Monadic.hs +++ b/vector/src/Data/Vector/Fusion/Bundle/Monadic.hs @@ -162,13 +162,15 @@ sized s sz = s { sSize = sz } -- | Length of a 'Bundle' length :: Monad m => Bundle m v a -> m Int {-# INLINE_FUSED length #-} -length Bundle{sSize = Exact n} = return n +-- FIXME: SIZE +-- length Bundle{sSize = Exact n} = return n length Bundle{sChunks = s} = S.foldl' (\n (Chunk k _) -> n+k) 0 s -- | Check if a 'Bundle' is empty null :: Monad m => Bundle m v a -> m Bool {-# INLINE_FUSED null #-} -null Bundle{sSize = Exact n} = return (n == 0) +-- FIXME: SIZE +-- null Bundle{sSize = Exact n} = return (n == 0) null Bundle{sChunks = s} = S.foldr (\(Chunk n _) z -> n == 0 && z) True s -- Construction @@ -177,12 +179,12 @@ null Bundle{sChunks = s} = S.foldr (\(Chunk n _) z -> n == 0 && z) True s -- | Empty 'Bundle' empty :: Monad m => Bundle m v a {-# INLINE_FUSED empty #-} -empty = fromStream S.empty (Exact 0) +empty = fromStream S.empty (exact 0) -- | Singleton 'Bundle' singleton :: Monad m => a -> Bundle m v a {-# INLINE_FUSED singleton #-} -singleton x = fromStream (S.singleton x) (Exact 1) +singleton x = fromStream (S.singleton x) (exact 1) -- | Replicate a value to a given length replicate :: Monad m => Int -> a -> Bundle m v a @@ -190,7 +192,7 @@ replicate :: Monad m => Int -> a -> Bundle m v a replicate n x = Bundle (S.replicate n x) (S.singleton $ Chunk len (\v -> stToPrim $ M.basicSet v x)) Nothing - (Exact len) + (exact len) where len = delay_inline max n 0 @@ -200,7 +202,7 @@ replicateM :: Monad m => Int -> m a -> Bundle m v a {-# INLINE_FUSED replicateM #-} -- NOTE: We delay inlining max here because GHC will create a join point for -- the call to newArray# otherwise which is not really nice. -replicateM n p = fromStream (S.replicateM n p) (Exact (delay_inline max n 0)) +replicateM n p = fromStream (S.replicateM n p) (exact (delay_inline max n 0)) generate :: Monad m => Int -> (Int -> a) -> Bundle m v a {-# INLINE generate #-} @@ -209,7 +211,7 @@ generate n f = generateM n (return . f) -- | Generate a stream from its indices generateM :: Monad m => Int -> (Int -> m a) -> Bundle m v a {-# INLINE_FUSED generateM #-} -generateM n f = fromStream (S.generateM n f) (Exact (delay_inline max n 0)) +generateM n f = fromStream (S.generateM n f) (exact (delay_inline max n 0)) -- | Prepend an element cons :: Monad m => a -> Bundle m v a -> Bundle m v a @@ -276,13 +278,13 @@ tail Bundle{sElems = s, sSize = sz} = fromStream (S.tail s) (sz-1) -- | The first @n@ elements take :: Monad m => Int -> Bundle m v a -> Bundle m v a {-# INLINE_FUSED take #-} -take n Bundle{sElems = s, sSize = sz} = fromStream (S.take n s) (smallerThan n sz) +take n Bundle{sElems = s, sSize = sz} = fromStream (S.take n s) (setUpperBound n sz) -- | All but the first @n@ elements drop :: Monad m => Int -> Bundle m v a -> Bundle m v a {-# INLINE_FUSED drop #-} drop n Bundle{sElems = s, sSize = sz} = - fromStream (S.drop n s) (clampedSubtract sz (Exact n)) + fromStream (S.drop n s) (sz - exact n) --FIXME: specialize? -- Mapping -- ------- @@ -440,9 +442,10 @@ eqBy eq x y | otherwise = S.eqBy eq (sElems x) (sElems y) where sizesAreDifferent :: Size -> Size -> Bool - sizesAreDifferent (Exact a) (Exact b) = a /= b - sizesAreDifferent (Exact a) (Max b) = a > b - sizesAreDifferent (Max a) (Exact b) = a < b + -- FIXME: SIZE + -- sizesAreDifferent (Exact a) (Exact b) = a /= b + -- sizesAreDifferent (Exact a) (Max b) = a > b + -- sizesAreDifferent (Max a) (Exact b) = a < b sizesAreDifferent _ _ = False -- | Lexicographically compare two 'Bundle's @@ -461,14 +464,14 @@ filter f = filterM (return . f) -- | Drop elements which do not satisfy the monadic predicate filterM :: Monad m => (a -> m Bool) -> Bundle m v a -> Bundle m v a {-# INLINE_FUSED filterM #-} -filterM f Bundle{sElems = s, sSize = n} = fromStream (S.filterM f s) (toMax n) +filterM f Bundle{sElems = s, sSize = n} = fromStream (S.filterM f s) (zeroLowerBound n) -- | Apply monadic function to each element and drop all Nothings -- -- @since 0.12.2.0 mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Bundle m v a -> Bundle m v b {-# INLINE_FUSED mapMaybeM #-} -mapMaybeM f Bundle{sElems = s, sSize = n} = fromStream (S.mapMaybeM f s) (toMax n) +mapMaybeM f Bundle{sElems = s, sSize = n} = fromStream (S.mapMaybeM f s) (zeroLowerBound n) -- | Longest prefix of elements that satisfy the predicate takeWhile :: Monad m => (a -> Bool) -> Bundle m v a -> Bundle m v a @@ -478,7 +481,7 @@ takeWhile f = takeWhileM (return . f) -- | Longest prefix of elements that satisfy the monadic predicate takeWhileM :: Monad m => (a -> m Bool) -> Bundle m v a -> Bundle m v a {-# INLINE_FUSED takeWhileM #-} -takeWhileM f Bundle{sElems = s, sSize = n} = fromStream (S.takeWhileM f s) (toMax n) +takeWhileM f Bundle{sElems = s, sSize = n} = fromStream (S.takeWhileM f s) (zeroLowerBound n) -- | Drop the longest prefix of elements that satisfy the predicate dropWhile :: Monad m => (a -> Bool) -> Bundle m v a -> Bundle m v a @@ -488,7 +491,7 @@ dropWhile f = dropWhileM (return . f) -- | Drop the longest prefix of elements that satisfy the monadic predicate dropWhileM :: Monad m => (a -> m Bool) -> Bundle m v a -> Bundle m v a {-# INLINE_FUSED dropWhileM #-} -dropWhileM f Bundle{sElems = s, sSize = n} = fromStream (S.dropWhileM f s) (toMax n) +dropWhileM f Bundle{sElems = s, sSize = n} = fromStream (S.dropWhileM f s) (zeroLowerBound n) -- Searching -- --------- @@ -630,7 +633,7 @@ concatMap f = concatMapM (return . f) concatMapM :: Monad m => (a -> m (Bundle m v b)) -> Bundle m v a -> Bundle m v b {-# INLINE_FUSED concatMapM #-} -concatMapM f Bundle{sElems = s} = fromStream (S.concatMapM (liftM sElems . f) s) Unknown +concatMapM f Bundle{sElems = s} = fromStream (S.concatMapM (liftM sElems . f) s) unknown -- | Create a 'Bundle' of values from a 'Bundle' of streamable things flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Size @@ -649,7 +652,7 @@ unfoldr f = unfoldrM (return . f) -- | Unfold with a monadic function unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Bundle m u a {-# INLINE_FUSED unfoldrM #-} -unfoldrM f s = fromStream (S.unfoldrM f s) Unknown +unfoldrM f s = fromStream (S.unfoldrM f s) unknown -- | Unfold at most @n@ elements unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Bundle m u a @@ -660,6 +663,7 @@ unfoldrN n f = unfoldrNM n (return . f) unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Bundle m u a {-# INLINE_FUSED unfoldrNM #-} unfoldrNM n f s = fromStream (S.unfoldrNM n f s) Unknown +unfoldrNM n f s = fromStream (S.unfoldrNM n f s) (maxSize (delay_inline max n 0)) -- | Unfold exactly @n@ elements -- @@ -673,14 +677,14 @@ unfoldrExactN n f = unfoldrExactNM n (return . f) -- @since 0.12.2.0 unfoldrExactNM :: Monad m => Int -> (s -> m (a, s)) -> s -> Bundle m u a {-# INLINE_FUSED unfoldrExactNM #-} -unfoldrExactNM n f s = fromStream (S.unfoldrExactNM n f s) (Max (delay_inline max n 0)) +unfoldrExactNM n f s = fromStream (S.unfoldrExactNM n f s) (maxSize (delay_inline max n 0)) -- | /O(n)/ Apply monadic function \(\max(n - 1, 0)\) times to an initial value, producing -- a monadic bundle of exact length \(\max(n, 0)\). Zeroth element will contain the initial -- value. iterateNM :: Monad m => Int -> (a -> m a) -> a -> Bundle m u a {-# INLINE_FUSED iterateNM #-} -iterateNM n f x0 = fromStream (S.iterateNM n f x0) (Exact (delay_inline max n 0)) +iterateNM n f x0 = fromStream (S.iterateNM n f x0) (exact (delay_inline max n 0)) -- | /O(n)/ Apply function \(\max(n - 1, 0)\) times to an initial value, producing a -- monadic bundle of exact length \(\max(n, 0)\). Zeroth element will contain the initial @@ -784,7 +788,7 @@ scanl1M' f Bundle{sElems = s, sSize = sz} = fromStream (S.scanl1M' f s) sz -- @x+y+y@ etc. enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Bundle m v a {-# INLINE_FUSED enumFromStepN #-} -enumFromStepN x y n = fromStream (S.enumFromStepN x y n) (Exact (delay_inline max n 0)) +enumFromStepN x y n = fromStream (S.enumFromStepN x y n) (exact (delay_inline max n 0)) -- | Enumerate values -- @@ -800,7 +804,7 @@ enumFromTo x y = fromList [x .. y] -- FIXME: add "too large" test for Int enumFromTo_small :: (Integral a, Monad m) => a -> a -> Bundle m v a {-# INLINE_FUSED enumFromTo_small #-} -enumFromTo_small x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact n) +enumFromTo_small x y = x `seq` y `seq` fromStream (Stream step (Just x)) (exact n) where n = delay_inline max (fromIntegral y - fromIntegral x + 1) 0 @@ -852,7 +856,7 @@ enumFromTo_small x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact enumFromTo_int :: forall m v. (HasCallStack, Monad m) => Int -> Int -> Bundle m v Int {-# INLINE_FUSED enumFromTo_int #-} -enumFromTo_int x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact (len x y)) +enumFromTo_int x y = x `seq` y `seq` fromStream (Stream step (Just x)) (exact (len x y)) where {-# INLINE [0] len #-} len :: HasCallStack => Int -> Int -> Int @@ -869,7 +873,7 @@ enumFromTo_int x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact (l enumFromTo_intlike :: forall m v a. (HasCallStack, Integral a, Monad m) => a -> a -> Bundle m v a {-# INLINE_FUSED enumFromTo_intlike #-} -enumFromTo_intlike x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact (len x y)) +enumFromTo_intlike x y = x `seq` y `seq` fromStream (Stream step (Just x)) (exact (len x y)) where {-# INLINE [0] len #-} len :: HasCallStack => a -> a -> Int @@ -907,7 +911,7 @@ enumFromTo_intlike x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exac enumFromTo_big_word :: forall m v a. (HasCallStack, Integral a, Monad m) => a -> a -> Bundle m v a {-# INLINE_FUSED enumFromTo_big_word #-} -enumFromTo_big_word x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact (len x y)) +enumFromTo_big_word x y = x `seq` y `seq` fromStream (Stream step (Just x)) (exact (len x y)) where {-# INLINE [0] len #-} len :: HasCallStack => a -> a -> Int @@ -951,7 +955,7 @@ enumFromTo_big_word x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exa -- FIXME: the "too large" test is totally wrong enumFromTo_big_int :: forall m v a. (HasCallStack, Integral a, Monad m) => a -> a -> Bundle m v a {-# INLINE_FUSED enumFromTo_big_int #-} -enumFromTo_big_int x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exact (len x y)) +enumFromTo_big_int x y = x `seq` y `seq` fromStream (Stream step (Just x)) (exact (len x y)) where {-# INLINE [0] len #-} len :: HasCallStack => a -> a -> Int @@ -980,7 +984,7 @@ enumFromTo_big_int x y = x `seq` y `seq` fromStream (Stream step (Just x)) (Exac enumFromTo_char :: Monad m => Char -> Char -> Bundle m v Char {-# INLINE_FUSED enumFromTo_char #-} -enumFromTo_char x y = x `seq` y `seq` fromStream (Stream step xn) (Exact n) +enumFromTo_char x y = x `seq` y `seq` fromStream (Stream step xn) (exact n) where xn = ord x yn = ord y @@ -1005,7 +1009,7 @@ enumFromTo_char x y = x `seq` y `seq` fromStream (Stream step xn) (Exact n) enumFromTo_double :: forall m v a. (HasCallStack, Monad m, Ord a, RealFrac a) => a -> a -> Bundle m v a {-# INLINE_FUSED enumFromTo_double #-} -enumFromTo_double n m = n `seq` m `seq` fromStream (Stream step ini) (Max (len n lim)) +enumFromTo_double n m = n `seq` m `seq` fromStream (Stream step ini) (maxSize (len n lim)) where lim = m + 1/2 -- important to float out @@ -1069,12 +1073,12 @@ toList = foldr (:) [] -- | Convert a list to a 'Bundle' fromList :: Monad m => [a] -> Bundle m v a {-# INLINE fromList #-} -fromList xs = unsafeFromList Unknown xs +fromList xs = unsafeFromList unknown xs -- | Convert the first @n@ elements of a list to a 'Bundle' fromListN :: Monad m => Int -> [a] -> Bundle m v a {-# INLINE_FUSED fromListN #-} -fromListN n xs = fromStream (S.fromListN n xs) (Max (delay_inline max n 0)) +fromListN n xs = fromStream (S.fromListN n xs) (exact (delay_inline max n 0)) -- | Convert a list to a 'Bundle' with the given 'Size' hint. unsafeFromList :: Monad m => Size -> [a] -> Bundle m v a @@ -1086,7 +1090,7 @@ fromVector :: (Monad m, Vector v a) => v a -> Bundle m v a fromVector v = v `seq` n `seq` Bundle (Stream step 0) (Stream vstep True) (Just v) - (Exact n) + (exact n) where n = basicLength v @@ -1105,7 +1109,7 @@ fromVectors :: forall m v a. (Monad m, Vector v a) => [v a] -> Bundle m v a fromVectors us = Bundle (Stream pstep (Left us)) (Stream vstep us) Nothing - (Exact n) + (exact n) where n = List.foldl' (\k v -> k + basicLength v) 0 us @@ -1134,7 +1138,7 @@ concatVectors Bundle{sElems = Stream step t} = Bundle (Stream pstep (Left t)) (Stream vstep t) Nothing - Unknown + unknown where pstep (Left s) = do r <- step s diff --git a/vector/src/Data/Vector/Fusion/Bundle/Size.hs b/vector/src/Data/Vector/Fusion/Bundle/Size.hs index da757adf..ec17b213 100644 --- a/vector/src/Data/Vector/Fusion/Bundle/Size.hs +++ b/vector/src/Data/Vector/Fusion/Bundle/Size.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE PatternSynonyms #-} -- | -- Module : Data.Vector.Fusion.Bundle.Size -- Copyright : (c) Roman Leshchinskiy 2008-2010 @@ -14,119 +15,107 @@ -- module Data.Vector.Fusion.Bundle.Size ( - Size(..), clampedSubtract, smaller, smallerThan, larger, toMax, upperBound, lowerBound + Size(..), exact, unknown, maxSize, + + {-clampedSubtract,-} smaller, {-smallerThan,-} larger + -- , toMax, upperBound, lowerBound + , zeroLowerBound + , setUpperBound ) where import Data.Vector.Fusion.Util ( delay_inline ) --- | Size hint -data Size = Exact {-# UNPACK #-} !Int -- ^ Exact size - | Max {-# UNPACK #-} !Int -- ^ Upper bound on the size - | Unknown -- ^ Unknown size - deriving( Eq, Show ) +-- | Hint about size of vector that will be produced from bundle. It's +-- an interval which bounds size of produced vector. +-- +-- Upper bound is treated specially. @maxBound::Int@ means unbounded +-- stream and any estimate that overflows @Int@ is treated as +-- unbounded. +-- +-- Note that it's possible to create vectors with length @maxBound@ +-- since unboxed vectors for () use O(1) memory. So we won't run +-- out of memory first. +data Size = Size + { lowerBound :: !Int + -- ^ Lower bound on size of vector. + , upperBound :: !Int + -- ^ Upper bound on size of vector. + } + deriving( Eq, Show ) + +-- | Hint for size that is known exactly +exact :: Int -> Size +exact n = Size n n + +-- | Unknown size +unknown :: Size +unknown = Size 0 maxBound + +-- | Hint for case when we have upper bound but no lower bound +maxSize :: Int -> Size +maxSize n = Size 0 n instance Num Size where - Exact m + Exact n = checkedAdd Exact m n - Exact m + Max n = checkedAdd Max m n - - Max m + Exact n = checkedAdd Max m n - Max m + Max n = checkedAdd Max m n - - _ + _ = Unknown + Size lA uA + Size lB uB = Size (checkedAdd lA lB) (saturatedAdd uA uB) + -- + Size lA uA - Size lB _ + | uA == maxBound = Size lR maxBound + | otherwise = Size lR (saturatedSub uA lB) + where + lR = saturatedSub lA lB + -- + fromInteger = exact . fromInteger + (*) = error "vector: internal error * for Bundle.size isn't defined" + abs = error "vector: internal error abs for Bundle.size isn't defined" + signum = error "vector: internal error signum for Bundle.size isn't defined" - Exact m - Exact n = checkedSubtract Exact m n - Exact m - Max _ = Max m - - Max m - Exact n = checkedSubtract Max m n - Max m - Max _ = Max m - Max m - Unknown = Max m - - _ - _ = Unknown +-- | Add two non-negative integers and in case of overflow return maxBound +saturatedAdd :: Int -> Int -> Int +saturatedAdd a b | n < 0 = maxBound + | otherwise = n + where n = a + b - fromInteger n = Exact (fromInteger n) +-- | Subtract two non-negative integers. If result is negative it's set to zero +saturatedSub :: Int -> Int -> Int +saturatedSub a b | a < b = 0 + | otherwise = a - b - (*) = error "vector: internal error * for Bundle.size isn't defined" - abs = error "vector: internal error abs for Bundle.size isn't defined" - signum = error "vector: internal error signum for Bundle.size isn't defined" +-- | Add two non-negative integers and throw error in case of overflow +checkedAdd :: Int -> Int -> Int {-# INLINE checkedAdd #-} -checkedAdd :: (Int -> Size) -> Int -> Int -> Size -checkedAdd con m n +checkedAdd m n -- Note: we assume m and n are >= 0. | r < m || r < n = error $ "Data.Vector.Fusion.Bundle.Size.checkedAdd: overflow: " ++ show r - | otherwise = con r + | otherwise = r where r = m + n -{-# INLINE checkedSubtract #-} -checkedSubtract :: (Int -> Size) -> Int -> Int -> Size -checkedSubtract con m n - | r < 0 = - error $ "Data.Vector.Fusion.Bundle.Size.checkedSubtract: underflow: " ++ show r - | otherwise = con r - where - r = m - n - --- | Subtract two sizes with clamping to 0, for drop-like things -{-# INLINE clampedSubtract #-} -clampedSubtract :: Size -> Size -> Size -clampedSubtract (Exact m) (Exact n) = Exact (max 0 (m - n)) -clampedSubtract (Max m) (Exact n) - | m <= n = Exact 0 - | otherwise = Max (m - n) -clampedSubtract (Exact m) (Max _) = Max m -clampedSubtract (Max m) (Max _) = Max m -clampedSubtract _ _ = Unknown -- | Minimum of two size hints smaller :: Size -> Size -> Size {-# INLINE smaller #-} -smaller (Exact m) (Exact n) = Exact (delay_inline min m n) -smaller (Exact m) (Max n) = Max (delay_inline min m n) -smaller (Exact m) Unknown = Max m -smaller (Max m) (Exact n) = Max (delay_inline min m n) -smaller (Max m) (Max n) = Max (delay_inline min m n) -smaller (Max m) Unknown = Max m -smaller Unknown (Exact n) = Max n -smaller Unknown (Max n) = Max n -smaller Unknown Unknown = Unknown - --- | Select a safe smaller than known size. -smallerThan :: Int -> Size -> Size -{-# INLINE smallerThan #-} -smallerThan m (Exact n) = Exact (delay_inline min m n) -smallerThan m (Max n) = Max (delay_inline min m n) -smallerThan _ Unknown = Unknown +smaller (Size lA uA) (Size lB uB) + = Size (min lA lB) (min uA uB) -- | Maximum of two size hints larger :: Size -> Size -> Size {-# INLINE larger #-} -larger (Exact m) (Exact n) = Exact (delay_inline max m n) -larger (Exact m) (Max n) | m >= n = Exact m - | otherwise = Max n -larger (Max m) (Exact n) | n >= m = Exact n - | otherwise = Max m -larger (Max m) (Max n) = Max (delay_inline max m n) -larger _ _ = Unknown - --- | Convert a size hint to an upper bound -toMax :: Size -> Size -toMax (Exact n) = Max n -toMax (Max n) = Max n -toMax Unknown = Unknown - --- | Compute the minimum size from a size hint -lowerBound :: Size -> Int -lowerBound (Exact n) = n -lowerBound _ = 0 - --- | Compute the maximum size from a size hint if possible -upperBound :: Size -> Maybe Int -upperBound (Exact n) = Just n -upperBound (Max n) = Just n -upperBound Unknown = Nothing +larger (Size lA uA) (Size lB uB) + = Size (max lA lB) (max uA uB) + + + +-- | Set lower bound of size hint to zero +zeroLowerBound :: Size -> Size +{-# INLINE zeroLowerBound #-} +zeroLowerBound (Size _ ub) = Size 0 ub +-- | Set upper bound of size to given value +setUpperBound :: Int -> Size -> Size +{-# INLINE setUpperBound #-} +setUpperBound n (Size lb ub) = Size (min n lb) (min n ub) diff --git a/vector/src/Data/Vector/Generic.hs b/vector/src/Data/Vector/Generic.hs index 35acd0fb..5bf8ba6c 100644 --- a/vector/src/Data/Vector/Generic.hs +++ b/vector/src/Data/Vector/Generic.hs @@ -1385,14 +1385,14 @@ unzip6 xs = (map (\(a, _, _, _, _, _) -> a) xs, -- | /O(n)/ Drop all elements that do not satisfy the predicate. filter :: Vector v a => (a -> Bool) -> v a -> v a {-# INLINE filter #-} -filter f = unstream . inplace (S.filter f) toMax . stream +filter f = unstream . inplace (S.filter f) zeroLowerBound . stream -- | /O(n)/ Drop all elements that do not satisfy the predicate which is applied to -- the values and their indices. ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a {-# INLINE ifilter #-} ifilter f = unstream - . inplace (S.map snd . S.filter (uncurry f) . S.indexed) toMax + . inplace (S.map snd . S.filter (uncurry f) . S.indexed) zeroLowerBound . stream -- | /O(n)/ Drop repeated adjacent elements. The first element in each group is returned. @@ -1407,18 +1407,18 @@ ifilter f = unstream -- [Arg 1 'a'] uniq :: (Vector v a, Eq a) => v a -> v a {-# INLINE uniq #-} -uniq = unstream . inplace S.uniq toMax . stream +uniq = unstream . inplace S.uniq zeroLowerBound . stream -- | /O(n)/ Map the values and collect the 'Just' results. mapMaybe :: (Vector v a, Vector v b) => (a -> Maybe b) -> v a -> v b {-# INLINE mapMaybe #-} -mapMaybe f = unstream . inplace (S.mapMaybe f) toMax . stream +mapMaybe f = unstream . inplace (S.mapMaybe f) zeroLowerBound . stream -- | /O(n)/ Map the indices/values and collect the 'Just' results. imapMaybe :: (Vector v a, Vector v b) => (Int -> a -> Maybe b) -> v a -> v b {-# INLINE imapMaybe #-} imapMaybe f = unstream - . inplace (S.mapMaybe (uncurry f) . S.indexed) toMax + . inplace (S.mapMaybe (uncurry f) . S.indexed) zeroLowerBound . stream @@ -1702,7 +1702,7 @@ findIndexR f v = fmap (length v - 1 -) . Bundle.findIndex f $ streamR v findIndices :: (Vector v a, Vector v Int) => (a -> Bool) -> v a -> v Int {-# INLINE findIndices #-} findIndices f = unstream - . inplace (S.map fst . S.filter (f . snd) . S.indexed) toMax + . inplace (S.map fst . S.filter (f . snd) . S.indexed) zeroLowerBound . stream -- | /O(n)/ Yield 'Just' the index of the first occurrence of the given element or @@ -2554,7 +2554,7 @@ unstream s = new (New.unstream s) -- | /O(1)/ Convert a vector to a 'Bundle', proceeding from right to left. streamR :: Vector v a => v a -> Bundle u a {-# INLINE_FUSED streamR #-} -streamR v = v `seq` n `seq` (Bundle.unfoldr get n `Bundle.sized` Exact n) +streamR v = v `seq` n `seq` (Bundle.unfoldr get n `Bundle.sized` exact n) where n = length v diff --git a/vector/src/Data/Vector/Generic/Mutable.hs b/vector/src/Data/Vector/Generic/Mutable.hs index 0f55e6f1..571a6fc0 100644 --- a/vector/src/Data/Vector/Generic/Mutable.hs +++ b/vector/src/Data/Vector/Generic/Mutable.hs @@ -104,32 +104,31 @@ import Data.Bits ( Bits(shiftR) ) -- Internal functions -- ------------------ -unsafeAppend1 :: (PrimMonad m, MVector v a) - => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a) +unsafeAppend1 + :: (PrimMonad m, MVector v a) + => Int -> v (PrimState m) a -> Int -> a -> m (v (PrimState m) a) {-# INLINE_INNER unsafeAppend1 #-} -- NOTE: The case distinction has to be on the outside because -- GHC creates a join point for the unsafeWrite even when everything -- is inlined. This is bad because with the join point, v isn't getting -- unboxed. -unsafeAppend1 v i x - | i < length v = do - unsafeWrite v i x - return v - | otherwise = do - v' <- enlarge v - checkIndex Internal i (length v') $ unsafeWrite v' i x - return v' +unsafeAppend1 upTo v i x + | i < length v = do unsafeWrite v i x + return v + | otherwise = do v' <- enlargeUpTo upTo v + checkIndex Internal i (length v') $ unsafeWrite v' i x + return v' unsafePrepend1 :: (PrimMonad m, MVector v a) - => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int) + => Int -> v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int) {-# INLINE_INNER unsafePrepend1 #-} -unsafePrepend1 v i x +unsafePrepend1 upTo v i x | i /= 0 = do let i' = i-1 unsafeWrite v i' x return (v, i') | otherwise = do - (v', j) <- enlargeFront v + (v', j) <- enlargeFrontUpTo upTo v let i' = j-1 checkIndex Internal i' (length v') $ unsafeWrite v' i' x return (v', i') @@ -213,37 +212,36 @@ unstream s = munstream (Bundle.lift s) munstream :: (PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE_FUSED munstream #-} -munstream s = case upperBound (MBundle.size s) of - Just n -> munstreamMax s n - Nothing -> munstreamUnknown s - -munstreamMax :: (PrimMonad m, MVector v a) - => MBundle m u a -> Int -> m (v (PrimState m) a) -{-# INLINE munstreamMax #-} -munstreamMax s n - = do - v <- checkLength Internal n $ unsafeNew n - let put i x = do - checkIndex Internal i n $ unsafeWrite v i x - return (i+1) - n' <- MBundle.foldM' put 0 s - return $ checkSlice Internal 0 n' n - $ unsafeSlice 0 n' v +munstream s = + case MBundle.size s of + Size lb ub | lb == ub -> munstreamExact s lb + _ -> munstreamUnknown s + +munstreamExact + :: (PrimMonad m, MVector v a) + => MBundle m u a -> Int -> m (v (PrimState m) a) +{-# INLINE munstreamExact #-} +munstreamExact s n = do + v <- checkLength Internal n $ unsafeNew n + let put i x = do checkIndex Internal i n $ unsafeWrite v i x + return (i+1) + n' <- MBundle.foldM' put 0 s + return $ checkSlice Internal 0 n' n + $ unsafeSlice 0 n' v munstreamUnknown :: (PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE munstreamUnknown #-} -munstreamUnknown s - = do - v <- unsafeNew 0 - (v', n) <- MBundle.foldM put (v, 0) s - return $ checkSlice Internal 0 n (length v') - $ unsafeSlice 0 n v' +munstreamUnknown s = do + v <- unsafeNew lb + (v', n) <- MBundle.foldM put (v, 0) s + return $ checkSlice Internal 0 n (length v') + $ unsafeSlice 0 n v' where + Size lb ub = MBundle.size s {-# INLINE_INNER put #-} - put (v,i) x = do - v' <- unsafeAppend1 v i x - return (v',i+1) + put (v,i) x = do v' <- unsafeAppend1 ub v i x + return (v',i+1) -- | Create a new mutable vector and fill it with elements from the 'Bundle'. @@ -261,14 +259,15 @@ vunstream s = vmunstream (Bundle.lift s) vmunstream :: (PrimMonad m, V.Vector v a) => MBundle m v a -> m (V.Mutable v (PrimState m) a) {-# INLINE_FUSED vmunstream #-} -vmunstream s = case upperBound (MBundle.size s) of - Just n -> vmunstreamMax s n - Nothing -> vmunstreamUnknown s +vmunstream s = + case MBundle.size s of + Size ub lb | lb == ub -> vmunstreamExact s ub + _ -> vmunstreamUnknown s -vmunstreamMax :: (PrimMonad m, V.Vector v a) +vmunstreamExact :: (PrimMonad m, V.Vector v a) => MBundle m v a -> Int -> m (V.Mutable v (PrimState m) a) -{-# INLINE vmunstreamMax #-} -vmunstreamMax s n +{-# INLINE vmunstreamExact #-} +vmunstreamExact s n = do v <- checkLength Internal n $ unsafeNew n let {-# INLINE_INNER copyChunk #-} @@ -282,21 +281,23 @@ vmunstreamMax s n $ unsafeSlice 0 n' v vmunstreamUnknown :: (PrimMonad m, V.Vector v a) - => MBundle m v a -> m (V.Mutable v (PrimState m) a) + => MBundle m v a -> m (V.Mutable v (PrimState m) a) {-# INLINE vmunstreamUnknown #-} vmunstreamUnknown s = do - v <- unsafeNew 0 + v <- unsafeNew lb (v', n) <- Stream.foldlM copyChunk (v,0) (MBundle.chunks s) return $ checkSlice Internal 0 n (length v') $ unsafeSlice 0 n v' where + lb = lowerBound $ MBundle.size s + ub = upperBound $ MBundle.size s {-# INLINE_INNER copyChunk #-} copyChunk (v,i) (Chunk n f) = do let j = i+n v' <- if basicLength v < j - then unsafeGrow v (delay_inline max (enlarge_delta v) (j - basicLength v)) + then unsafeGrow v $ delay_inline enlarge_delta' v n ub else return v checkSlice Internal i n (length v') $ f (basicUnsafeSlice i n v') return (v',j) @@ -317,14 +318,15 @@ unstreamR s = munstreamR (Bundle.lift s) munstreamR :: (PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE_FUSED munstreamR #-} -munstreamR s = case upperBound (MBundle.size s) of - Just n -> munstreamRMax s n - Nothing -> munstreamRUnknown s +munstreamR s = + case MBundle.size s of + Size lb ub | lb == ub -> munstreamRExact s lb + _ -> munstreamRUnknown s -munstreamRMax :: (PrimMonad m, MVector v a) +munstreamRExact :: (PrimMonad m, MVector v a) => MBundle m u a -> Int -> m (v (PrimState m) a) -{-# INLINE munstreamRMax #-} -munstreamRMax s n +{-# INLINE munstreamRExact #-} +munstreamRExact s n = do v <- checkLength Internal n $ unsafeNew n let put i x = do @@ -341,14 +343,15 @@ munstreamRUnknown :: (HasCallStack, PrimMonad m, MVector v a) {-# INLINE munstreamRUnknown #-} munstreamRUnknown s = do - v <- unsafeNew 0 + v <- unsafeNew $ lowerBound $ MBundle.size s (v', i) <- MBundle.foldM put (v, 0) s let n = length v' return $ checkSlice Internal i (n-i) n $ unsafeSlice i (n-i) v' where + ub = upperBound $ MBundle.size s {-# INLINE_INNER put #-} - put (v,i) x = unsafePrepend1 v i x + put (v,i) x = unsafePrepend1 ub v i x -- Length -- ------ @@ -560,6 +563,27 @@ growFront v by = checkLength Bounds by enlarge_delta :: MVector v a => v s a -> Int enlarge_delta v = max (length v) 1 +enlarge_delta' :: MVector v a + => v s a + -> Int -- Minimal increment + -> Int -- Upper bound + -> Int +enlarge_delta' v incMin upperBnd + -- We're trying to grow vector which already reached upper bound. + -- + -- This could happen if vector upper limit is not correnct or if we + -- reached maxBound for stream that in fact larger than that. This + -- is actually possible without running out of memory on 32bit + -- platforms for unboxed vectors of (). + | n >= upperBnd = error "FIXME: we overflowed vector" + | otherwise = d1 + where + n = length v + -- We double vector and want to increment length by at least n + d0 = n `max` incMin + -- At the same time we want to avoid overshooting upperLimit + d1 = d0 `min` (upperBnd - n) + -- | Grow a vector logarithmically. enlarge :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a) @@ -571,6 +595,21 @@ enlarge v = stToPrim $ do where by = enlarge_delta v +-- | Grow a vector logarithmically. +enlargeUpTo :: (PrimMonad m, MVector v a) + => Int -> v (PrimState m) a -> m (v (PrimState m) a) +{-# INLINE enlargeUpTo #-} +enlargeUpTo upTo v = stToPrim $ do + -- FIXME: elaborate dealing with overflow in size + vnew <- unsafeGrow v by + basicInitialize $ basicUnsafeSlice n by vnew + return vnew + where + n = length v + delta = enlarge_delta v + by | delta + n <= upTo = delta + | otherwise = upTo - n + enlargeFront :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a, Int) {-# INLINE enlargeFront #-} @@ -581,6 +620,20 @@ enlargeFront v = stToPrim $ do where by = enlarge_delta v +enlargeFrontUpTo :: (PrimMonad m, MVector v a) + => Int -> v (PrimState m) a -> m (v (PrimState m) a, Int) +{-# INLINE enlargeFrontUpTo #-} +enlargeFrontUpTo upTo v = stToPrim $ do + v' <- unsafeGrowFront v by + basicInitialize $ basicUnsafeSlice 0 by v' + return (v', by) + where + n = length v + delta = enlarge_delta v + -- FIXME: elaborate dealing with overflow in size + by | delta + n <= upTo = delta + | otherwise = upTo - n + -- | Grow a vector by allocating a new mutable vector of the same size plus the -- the given number of elements and copying all the data over to the new vector, -- starting at its beginning. The newly allocated memory is not initialized and @@ -1075,16 +1128,16 @@ unstablePartition f !v = from_left 0 (length v) unstablePartitionBundle :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE unstablePartitionBundle #-} -unstablePartitionBundle f s - = case upperBound (Bundle.size s) of - Just n -> unstablePartitionMax f s n - Nothing -> partitionUnknown f s +unstablePartitionBundle f s = + case MBundle.size s of + Size lb ub | lb == ub -> unstablePartitionExact f s ub + _ -> partitionUnknown f s -unstablePartitionMax :: (PrimMonad m, MVector v a) +unstablePartitionExact :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> Int -> m (v (PrimState m) a, v (PrimState m) a) -{-# INLINE unstablePartitionMax #-} -unstablePartitionMax f s n +{-# INLINE unstablePartitionExact #-} +unstablePartitionExact f s n = do v <- checkLength Internal n $ unsafeNew n let {-# INLINE_INNER put #-} @@ -1102,36 +1155,8 @@ unstablePartitionMax f s n partitionBundle :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE partitionBundle #-} -partitionBundle f s - = case upperBound (Bundle.size s) of - Just n -> partitionMax f s n - Nothing -> partitionUnknown f s - -partitionMax :: (PrimMonad m, MVector v a) - => (a -> Bool) -> Bundle u a -> Int -> m (v (PrimState m) a, v (PrimState m) a) -{-# INLINE partitionMax #-} -partitionMax f s n - = do - v <- checkLength Internal n $ unsafeNew n - - let {-# INLINE_INNER put #-} - put (i,j) x - | f x = do - unsafeWrite v i x - return (i+1,j) - - | otherwise = let j' = j-1 in - do - unsafeWrite v j' x - return (i,j') - - (i,j) <- Bundle.foldM' put (0,n) s - check Internal "invalid indices" (i <= j) - $ return () - let l = unsafeSlice 0 i v - r = unsafeSlice j (n-j) v - reverse r - return (l,r) +-- FIXME: +partitionBundle f s = partitionUnknown f s partitionUnknown :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) @@ -1145,6 +1170,7 @@ partitionUnknown f s $ checkSlice Internal 0 n2 (length v2') $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2') where + ub = upperBound $ MBundle.size s -- NOTE: The case distinction has to be on the outside because -- GHC creates a join point for the unsafeWrite even when everything -- is inlined. This is bad because with the join point, v isn't getting @@ -1152,40 +1178,17 @@ partitionUnknown f s {-# INLINE_INNER put #-} put (v1, i1, v2, i2) x | f x = do - v1' <- unsafeAppend1 v1 i1 x + v1' <- unsafeAppend1 ub v1 i1 x return (v1', i1+1, v2, i2) | otherwise = do - v2' <- unsafeAppend1 v2 i2 x + v2' <- unsafeAppend1 ub v2 i2 x return (v1, i1, v2', i2+1) partitionWithBundle :: (PrimMonad m, MVector v a, MVector v b, MVector v c) => (a -> Either b c) -> Bundle u a -> m (v (PrimState m) b, v (PrimState m) c) {-# INLINE partitionWithBundle #-} -partitionWithBundle f s - = case upperBound (Bundle.size s) of - Just n -> partitionWithMax f s n - Nothing -> partitionWithUnknown f s - -partitionWithMax :: (PrimMonad m, MVector v a, MVector v b, MVector v c) - => (a -> Either b c) -> Bundle u a -> Int -> m (v (PrimState m) b, v (PrimState m) c) -{-# INLINE partitionWithMax #-} -partitionWithMax f s n - = do - v1 <- unsafeNew n - v2 <- unsafeNew n - let {-# INLINE_INNER put #-} - put (i1, i2) x = case f x of - Left b -> do - unsafeWrite v1 i1 b - return (i1+1, i2) - Right c -> do - unsafeWrite v2 i2 c - return (i1, i2+1) - (n1, n2) <- Bundle.foldM' put (0, 0) s - checkSlice Internal 0 n1 (length v1) - $ checkSlice Internal 0 n2 (length v2) - $ return (unsafeSlice 0 n1 v1, unsafeSlice 0 n2 v2) +partitionWithBundle f s = partitionWithUnknown f s partitionWithUnknown :: forall m v u a b c. (PrimMonad m, MVector v a, MVector v b, MVector v c) @@ -1200,16 +1203,17 @@ partitionWithUnknown f s $ checkSlice Internal 0 n2 (length v2') $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2') where + ub = upperBound $ MBundle.size s put :: (v (PrimState m) b, Int, v (PrimState m) c, Int) -> a -> m (v (PrimState m) b, Int, v (PrimState m) c, Int) {-# INLINE_INNER put #-} put (v1, i1, v2, i2) x = case f x of Left b -> do - v1' <- unsafeAppend1 v1 i1 b + v1' <- unsafeAppend1 ub v1 i1 b return (v1', i1+1, v2, i2) Right c -> do - v2' <- unsafeAppend1 v2 i2 c + v2' <- unsafeAppend1 ub v2 i2 c return (v1, i1, v2', i2+1) -- Modifying vectors