{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}

-- | Zero based arrays.
--
-- Note that no bounds checking are performed.
module Data.HashMap.Array
    ( Array
    , MArray

      -- * Creation
    , new
    , new_
    , singleton
    , singletonM
    , pair

      -- * Basic interface
    , length
    , lengthM
    , read
    , write
    , index
    , indexM
    , index#
    , update
    , updateWith'
    , unsafeUpdateM
    , insert
    , insertM
    , delete
    , sameArray1
    , trim

    , unsafeFreeze
    , unsafeThaw
    , unsafeSameArray
    , run
    , run2
    , copy
    , copyM

      -- * Folds
    , foldl'
    , foldr

    , thaw
    , map
    , map'
    , traverse
    , traverse'
    , toList
    , fromList
    ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative (Applicative (..), (<$>))
#endif
import Control.Applicative (liftA2)
import Control.DeepSeq
import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#)
import GHC.ST (ST(..))
import Control.Monad.ST (stToIO)

#if __GLASGOW_HASKELL__ >= 709
import Prelude hiding (filter, foldr, length, map, read, traverse)
#else
import Prelude hiding (filter, foldr, length, map, read)
#endif

#if __GLASGOW_HASKELL__ >= 710
import GHC.Exts (SmallArray#, newSmallArray#, readSmallArray#, writeSmallArray#,
                 indexSmallArray#, unsafeFreezeSmallArray#, unsafeThawSmallArray#,
                 SmallMutableArray#, sizeofSmallArray#, copySmallArray#, thawSmallArray#,
                 sizeofSmallMutableArray#, copySmallMutableArray#, cloneSmallMutableArray#)

#else
import GHC.Exts (Array#, newArray#, readArray#, writeArray#,
                 indexArray#, unsafeFreezeArray#, unsafeThawArray#,
                 MutableArray#, sizeofArray#, copyArray#, thawArray#,
                 sizeofMutableArray#, copyMutableArray#, cloneMutableArray#)
#endif

#if defined(ASSERTS)
import qualified Prelude
#endif

import Data.HashMap.Unsafe (runST)
import Control.Monad ((>=>))


#if __GLASGOW_HASKELL__ >= 710
type Array# a = SmallArray# a
type MutableArray# a = SmallMutableArray# a

newArray# :: Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newArray# :: Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newArray# = Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
forall a d.
Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newSmallArray#

unsafeFreezeArray# :: SmallMutableArray# d a
                   -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeArray# :: SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeArray# = SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
forall d a.
SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeSmallArray#

readArray# :: SmallMutableArray# d a
           -> Int# -> State# d -> (# State# d, a #)
readArray# :: SmallMutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# = SmallMutableArray# d a -> Int# -> State# d -> (# State# d, a #)
forall d a.
SmallMutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readSmallArray#

writeArray# :: SmallMutableArray# d a
            -> Int# -> a -> State# d -> State# d
writeArray# :: SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
writeArray# = SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
forall d a.
SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
writeSmallArray#

indexArray# :: SmallArray# a -> Int# -> (# a #)
indexArray# :: SmallArray# a -> Int# -> (# a #)
indexArray# = SmallArray# a -> Int# -> (# a #)
forall a. SmallArray# a -> Int# -> (# a #)
indexSmallArray#

unsafeThawArray# :: SmallArray# a
                 -> State# d -> (# State# d, SmallMutableArray# d a #)
unsafeThawArray# :: SmallArray# a -> State# d -> (# State# d, SmallMutableArray# d a #)
unsafeThawArray# = SmallArray# a -> State# d -> (# State# d, SmallMutableArray# d a #)
forall a d.
SmallArray# a -> State# d -> (# State# d, SmallMutableArray# d a #)
unsafeThawSmallArray#

sizeofArray# :: SmallArray# a -> Int#
sizeofArray# :: SmallArray# a -> Int#
sizeofArray# = SmallArray# a -> Int#
forall a. SmallArray# a -> Int#
sizeofSmallArray#

copyArray# :: SmallArray# a
           -> Int#
           -> SmallMutableArray# d a
           -> Int#
           -> Int#
           -> State# d
           -> State# d
copyArray# :: SmallArray# a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copyArray# = SmallArray# a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
forall a d.
SmallArray# a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copySmallArray#

cloneMutableArray# :: SmallMutableArray# s a
                   -> Int#
                   -> Int#
                   -> State# s
                   -> (# State# s, SmallMutableArray# s a #)
cloneMutableArray# :: SmallMutableArray# s a
-> Int#
-> Int#
-> State# s
-> (# State# s, SmallMutableArray# s a #)
cloneMutableArray# = SmallMutableArray# s a
-> Int#
-> Int#
-> State# s
-> (# State# s, SmallMutableArray# s a #)
forall d a.
SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
cloneSmallMutableArray#

thawArray# :: SmallArray# a
           -> Int#
           -> Int#
           -> State# d
           -> (# State# d, SmallMutableArray# d a #)
thawArray# :: SmallArray# a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
thawArray# = SmallArray# a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
forall a d.
SmallArray# a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
thawSmallArray#

sizeofMutableArray# :: SmallMutableArray# s a -> Int#
sizeofMutableArray# :: SmallMutableArray# s a -> Int#
sizeofMutableArray# = SmallMutableArray# s a -> Int#
forall d a. SmallMutableArray# d a -> Int#
sizeofSmallMutableArray#

copyMutableArray# :: SmallMutableArray# d a
                  -> Int#
                  -> SmallMutableArray# d a
                  -> Int#
                  -> Int#
                  -> State# d
                  -> State# d
copyMutableArray# :: SmallMutableArray# d a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableArray# = SmallMutableArray# d a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
forall d a.
SmallMutableArray# d a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copySmallMutableArray#
#endif

------------------------------------------------------------------------

#if defined(ASSERTS)
-- This fugly hack is brought by GHC's apparent reluctance to deal
-- with MagicHash and UnboxedTuples when inferring types. Eek!
# define CHECK_BOUNDS(_func_,_len_,_k_) \
if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.HashMap.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else
# define CHECK_OP(_func_,_op_,_lhs_,_rhs_) \
if not ((_lhs_) _op_ (_rhs_)) then error ("Data.HashMap.Array." ++ (_func_) ++ ": Check failed: _lhs_ _op_ _rhs_ (" ++ show (_lhs_) ++ " vs. " ++ show (_rhs_) ++ ")") else
# define CHECK_GT(_func_,_lhs_,_rhs_) CHECK_OP(_func_,>,_lhs_,_rhs_)
# define CHECK_LE(_func_,_lhs_,_rhs_) CHECK_OP(_func_,<=,_lhs_,_rhs_)
# define CHECK_EQ(_func_,_lhs_,_rhs_) CHECK_OP(_func_,==,_lhs_,_rhs_)
#else
# define CHECK_BOUNDS(_func_,_len_,_k_)
# define CHECK_OP(_func_,_op_,_lhs_,_rhs_)
# define CHECK_GT(_func_,_lhs_,_rhs_)
# define CHECK_LE(_func_,_lhs_,_rhs_)
# define CHECK_EQ(_func_,_lhs_,_rhs_)
#endif

data Array a = Array {
      Array a -> Array# a
unArray :: !(Array# a)
    }

instance Show a => Show (Array a) where
    show :: Array a -> String
show = [a] -> String
forall a. Show a => a -> String
show ([a] -> String) -> (Array a -> [a]) -> Array a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array a -> [a]
forall a. Array a -> [a]
toList

-- Determines whether two arrays have the same memory address.
-- This is more reliable than testing pointer equality on the
-- Array wrappers, but it's still slightly bogus.
unsafeSameArray :: Array a -> Array b -> Bool
unsafeSameArray :: Array a -> Array b -> Bool
unsafeSameArray (Array xs :: Array# a
xs) (Array ys :: Array# b
ys) =
  Int# -> Bool
forall a. Int# -> a
tagToEnum# ((Any -> Any -> Int#) -> Array# a -> Array# b -> Int#
unsafeCoerce# Any -> Any -> Int#
forall a. a -> a -> Int#
reallyUnsafePtrEquality# Array# a
xs Array# b
ys)

sameArray1 :: (a -> b -> Bool) -> Array a -> Array b -> Bool
sameArray1 :: (a -> b -> Bool) -> Array a -> Array b -> Bool
sameArray1 eq :: a -> b -> Bool
eq !Array a
xs0 !Array b
ys0
  | Int
lenxs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
lenys = Bool
False
  | Bool
otherwise = Int -> Array a -> Array b -> Bool
go 0 Array a
xs0 Array b
ys0
  where
    go :: Int -> Array a -> Array b -> Bool
go !Int
k !Array a
xs !Array b
ys
      | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
lenxs = Bool
True
      | (# x :: a
x #) <- Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
index# Array a
xs Int
k
      , (# y :: b
y #) <- Array b -> Int -> (# b #)
forall a. Array a -> Int -> (# a #)
index# Array b
ys Int
k
      = a -> b -> Bool
eq a
x b
y Bool -> Bool -> Bool
&& Int -> Array a -> Array b -> Bool
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) Array a
xs Array b
ys

    !lenxs :: Int
lenxs = Array a -> Int
forall a. Array a -> Int
length Array a
xs0
    !lenys :: Int
lenys = Array b -> Int
forall a. Array a -> Int
length Array b
ys0

length :: Array a -> Int
length :: Array a -> Int
length ary :: Array a
ary = Int# -> Int
I# (SmallArray# a -> Int#
forall a. SmallArray# a -> Int#
sizeofArray# (Array a -> SmallArray# a
forall a. Array a -> Array# a
unArray Array a
ary))
{-# INLINE length #-}

-- | Smart constructor
array :: Array# a -> Int -> Array a
array :: Array# a -> Int -> Array a
array ary :: Array# a
ary _n :: Int
_n = Array# a -> Array a
forall a. Array# a -> Array a
Array Array# a
ary
{-# INLINE array #-}

data MArray s a = MArray {
      MArray s a -> MutableArray# s a
unMArray :: !(MutableArray# s a)
    }

lengthM :: MArray s a -> Int
lengthM :: MArray s a -> Int
lengthM mary :: MArray s a
mary = Int# -> Int
I# (SmallMutableArray# s a -> Int#
forall d a. SmallMutableArray# d a -> Int#
sizeofMutableArray# (MArray s a -> SmallMutableArray# s a
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s a
mary))
{-# INLINE lengthM #-}

-- | Smart constructor
marray :: MutableArray# s a -> Int -> MArray s a
marray :: MutableArray# s a -> Int -> MArray s a
marray mary :: MutableArray# s a
mary _n :: Int
_n = MutableArray# s a -> MArray s a
forall s a. MutableArray# s a -> MArray s a
MArray MutableArray# s a
mary
{-# INLINE marray #-}

------------------------------------------------------------------------

instance NFData a => NFData (Array a) where
    rnf :: Array a -> ()
rnf = Array a -> ()
forall a. NFData a => Array a -> ()
rnfArray

rnfArray :: NFData a => Array a -> ()
rnfArray :: Array a -> ()
rnfArray ary0 :: Array a
ary0 = Array a -> Int -> Int -> ()
forall a. NFData a => Array a -> Int -> Int -> ()
go Array a
ary0 Int
n0 0
  where
    n0 :: Int
n0 = Array a -> Int
forall a. Array a -> Int
length Array a
ary0
    go :: Array a -> Int -> Int -> ()
go !Array a
ary !Int
n !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = ()
        | (# x :: a
x #) <- Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
index# Array a
ary Int
i
        = a -> ()
forall a. NFData a => a -> ()
rnf a
x () -> () -> ()
forall a a. a -> a -> a
`seq` Array a -> Int -> Int -> ()
go Array a
ary Int
n (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
-- We use index# just in case GHC can't see that the
-- relevant rnf is strict, or in case it actually isn't.
{-# INLINE rnfArray #-}

-- | Create a new mutable array of specified size, in the specified
-- state thread, with each element containing the specified initial
-- value.
new :: Int -> a -> ST s (MArray s a)
new :: Int -> a -> ST s (MArray s a)
new n :: Int
n@(I# n# :: Int#
n#) b :: a
b =
    CHECK_GT("new",n,(0 :: Int))
    STRep s (MArray s a) -> ST s (MArray s a)
forall s a. STRep s a -> ST s a
ST (STRep s (MArray s a) -> ST s (MArray s a))
-> STRep s (MArray s a) -> ST s (MArray s a)
forall a b. (a -> b) -> a -> b
$ \s :: State# s
s ->
        case Int# -> a -> State# s -> (# State# s, SmallMutableArray# s a #)
forall a d.
Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #)
newArray# Int#
n# a
b State# s
s of
            (# s' :: State# s
s', ary :: SmallMutableArray# s a
ary #) -> (# State# s
s', SmallMutableArray# s a -> Int -> MArray s a
forall s a. MutableArray# s a -> Int -> MArray s a
marray SmallMutableArray# s a
ary Int
n #)
{-# INLINE new #-}

new_ :: Int -> ST s (MArray s a)
new_ :: Int -> ST s (MArray s a)
new_ n :: Int
n = Int -> a -> ST s (MArray s a)
forall a s. Int -> a -> ST s (MArray s a)
new Int
n a
forall a. a
undefinedElem

singleton :: a -> Array a
singleton :: a -> Array a
singleton x :: a
x = (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST (a -> ST s (Array a)
forall a s. a -> ST s (Array a)
singletonM a
x)
{-# INLINE singleton #-}

singletonM :: a -> ST s (Array a)
singletonM :: a -> ST s (Array a)
singletonM x :: a
x = Int -> a -> ST s (MArray s a)
forall a s. Int -> a -> ST s (MArray s a)
new 1 a
x ST s (MArray s a)
-> (MArray s a -> ST s (Array a)) -> ST s (Array a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray s a -> ST s (Array a)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze
{-# INLINE singletonM #-}

pair :: a -> a -> Array a
pair :: a -> a -> Array a
pair x :: a
x y :: a
y = (forall s. ST s (MArray s a)) -> Array a
forall e. (forall s. ST s (MArray s e)) -> Array e
run ((forall s. ST s (MArray s a)) -> Array a)
-> (forall s. ST s (MArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    MArray s a
ary <- Int -> a -> ST s (MArray s a)
forall a s. Int -> a -> ST s (MArray s a)
new 2 a
x
    MArray s a -> Int -> a -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s a
ary 1 a
y
    MArray s a -> ST s (MArray s a)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray s a
ary
{-# INLINE pair #-}

read :: MArray s a -> Int -> ST s a
read :: MArray s a -> Int -> ST s a
read ary :: MArray s a
ary _i :: Int
_i@(I# i# :: Int#
i#) = STRep s a -> ST s a
forall s a. STRep s a -> ST s a
ST (STRep s a -> ST s a) -> STRep s a -> ST s a
forall a b. (a -> b) -> a -> b
$ \ s :: State# s
s ->
    CHECK_BOUNDS("read", lengthM ary, _i)
        SmallMutableArray# s a -> Int# -> STRep s a
forall d a.
SmallMutableArray# d a -> Int# -> State# d -> (# State# d, a #)
readArray# (MArray s a -> SmallMutableArray# s a
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s a
ary) Int#
i# State# s
s
{-# INLINE read #-}

write :: MArray s a -> Int -> a -> ST s ()
write :: MArray s a -> Int -> a -> ST s ()
write ary :: MArray s a
ary _i :: Int
_i@(I# i# :: Int#
i#) b :: a
b = STRep s () -> ST s ()
forall s a. STRep s a -> ST s a
ST (STRep s () -> ST s ()) -> STRep s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ s :: State# s
s ->
    CHECK_BOUNDS("write", lengthM ary, _i)
        case SmallMutableArray# s a -> Int# -> a -> State# s -> State# s
forall d a.
SmallMutableArray# d a -> Int# -> a -> State# d -> State# d
writeArray# (MArray s a -> SmallMutableArray# s a
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s a
ary) Int#
i# a
b State# s
s of
            s' :: State# s
s' -> (# State# s
s' , () #)
{-# INLINE write #-}

index :: Array a -> Int -> a
index :: Array a -> Int -> a
index ary :: Array a
ary _i :: Int
_i@(I# i# :: Int#
i#) =
    CHECK_BOUNDS("index", length ary, _i)
        case SmallArray# a -> Int# -> (# a #)
forall a. SmallArray# a -> Int# -> (# a #)
indexArray# (Array a -> SmallArray# a
forall a. Array a -> Array# a
unArray Array a
ary) Int#
i# of (# b :: a
b #) -> a
b
{-# INLINE index #-}

index# :: Array a -> Int -> (# a #)
index# :: Array a -> Int -> (# a #)
index# ary :: Array a
ary _i :: Int
_i@(I# i# :: Int#
i#) =
    CHECK_BOUNDS("index#", length ary, _i)
        SmallArray# a -> Int# -> (# a #)
forall a. SmallArray# a -> Int# -> (# a #)
indexArray# (Array a -> SmallArray# a
forall a. Array a -> Array# a
unArray Array a
ary) Int#
i#
{-# INLINE index# #-}

indexM :: Array a -> Int -> ST s a
indexM :: Array a -> Int -> ST s a
indexM ary :: Array a
ary _i :: Int
_i@(I# i# :: Int#
i#) =
    CHECK_BOUNDS("indexM", length ary, _i)
        case SmallArray# a -> Int# -> (# a #)
forall a. SmallArray# a -> Int# -> (# a #)
indexArray# (Array a -> SmallArray# a
forall a. Array a -> Array# a
unArray Array a
ary) Int#
i# of (# b :: a
b #) -> a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return a
b
{-# INLINE indexM #-}

unsafeFreeze :: MArray s a -> ST s (Array a)
unsafeFreeze :: MArray s a -> ST s (Array a)
unsafeFreeze mary :: MArray s a
mary
    = STRep s (Array a) -> ST s (Array a)
forall s a. STRep s a -> ST s a
ST (STRep s (Array a) -> ST s (Array a))
-> STRep s (Array a) -> ST s (Array a)
forall a b. (a -> b) -> a -> b
$ \s :: State# s
s -> case SmallMutableArray# s a -> State# s -> (# State# s, SmallArray# a #)
forall d a.
SmallMutableArray# d a -> State# d -> (# State# d, SmallArray# a #)
unsafeFreezeArray# (MArray s a -> SmallMutableArray# s a
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s a
mary) State# s
s of
                   (# s' :: State# s
s', ary :: SmallArray# a
ary #) -> (# State# s
s', SmallArray# a -> Int -> Array a
forall a. Array# a -> Int -> Array a
array SmallArray# a
ary (MArray s a -> Int
forall s a. MArray s a -> Int
lengthM MArray s a
mary) #)
{-# INLINE unsafeFreeze #-}

unsafeThaw :: Array a -> ST s (MArray s a)
unsafeThaw :: Array a -> ST s (MArray s a)
unsafeThaw ary :: Array a
ary
    = STRep s (MArray s a) -> ST s (MArray s a)
forall s a. STRep s a -> ST s a
ST (STRep s (MArray s a) -> ST s (MArray s a))
-> STRep s (MArray s a) -> ST s (MArray s a)
forall a b. (a -> b) -> a -> b
$ \s :: State# s
s -> case SmallArray# a -> State# s -> (# State# s, SmallMutableArray# s a #)
forall a d.
SmallArray# a -> State# d -> (# State# d, SmallMutableArray# d a #)
unsafeThawArray# (Array a -> SmallArray# a
forall a. Array a -> Array# a
unArray Array a
ary) State# s
s of
                   (# s' :: State# s
s', mary :: SmallMutableArray# s a
mary #) -> (# State# s
s', SmallMutableArray# s a -> Int -> MArray s a
forall s a. MutableArray# s a -> Int -> MArray s a
marray SmallMutableArray# s a
mary (Array a -> Int
forall a. Array a -> Int
length Array a
ary) #)
{-# INLINE unsafeThaw #-}

run :: (forall s . ST s (MArray s e)) -> Array e
run :: (forall s. ST s (MArray s e)) -> Array e
run act :: forall s. ST s (MArray s e)
act = (forall s. ST s (Array e)) -> Array e
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array e)) -> Array e)
-> (forall s. ST s (Array e)) -> Array e
forall a b. (a -> b) -> a -> b
$ ST s (MArray s e)
forall s. ST s (MArray s e)
act ST s (MArray s e)
-> (MArray s e -> ST s (Array e)) -> ST s (Array e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray s e -> ST s (Array e)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze
{-# INLINE run #-}

run2 :: (forall s. ST s (MArray s e, a)) -> (Array e, a)
run2 :: (forall s. ST s (MArray s e, a)) -> (Array e, a)
run2 k :: forall s. ST s (MArray s e, a)
k = (forall s. ST s (Array e, a)) -> (Array e, a)
forall a. (forall s. ST s a) -> a
runST (do
                 (marr :: MArray s e
marr,b :: a
b) <- ST s (MArray s e, a)
forall s. ST s (MArray s e, a)
k
                 Array e
arr <- MArray s e -> ST s (Array e)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze MArray s e
marr
                 (Array e, a) -> ST s (Array e, a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Array e
arr,a
b))

-- | Unsafely copy the elements of an array. Array bounds are not checked.
copy :: Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
copy :: Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
copy !Array e
src !_sidx :: Int
_sidx@(I# sidx# :: Int#
sidx#) !MArray s e
dst !_didx :: Int
_didx@(I# didx# :: Int#
didx#) _n :: Int
_n@(I# n# :: Int#
n#) =
    CHECK_LE("copy", _sidx + _n, length src)
    CHECK_LE("copy", _didx + _n, lengthM dst)
        STRep s () -> ST s ()
forall s a. STRep s a -> ST s a
ST (STRep s () -> ST s ()) -> STRep s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ s# :: State# s
s# ->
        case SmallArray# e
-> Int#
-> SmallMutableArray# s e
-> Int#
-> Int#
-> State# s
-> State# s
forall a d.
SmallArray# a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copyArray# (Array e -> SmallArray# e
forall a. Array a -> Array# a
unArray Array e
src) Int#
sidx# (MArray s e -> SmallMutableArray# s e
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s e
dst) Int#
didx# Int#
n# State# s
s# of
            s2 :: State# s
s2 -> (# State# s
s2, () #)

-- | Unsafely copy the elements of an array. Array bounds are not checked.
copyM :: MArray s e -> Int -> MArray s e -> Int -> Int -> ST s ()
copyM :: MArray s e -> Int -> MArray s e -> Int -> Int -> ST s ()
copyM !MArray s e
src !_sidx :: Int
_sidx@(I# sidx# :: Int#
sidx#) !MArray s e
dst !_didx :: Int
_didx@(I# didx# :: Int#
didx#) _n :: Int
_n@(I# n# :: Int#
n#) =
    CHECK_BOUNDS("copyM: src", lengthM src, _sidx + _n - 1)
    CHECK_BOUNDS("copyM: dst", lengthM dst, _didx + _n - 1)
    STRep s () -> ST s ()
forall s a. STRep s a -> ST s a
ST (STRep s () -> ST s ()) -> STRep s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ s# :: State# s
s# ->
    case SmallMutableArray# s e
-> Int#
-> SmallMutableArray# s e
-> Int#
-> Int#
-> State# s
-> State# s
forall d a.
SmallMutableArray# d a
-> Int#
-> SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> State# d
copyMutableArray# (MArray s e -> SmallMutableArray# s e
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s e
src) Int#
sidx# (MArray s e -> SmallMutableArray# s e
forall s a. MArray s a -> MutableArray# s a
unMArray MArray s e
dst) Int#
didx# Int#
n# State# s
s# of
        s2 :: State# s
s2 -> (# State# s
s2, () #)

cloneM :: MArray s a -> Int -> Int -> ST s (MArray s a)
cloneM :: MArray s a -> Int -> Int -> ST s (MArray s a)
cloneM _mary :: MArray s a
_mary@(MArray mary# :: MutableArray# s a
mary#) _off :: Int
_off@(I# off# :: Int#
off#) _len :: Int
_len@(I# len# :: Int#
len#) =
    CHECK_BOUNDS("cloneM_off", lengthM _mary, _off - 1)
    CHECK_BOUNDS("cloneM_end", lengthM _mary, _off + _len - 1)
    STRep s (MArray s a) -> ST s (MArray s a)
forall s a. STRep s a -> ST s a
ST (STRep s (MArray s a) -> ST s (MArray s a))
-> STRep s (MArray s a) -> ST s (MArray s a)
forall a b. (a -> b) -> a -> b
$ \ s :: State# s
s ->
    case MutableArray# s a
-> Int# -> Int# -> State# s -> (# State# s, MutableArray# s a #)
forall d a.
SmallMutableArray# d a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
cloneMutableArray# MutableArray# s a
mary# Int#
off# Int#
len# State# s
s of
      (# s' :: State# s
s', mary'# :: MutableArray# s a
mary'# #) -> (# State# s
s', MutableArray# s a -> MArray s a
forall s a. MutableArray# s a -> MArray s a
MArray MutableArray# s a
mary'# #)

-- | Create a new array of the @n@ first elements of @mary@.
trim :: MArray s a -> Int -> ST s (Array a)
trim :: MArray s a -> Int -> ST s (Array a)
trim mary :: MArray s a
mary n :: Int
n = MArray s a -> Int -> Int -> ST s (MArray s a)
forall s a. MArray s a -> Int -> Int -> ST s (MArray s a)
cloneM MArray s a
mary 0 Int
n ST s (MArray s a)
-> (MArray s a -> ST s (Array a)) -> ST s (Array a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MArray s a -> ST s (Array a)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze
{-# INLINE trim #-}

-- | /O(n)/ Insert an element at the given position in this array,
-- increasing its size by one.
insert :: Array e -> Int -> e -> Array e
insert :: Array e -> Int -> e -> Array e
insert ary :: Array e
ary idx :: Int
idx b :: e
b = (forall s. ST s (Array e)) -> Array e
forall a. (forall s. ST s a) -> a
runST (Array e -> Int -> e -> ST s (Array e)
forall e s. Array e -> Int -> e -> ST s (Array e)
insertM Array e
ary Int
idx e
b)
{-# INLINE insert #-}

-- | /O(n)/ Insert an element at the given position in this array,
-- increasing its size by one.
insertM :: Array e -> Int -> e -> ST s (Array e)
insertM :: Array e -> Int -> e -> ST s (Array e)
insertM ary :: Array e
ary idx :: Int
idx b :: e
b =
    CHECK_BOUNDS("insertM", count + 1, idx)
        do MArray s e
mary <- Int -> ST s (MArray s e)
forall s a. Int -> ST s (MArray s a)
new_ (Int
countInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)
           Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
forall e s. Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
copy Array e
ary 0 MArray s e
mary 0 Int
idx
           MArray s e -> Int -> e -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s e
mary Int
idx e
b
           Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
forall e s. Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
copy Array e
ary Int
idx MArray s e
mary (Int
idxInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) (Int
countInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
idx)
           MArray s e -> ST s (Array e)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze MArray s e
mary
  where !count :: Int
count = Array e -> Int
forall a. Array a -> Int
length Array e
ary
{-# INLINE insertM #-}

-- | /O(n)/ Update the element at the given position in this array.
update :: Array e -> Int -> e -> Array e
update :: Array e -> Int -> e -> Array e
update ary :: Array e
ary idx :: Int
idx b :: e
b = (forall s. ST s (Array e)) -> Array e
forall a. (forall s. ST s a) -> a
runST (Array e -> Int -> e -> ST s (Array e)
forall e s. Array e -> Int -> e -> ST s (Array e)
updateM Array e
ary Int
idx e
b)
{-# INLINE update #-}

-- | /O(n)/ Update the element at the given position in this array.
updateM :: Array e -> Int -> e -> ST s (Array e)
updateM :: Array e -> Int -> e -> ST s (Array e)
updateM ary :: Array e
ary idx :: Int
idx b :: e
b =
    CHECK_BOUNDS("updateM", count, idx)
        do MArray s e
mary <- Array e -> Int -> Int -> ST s (MArray s e)
forall e s. Array e -> Int -> Int -> ST s (MArray s e)
thaw Array e
ary 0 Int
count
           MArray s e -> Int -> e -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s e
mary Int
idx e
b
           MArray s e -> ST s (Array e)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze MArray s e
mary
  where !count :: Int
count = Array e -> Int
forall a. Array a -> Int
length Array e
ary
{-# INLINE updateM #-}

-- | /O(n)/ Update the element at the given positio in this array, by
-- applying a function to it.  Evaluates the element to WHNF before
-- inserting it into the array.
updateWith' :: Array e -> Int -> (e -> e) -> Array e
updateWith' :: Array e -> Int -> (e -> e) -> Array e
updateWith' ary :: Array e
ary idx :: Int
idx f :: e -> e
f
  | (# x :: e
x #) <- Array e -> Int -> (# e #)
forall a. Array a -> Int -> (# a #)
index# Array e
ary Int
idx
  = Array e -> Int -> e -> Array e
forall e. Array e -> Int -> e -> Array e
update Array e
ary Int
idx (e -> Array e) -> e -> Array e
forall a b. (a -> b) -> a -> b
$! e -> e
f e
x
{-# INLINE updateWith' #-}

-- | /O(1)/ Update the element at the given position in this array,
-- without copying.
unsafeUpdateM :: Array e -> Int -> e -> ST s ()
unsafeUpdateM :: Array e -> Int -> e -> ST s ()
unsafeUpdateM ary :: Array e
ary idx :: Int
idx b :: e
b =
    CHECK_BOUNDS("unsafeUpdateM", length ary, idx)
        do MArray s e
mary <- Array e -> ST s (MArray s e)
forall a s. Array a -> ST s (MArray s a)
unsafeThaw Array e
ary
           MArray s e -> Int -> e -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s e
mary Int
idx e
b
           Array e
_ <- MArray s e -> ST s (Array e)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze MArray s e
mary
           () -> ST s ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE unsafeUpdateM #-}

foldl' :: (b -> a -> b) -> b -> Array a -> b
foldl' :: (b -> a -> b) -> b -> Array a -> b
foldl' f :: b -> a -> b
f = \ z0 :: b
z0 ary0 :: Array a
ary0 -> Array a -> Int -> Int -> b -> b
go Array a
ary0 (Array a -> Int
forall a. Array a -> Int
length Array a
ary0) 0 b
z0
  where
    go :: Array a -> Int -> Int -> b -> b
go ary :: Array a
ary n :: Int
n i :: Int
i !b
z
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = b
z
        | Bool
otherwise
        = case Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
index# Array a
ary Int
i of
            (# x :: a
x #) -> Array a -> Int -> Int -> b -> b
go Array a
ary Int
n (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) (b -> a -> b
f b
z a
x)
{-# INLINE foldl' #-}

foldr :: (a -> b -> b) -> b -> Array a -> b
foldr :: (a -> b -> b) -> b -> Array a -> b
foldr f :: a -> b -> b
f = \ z0 :: b
z0 ary0 :: Array a
ary0 -> Array a -> Int -> Int -> b -> b
go Array a
ary0 (Array a -> Int
forall a. Array a -> Int
length Array a
ary0) 0 b
z0
  where
    go :: Array a -> Int -> Int -> b -> b
go ary :: Array a
ary n :: Int
n i :: Int
i z :: b
z
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = b
z
        | Bool
otherwise
        = case Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
index# Array a
ary Int
i of
            (# x :: a
x #) -> a -> b -> b
f a
x (Array a -> Int -> Int -> b -> b
go Array a
ary Int
n (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) b
z)
{-# INLINE foldr #-}

undefinedElem :: a
undefinedElem :: a
undefinedElem = String -> a
forall a. HasCallStack => String -> a
error "Data.HashMap.Array: Undefined element"
{-# NOINLINE undefinedElem #-}

thaw :: Array e -> Int -> Int -> ST s (MArray s e)
thaw :: Array e -> Int -> Int -> ST s (MArray s e)
thaw !Array e
ary !_o :: Int
_o@(I# o# :: Int#
o#) !n :: Int
n@(I# n# :: Int#
n#) =
    CHECK_LE("thaw", _o + n, length ary)
        STRep s (MArray s e) -> ST s (MArray s e)
forall s a. STRep s a -> ST s a
ST (STRep s (MArray s e) -> ST s (MArray s e))
-> STRep s (MArray s e) -> ST s (MArray s e)
forall a b. (a -> b) -> a -> b
$ \ s :: State# s
s -> case SmallArray# e
-> Int#
-> Int#
-> State# s
-> (# State# s, SmallMutableArray# s e #)
forall a d.
SmallArray# a
-> Int#
-> Int#
-> State# d
-> (# State# d, SmallMutableArray# d a #)
thawArray# (Array e -> SmallArray# e
forall a. Array a -> Array# a
unArray Array e
ary) Int#
o# Int#
n# State# s
s of
            (# s2 :: State# s
s2, mary# :: SmallMutableArray# s e
mary# #) -> (# State# s
s2, SmallMutableArray# s e -> Int -> MArray s e
forall s a. MutableArray# s a -> Int -> MArray s a
marray SmallMutableArray# s e
mary# Int
n #)
{-# INLINE thaw #-}

-- | /O(n)/ Delete an element at the given position in this array,
-- decreasing its size by one.
delete :: Array e -> Int -> Array e
delete :: Array e -> Int -> Array e
delete ary :: Array e
ary idx :: Int
idx = (forall s. ST s (Array e)) -> Array e
forall a. (forall s. ST s a) -> a
runST (Array e -> Int -> ST s (Array e)
forall e s. Array e -> Int -> ST s (Array e)
deleteM Array e
ary Int
idx)
{-# INLINE delete #-}

-- | /O(n)/ Delete an element at the given position in this array,
-- decreasing its size by one.
deleteM :: Array e -> Int -> ST s (Array e)
deleteM :: Array e -> Int -> ST s (Array e)
deleteM ary :: Array e
ary idx :: Int
idx = do
    CHECK_BOUNDS("deleteM", count, idx)
        do MArray s e
mary <- Int -> ST s (MArray s e)
forall s a. Int -> ST s (MArray s a)
new_ (Int
countInt -> Int -> Int
forall a. Num a => a -> a -> a
-1)
           Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
forall e s. Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
copy Array e
ary 0 MArray s e
mary 0 Int
idx
           Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
forall e s. Array e -> Int -> MArray s e -> Int -> Int -> ST s ()
copy Array e
ary (Int
idxInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) MArray s e
mary Int
idx (Int
countInt -> Int -> Int
forall a. Num a => a -> a -> a
-(Int
idxInt -> Int -> Int
forall a. Num a => a -> a -> a
+1))
           MArray s e -> ST s (Array e)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze MArray s e
mary
  where !count :: Int
count = Array e -> Int
forall a. Array a -> Int
length Array e
ary
{-# INLINE deleteM #-}

map :: (a -> b) -> Array a -> Array b
map :: (a -> b) -> Array a -> Array b
map f :: a -> b
f = \ ary :: Array a
ary ->
    let !n :: Int
n = Array a -> Int
forall a. Array a -> Int
length Array a
ary
    in (forall s. ST s (MArray s b)) -> Array b
forall e. (forall s. ST s (MArray s e)) -> Array e
run ((forall s. ST s (MArray s b)) -> Array b)
-> (forall s. ST s (MArray s b)) -> Array b
forall a b. (a -> b) -> a -> b
$ do
        MArray s b
mary <- Int -> ST s (MArray s b)
forall s a. Int -> ST s (MArray s a)
new_ Int
n
        Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
forall s. Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
go Array a
ary MArray s b
mary 0 Int
n
  where
    go :: Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
go ary :: Array a
ary mary :: MArray s b
mary i :: Int
i n :: Int
n
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = MArray s b -> ST s (MArray s b)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray s b
mary
        | Bool
otherwise = do
             a
x <- Array a -> Int -> ST s a
forall a s. Array a -> Int -> ST s a
indexM Array a
ary Int
i
             MArray s b -> Int -> b -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s b
mary Int
i (b -> ST s ()) -> b -> ST s ()
forall a b. (a -> b) -> a -> b
$ a -> b
f a
x
             Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
go Array a
ary MArray s b
mary (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
n
{-# INLINE map #-}

-- | Strict version of 'map'.
map' :: (a -> b) -> Array a -> Array b
map' :: (a -> b) -> Array a -> Array b
map' f :: a -> b
f = \ ary :: Array a
ary ->
    let !n :: Int
n = Array a -> Int
forall a. Array a -> Int
length Array a
ary
    in (forall s. ST s (MArray s b)) -> Array b
forall e. (forall s. ST s (MArray s e)) -> Array e
run ((forall s. ST s (MArray s b)) -> Array b)
-> (forall s. ST s (MArray s b)) -> Array b
forall a b. (a -> b) -> a -> b
$ do
        MArray s b
mary <- Int -> ST s (MArray s b)
forall s a. Int -> ST s (MArray s a)
new_ Int
n
        Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
forall s. Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
go Array a
ary MArray s b
mary 0 Int
n
  where
    go :: Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
go ary :: Array a
ary mary :: MArray s b
mary i :: Int
i n :: Int
n
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = MArray s b -> ST s (MArray s b)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray s b
mary
        | Bool
otherwise = do
             a
x <- Array a -> Int -> ST s a
forall a s. Array a -> Int -> ST s a
indexM Array a
ary Int
i
             MArray s b -> Int -> b -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s b
mary Int
i (b -> ST s ()) -> b -> ST s ()
forall a b. (a -> b) -> a -> b
$! a -> b
f a
x
             Array a -> MArray s b -> Int -> Int -> ST s (MArray s b)
go Array a
ary MArray s b
mary (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Int
n
{-# INLINE map' #-}

fromList :: Int -> [a] -> Array a
fromList :: Int -> [a] -> Array a
fromList n :: Int
n xs0 :: [a]
xs0 =
    CHECK_EQ("fromList", n, Prelude.length xs0)
        (forall s. ST s (MArray s a)) -> Array a
forall e. (forall s. ST s (MArray s e)) -> Array e
run ((forall s. ST s (MArray s a)) -> Array a)
-> (forall s. ST s (MArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
            MArray s a
mary <- Int -> ST s (MArray s a)
forall s a. Int -> ST s (MArray s a)
new_ Int
n
            [a] -> MArray s a -> Int -> ST s (MArray s a)
forall a s. [a] -> MArray s a -> Int -> ST s (MArray s a)
go [a]
xs0 MArray s a
mary 0
  where
    go :: [a] -> MArray s a -> Int -> ST s (MArray s a)
go [] !MArray s a
mary !Int
_   = MArray s a -> ST s (MArray s a)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray s a
mary
    go (x :: a
x:xs :: [a]
xs) mary :: MArray s a
mary i :: Int
i = do MArray s a -> Int -> a -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s a
mary Int
i a
x
                          [a] -> MArray s a -> Int -> ST s (MArray s a)
go [a]
xs MArray s a
mary (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+1)

toList :: Array a -> [a]
toList :: Array a -> [a]
toList = (a -> [a] -> [a]) -> [a] -> Array a -> [a]
forall a b. (a -> b -> b) -> b -> Array a -> b
foldr (:) []

newtype STA a = STA {STA a -> forall s. MutableArray# s a -> ST s (Array a)
_runSTA :: forall s. MutableArray# s a -> ST s (Array a)}

runSTA :: Int -> STA a -> Array a
runSTA :: Int -> STA a -> Array a
runSTA !Int
n (STA m :: forall s. MutableArray# s a -> ST s (Array a)
m) = (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array a)) -> Array a)
-> (forall s. ST s (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$ Int -> ST s (MArray s a)
forall s a. Int -> ST s (MArray s a)
new_ Int
n ST s (MArray s a)
-> (MArray s a -> ST s (Array a)) -> ST s (Array a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ (MArray ar :: MutableArray# s a
ar) -> MutableArray# s a -> ST s (Array a)
forall s. MutableArray# s a -> ST s (Array a)
m MutableArray# s a
ar

traverse :: Applicative f => (a -> f b) -> Array a -> f (Array b)
traverse :: (a -> f b) -> Array a -> f (Array b)
traverse f :: a -> f b
f = \ !Array a
ary ->
  let
    !len :: Int
len = Array a -> Int
forall a. Array a -> Int
length Array a
ary
    go :: Int -> f (STA b)
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = STA b -> f (STA b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (STA b -> f (STA b)) -> STA b -> f (STA b)
forall a b. (a -> b) -> a -> b
$ (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a. (forall s. MutableArray# s a -> ST s (Array a)) -> STA a
STA ((forall s. MutableArray# s b -> ST s (Array b)) -> STA b)
-> (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a b. (a -> b) -> a -> b
$ \mary :: MutableArray# s b
mary -> MArray s b -> ST s (Array b)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze (MutableArray# s b -> MArray s b
forall s a. MutableArray# s a -> MArray s a
MArray MutableArray# s b
mary)
      | (# x :: a
x #) <- Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
index# Array a
ary Int
i
      = (b -> STA b -> STA b) -> f b -> f (STA b) -> f (STA b)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\b :: b
b (STA m :: forall s. MutableArray# s b -> ST s (Array b)
m) -> (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a. (forall s. MutableArray# s a -> ST s (Array a)) -> STA a
STA ((forall s. MutableArray# s b -> ST s (Array b)) -> STA b)
-> (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a b. (a -> b) -> a -> b
$ \mary :: MutableArray# s b
mary ->
                  MArray s b -> Int -> b -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write (MutableArray# s b -> MArray s b
forall s a. MutableArray# s a -> MArray s a
MArray MutableArray# s b
mary) Int
i b
b ST s () -> ST s (Array b) -> ST s (Array b)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MutableArray# s b -> ST s (Array b)
forall s. MutableArray# s b -> ST s (Array b)
m MutableArray# s b
mary)
               (a -> f b
f a
x) (Int -> f (STA b)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1))
  in Int -> STA b -> Array b
forall a. Int -> STA a -> Array a
runSTA Int
len (STA b -> Array b) -> f (STA b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA b)
go 0
{-# INLINE [1] traverse #-}

-- TODO: Would it be better to just use a lazy traversal
-- and then force the elements of the result? My guess is
-- yes.
traverse' :: Applicative f => (a -> f b) -> Array a -> f (Array b)
traverse' :: (a -> f b) -> Array a -> f (Array b)
traverse' f :: a -> f b
f = \ !Array a
ary ->
  let
    !len :: Int
len = Array a -> Int
forall a. Array a -> Int
length Array a
ary
    go :: Int -> f (STA b)
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = STA b -> f (STA b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (STA b -> f (STA b)) -> STA b -> f (STA b)
forall a b. (a -> b) -> a -> b
$ (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a. (forall s. MutableArray# s a -> ST s (Array a)) -> STA a
STA ((forall s. MutableArray# s b -> ST s (Array b)) -> STA b)
-> (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a b. (a -> b) -> a -> b
$ \mary :: MutableArray# s b
mary -> MArray s b -> ST s (Array b)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze (MutableArray# s b -> MArray s b
forall s a. MutableArray# s a -> MArray s a
MArray MutableArray# s b
mary)
      | (# x :: a
x #) <- Array a -> Int -> (# a #)
forall a. Array a -> Int -> (# a #)
index# Array a
ary Int
i
      = (b -> STA b -> STA b) -> f b -> f (STA b) -> f (STA b)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\ !b
b (STA m :: forall s. MutableArray# s b -> ST s (Array b)
m) -> (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a. (forall s. MutableArray# s a -> ST s (Array a)) -> STA a
STA ((forall s. MutableArray# s b -> ST s (Array b)) -> STA b)
-> (forall s. MutableArray# s b -> ST s (Array b)) -> STA b
forall a b. (a -> b) -> a -> b
$ \mary :: MutableArray# s b
mary ->
                    MArray s b -> Int -> b -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write (MutableArray# s b -> MArray s b
forall s a. MutableArray# s a -> MArray s a
MArray MutableArray# s b
mary) Int
i b
b ST s () -> ST s (Array b) -> ST s (Array b)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> MutableArray# s b -> ST s (Array b)
forall s. MutableArray# s b -> ST s (Array b)
m MutableArray# s b
mary)
               (a -> f b
f a
x) (Int -> f (STA b)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1))
  in Int -> STA b -> Array b
forall a. Int -> STA a -> Array a
runSTA Int
len (STA b -> Array b) -> f (STA b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA b)
go 0
{-# INLINE [1] traverse' #-}

-- Traversing in ST, we don't need to get fancy; we
-- can just do it directly.
traverseST :: (a -> ST s b) -> Array a -> ST s (Array b)
traverseST :: (a -> ST s b) -> Array a -> ST s (Array b)
traverseST f :: a -> ST s b
f = \ ary0 :: Array a
ary0 ->
  let
    !len :: Int
len = Array a -> Int
forall a. Array a -> Int
length Array a
ary0
    go :: Int -> MArray s b -> ST s (MArray s b)
go k :: Int
k !MArray s b
mary
      | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = MArray s b -> ST s (MArray s b)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray s b
mary
      | Bool
otherwise = do
          a
x <- Array a -> Int -> ST s a
forall a s. Array a -> Int -> ST s a
indexM Array a
ary0 Int
k
          b
y <- a -> ST s b
f a
x
          MArray s b -> Int -> b -> ST s ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray s b
mary Int
k b
y
          Int -> MArray s b -> ST s (MArray s b)
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) MArray s b
mary
  in Int -> ST s (MArray s b)
forall s a. Int -> ST s (MArray s a)
new_ Int
len ST s (MArray s b)
-> (MArray s b -> ST s (Array b)) -> ST s (Array b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> MArray s b -> ST s (MArray s b)
go 0 (MArray s b -> ST s (MArray s b))
-> (MArray s b -> ST s (Array b)) -> MArray s b -> ST s (Array b)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> MArray s b -> ST s (Array b)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze)
{-# INLINE traverseST #-}

traverseIO :: (a -> IO b) -> Array a -> IO (Array b)
traverseIO :: (a -> IO b) -> Array a -> IO (Array b)
traverseIO f :: a -> IO b
f = \ ary0 :: Array a
ary0 ->
  let
    !len :: Int
len = Array a -> Int
forall a. Array a -> Int
length Array a
ary0
    go :: Int -> MArray RealWorld b -> IO (MArray RealWorld b)
go k :: Int
k !MArray RealWorld b
mary
      | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = MArray RealWorld b -> IO (MArray RealWorld b)
forall (m :: * -> *) a. Monad m => a -> m a
return MArray RealWorld b
mary
      | Bool
otherwise = do
          a
x <- ST RealWorld a -> IO a
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld a -> IO a) -> ST RealWorld a -> IO a
forall a b. (a -> b) -> a -> b
$ Array a -> Int -> ST RealWorld a
forall a s. Array a -> Int -> ST s a
indexM Array a
ary0 Int
k
          b
y <- a -> IO b
f a
x
          ST RealWorld () -> IO ()
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld () -> IO ()) -> ST RealWorld () -> IO ()
forall a b. (a -> b) -> a -> b
$ MArray RealWorld b -> Int -> b -> ST RealWorld ()
forall s a. MArray s a -> Int -> a -> ST s ()
write MArray RealWorld b
mary Int
k b
y
          Int -> MArray RealWorld b -> IO (MArray RealWorld b)
go (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) MArray RealWorld b
mary
  in ST RealWorld (MArray RealWorld b) -> IO (MArray RealWorld b)
forall a. ST RealWorld a -> IO a
stToIO (Int -> ST RealWorld (MArray RealWorld b)
forall s a. Int -> ST s (MArray s a)
new_ Int
len) IO (MArray RealWorld b)
-> (MArray RealWorld b -> IO (Array b)) -> IO (Array b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Int -> MArray RealWorld b -> IO (MArray RealWorld b)
go 0 (MArray RealWorld b -> IO (MArray RealWorld b))
-> (MArray RealWorld b -> IO (Array b))
-> MArray RealWorld b
-> IO (Array b)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> ST RealWorld (Array b) -> IO (Array b)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (Array b) -> IO (Array b))
-> (MArray RealWorld b -> ST RealWorld (Array b))
-> MArray RealWorld b
-> IO (Array b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MArray RealWorld b -> ST RealWorld (Array b)
forall s a. MArray s a -> ST s (Array a)
unsafeFreeze)
{-# INLINE traverseIO #-}


-- Why don't we have similar RULES for traverse'? The efficient
-- way to traverse strictly in IO or ST is to force results as
-- they come in, which leads to different semantics. In particular,
-- we need to ensure that
--
--  traverse' (\x -> print x *> pure undefined) xs
--
-- will actually print all the values and then return undefined.
-- We could add a strict mapMWithIndex, operating in an arbitrary
-- Monad, that supported such rules, but we don't have that right now.
{-# RULES
"traverse/ST" forall f. traverse f = traverseST f
"traverse/IO" forall f. traverse f = traverseIO f
 #-}