Hint to GHC that indices are to be used strictly

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')
 

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,33 +701,37 @@ 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.
 --
 -- @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 ()