{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedFFITypes #-}

{-# OPTIONS_GHC -fno-warn-unused-matches #-}
-- |
-- Module      : Data.Text.Array
-- Copyright   : (c) 2009, 2010, 2011 Bryan O'Sullivan
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Portability : portable
--
-- Packed, unboxed, heap-resident arrays.  Suitable for performance
-- critical use, both in terms of large data quantities and high
-- speed.
--
-- This module is intended to be imported @qualified@, to avoid name
-- clashes with "Prelude" functions, e.g.
--
-- > import qualified Data.Text.Array as A
--
-- The names in this module resemble those in the 'Data.Array' family
-- of modules, but are shorter due to the assumption of qualified
-- naming.
module Data.Text.Array
    (
    -- * Types
      Array
    , pattern ByteArray
    , MArray
    , pattern MutableByteArray
    -- * Functions
    , resizeM
    , shrinkM
    , copyM
    , copyI
    , copyFromPointer
    , copyToPointer
    , empty
    , equal
    , compare
    , run
    , run2
    , toList
    , unsafeFreeze
    , unsafeIndex
    , new
    , newPinned
    , newFilled
    , unsafeWrite
    , tile
    , getSizeofMArray
    ) where

#if defined(ASSERTS)
import GHC.Stack (HasCallStack)
#endif
#if !MIN_VERSION_base(4,11,0)
import Foreign.C.Types (CInt(..))
#endif
import GHC.Exts hiding (toList)
import GHC.ST (ST(..), runST)
import GHC.Word (Word8(..))
import qualified Prelude
import Prelude hiding (length, read, compare)
import Data.Array.Byte (ByteArray(..), MutableByteArray(..))

-- | Immutable array type.
type Array = ByteArray

-- | Mutable array type, for use in the ST monad.
type MArray = MutableByteArray

-- | Create an uninitialized mutable array.
new :: forall s. Int -> ST s (MArray s)
new :: forall s. Int -> ST s (MArray s)
new (I# Int#
len#)
#if defined(ASSERTS)
  | I# len# < 0 = error "Data.Text.Array.new: size overflow"
#endif
  | Bool
otherwise = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
len# State# s
s1# of
      (# State# s
s2#, MutableByteArray# s
marr# #) -> (# State# s
s2#, forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# s
marr# #)
{-# INLINE new #-}

-- | Create an uninitialized mutable pinned array.
--
-- @since 2.0
newPinned :: forall s. Int -> ST s (MArray s)
newPinned :: forall s. Int -> ST s (MArray s)
newPinned (I# Int#
len#)
#if defined(ASSERTS)
  | I# len# < 0 = error "Data.Text.Array.newPinned: size overflow"
#endif
  | Bool
otherwise = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newPinnedByteArray# Int#
len# State# s
s1# of
      (# State# s
s2#, MutableByteArray# s
marr# #) -> (# State# s
s2#, forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# s
marr# #)
{-# INLINE newPinned #-}

-- | @since 2.0
newFilled :: Int -> Int -> ST s (MArray s)
newFilled :: forall s. Int -> Int -> ST s (MArray s)
newFilled (I# Int#
len#) (I# Int#
c#) = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
  case forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# Int#
len# State# s
s1# of
    (# State# s
s2#, MutableByteArray# s
marr# #) -> case forall d.
MutableByteArray# d -> Int# -> Int# -> Int# -> State# d -> State# d
setByteArray# MutableByteArray# s
marr# Int#
0# Int#
len# Int#
c# State# s
s2# of
      State# s
s3# -> (# State# s
s3#, forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# s
marr# #)
{-# INLINE newFilled #-}

-- | @since 2.0
tile :: MArray s -> Int -> ST s ()
tile :: forall s. MArray s -> Int -> ST s ()
tile MArray s
marr Int
tileLen = do
  Int
totalLen <- forall s. MArray s -> ST s Int
getSizeofMArray MArray s
marr
  let go :: Int -> ST s ()
go Int
l
        | Int
2 forall a. Num a => a -> a -> a
* Int
l forall a. Ord a => a -> a -> Bool
> Int
totalLen = forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
copyM MArray s
marr Int
l MArray s
marr Int
0 (Int
totalLen forall a. Num a => a -> a -> a
- Int
l)
        | Bool
otherwise = forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
copyM MArray s
marr Int
l MArray s
marr Int
0 Int
l forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> ST s ()
go (Int
2 forall a. Num a => a -> a -> a
* Int
l)
  Int -> ST s ()
go Int
tileLen
{-# INLINE tile #-}

-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
unsafeFreeze :: MArray s -> ST s Array
unsafeFreeze :: forall s. MArray s -> ST s Array
unsafeFreeze (MutableByteArray MutableByteArray# s
marr) = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d.
MutableByteArray# d -> State# d -> (# State# d, ByteArray# #)
unsafeFreezeByteArray# MutableByteArray# s
marr State# s
s1# of
        (# State# s
s2#, ByteArray#
ba# #) -> (# State# s
s2#, ByteArray# -> Array
ByteArray ByteArray#
ba# #)
{-# INLINE unsafeFreeze #-}

-- | Unchecked read of an immutable array.  May return garbage or
-- crash on an out-of-bounds access.
unsafeIndex ::
#if defined(ASSERTS)
  HasCallStack =>
#endif
  Array -> Int -> Word8
unsafeIndex :: Array -> Int -> Word8
unsafeIndex (ByteArray ByteArray#
arr) i :: Int
i@(I# Int#
i#) =
#if defined(ASSERTS)
  let word8len = I# (sizeofByteArray# arr) in
  if i < 0 || i >= word8len
  then error ("Data.Text.Array.unsafeIndex: bounds error, offset " ++ show i ++ ", length " ++ show word8len)
  else
#endif
  case ByteArray# -> Int# -> Word8#
indexWord8Array# ByteArray#
arr Int#
i# of Word8#
r# -> (Word8# -> Word8
W8# Word8#
r#)
{-# INLINE unsafeIndex #-}

-- | @since 2.0
getSizeofMArray :: MArray s -> ST s Int
getSizeofMArray :: forall s. MArray s -> ST s Int
getSizeofMArray (MutableByteArray MutableByteArray# s
marr) = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s0# ->
  -- Cannot simply use (deprecated) 'sizeofMutableByteArray#', because it is
  -- unsafe in the presence of 'shrinkMutableByteArray#' and 'resizeMutableByteArray#'.
  case forall d. MutableByteArray# d -> State# d -> (# State# d, Int# #)
getSizeofMutableByteArray# MutableByteArray# s
marr State# s
s0# of
    (# State# s
s1#, Int#
word8len# #) -> (# State# s
s1#, Int# -> Int
I# Int#
word8len# #)

#if defined(ASSERTS)
checkBoundsM :: HasCallStack => MArray s -> Int -> Int -> ST s ()
checkBoundsM ma i elSize = do
  len <- getSizeofMArray ma
  if i < 0 || i + elSize > len
    then error ("bounds error, offset " ++ show i ++ ", length " ++ show len)
    else return ()
#endif

-- | Unchecked write of a mutable array.  May return garbage or crash
-- on an out-of-bounds access.
unsafeWrite ::
#if defined(ASSERTS)
  HasCallStack =>
#endif
  MArray s -> Int -> Word8 -> ST s ()
unsafeWrite :: forall s. MArray s -> Int -> Word8 -> ST s ()
unsafeWrite ma :: MArray s
ma@(MutableByteArray MutableByteArray# s
marr) i :: Int
i@(I# Int#
i#) (W8# Word8#
e#) =
#if defined(ASSERTS)
  checkBoundsM ma i 1 >>
#endif
  (forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# -> case forall d.
MutableByteArray# d -> Int# -> Word8# -> State# d -> State# d
writeWord8Array# MutableByteArray# s
marr Int#
i# Word8#
e# State# s
s1# of
    State# s
s2# -> (# State# s
s2#, () #))
{-# INLINE unsafeWrite #-}

-- | Convert an immutable array to a list.
toList :: Array -> Int -> Int -> [Word8]
toList :: Array -> Int -> Int -> [Word8]
toList Array
ary Int
off Int
len = Int -> [Word8]
loop Int
0
    where loop :: Int -> [Word8]
loop Int
i | Int
i forall a. Ord a => a -> a -> Bool
< Int
len   = Array -> Int -> Word8
unsafeIndex Array
ary (Int
offforall a. Num a => a -> a -> a
+Int
i) forall a. a -> [a] -> [a]
: Int -> [Word8]
loop (Int
iforall a. Num a => a -> a -> a
+Int
1)
                 | Bool
otherwise = []

-- | An empty immutable array.
empty :: Array
empty :: Array
empty = forall a. (forall s. ST s a) -> a
runST (forall s. Int -> ST s (MArray s)
new Int
0 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s. MArray s -> ST s Array
unsafeFreeze)

-- | Run an action in the ST monad and return an immutable array of
-- its result.
run :: (forall s. ST s (MArray s)) -> Array
run :: (forall s. ST s (MArray s)) -> Array
run forall s. ST s (MArray s)
k = forall a. (forall s. ST s a) -> a
runST (forall s. ST s (MArray s)
k forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s. MArray s -> ST s Array
unsafeFreeze)

-- | Run an action in the ST monad and return an immutable array of
-- its result paired with whatever else the action returns.
run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
run2 :: forall a. (forall s. ST s (MArray s, a)) -> (Array, a)
run2 forall s. ST s (MArray s, a)
k = forall a. (forall s. ST s a) -> a
runST (do
                 (MArray s
marr,a
b) <- forall s. ST s (MArray s, a)
k
                 Array
arr <- forall s. MArray s -> ST s Array
unsafeFreeze MArray s
marr
                 forall (m :: * -> *) a. Monad m => a -> m a
return (Array
arr,a
b))
{-# INLINE run2 #-}

-- | @since 2.0
resizeM :: MArray s -> Int -> ST s (MArray s)
resizeM :: forall s. MArray s -> Int -> ST s (MArray s)
resizeM (MutableByteArray MutableByteArray# s
ma) i :: Int
i@(I# Int#
i#) = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
  case forall d.
MutableByteArray# d
-> Int# -> State# d -> (# State# d, MutableByteArray# d #)
resizeMutableByteArray# MutableByteArray# s
ma Int#
i# State# s
s1# of
    (# State# s
s2#, MutableByteArray# s
newArr #) -> (# State# s
s2#, forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# s
newArr #)
{-# INLINE resizeM #-}

-- | @since 2.0
shrinkM ::
#if defined(ASSERTS)
  HasCallStack =>
#endif
  MArray s -> Int -> ST s ()
shrinkM :: forall s. MArray s -> Int -> ST s ()
shrinkM (MutableByteArray MutableByteArray# s
marr) i :: Int
i@(I# Int#
newSize) = do
#if defined(ASSERTS)
  oldSize <- getSizeofMArray (MutableByteArray marr)
  if I# newSize > oldSize
    then error $ "shrinkM: shrink cannot grow " ++ show oldSize ++ " to " ++ show (I# newSize)
    else return ()
#endif
  forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d. MutableByteArray# d -> Int# -> State# d -> State# d
shrinkMutableByteArray# MutableByteArray# s
marr Int#
newSize State# s
s1# of
      State# s
s2# -> (# State# s
s2#, () #)
{-# INLINE shrinkM #-}

-- | Copy some elements of a mutable array.
copyM :: MArray s               -- ^ Destination
      -> Int                    -- ^ Destination offset
      -> MArray s               -- ^ Source
      -> Int                    -- ^ Source offset
      -> Int                    -- ^ Count
      -> ST s ()
copyM :: forall s. MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
copyM dst :: MArray s
dst@(MutableByteArray MutableByteArray# s
dst#) dstOff :: Int
dstOff@(I# Int#
dstOff#) src :: MArray s
src@(MutableByteArray MutableByteArray# s
src#) srcOff :: Int
srcOff@(I# Int#
srcOff#) count :: Int
count@(I# Int#
count#)
#if defined(ASSERTS)
  | count < 0 = error $
    "copyM: count must be >= 0, but got " ++ show count
#endif
    | Bool
otherwise = do
#if defined(ASSERTS)
    srcLen <- getSizeofMArray src
    dstLen <- getSizeofMArray dst
    if srcOff + count > srcLen
      then error "copyM: source is too short"
      else return ()
    if dstOff + count > dstLen
      then error "copyM: destination is too short"
      else return ()
#endif
    forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# -> case forall d.
MutableByteArray# d
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableByteArray# MutableByteArray# s
src# Int#
srcOff# MutableByteArray# s
dst# Int#
dstOff# Int#
count# State# s
s1# of
      State# s
s2# -> (# State# s
s2#, () #)
{-# INLINE copyM #-}

-- | Copy some elements of an immutable array.
copyI :: Int                    -- ^ Count
      -> MArray s               -- ^ Destination
      -> Int                    -- ^ Destination offset
      -> Array                  -- ^ Source
      -> Int                    -- ^ Source offset
      -> ST s ()
copyI :: forall s. Int -> MArray s -> Int -> Array -> Int -> ST s ()
copyI count :: Int
count@(I# Int#
count#) (MutableByteArray MutableByteArray# s
dst#) dstOff :: Int
dstOff@(I# Int#
dstOff#) (ByteArray ByteArray#
src#) (I# Int#
srcOff#)
#if defined(ASSERTS)
  | count < 0 = error $
    "copyI: count must be >= 0, but got " ++ show count
#endif
  | Bool
otherwise = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d.
ByteArray#
-> Int#
-> MutableByteArray# d
-> Int#
-> Int#
-> State# d
-> State# d
copyByteArray# ByteArray#
src# Int#
srcOff# MutableByteArray# s
dst# Int#
dstOff# Int#
count# State# s
s1# of
      State# s
s2# -> (# State# s
s2#, () #)
{-# INLINE copyI #-}

-- | Copy from pointer.
--
-- @since 2.0
copyFromPointer
  :: MArray s               -- ^ Destination
  -> Int                    -- ^ Destination offset
  -> Ptr Word8              -- ^ Source
  -> Int                    -- ^ Count
  -> ST s ()
copyFromPointer :: forall s. MArray s -> Int -> Ptr Word8 -> Int -> ST s ()
copyFromPointer (MutableByteArray MutableByteArray# s
dst#) dstOff :: Int
dstOff@(I# Int#
dstOff#) (Ptr Addr#
src#) count :: Int
count@(I# Int#
count#)
#if defined(ASSERTS)
  | count < 0 = error $
    "copyFromPointer: count must be >= 0, but got " ++ show count
#endif
  | Bool
otherwise = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d.
Addr#
-> MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
copyAddrToByteArray# Addr#
src# MutableByteArray# s
dst# Int#
dstOff# Int#
count# State# s
s1# of
      State# s
s2# -> (# State# s
s2#, () #)
{-# INLINE copyFromPointer #-}

-- | Copy to pointer.
--
-- @since 2.0
copyToPointer
  :: Array                  -- ^ Source
  -> Int                    -- ^ Source offset
  -> Ptr Word8              -- ^ Destination
  -> Int                    -- ^ Count
  -> ST s ()
copyToPointer :: forall s. Array -> Int -> Ptr Word8 -> Int -> ST s ()
copyToPointer (ByteArray ByteArray#
src#) srcOff :: Int
srcOff@(I# Int#
srcOff#) (Ptr Addr#
dst#) count :: Int
count@(I# Int#
count#)
#if defined(ASSERTS)
  | count < 0 = error $
    "copyToPointer: count must be >= 0, but got " ++ show count
#endif
  | Bool
otherwise = forall s a. STRep s a -> ST s a
ST forall a b. (a -> b) -> a -> b
$ \State# s
s1# ->
    case forall d.
ByteArray# -> Int# -> Addr# -> Int# -> State# d -> State# d
copyByteArrayToAddr# ByteArray#
src# Int#
srcOff# Addr#
dst# Int#
count# State# s
s1# of
      State# s
s2# -> (# State# s
s2#, () #)
{-# INLINE copyToPointer #-}

-- | Compare portions of two arrays for equality.  No bounds checking
-- is performed.
equal :: Array -> Int -> Array -> Int -> Int -> Bool
equal :: Array -> Int -> Array -> Int -> Int -> Bool
equal Array
src1 Int
off1 Array
src2 Int
off2 Int
count = Array -> Int -> Array -> Int -> Int -> Int
compareInternal Array
src1 Int
off1 Array
src2 Int
off2 Int
count forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE equal #-}

-- | Compare portions of two arrays. No bounds checking is performed.
--
-- @since 2.0
compare :: Array -> Int -> Array -> Int -> Int -> Ordering
compare :: Array -> Int -> Array -> Int -> Int -> Ordering
compare Array
src1 Int
off1 Array
src2 Int
off2 Int
count = Array -> Int -> Array -> Int -> Int -> Int
compareInternal Array
src1 Int
off1 Array
src2 Int
off2 Int
count forall a. Ord a => a -> a -> Ordering
`Prelude.compare` Int
0
{-# INLINE compare #-}

compareInternal
      :: Array                  -- ^ First
      -> Int                    -- ^ Offset into first
      -> Array                  -- ^ Second
      -> Int                    -- ^ Offset into second
      -> Int                    -- ^ Count
      -> Int
compareInternal :: Array -> Int -> Array -> Int -> Int -> Int
compareInternal (ByteArray ByteArray#
src1#) (I# Int#
off1#) (ByteArray ByteArray#
src2#) (I# Int#
off2#) (I# Int#
count#) = Int
i
  where
#if MIN_VERSION_base(4,11,0)
    i :: Int
i = Int# -> Int
I# (ByteArray# -> Int# -> ByteArray# -> Int# -> Int# -> Int#
compareByteArrays# ByteArray#
src1# Int#
off1# ByteArray#
src2# Int#
off2# Int#
count#)
#else
    i = fromIntegral (memcmp src1# off1# src2# off2# count#)

foreign import ccall unsafe "_hs_text_memcmp2" memcmp
    :: ByteArray# -> Int# -> ByteArray# -> Int# -> Int# -> CInt
#endif
{-# INLINE compareInternal #-}