Introduce basicUnsafeIndexM#

vector/internal/GenUnboxTuple.hs

@@ -1,7 +1,12 @@
+#!/usr/bin/env cabal
+{- cabal:
+build-depends: base, pretty
+-}
 {-# LANGUAGE ParallelListComp #-}
 
 module Main where
 
+import Prelude hiding ((<>))
 import Text.PrettyPrint
 
 import System.Environment ( getArgs )
@@ -235,6 +240,6 @@ generate n =
                       ,("basicUnsafeThaw",        gen_unsafeThaw)
                       ,("basicLength",            gen_length "V")
                       ,("basicUnsafeSlice",       gen_unsafeSlice "G" "V")
-                      ,("basicUnsafeIndexM",      gen_basicUnsafeIndexM)
+                      ,("basicUnsafeIndexM#",     gen_basicUnsafeIndexM)
                       ,("basicUnsafeCopy",        gen_unsafeCopy "V" qG)
                       ,("elemseq",                gen_elemseq)]

vector/src/Data/Vector/Fusion/Bundle/Monadic.hs

@@ -4,6 +4,7 @@
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -98,10 +99,10 @@ import Control.Monad.Primitive
 
 import qualified Data.List as List
 import Data.Char      ( ord )
-import GHC.Base       ( unsafeChr )
+import GHC.Base       ( unsafeChr, Int(..) )
 import Control.Monad  ( liftM )
 import Prelude
