diff --git a/vector/src/Data/Vector/Generic.hs b/vector/src/Data/Vector/Generic.hs index ec847c20..37380007 100644 --- a/vector/src/Data/Vector/Generic.hs +++ b/vector/src/Data/Vector/Generic.hs @@ -228,19 +228,35 @@ null = Bundle.null . stream -- Indexing -- -------- +-- NOTE: [Strict indexing] +-- ~~~~~~~~~~~~~~~~~~~~~~~ +-- +-- Why index parameters are strict in indexing ((!), (!?)) functions +-- and functions for accessing elements in mutable arrays ('unsafeRead', +-- 'unsafeWrite', 'unsafeModify'), and slice functions? +-- +-- These function call class methods ('basicUnsafeIndexM', +-- 'basicUnsafeRead', etc) and, unless (!) was already specialised to +-- a specific v, GHC has no clue that i is most certainly to be used +-- eagerly. Bang before i hints this vital for optimizer information. + + infixl 9 ! -- | O(1) Indexing. (!) :: (HasCallStack, Vector v a) => v a -> Int -> a {-# INLINE_FUSED (!) #-} -(!) v i = checkIndex Bounds i (length v) $ unBox (basicUnsafeIndexM v i) +-- See NOTE: [Strict indexing] +(!) v !i = checkIndex Bounds i (length v) $ unBox (basicUnsafeIndexM v i) infixl 9 !? -- | O(1) Safe indexing. (!?) :: Vector v a => v a -> Int -> Maybe a {-# INLINE_FUSED (!?) #-} +-- See NOTE: [Strict indexing] -- Use basicUnsafeIndexM @Box to perform the indexing eagerly. -v !? i | i `inRange` length v = case basicUnsafeIndexM v i of Box a -> Just a - | otherwise = Nothing +v !? (!i) + | i `inRange` length v = case basicUnsafeIndexM v i of Box a -> Just a + | otherwise = Nothing -- | /O(1)/ First element. @@ -256,7 +272,8 @@ last v = v ! (length v - 1) -- | /O(1)/ Unsafe indexing without bounds checking. unsafeIndex :: Vector v a => v a -> Int -> a {-# INLINE_FUSED unsafeIndex #-} -unsafeIndex v i = checkIndex Unsafe i (length v) $ unBox (basicUnsafeIndexM v i) +-- See NOTE: [Strict indexing] +unsafeIndex v !i = checkIndex Unsafe i (length v) $ unBox (basicUnsafeIndexM v i) -- | /O(1)/ First element, without checking if the vector is empty. unsafeHead :: Vector v a => v a -> a @@ -316,7 +333,7 @@ unsafeLast v = unsafeIndex v (length v - 1) -- element) is evaluated eagerly. indexM :: (HasCallStack, Vector v a, Monad m) => v a -> Int -> m a {-# INLINE_FUSED indexM #-} -indexM v i = checkIndex Bounds i (length v) $ liftBox $ basicUnsafeIndexM v i +indexM v !i = checkIndex Bounds i (length v) $ liftBox $ basicUnsafeIndexM v i -- | /O(1)/ First element of a vector in a monad. See 'indexM' for an -- explanation of why this is useful. @@ -334,7 +351,7 @@ lastM v = indexM v (length v - 1) -- explanation of why this is useful. unsafeIndexM :: (Vector v a, Monad m) => v a -> Int -> m a {-# INLINE_FUSED unsafeIndexM #-} -unsafeIndexM v i = checkIndex Unsafe i (length v) +unsafeIndexM v !i = checkIndex Unsafe i (length v) $ liftBox $ basicUnsafeIndexM v i @@ -452,7 +469,8 @@ unsafeSlice :: Vector v a => Int -- ^ @i@ starting index -> v a -> v a {-# INLINE_FUSED unsafeSlice #-} -unsafeSlice i n v = checkSlice Unsafe i n (length v) $ basicUnsafeSlice i n v +-- See NOTE: [Strict indexing] +unsafeSlice !i !n v = checkSlice Unsafe i n (length v) $ basicUnsafeSlice i n v -- | /O(1)/ Yield all but the last element without copying. The vector may not -- be empty, but this is not checked. @@ -993,7 +1011,7 @@ backpermute v is = seq v -- NOTE: we do it this way to avoid triggering LiberateCase on n in -- polymorphic code index :: HasCallStack => Int -> Box a - index i = checkIndex Bounds i n $ basicUnsafeIndexM v i + index !i = checkIndex Bounds i n $ basicUnsafeIndexM v i -- | Same as 'backpermute', but without bounds checking. unsafeBackpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a @@ -1010,7 +1028,7 @@ unsafeBackpermute v is = seq v {-# INLINE index #-} -- NOTE: we do it this way to avoid triggering LiberateCase on n in -- polymorphic code - index i = checkIndex Unsafe i n $ basicUnsafeIndexM v i + index !i = checkIndex Unsafe i n $ basicUnsafeIndexM v i -- Safe destructive updates -- ------------------------ @@ -2534,7 +2552,7 @@ streamR v = v `seq` n `seq` (Bundle.unfoldr get n `Bundle.sized` Exact n) {-# INLINE get #-} get 0 = Nothing - get i = let i' = i-1 + get i = let !i' = i-1 in case basicUnsafeIndexM v i' of Box x -> Just (x, i') diff --git a/vector/src/Data/Vector/Generic/Mutable.hs b/vector/src/Data/Vector/Generic/Mutable.hs index 67e660b7..71a8a185 100644 --- a/vector/src/Data/Vector/Generic/Mutable.hs +++ b/vector/src/Data/Vector/Generic/Mutable.hs @@ -425,8 +425,9 @@ unsafeSlice :: MVector v a => Int -- ^ starting index -> v s a -> v s a {-# INLINE unsafeSlice #-} -unsafeSlice i n v = checkSlice Unsafe i n (length v) - $ basicUnsafeSlice i n v +-- See NOTE: [Strict indexing] in D.V.Generic +unsafeSlice !i !n v = checkSlice Unsafe i n (length v) + $ basicUnsafeSlice i n v -- | Same as 'init', but doesn't do range checks. unsafeInit :: MVector v a => v s a -> v s a @@ -700,24 +701,27 @@ exchange v i x = checkIndex Bounds i (length v) $ unsafeExchange v i x -- | Yield the element at the given position. No bounds checks are performed. unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a {-# INLINE unsafeRead #-} -unsafeRead v i = checkIndex Unsafe i (length v) - $ stToPrim - $ basicUnsafeRead v i +-- See NOTE: [Strict indexing] in D.V.Generic +unsafeRead v !i = checkIndex Unsafe i (length v) + $ stToPrim + $ basicUnsafeRead v i -- | Replace the element at the given position. No bounds checks are performed. unsafeWrite :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m () {-# INLINE unsafeWrite #-} -unsafeWrite v i x = checkIndex Unsafe i (length v) - $ stToPrim - $ basicUnsafeWrite v i x +-- See NOTE: [Strict indexing] in D.V.Generic +unsafeWrite v !i x = checkIndex Unsafe i (length v) + $ stToPrim + $ basicUnsafeWrite v i x -- | Modify the element at the given position. No bounds checks are performed. unsafeModify :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> a) -> Int -> m () {-# INLINE unsafeModify #-} -unsafeModify v f i = checkIndex Unsafe i (length v) - $ stToPrim - $ basicUnsafeRead v i >>= \x -> - basicUnsafeWrite v i (f x) +-- See NOTE: [Strict indexing] in D.V.Generic +unsafeModify v f !i = checkIndex Unsafe i (length v) + $ stToPrim + $ basicUnsafeRead v i >>= \x -> + basicUnsafeWrite v i (f x) -- | Modify the element at the given position using a monadic -- function. No bounds checks are performed. @@ -725,8 +729,9 @@ unsafeModify v f i = checkIndex Unsafe i (length v) -- @since 0.12.3.0 unsafeModifyM :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> m a) -> Int -> m () {-# INLINE unsafeModifyM #-} -unsafeModifyM v f i = checkIndex Unsafe i (length v) - $ stToPrim . basicUnsafeWrite v i =<< f =<< stToPrim (basicUnsafeRead v i) +-- See NOTE: [Strict indexing] in D.V.Generic +unsafeModifyM v f !i = checkIndex Unsafe i (length v) + $ stToPrim . basicUnsafeWrite v i =<< f =<< stToPrim (basicUnsafeRead v i) -- | Swap the elements at the given positions. No bounds checks are performed. unsafeSwap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m ()