Skip to content

Commit

Permalink
Hint to GHC that indices are to be used strictly
Browse files Browse the repository at this point in the history
'basicUnsafe{Read/Write/IndexM}' are class members and, unless a call site was
already specialised to a specific vector instance, GHC has no clue that the index
is most certainly to be used eagerly. Bang before the index provides
this vital for optimizer information.
  • Loading branch information
Bodigrim committed Mar 11, 2024
1 parent 197eeb4 commit 5de58ff
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 13 deletions.
23 changes: 14 additions & 9 deletions vector/src/Data/Vector/Generic.hs
Original file line number Diff line number Diff line change
Expand Up @@ -232,15 +232,20 @@ 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)
(!) v !i = checkIndex Bounds i (length v) $ unBox (basicUnsafeIndexM v i)
-- Why do we need ! before i?
-- The reason is that 'basicUnsafeIndexM' is a class member 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) Safe indexing.
(!?) :: Vector v a => v a -> Int -> Maybe a
{-# INLINE_FUSED (!?) #-}
-- 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.
Expand All @@ -256,7 +261,7 @@ 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)
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
Expand Down Expand Up @@ -316,7 +321,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.
Expand All @@ -334,7 +339,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

Expand Down Expand Up @@ -993,7 +998,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
Expand All @@ -1010,7 +1015,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
-- ------------------------
Expand Down Expand Up @@ -2534,7 +2539,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')

Expand Down
12 changes: 8 additions & 4 deletions vector/src/Data/Vector/Generic/Mutable.hs
Original file line number Diff line number Diff line change
Expand Up @@ -700,21 +700,25 @@ 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)
unsafeRead v !i = checkIndex Unsafe i (length v)
$ stToPrim
$ basicUnsafeRead v i
-- Why do we need ! before i?
-- The reason is that 'basicUnsafeRead' is a class member and, unless 'unsafeRead' 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.

-- | 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)
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)
unsafeModify v f !i = checkIndex Unsafe i (length v)
$ stToPrim
$ basicUnsafeRead v i >>= \x ->
basicUnsafeWrite v i (f x)
Expand All @@ -725,7 +729,7 @@ 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)
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.
Expand Down

0 comments on commit 5de58ff

Please sign in to comment.