@@ -57,9 +57,10 @@ if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.Text.Array." ++ (_func_) ++ "
5757#if defined(ASSERTS)
5858import Control.Exception (assert)
5959#endif
60+import Data.Bits ((.&.))
6061import Data.Text.UnsafeShift (shiftL, shiftR)
6162import GHC.Base (ByteArray#, MutableByteArray#, Int(..),
62− indexWord16Array#, newByteArray#,
63+ indexWord16Array#, indexWordArray#, newByteArray#,
6364 readWord16Array#, readWordArray#, unsafeCoerce#,
6465 writeWord16Array#, writeWordArray#)
6566import GHC.ST (ST(..), runST)
@@ -125,6 +126,14 @@ unsafeIndex (Array len ba#) i@(I# i#) =
125126 case indexWord16Array# ba# i# of r# -> (W16# r#)
126127{-# INLINE unsafeIndex #-}
127128
129+-- | Unchecked read of an immutable array. May return garbage or
130+-- crash on an out-of-bounds access.
131+unsafeIndexWord :: Array -> Int -> Word
132+unsafeIndexWord (Array len ba#) i@(I# i#) =
133+ CHECK_BOUNDS("unsafeIndexWord",len,i)
134+ case indexWordArray# ba# i# of r# -> (W# r#)
135+{-# INLINE unsafeIndexWord #-}
136+
128137-- | Unchecked read of a mutable array. May return garbage or
129138-- crash on an out-of-bounds access.
130139unsafeRead :: MArray s -> Int -> ST s Word16
@@ -185,6 +194,15 @@ run2 k = runST (do
185194 arr <- unsafeFreeze marr
186195 return (arr,b))
187196
197+-- | The amount to divide or multiply by to switch between units of
198+-- 'Word16' and units of 'Word'.
199+wordFactor :: Int
200+wordFactor = SIZEOF_HSWORD `shiftR` 1
201+
202+-- | Indicate whether an offset is word-aligned.
203+wordAligned :: Int -> Bool
204+wordAligned i = i .&. (wordFactor - 1) == 0
205+
188206-- | Copy an array in its entirety. The destination array must be at
189207-- least as big as the source.
190208copy :: MArray s -- ^ source array
@@ -194,9 +212,9 @@ copy dest@(MArray dlen _) src@(MArray slen _)
194212 | dlen >= slen = fast_loop 0
195213 | otherwise = fail "Data.Text.Array.copy: array too small"
196214 where
197− nwds = slen `div` (SIZEOF_HSWORD `shiftR` 1)
215+ nwds = slen `div` wordFactor
198216 fast_loop !i
199− | i >= nwds = copy_loop (i * (SIZEOF_HSWORD `shiftR` 1))
217+ | i >= nwds = copy_loop (i * wordFactor)
200218 | otherwise = do unsafeReadWord src i >>= unsafeWriteWord dest i
201219 fast_loop (i+1)
202220 copy_loop !i
@@ -205,7 +223,12 @@ copy dest@(MArray dlen _) src@(MArray slen _)
205223 copy_loop (i+1)
206224
207225-- | Copy some elements of a mutable array.
208−partialCopyM :: MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
226+partialCopyM :: MArray s -- ^ Destination
227+ -> Int -- ^ Destination offset
228+ -> MArray s -- ^ Source
229+ -> Int -- ^ Source offset
230+ -> Int -- ^ Count
231+ -> ST s ()
209232partialCopyM dest didx src sidx count =
210233#if defined(ASSERTS)
211234 assert (sidx + count <= length src) .
@@ -219,9 +242,23 @@ partialCopyM dest didx src sidx count =
219242 copy_loop (i+1) (j+1) (c+1)
220243
221244-- | Copy some elements of an immutable array.
222−partialCopyI :: MArray s -> Int -> Array -> Int -> Int -> ST s ()
223−partialCopyI dest i0 src j0 top = go i0 j0
245+partialCopyI :: MArray s -- ^ Destination
246+ -> Int -- ^ Destination offset
247+ -> Array -- ^ Source
248+ -> Int -- ^ Source offset
249+ -> Int -- ^ First offset in source /not/ to
250+ -- copy (i.e. /not/ length)
251+ -> ST s ()
252+partialCopyI dest i0 src j0 top
253+ | wordAligned i0 && wordAligned j0 = fast (i0 `div` wordFactor) (j0 `div` wordFactor)
254+ | otherwise = slow i0 j0
224255 where
225− go !i !j | i >= top = return ()
226− | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
227− go (i+1) (j+1)
256+ topwds = top `div` wordFactor
257+ fast !i !j
258+ | i >= topwds = slow (i * wordFactor) (j * wordFactor)
259+ | otherwise = do unsafeWriteWord dest i (src `unsafeIndexWord` j)
260+ fast (i+1) (j+1)
261+ slow !i !j
262+ | i >= top = return ()
263+ | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
264+ slow (i+1) (j+1)