1 {-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, RecordWildCards,
    2     UnboxedTuples #-}
    3 {-# OPTIONS_GHC -fno-warn-unused-matches #-}
    4 -- |
    5 -- Module      : Data.Text.Array
    6 -- Copyright   : (c) 2009, 2010 Bryan O'Sullivan
    7 --
    8 -- License     : BSD-style
    9 -- Maintainer  : bos@serpentine.com, rtomharper@googlemail.com,
   10 --               duncan@haskell.org
   11 -- Stability   : experimental
   12 -- Portability : portable
   13 --
   14 -- Packed, unboxed, heap-resident arrays.  Suitable for performance
   15 -- critical use, both in terms of large data quantities and high
   16 -- speed.
   17 --
   18 -- This module is intended to be imported @qualified@, to avoid name
   19 -- clashes with "Prelude" functions, e.g.
   20 --
   21 -- > import qualified Data.Text.Array as A
   22 --
   23 -- The names in this module resemble those in the 'Data.Array' family
   24 -- of modules, but are shorter due to the assumption of qualifid
   25 -- naming.
   26 module Data.Text.Array
   27     (
   28     -- * Types
   29       Array
   30     , MArray
   31 
   32     -- * Functions
   33     , copyM
   34     , copyI
   35     , empty
   36 #if defined(ASSERTS)
   37     , length
   38 #endif
   39     , run
   40     , run2
   41     , toList
   42     , unsafeFreeze
   43     , unsafeIndex
   44     , unsafeNew
   45     , unsafeWrite
   46     ) where
   47 
   48 #if defined(ASSERTS)
   49 -- This fugly hack is brought by GHC's apparent reluctance to deal
   50 -- with MagicHash and UnboxedTuples when inferring types. Eek!
   51 # define CHECK_BOUNDS(_func_,_len_,_k_) \
   52 if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.Text.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else
   53 #else
   54 # define CHECK_BOUNDS(_func_,_len_,_k_)
   55 #endif
   56 
   57 #include "MachDeps.h"
   58 
   59 #if defined(ASSERTS)
   60 import Control.Exception (assert)
   61 #endif
   62 import Data.Bits ((.&.))
   63 import Data.Text.UnsafeShift (shiftL, shiftR)
   64 import GHC.Base (ByteArray#, MutableByteArray#, Int(..),
   65                  indexWord16Array#, indexWordArray#, newByteArray#,
   66                  readWord16Array#, readWordArray#, unsafeCoerce#,
   67                  writeWord16Array#, writeWordArray#)
   68 import GHC.ST (ST(..), runST)
   69 import GHC.Word (Word16(..), Word(..))
   70 import Prelude hiding (length, read)
   71 
   72 -- | Immutable array type.
   73 data Array = Array {
   74       aBA :: ByteArray#
   75 #if defined(ASSERTS)
   76     , aLen :: {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
   77 #endif
   78     }
   79 
   80 -- | Mutable array type, for use in the ST monad.
   81 data MArray s = MArray {
   82       maBA :: MutableByteArray# s
   83 #if defined(ASSERTS)
   84     , maLen :: {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
   85 #endif
   86     }
   87 
   88 #if defined(ASSERTS)
   89 -- | Operations supported by all arrays.
   90 class IArray a where
   91     -- | Return the length of an array.
   92     length :: a -> Int
   93 
   94 instance IArray Array where
   95     length = aLen
   96     {-# INLINE length #-}
   97 
   98 instance IArray (MArray s) where
   99     length = maLen
  100     {-# INLINE length #-}
  101 #endif
  102 
  103 -- | Create an uninitialized mutable array.
  104 unsafeNew :: forall s. Int -> ST s (MArray s)
  105 -- entered 467,933 timesunsafeNew n =
  106 #if defined(ASSERTS)
  107     assert (n >= 0) .
  108 #endif
  109     ST $ \s1# ->
  110     case bytesInArray n of
  111       len@(I# len#) ->
  112 #if defined(ASSERTS)
  113          if len < 0 then error (show ("unsafeNew",len)) else
  114 #endif
  115          case newByteArray# len# s1# of
  116            (# s2#, marr# #) -> (# s2#, MArray marr#
  117 #if defined(ASSERTS)
  118                                   n
  119 #endif
  120                                   #)
  121 {-# INLINE unsafeNew #-}
  122 
  123 -- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
  124 unsafeFreeze :: MArray s -> ST s Array
  125 -- entered 226,726 timesunsafeFreeze MArray{..} = ST $ \s# ->
  126                           (# s#, Array (unsafeCoerce# maBA)
  127 #if defined(ASSERTS)
  128                              maLen
  129 #endif
  130                              #)
  131 {-# INLINE unsafeFreeze #-}
  132 
  133 -- | Indicate how many bytes would be used for an array of the given
  134 -- size.
  135 bytesInArray :: Int -> Int
  136 -- entered 467,933 timesbytesInArray n = n `shiftL` 1
  137 {-# INLINE bytesInArray #-}
  138 
  139 -- | Unchecked read of an immutable array.  May return garbage or
  140 -- crash on an out-of-bounds access.
  141 unsafeIndex :: Array -> Int -> Word16
  142 -- entered 9,801,934 timesunsafeIndex Array{..} i@(I# i#) =
  143   CHECK_BOUNDS("unsafeIndex",aLen,i)
  144     case indexWord16Array# aBA i# of r# -> (W16# r#)
  145 {-# INLINE unsafeIndex #-}
  146 
  147 -- | Unchecked read of an immutable array.  May return garbage or
  148 -- crash on an out-of-bounds access.
  149 unsafeIndexWord :: Array -> Int -> Word
  150 -- entered 19,281 timesunsafeIndexWord Array{..} i@(I# i#) =
  151   CHECK_BOUNDS("unsafeIndexWord",aLen`div`wordFactor,i)
  152     case indexWordArray# aBA i# of r# -> (W# r#)
  153 {-# INLINE unsafeIndexWord #-}
  154 
  155 -- | Unchecked read of a mutable array.  May return garbage or
  156 -- crash on an out-of-bounds access.
  157 unsafeRead :: MArray s -> Int -> ST s Word16
  158 -- entered 89,481 timesunsafeRead MArray{..} i@(I# i#) = ST $ \s# ->
  159   CHECK_BOUNDS("unsafeRead",maLen,i)
  160   case readWord16Array# maBA i# s# of
  161     (# s2#, r# #) -> (# s2#, W16# r# #)
  162 {-# INLINE unsafeRead #-}
  163 
  164 -- | Unchecked write of a mutable array.  May return garbage or crash
  165 -- on an out-of-bounds access.
  166 unsafeWrite :: MArray s -> Int -> Word16 -> ST s ()
  167 -- entered 4,142,519 timesunsafeWrite MArray{..} i@(I# i#) (W16# e#) = ST $ \s1# ->
  168   CHECK_BOUNDS("unsafeWrite",maLen,i)
  169   case writeWord16Array# maBA i# e# s1# of
  170     s2# -> (# s2#, () #)
  171 {-# INLINE unsafeWrite #-}
  172 
  173 -- | Unchecked read of a mutable array.  May return garbage or
  174 -- crash on an out-of-bounds access.
  175 unsafeReadWord :: MArray s -> Int -> ST s Word
  176 -- entered 323,420 timesunsafeReadWord MArray{..} i@(I# i#) = ST $ \s# ->
  177   CHECK_BOUNDS("unsafeRead64",maLen`div`wordFactor,i)
  178   case readWordArray# maBA i# s# of
  179     (# s2#, r# #) -> (# s2#, W# r# #)
  180 {-# INLINE unsafeReadWord #-}
  181 
  182 -- | Unchecked write of a mutable array.  May return garbage or crash
  183 -- on an out-of-bounds access.
  184 unsafeWriteWord :: MArray s -> Int -> Word -> ST s ()
  185 -- entered 342,323 timesunsafeWriteWord MArray{..} i@(I# i#) (W# e#) = ST $ \s1# ->
  186   CHECK_BOUNDS("unsafeWriteWord",maLen`div`wordFactor,i)
  187   case writeWordArray# maBA i# e# s1# of
  188     s2# -> (# s2#, () #)
  189 {-# INLINE unsafeWriteWord #-}
  190 
  191 -- | Convert an immutable array to a list.
  192 toList :: Array -> Int -> Int -> [Word16]
  193 -- never enteredtoList ary off len = loop 0
  194     where loop i | i < len   = unsafeIndex ary (off+i) : loop (i+1)
  195                  | otherwise = []
  196 
  197 -- | An empty immutable array.
  198 empty :: Array
  199 -- entered onceempty = runST (unsafeNew 0 >>= unsafeFreeze)
  200 
  201 -- | Run an action in the ST monad and return an immutable array of
  202 -- its result.
  203 run :: (forall s. ST s (MArray s)) -> Array
  204 -- entered 4261 timesrun k = runST (k >>= unsafeFreeze)
  205 
  206 -- | Run an action in the ST monad and return an immutable array of
  207 -- its result paired with whatever else the action returns.
  208 run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
  209 -- entered 221,540 timesrun2 k = runST (do
  210                  (marr,b) <- k
  211                  arr <- unsafeFreeze marr
  212                  return (arr,b))
  213 
  214 -- | The amount to divide or multiply by to switch between units of
  215 -- 'Word16' and units of 'Word'.
  216 wordFactor :: Int
  217 -- entered 344,681 timeswordFactor = SIZEOF_HSWORD `shiftR` 1
  218 
  219 -- | Indicate whether an offset is word-aligned.
  220 wordAligned :: Int -> Bool
  221 -- entered 58,910 timeswordAligned i = i .&. (wordFactor - 1) == 0
  222 
  223 -- | Copy some elements of a mutable array.
  224 copyM :: MArray s               -- ^ Destination
  225       -> Int                    -- ^ Destination offset
  226       -> MArray s               -- ^ Source
  227       -> Int                    -- ^ Source offset
  228       -> Int                    -- ^ Count
  229       -> ST s ()
  230 -- entered 323,193 timescopyM dest didx src sidx count =
  231 #if defined(ASSERTS)
  232     assert (sidx + count <= length src) .
  233     assert (didx + count <= length dest) $
  234 #endif
  235     if srem == 0 && drem == 0
  236     then fast_loop 0
  237     else slow_loop 0
  238     where
  239       (swidx,srem) = sidx `divMod` wordFactor
  240       (dwidx,drem) = didx `divMod` wordFactor
  241       nwds         = count `div` wordFactor
  242       fast_loop !i
  243           | i >= nwds = slow_loop (i * wordFactor)
  244           | otherwise = do w <- unsafeReadWord src (swidx+i)
  245                            unsafeWriteWord dest (dwidx+i) w
  246                            fast_loop (i+1)
  247       slow_loop !i
  248           | i >= count= return ()
  249           | otherwise = do unsafeRead src (sidx+i) >>= unsafeWrite dest (didx+i)
  250                            slow_loop (i+1)
  251 
  252 -- | Copy some elements of an immutable array.
  253 copyI :: MArray s               -- ^ Destination
  254       -> Int                    -- ^ Destination offset
  255       -> Array                  -- ^ Source
  256       -> Int                    -- ^ Source offset
  257       -> Int                    -- ^ First offset in source /not/ to
  258                                 -- copy (i.e. /not/ length)
  259       -> ST s ()
  260 -- entered 58,910 timescopyI dest i0 src j0 top
  261     | wordAligned i0 && wordAligned j0 = fast (i0 `div` wordFactor) (j0 `div` wordFactor)
  262     | otherwise = slow i0 j0
  263   where
  264     topwds = top `div` wordFactor
  265     fast !i !j
  266         | i >= topwds = slow (i * wordFactor) (j * wordFactor)
  267         | otherwise   = do unsafeWriteWord dest i (src `unsafeIndexWord` j)
  268                            fast (i+1) (j+1)
  269     slow !i !j
  270         | i >= top  = return ()
  271         | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
  272                          slow (i+1) (j+1)