diff --git a/vector/src/Data/Vector/Generic.hs b/vector/src/Data/Vector/Generic.hs index 321e173b..4cd67f71 100644 --- a/vector/src/Data/Vector/Generic.hs +++ b/vector/src/Data/Vector/Generic.hs @@ -245,8 +245,10 @@ infixl 9 !? -- | O(1) Safe indexing. (!?) :: Vector v a => v a -> Int -> Maybe a {-# INLINE_FUSED (!?) #-} -v !? i | i < 0 || i >= length v = Nothing - | otherwise = Just $ unsafeIndex v i +-- 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 + -- | /O(1)/ First element. head :: Vector v a => v a -> a diff --git a/vector/src/Data/Vector/Internal/Check.hs b/vector/src/Data/Vector/Internal/Check.hs index 8950619b..4a9f020b 100644 --- a/vector/src/Data/Vector/Internal/Check.hs +++ b/vector/src/Data/Vector/Internal/Check.hs @@ -18,7 +18,8 @@ module Data.Vector.Internal.Check ( Checks(..), doChecks, internalError, - check, checkIndex, checkLength, checkSlice + check, checkIndex, checkLength, checkSlice, + inRange ) where import GHC.Base( Int(..) ) @@ -112,7 +113,7 @@ checkIndex_msg# i# n# = "index out of bounds " ++ show (I# i#, I# n#) checkIndex :: HasCallStack => Checks -> Int -> Int -> a -> a {-# INLINE checkIndex #-} checkIndex kind i n x - = check kind (checkIndex_msg i n) (i >= 0 && i String @@ -141,3 +142,8 @@ checkSlice :: HasCallStack => Checks -> Int -> Int -> Int -> a -> a checkSlice kind i m n x = check kind (checkSlice_msg i m n) (i >= 0 && m >= 0 && m <= n - i) x +-- Lengths are never negative, so we can check @0 <= i < length v@ +-- using one unsigned comparison. +inRange :: Int -> Int -> Bool +{-# INLINE inRange #-} +inRange i n = (fromIntegral i :: Word) < (fromIntegral n :: Word)