-  ( Eq, Ord, Num, Enum, Functor, Monad, Bool(..), Ordering, Char, Int, Word, Integer, Float, Double, Maybe(..), Either(..), Integral, RealFrac
+  ( Eq, Ord, Num, Enum, Functor, Monad, Bool(..), Ordering, Char, Word, Integer, Float, Double, Maybe(..), Either(..), Integral, RealFrac
   , return, fmap, otherwise, id, const, seq, max, maxBound, fromIntegral, truncate
   , (+), (-), (<), (<=), (>), (>=), (==), (/=), (&&), (.), ($), (<$), (/) )
 
@@ -1079,9 +1080,10 @@ fromVector v = v `seq` n `seq` Bundle (Stream step 0)
     n = basicLength v
 
     {-# INLINE step #-}
-    step i | i >= n = return Done
-           | otherwise = case basicUnsafeIndexM v i of
-                           Box x -> return $ Yield x (i+1)
+    step (I# i)
+      | I# i >= n = return Done
+      | otherwise = case basicUnsafeIndexM# v i of
+                           Box x -> return $ Yield x (I# i + 1)
 
 
     {-# INLINE vstep #-}
@@ -1100,10 +1102,10 @@ fromVectors us = Bundle (Stream pstep (Left us))
     pstep (Left []) = return Done
     pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs)))
 
-    pstep (Right (v,i,vs))
-      | i >= basicLength v = return $ Skip (Left vs)
-      | otherwise          = case basicUnsafeIndexM v i of
-                               Box x -> return $ Yield x (Right (v,i+1,vs))
+    pstep (Right (v, I# i, vs))
+      | I# i >= basicLength v = return $ Skip (Left vs)
+      | otherwise          = case basicUnsafeIndexM# v i of
+                               Box x -> return $ Yield x (Right (v, I# i + 1, vs))
 
     -- FIXME: work around bug in GHC 7.6.1
     vstep :: HasCallStack => [v a] -> m (Step [v a] (Chunk v a))
@@ -1131,10 +1133,10 @@ concatVectors Bundle{sElems = Stream step t}
         Skip    s' -> return (Skip (Left s'))
         Done       -> return Done
 
-    pstep (Right (v,i,s))
-      | i >= basicLength v = return (Skip (Left s))
-      | otherwise          = case basicUnsafeIndexM v i of
-                               Box x -> return (Yield x (Right (v,i+1,s)))
+    pstep (Right (v, I# i, s))
+      | I# i >= basicLength v = return (Skip (Left s))
+      | otherwise          = case basicUnsafeIndexM# v i of
+                               Box x -> return (Yield x (Right (v, I# i + 1,s)))
 
 
     vstep s = do

vector/src/Data/Vector/Generic.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -218,6 +219,7 @@ import Data.Typeable ( Typeable, gcast1 )
 import Data.Data ( Data, DataType, Constr, Fixity(Prefix),
                    mkDataType, mkConstr, constrIndex, mkNoRepType )
 import qualified Data.Traversable as T (Traversable(mapM))
+import GHC.Exts (Int(..))
 
 -- Length information
 -- ------------------
@@ -253,16 +255,16 @@ infixl 9 !
 (!) :: (HasCallStack, Vector v a) => v a -> Int -> a
 {-# INLINE_FUSED (!) #-}
 -- See NOTE: [Strict indexing]
-(!) v !i = checkIndex Bounds i (length v) $ unBox (basicUnsafeIndexM v i)
+(!) v (I# i) = checkIndex Bounds (I# 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
+-- Use basicUnsafeIndexM# @Box to perform the indexing eagerly.
+v !? (I# i)
+  | I# i `inRange` length v = case basicUnsafeIndexM# v i of Box a -> Just a
   | otherwise            = Nothing
 
 
@@ -280,7 +282,7 @@ last v = v ! (length v - 1)
 unsafeIndex :: Vector v a => v a -> Int -> a
 {-# INLINE_FUSED unsafeIndex #-}
 -- See NOTE: [Strict indexing]
-unsafeIndex v !i = checkIndex Unsafe i (length v) $ unBox (basicUnsafeIndexM v i)
+unsafeIndex v (I# i) = checkIndex Unsafe (I# 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
@@ -340,7 +342,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# i) = checkIndex Bounds (I# 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.
@@ -358,9 +360,9 @@ 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# i) = checkIndex Unsafe (I# i) (length v)
                  $ liftBox
-                 $ basicUnsafeIndexM v i
+                 $ basicUnsafeIndexM# v i
 
 -- | /O(1)/ First element in a monad, without checking for empty vectors.
 -- See 'indexM' for an explanation of why this is useful.
@@ -1019,7 +1021,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# i) = checkIndex Bounds (I# 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
@@ -1036,7 +1038,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# i) = checkIndex Unsafe (I# i) n $ basicUnsafeIndexM# v i
 
 -- Safe destructive updates
 -- ------------------------
@@ -2581,9 +2583,9 @@ 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') = i-1
             in
-            case basicUnsafeIndexM v i' of Box x -> Just (x, i')
+            case basicUnsafeIndexM# v i' of Box x -> Just (x, I# i')
 
 -- | /O(n)/ Construct a vector from a 'Bundle', proceeding from right to left.
 unstreamR :: Vector v a => Bundle v a -> v a

vector/src/Data/Vector/Generic/Base.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
@@ -34,6 +35,7 @@ import qualified Data.Primitive.PrimArray as Prim
 
 import Control.Monad.ST
 import Data.Kind (Type)
+import GHC.Exts (Int(..), Int#)
 
 -- | @Mutable v s a@ is the mutable version of the immutable vector type @v a@ with
 -- the state token @s@. It is injective on GHC 8 and newer.
@@ -112,7 +114,13 @@ class MVector (Mutable v) a => Vector v a where
   --
   -- which does not have this problem, because indexing (but not the returned
   -- element!) is evaluated immediately.
-  basicUnsafeIndexM  :: v a -> Int -> Box a
+  basicUnsafeIndexM :: v a -> Int -> Box a
+  basicUnsafeIndexM xs (I# i) = basicUnsafeIndexM# xs i
+  {-# INLINE basicUnsafeIndexM #-}
+
+  basicUnsafeIndexM# :: v a -> Int# -> Box a
+  basicUnsafeIndexM# xs i = basicUnsafeIndexM xs (I# i)
+  {-# INLINE basicUnsafeIndexM# #-}
 
   -- |  /Assumed complexity: O(n)/
   --
@@ -131,11 +139,12 @@ class MVector (Mutable v) a => Vector v a where
     where
       !n = basicLength src
 
-      do_copy i | i < n = do
-                            x <- liftBox $ basicUnsafeIndexM src i
-                            M.basicUnsafeWrite dst i x
-                            do_copy (i+1)
-                | otherwise = return ()
+      do_copy (I# i)
+        | I# i < n = do
+          x <- liftBox $ basicUnsafeIndexM# src i
+          M.basicUnsafeWrite dst (I# i) x
+          do_copy (I# i + 1)
+        | otherwise = return ()
 
   -- | Evaluate @a@ as far as storing it in a vector would and yield @b@.
   -- The @v a@ argument only fixes the type and is not touched. This method is
@@ -152,4 +161,4 @@ class MVector (Mutable v) a => Vector v a where
   elemseq _ = \_ x -> x
 
   {-# MINIMAL basicUnsafeFreeze, basicUnsafeThaw, basicLength,
-              basicUnsafeSlice, basicUnsafeIndexM #-}
+              basicUnsafeSlice, (basicUnsafeIndexM | basicUnsafeIndexM#) #-}

vector/src/Data/Vector/Primitive.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE RoleAnnotations #-}
@@ -1721,12 +1722,12 @@ toList :: Prim a => Vector a -> [a]
 {-# INLINE toList #-}
 toList = G.toList
 
--- | /O(n)/ Convert a list to a vector. During the operation, the 
--- vector’s capacity will be doubling until the list's contents are 
--- in the vector. Depending on the list’s size, up to half of the vector’s 
--- capacity might be empty. If you’d rather avoid this, you can use 
--- 'fromListN', which will provide the exact space the list requires but will 
--- prevent list fusion, or @'force' . 'fromList'@, which will create the 
+-- | /O(n)/ Convert a list to a vector. During the operation, the
+-- vector’s capacity will be doubling until the list's contents are
+-- in the vector. Depending on the list’s size, up to half of the vector’s
+-- capacity might be empty. If you’d rather avoid this, you can use
+-- 'fromListN', which will provide the exact space the list requires but will
+-- prevent list fusion, or @'force' . 'fromList'@, which will create the
 -- vector and then copy it without the superfluous space.
 --
 -- @since 0.4

vector/src/Data/Vector/Primitive/Unsafe.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE RoleAnnotations #-}
@@ -103,8 +104,8 @@ instance Prim a => G.Vector Vector a where
   {-# INLINE basicUnsafeSlice #-}
   basicUnsafeSlice j n (UnsafeVector i _ arr) = UnsafeVector (i+j) n arr
 
-  {-# INLINE basicUnsafeIndexM #-}
-  basicUnsafeIndexM (UnsafeVector i _ arr) j = return $! indexByteArray arr (i+j)
+  {-# INLINE basicUnsafeIndexM# #-}
+  basicUnsafeIndexM# (UnsafeVector i _ arr) j = return $! indexByteArray arr (i + Exts.I# j)
 
   {-# INLINE basicUnsafeCopy #-}
   basicUnsafeCopy (UnsafeMVector i n dst) (UnsafeVector j _ src)

vector/src/Data/Vector/Storable.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE RoleAnnotations #-}
@@ -1760,12 +1761,12 @@ toList :: Storable a => Vector a -> [a]
 {-# INLINE toList #-}
 toList = G.toList
 
--- | /O(n)/ Convert a list to a vector. During the operation, the 
--- vector’s capacity will be doubling until the list's contents are 
--- in the vector. Depending on the list’s size, up to half of the vector’s 
--- capacity might be empty. If you’d rather avoid this, you can use 
--- 'fromListN', which will provide the exact space the list requires but will 
--- prevent list fusion, or @'force' . 'fromList'@, which will create the 
+-- | /O(n)/ Convert a list to a vector. During the operation, the
+-- vector’s capacity will be doubling until the list's contents are
+-- in the vector. Depending on the list’s size, up to half of the vector’s
+-- capacity might be empty. If you’d rather avoid this, you can use
+-- 'fromListN', which will provide the exact space the list requires but will
+-- prevent list fusion, or @'force' . 'fromList'@, which will create the
 -- vector and then copy it without the superfluous space.
 --
 -- @since 0.4

vector/src/Data/Vector/Storable/Unsafe.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE RoleAnnotations #-}
@@ -120,12 +121,12 @@ instance Storable a => G.Vector Vector a where
   {-# INLINE basicUnsafeSlice #-}
   basicUnsafeSlice i n (UnsafeVector _ fp) = UnsafeVector n (updPtr (`advancePtr` i) fp)
 
-  {-# INLINE basicUnsafeIndexM #-}
-  basicUnsafeIndexM (UnsafeVector _ fp) i
+  {-# INLINE basicUnsafeIndexM# #-}
+  basicUnsafeIndexM# (UnsafeVector _ fp) i
     = return
     . unsafeInlineIO
     $ unsafeWithForeignPtr fp $ \p ->
-      peekElemOff p i
+      peekElemOff p (Exts.I# i)
 
   {-# INLINE basicUnsafeCopy #-}
   basicUnsafeCopy (UnsafeMVector n fp) (UnsafeVector _ fq)

vector/src/Data/Vector/Strict/Unsafe.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE TypeFamilies #-}
@@ -119,8 +120,8 @@ instance G.Vector Vector a where
   basicLength = coerce (G.basicLength @V.Vector @a)
   {-# INLINE basicUnsafeSlice #-}
   basicUnsafeSlice = coerce (G.basicUnsafeSlice @V.Vector @a)
-  {-# INLINE basicUnsafeIndexM #-}
-  basicUnsafeIndexM = coerce (G.basicUnsafeIndexM @V.Vector @a)
+  {-# INLINE basicUnsafeIndexM# #-}
+  basicUnsafeIndexM# = coerce (G.basicUnsafeIndexM# @V.Vector @a)
   {-# INLINE basicUnsafeCopy #-}
   basicUnsafeCopy = coerce (G.basicUnsafeCopy @V.Vector @a)
   {-# INLINE elemseq #-}

vector/src/Data/Vector/Unboxed.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE TypeFamilies #-}
 -- |