{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PackageImports #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Primitive.Array
  ( Array(..)
  , MutableArray(..)
  , newArray
  , readArray
  , writeArray
  , indexArray
  , indexArrayM
  , indexArray##
  , freezeArray
  , thawArray
  , A.runArray
  , unsafeFreezeArray
  , A.unsafeThawArray
  , A.sameMutableArray
  , copyArray
  , copyMutableArray
  , cloneArray
  , cloneMutableArray
  , A.sizeofArray
  , A.sizeofMutableArray
  , A.fromListN
  , A.fromList
  , A.arrayFromListN
  , A.arrayFromList
  , A.mapArray'
  , A.traverseArrayP
  ) where

import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Exception (throw, ArrayException(..), Exception, toException)
import qualified Data.List as L
import "primitive" Data.Primitive.Array (Array, MutableArray)
import qualified "primitive" Data.Primitive.Array as A
import GHC.Exts (raise#)
import GHC.Stack

check :: HasCallStack => String -> Bool -> a -> a
check :: String -> Bool -> a -> a
check String
_      Bool
True  a
x = a
x
check String
errMsg Bool
False a
_ = ArrayException -> a
forall a e. Exception e => e -> a
throw (String -> ArrayException
IndexOutOfBounds (String -> ArrayException) -> String -> ArrayException
forall a b. (a -> b) -> a -> b
$ String
"Data.Primitive.Array." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
errMsg String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ CallStack -> String
prettyCallStack CallStack
HasCallStack => CallStack
callStack)

checkUnary :: HasCallStack => String -> Bool -> (# a #) -> (# a #)
checkUnary :: String -> Bool -> (# a #) -> (# a #)
checkUnary String
_      Bool
True  (# a #)
x = (# a #)
x
checkUnary String
errMsg Bool
False (# a #)
_ = ArrayException -> (# a #)
forall e a. Exception e => e -> (# a #)
throwUnary (String -> ArrayException
IndexOutOfBounds (String -> ArrayException) -> String -> ArrayException
forall a b. (a -> b) -> a -> b
$ String
"Data.Primitive.Array." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
errMsg String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ CallStack -> String
prettyCallStack CallStack
HasCallStack => CallStack
callStack)

throwUnary :: Exception e => e -> (# a #)
throwUnary :: e -> (# a #)
throwUnary e
e = SomeException -> (# a #)
forall b a. b -> a
raise# (e -> SomeException
forall e. Exception e => e -> SomeException
toException e
e)

newArray :: (HasCallStack, PrimMonad m) => Int -> a -> m (MutableArray (PrimState m) a)
newArray :: Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
n a
x = String
-> Bool
-> m (MutableArray (PrimState m) a)
-> m (MutableArray (PrimState m) a)
forall a. HasCallStack => String -> Bool -> a -> a
check String
"newArray: negative size" (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (Int -> a -> m (MutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
A.newArray Int
n a
x)

readArray :: (HasCallStack, PrimMonad m) => MutableArray (PrimState m) a -> Int -> m a
readArray :: MutableArray (PrimState m) a -> Int -> m a
readArray MutableArray (PrimState m) a
marr Int
i = do
  let siz :: Int
siz = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr
  String -> Bool -> m a -> m a
forall a. HasCallStack => String -> Bool -> a -> a
check String
"readArray: index out of bounds" (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
siz) (MutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
A.readArray MutableArray (PrimState m) a
marr Int
i)

writeArray :: (HasCallStack, PrimMonad m) => MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray :: MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray (PrimState m) a
marr Int
i a
x = do
  let siz :: Int
siz = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr
  String -> Bool -> m () -> m ()
forall a. HasCallStack => String -> Bool -> a -> a
check String
"writeArray: index out of bounds" (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
siz) (MutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
A.writeArray MutableArray (PrimState m) a
marr Int
i a
x)

indexArray :: HasCallStack => Array a -> Int -> a
indexArray :: Array a -> Int -> a
indexArray Array a
arr Int
i = String -> Bool -> a -> a
forall a. HasCallStack => String -> Bool -> a -> a
check String
"indexArray: index out of bounds"
  (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array a -> Int
forall a. Array a -> Int
A.sizeofArray Array a
arr)
  (Array a -> Int -> a
forall a. Array a -> Int -> a
A.indexArray Array a
arr Int
i)

indexArrayM :: (HasCallStack, Monad m) => Array a -> Int -> m a
indexArrayM :: Array a -> Int -> m a
indexArrayM Array a
arr Int
i = String -> Bool -> m a -> m a
forall a. HasCallStack => String -> Bool -> a -> a
check String
"indexArrayM: index out of bounds"
  (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array a -> Int
forall a. Array a -> Int
A.sizeofArray Array a
arr)
  (Array a -> Int -> m a
forall (m :: * -> *) a. Monad m => Array a -> Int -> m a
A.indexArrayM Array a
arr Int
i)

indexArray## :: HasCallStack => Array a -> Int -> (# a #)
indexArray## :: Array a -> Int -> (# a #)
indexArray## Array a
arr Int
i = String -> Bool -> (# a #) -> (# a #)
forall a. HasCallStack => String -> Bool -> (# a #) -> (# a #)
checkUnary String
"indexArray##: index out of bounds"
  (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array a -> Int
forall a. Array a -> Int
A.sizeofArray Array a
arr)
  (Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
A.indexArray## Array a
arr Int
i)

{-# NOINLINE errorUnsafeFreeze #-}
errorUnsafeFreeze :: a
errorUnsafeFreeze :: a
errorUnsafeFreeze =
  String -> a
forall a. HasCallStack => String -> a
error String
"Data.Primitive.Array.unsafeFreeze:\nAttempted to read from an array after unsafely freezing it."

-- | This installs error thunks in the argument array so that
-- any attempt to use it after an unsafeFreeze will fail.
unsafeFreezeArray
  :: (HasCallStack, PrimMonad m)
  => MutableArray (PrimState m) a
  -> m (Array a)
unsafeFreezeArray :: MutableArray (PrimState m) a -> m (Array a)
unsafeFreezeArray MutableArray (PrimState m) a
marr = do
  let sz :: Int
sz = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr
  Array a
arr <- MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
A.freezeArray MutableArray (PrimState m) a
marr Int
0 Int
sz
  let go :: Int -> m ()
go !Int
ix = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then MutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
A.writeArray MutableArray (PrimState m) a
marr Int
ix a
forall a. a
errorUnsafeFreeze m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        else () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  Int -> m ()
go Int
0
  Array a -> m (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return Array a
arr

freezeArray
  :: (HasCallStack, PrimMonad m)
  => MutableArray (PrimState m) a -- ^ source
  -> Int                          -- ^ offset
  -> Int                          -- ^ length
  -> m (Array a)
freezeArray :: MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
freezeArray MutableArray (PrimState m) a
marr Int
s Int
l = String -> Bool -> m (Array a) -> m (Array a)
forall a. HasCallStack => String -> Bool -> a -> a
check String
"freezeArray: index range of out bounds"
  (Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr)
  (MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
A.freezeArray MutableArray (PrimState m) a
marr Int
s Int
l)

thawArray
  :: (HasCallStack, PrimMonad m)
  => Array a -- ^ source
  -> Int     -- ^ offset
  -> Int     -- ^ length
  -> m (MutableArray (PrimState m) a)
thawArray :: Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
thawArray Array a
arr Int
s Int
l = String
-> Bool
-> m (MutableArray (PrimState m) a)
-> m (MutableArray (PrimState m) a)
forall a. HasCallStack => String -> Bool -> a -> a
check String
"thawArray: index range of out bounds"
  (Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a -> Int
forall a. Array a -> Int
A.sizeofArray Array a
arr)
  (Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
A.thawArray Array a
arr Int
s Int
l)

copyArray
  :: (HasCallStack, PrimMonad m)
  => MutableArray (PrimState m) a -- ^ destination array
  -> Int                          -- ^ offset into destination array
  -> Array a                      -- ^ source array
  -> Int                          -- ^ offset into source array
  -> Int                          -- ^ number of elements to copy
  -> m ()
copyArray :: MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
copyArray MutableArray (PrimState m) a
marr Int
s1 Array a
arr Int
s2 Int
l = do
  let siz :: Int
siz = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr
  String -> Bool -> m () -> m ()
forall a. HasCallStack => String -> Bool -> a -> a
check String
"copyArray: index range of out bounds"
    (Int
s1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
siz Bool -> Bool -> Bool
&& Int
s2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a -> Int
forall a. Array a -> Int
A.sizeofArray Array a
arr)
    (MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Array a -> Int -> Int -> m ()
A.copyArray MutableArray (PrimState m) a
marr Int
s1 Array a
arr Int
s2 Int
l)

copyMutableArray
  :: (HasCallStack, PrimMonad m)
  => MutableArray (PrimState m) a -- ^ destination array
  -> Int                          -- ^ offset into destination array
  -> MutableArray (PrimState m) a -- ^ source array
  -> Int                          -- ^ offset into source array
  -> Int                          -- ^ number of elements to copy
  -> m ()
copyMutableArray :: MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
copyMutableArray MutableArray (PrimState m) a
marr1 Int
s1 MutableArray (PrimState m) a
marr2 Int
s2 Int
l = do
  let siz1 :: Int
siz1 = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr1
  let siz2 :: Int
siz2 = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr2
  let explain :: String
explain = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat
        [ String
"[dst size: "
        , Int -> String
forall a. Show a => a -> String
show Int
siz1
        , String
", dst off: "
        , Int -> String
forall a. Show a => a -> String
show Int
s1
        , String
", src size: "
        , Int -> String
forall a. Show a => a -> String
show Int
siz2
        , String
", src off: "
        , Int -> String
forall a. Show a => a -> String
show Int
s2
        , String
", copy size: "
        , Int -> String
forall a. Show a => a -> String
show Int
l
        , String
"]"
        ]
  String -> Bool -> m () -> m ()
forall a. HasCallStack => String -> Bool -> a -> a
check (String
"copyMutableArray: index range of out bounds " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
explain)
    (Int
s1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
siz1 Bool -> Bool -> Bool
&& Int
s2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
siz2)
    (MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
A.copyMutableArray MutableArray (PrimState m) a
marr1 Int
s1 MutableArray (PrimState m) a
marr2 Int
s2 Int
l)

cloneArray
  :: HasCallStack
  => Array a -- ^ source array
  -> Int     -- ^ offset into source array
  -> Int     -- ^ number of elements to copy
  -> Array a
cloneArray :: Array a -> Int -> Int -> Array a
cloneArray Array a
arr Int
s Int
l = String -> Bool -> Array a -> Array a
forall a. HasCallStack => String -> Bool -> a -> a
check String
"cloneArray: index range of out bounds"
  (Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Array a -> Int
forall a. Array a -> Int
A.sizeofArray Array a
arr)
  (Array a -> Int -> Int -> Array a
forall a. Array a -> Int -> Int -> Array a
A.cloneArray Array a
arr Int
s Int
l)

cloneMutableArray
  :: (HasCallStack, PrimMonad m)
  => MutableArray (PrimState m) a -- ^ source array
  -> Int                          -- ^ offset into source array
  -> Int                          -- ^ number of elements to copy
  -> m (MutableArray (PrimState m) a)
cloneMutableArray :: MutableArray (PrimState m) a
-> Int -> Int -> m (MutableArray (PrimState m) a)
cloneMutableArray MutableArray (PrimState m) a
marr Int
s Int
l = String
-> Bool
-> m (MutableArray (PrimState m) a)
-> m (MutableArray (PrimState m) a)
forall a. HasCallStack => String -> Bool -> a -> a
check String
"cloneMutableArray: index range of out bounds"
  (Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray (PrimState m) a
marr)
  (MutableArray (PrimState m) a
-> Int -> Int -> m (MutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Int -> m (MutableArray (PrimState m) a)
A.cloneMutableArray MutableArray (PrimState m) a
marr Int
s Int
l)