{-# language BangPatterns #-}
{-# language FlexibleInstances #-}
{-# language LambdaCase #-}
{-# language MagicHash #-}
{-# language RankNTypes #-}
{-# language ScopedTypeVariables #-}
{-# language TypeFamilies #-}
{-# language TypeFamilyDependencies #-}
{-# language UnboxedTuples #-}

-- | The contiguous package presents a common API to a number of contiguous
-- array types and their mutable counterparts. This is enabled with the
-- 'Contiguous' typeclass, which parameterises over a contiguous array type and
-- defines the core operations. However, the stable part of the interface is
-- contained in this module, which combines those primitives into common,
-- efficient array algorithms suitable for replacing pointer-heavy list
-- manipulations.
module Data.Primitive.Contiguous
  (
    -- * Accessors
    -- ** Length Information
    size
  , sizeMut
  , null
    -- ** Indexing
  , index
  , index#
  , read
    -- ** Monadic indexing
  , indexM

    -- * Construction
    -- ** Initialisation
  , empty
  , new
  , singleton
  , doubleton
  , tripleton
  , quadrupleton
  , replicate
  , replicateMut
  , generate
  , generateM
  , generateMutable
  , iterateN
  , iterateMutableN
  , write
    -- ** Running
  , run
    -- ** Monadic initialisation
  , replicateMutM
  , generateMutableM
  , iterateMutableNM
  , create
  , createT
    -- ** Unfolding
  , unfoldr
  , unfoldrN
  , unfoldrMutable
    -- ** Enumeration
  , enumFromN
  , enumFromMutableN
    -- ** Concatenation
  , append
    -- ** Splitting and Splicing
  , insertAt

    -- * Slicing
  , Slice
  , MutableSlice
  , slice
  , sliceMut
  , toSlice
  , toSliceMut

    -- * Modifying arrays
  , replaceAt
  , modifyAt
  , modifyAt'
  , modifyAtF
  , modifyAtF'
  , deleteAt
    -- ** Permutations
  , reverse
  , reverseMutable
  , reverseSlice

    -- ** Resizing
  , resize
  , shrink
  , unsafeShrinkAndFreeze

    -- * Elementwise operations
    -- ** Mapping
  , map
  , map'
  , mapMutable
  , mapMutable'
  , imap
  , imap'
  , imapMutable
  , imapMutable'
  , modify
  , modify'
  , mapMaybe

    -- ** Zipping
  , zip
  , zipWith
  , izipWith

    -- ** Specific elements
  , swap

    -- * Working with predicates
    -- ** Filtering
  , filter
  , ifilter
  , catMaybes
  , lefts
  , rights
  , partitionEithers
    -- ** Searching
  , find
  , findIndex
  , elem
  , maximum
  , minimum
  , maximumBy
  , minimumBy
    -- ** Comparing for equality
  , equals
  , equalsMut
  , same

    -- * Folds
  , foldl
  , foldl'
  , foldr
  , foldr'
  , foldMap
  , foldMap'
  , foldlMap'
  , ifoldl'
  , ifoldr
  , ifoldr'
  , ifoldlMap'
  , ifoldlMap1'
  , foldlM'
  , ifoldlM'
  , asum
  , all
  , any
    -- ** Zipping Folds
  , foldrZipWith
  , ifoldrZipWith
  , foldlZipWithM'
  , ifoldlZipWithM'

    -- * Traversals
  , traverse
  , traverse_
  , itraverse
  , itraverse_
  , traverseP
  , itraverseP
  , mapM
  , forM
  , mapM_
  , forM_
  , for
  , for_
  , sequence
  , sequence_

    -- * Typeclass method defaults
  , (<$)
  , ap

    -- * Prefix sums (scans)
  , scanl
  , scanl'
  , iscanl
  , iscanl'
  , prescanl
  , prescanl'
  , iprescanl
  , iprescanl'
  --, postscanl
  --, ipostscanl

  , mapAccum'
  , mapAccumLM'

    -- * Conversions
    -- ** Lists
  , fromList
  , fromListN
  , fromListMutable
  , fromListMutableN
  , unsafeFromListN
  , unsafeFromListReverseN
  , unsafeFromListReverseMutableN
  , toList
  , toListMutable
    -- ** Other array types
  , convert
  , lift
  , liftMut
  , unlift
  , unliftMut
    -- ** Between mutable and immutable variants
  , clone
  , cloneMut
  , copy
  , copyMut
  , freeze
  , thaw
  , unsafeFreeze

    -- * Hashing
  , liftHashWithSalt

    -- * Forcing an array and its contents
  , rnf

    -- * Classes
  , Contiguous(Mutable,Element,Sliced,MutableSliced)
  , ContiguousU
  , Always

    -- * Re-Exports
  , Array
  , MutableArray
  , SmallArray
  , SmallMutableArray
  , PrimArray
  , MutablePrimArray
  , UnliftedArray
  , MutableUnliftedArray
  ) where

import Control.Monad.Primitive
import Data.Primitive hiding (fromList,fromListN)
import Data.Primitive.Unlifted.Array
import Prelude hiding (map,all,any,foldr,foldMap,traverse,read,filter,replicate,null,reverse,foldl,foldr,zip,zipWith,scanl,(<$),elem,maximum,minimum,mapM,mapM_,sequence,sequence_)

import Control.Applicative (liftA2)
import Control.Monad (when)
import Control.Monad.ST (runST,ST)
import Data.Bits (xor)
import Data.Coerce (coerce)
import Data.Primitive.Contiguous.Class (Contiguous(..), ContiguousU(..), Slice, MutableSlice, Always)
import Data.Semigroup (First(..))
import Data.Word (Word8)
import GHC.Base (build)
import GHC.Exts (MutableArrayArray#,unsafeCoerce#,sameMutableArrayArray#,isTrue#,dataToTag#,Int(..))

import qualified Control.Applicative as A
import qualified Prelude


-- | Append two arrays.
append :: (Contiguous arr, Element arr a) => arr a -> arr a -> arr a
append :: arr a -> arr a -> arr a
append !arr a
a !arr a
b = (forall s. ST s (arr a)) -> arr a
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr a)) -> arr a)
-> (forall s. ST s (arr a)) -> arr a
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr s a
m <- Int -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
b)
  Mutable arr (PrimState (ST s)) a -> Int -> Sliced arr a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> Sliced arr b -> m ()
copy Mutable arr s a
Mutable arr (PrimState (ST s)) a
m Int
0 (arr a -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Sliced arr a
toSlice arr a
a)
  Mutable arr (PrimState (ST s)) a -> Int -> Sliced arr a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> Sliced arr b -> m ()
copy Mutable arr s a
Mutable arr (PrimState (ST s)) a
m (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
a) (arr a -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Sliced arr a
toSlice arr a
b)
  Mutable arr (PrimState (ST s)) a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr s a
Mutable arr (PrimState (ST s)) a
m
{-# inline append #-}

-- | Delete the element at the given position.
deleteAt :: (Contiguous arr, Element arr a) => arr a -> Int -> arr a
deleteAt :: arr a -> Int -> arr a
deleteAt arr a
src Int
i = (forall s. ST s (arr a)) -> arr a
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr a)) -> arr a)
-> (forall s. ST s (arr a)) -> arr a
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr s a
dst <- Sliced arr a -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Sliced arr b -> m (Mutable arr (PrimState m) b)
thaw (arr a -> Int -> Int -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> Int -> Sliced arr a
slice arr a
src Int
0 (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
src Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
  let !i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  Mutable arr (PrimState (ST s)) a -> Int -> Sliced arr a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> Sliced arr b -> m ()
copy Mutable arr s a
Mutable arr (PrimState (ST s)) a
dst Int
i (arr a -> Int -> Int -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> Int -> Sliced arr a
slice arr a
src Int
i' (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
src Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i'))
  Mutable arr (PrimState (ST s)) a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr s a
Mutable arr (PrimState (ST s)) a
dst
{-# inline deleteAt #-}

-- | Create a copy of an array except the element at the index is replaced with
--   the given value.
replaceAt :: (Contiguous arr, Element arr a) => arr a -> Int -> a -> arr a
replaceAt :: arr a -> Int -> a -> arr a
replaceAt arr a
src Int
i a
x = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr s a)) -> arr a)
-> (forall s. ST s (Mutable arr s a)) -> arr a
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr s a
dst <- Sliced arr a -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Sliced arr b -> m (Mutable arr (PrimState m) b)
thaw (arr a -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Sliced arr a
toSlice arr a
src)
  Mutable arr (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr s a
Mutable arr (PrimState (ST s)) a
dst Int
i a
x
  Mutable arr s a -> ST s (Mutable arr s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr s a
dst
{-# inline replaceAt #-}

modifyAt :: (Contiguous arr, Element arr a)
  => (a -> a) -> arr a -> Int -> arr a
modifyAt :: (a -> a) -> arr a -> Int -> arr a
modifyAt a -> a
f arr a
src Int
i = arr a -> Int -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> a -> arr a
replaceAt arr a
src Int
i (a -> arr a) -> a -> arr a
forall a b. (a -> b) -> a -> b
$ a -> a
f (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
src Int
i)
{-# inline modifyAt #-}

-- | Variant of modifyAt that forces the result before installing it in the
-- array.
modifyAt' :: (Contiguous arr, Element arr a)
  => (a -> a) -> arr a -> Int -> arr a
modifyAt' :: (a -> a) -> arr a -> Int -> arr a
modifyAt' a -> a
f arr a
src Int
i = arr a -> Int -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> a -> arr a
replaceAt arr a
src Int
i (a -> arr a) -> a -> arr a
forall a b. (a -> b) -> a -> b
$! a -> a
f (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
src Int
i)
{-# inline modifyAt' #-}

modifyAtF :: (Contiguous arr, Element arr a, Functor f)
  => (a -> f a) -> arr a -> Int -> f (arr a)
modifyAtF :: (a -> f a) -> arr a -> Int -> f (arr a)
modifyAtF a -> f a
f arr a
src Int
i = arr a -> Int -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> a -> arr a
replaceAt arr a
src Int
i (a -> arr a) -> f a -> f (arr a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
src Int
i)
{-# inline modifyAtF #-}

-- | Variant of modifyAtF that forces the result before installing it in the
-- array. Note that this requires 'Monad' rather than 'Functor'.
modifyAtF' :: (Contiguous arr, Element arr a, Monad f)
  => (a -> f a) -> arr a -> Int -> f (arr a)
modifyAtF' :: (a -> f a) -> arr a -> Int -> f (arr a)
modifyAtF' a -> f a
f arr a
src Int
i = do
  !a
r <- a -> f a
f (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
src Int
i)
  let !dst :: arr a
dst = arr a -> Int -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> a -> arr a
replaceAt arr a
src Int
i a
r
  arr a -> f (arr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure arr a
dst
{-# inline modifyAtF' #-}

-- | Map over the elements of an array with the index.
imap :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c)
  => (Int -> b -> c) -> arr1 b -> arr2 c
imap :: (Int -> b -> c) -> arr1 b -> arr2 c
imap Int -> b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr2 c)) -> arr2 c)
-> (forall s. ST s (arr2 c)) -> arr2 c
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr2 s c
mb <- Int -> ST s (Mutable arr2 (PrimState (ST s)) c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go :: Int -> ST s ()
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            b
x <- arr1 b -> Int -> ST s b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 b
a Int
i
            Mutable arr2 (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb Int
i (Int -> b -> c
f Int
i b
x)
            Int -> ST s ()
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
  Int -> ST s ()
go Int
0
  Mutable arr2 (PrimState (ST s)) c -> ST s (arr2 c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb
{-# inline imap #-}

-- | Map strictly over the elements of an array with the index.
--
--   Note that because a new array must be created, the resulting
--   array type can be /different/ than the original.
imap' :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c)
  => (Int -> b -> c) -> arr1 b -> arr2 c
imap' :: (Int -> b -> c) -> arr1 b -> arr2 c
imap' Int -> b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr2 c)) -> arr2 c)
-> (forall s. ST s (arr2 c)) -> arr2 c
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr2 s c
mb <- Int -> ST s (Mutable arr2 (PrimState (ST s)) c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go :: Int -> ST s ()
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            b
x <- arr1 b -> Int -> ST s b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 b
a Int
i
            let !b :: c
b = Int -> b -> c
f Int
i b
x
            Mutable arr2 (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb Int
i c
b
            Int -> ST s ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
  Mutable arr2 (PrimState (ST s)) c -> ST s (arr2 c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb
{-# inline imap' #-}

-- | Map over the elements of an array.
--
--   Note that because a new array must be created, the resulting
--   array type can be /different/ than the original.
map :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c)
  => (b -> c) -> arr1 b -> arr2 c
map :: (b -> c) -> arr1 b -> arr2 c
map b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr2 c)) -> arr2 c)
-> (forall s. ST s (arr2 c)) -> arr2 c
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr2 s c
mb <- Int -> ST s (Mutable arr2 (PrimState (ST s)) c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go :: Int -> ST s ()
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            b
x <- arr1 b -> Int -> ST s b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 b
a Int
i
            Mutable arr2 (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb Int
i (b -> c
f b
x)
            Int -> ST s ()
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
  Int -> ST s ()
go Int
0
  Mutable arr2 (PrimState (ST s)) c -> ST s (arr2 c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb
{-# inline map #-}

-- | Map strictly over the elements of an array.
--
--   Note that because a new array must be created, the resulting
--   array type can be /different/ than the original.
map' :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 c)
  => (b -> c) -> arr1 b -> arr2 c
map' :: (b -> c) -> arr1 b -> arr2 c
map' b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr2 c)) -> arr2 c)
-> (forall s. ST s (arr2 c)) -> arr2 c
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr2 s c
mb <- Int -> ST s (Mutable arr2 (PrimState (ST s)) c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go :: Int -> ST s ()
go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            b
x <- arr1 b -> Int -> ST s b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 b
a Int
i
            let !b :: c
b = b -> c
f b
x
            Mutable arr2 (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb Int
i c
b
            Int -> ST s ()
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
  Int -> ST s ()
go Int
0
  Mutable arr2 (PrimState (ST s)) c -> ST s (arr2 c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
mb
{-# inline map' #-}

-- | Convert one type of array into another.
convert :: (Contiguous arr1, Element arr1 b, Contiguous arr2, Element arr2 b)
  => arr1 b -> arr2 b
convert :: arr1 b -> arr2 b
convert arr1 b
a = (b -> b) -> arr1 b -> arr2 b
forall (arr1 :: * -> *) b (arr2 :: * -> *) c.
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 c) =>
(b -> c) -> arr1 b -> arr2 c
map b -> b
forall a. a -> a
id arr1 b
a
{-# inline convert #-}

-- | Right fold over the element of an array.
foldr :: (Contiguous arr, Element arr a) => (a -> b -> b) -> b -> arr a -> b
{-# inline foldr #-}
foldr :: (a -> b -> b) -> b -> arr a -> b
foldr a -> b -> b
f b
z = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b
go !Int
ix = if Int
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
ix
        then case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> a -> b -> b
f a
x (Int -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        else b
z
  in Int -> b
go Int
0

-- | Right fold over the element of an array, lazy in the accumulator,
-- provides index to the step function.
ifoldr :: (Contiguous arr, Element arr a) => (Int -> a -> b -> b) -> b -> arr a -> b
{-# inline ifoldr #-}
ifoldr :: (Int -> a -> b -> b) -> b -> arr a -> b
ifoldr Int -> a -> b -> b
f b
z = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b
go !Int
ix = if Int
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
ix
        then case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> a -> b -> b
f Int
ix a
x (Int -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        else b
z
  in Int -> b
go Int
0

-- | Strict right fold over the elements of an array.
foldr' :: (Contiguous arr, Element arr a) => (a -> b -> b) -> b -> arr a -> b
foldr' :: (a -> b -> b) -> b -> arr a -> b
foldr' a -> b -> b
f !b
z = \arr a
arr ->
  let go :: Int -> b -> b
go !Int
ix !b
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1
        then b
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> b -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (a -> b -> b
f a
x b
acc)
  in Int -> b -> b
go (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) b
z
{-# inline foldr' #-}

-- | Left fold over the elements of an array.
foldl :: (Contiguous arr, Element arr a) => (b -> a -> b) -> b -> arr a -> b
foldl :: (b -> a -> b) -> b -> arr a -> b
foldl b -> a -> b
f b
z = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b -> b
go !Int
ix b
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then b
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> b -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b -> a -> b
f b
acc a
x)
  in Int -> b -> b
go Int
0 b
z
{-# inline foldl #-}

-- | Strict left fold over the elements of an array.
foldl' :: (Contiguous arr, Element arr a) => (b -> a -> b) -> b -> arr a -> b
foldl' :: (b -> a -> b) -> b -> arr a -> b
foldl' b -> a -> b
f !b
z = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b -> b
go !Int
ix !b
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then b
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> b -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b -> a -> b
f b
acc a
x)
  in Int -> b -> b
go Int
0 b
z
{-# inline foldl' #-}

-- | Strict left fold over the elements of an array, where the accumulating
--   function cares about the index of the element.
ifoldl' :: (Contiguous arr, Element arr a)
  => (b -> Int -> a -> b) -> b -> arr a -> b
ifoldl' :: (b -> Int -> a -> b) -> b -> arr a -> b
ifoldl' b -> Int -> a -> b
f !b
z = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b -> b
go !Int
ix !b
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then b
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> b -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b -> Int -> a -> b
f b
acc Int
ix a
x)
  in Int -> b -> b
go Int
0 b
z
{-# inline ifoldl' #-}

-- | Strict right fold over the elements of an array, where the accumulating
--   function cares about the index of the element.
ifoldr' :: (Contiguous arr, Element arr a)
  => (Int -> a -> b -> b) -> b -> arr a -> b
ifoldr' :: (Int -> a -> b -> b) -> b -> arr a -> b
ifoldr' Int -> a -> b -> b
f !b
z = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b -> b
go !Int
ix !b
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (-Int
1)
        then b
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> b -> b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> a -> b -> b
f Int
ix a
x b
acc)
  in Int -> b -> b
go (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) b
z
{-# inline ifoldr' #-}

-- | Monoidal fold over the element of an array.
foldMap :: (Contiguous arr, Element arr a, Monoid m) => (a -> m) -> arr a -> m
foldMap :: (a -> m) -> arr a -> m
foldMap a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> m
go !Int
ix = if Int
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
ix
        then case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> m -> m -> m
forall a. Monoid a => a -> a -> a
mappend (a -> m
f a
x) (Int -> m
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        else m
forall a. Monoid a => a
mempty
  in Int -> m
go Int
0
{-# inline foldMap #-}

-- | Strict monoidal fold over the elements of an array.
foldMap' :: (Contiguous arr, Element arr a, Monoid m)
  => (a -> m) -> arr a -> m
foldMap' :: (a -> m) -> arr a -> m
foldMap' a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> m -> m
go !Int
ix !m
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then m
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
          of (# a
x #) -> Int -> m -> m
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (m -> m -> m
forall a. Monoid a => a -> a -> a
mappend m
acc (a -> m
f a
x))
  in Int -> m -> m
go Int
0 m
forall a. Monoid a => a
mempty
{-# inline foldMap' #-}

-- | Strict left monoidal fold over the elements of an array.
foldlMap' :: (Contiguous arr, Element arr a, Monoid m)
  => (a -> m) -> arr a -> m
foldlMap' :: (a -> m) -> arr a -> m
foldlMap' = (a -> m) -> arr a -> m
forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Monoid m) =>
(a -> m) -> arr a -> m
foldMap'
{-# inline foldlMap' #-}

-- | Strict monoidal fold over the elements of an array.
ifoldlMap' :: (Contiguous arr, Element arr a, Monoid m)
  => (Int -> a -> m)
  -> arr a
  -> m
ifoldlMap' :: (Int -> a -> m) -> arr a -> m
ifoldlMap' Int -> a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> m -> m
go !Int
ix !m
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then m
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> m -> m
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (m -> m -> m
forall a. Monoid a => a -> a -> a
mappend m
acc (Int -> a -> m
f Int
ix a
x))
  in Int -> m -> m
go Int
0 m
forall a. Monoid a => a
mempty
{-# inline ifoldlMap' #-}

-- | Strict monoidal fold over the elements of an array.
ifoldlMap1' :: (Contiguous arr, Element arr a, Semigroup m)
  => (Int -> a -> m)
  -> arr a
  -> m
ifoldlMap1' :: (Int -> a -> m) -> arr a -> m
ifoldlMap1' Int -> a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> m -> m
go !Int
ix !m
acc = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then m
acc
        else case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          (# a
x #) -> Int -> m -> m
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (m
acc m -> m -> m
forall a. Semigroup a => a -> a -> a
<> Int -> a -> m
f Int
ix a
x)
      !(# a
e0 #) = arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
0
  in Int -> m -> m
go Int
1 (Int -> a -> m
f Int
0 a
e0)
{-# inline ifoldlMap1' #-}

-- | Strict left monadic fold over the elements of an array.
foldlM' :: (Contiguous arr, Element arr a, Monad m)
  => (b -> a -> m b) -> b -> arr a -> m b
foldlM' :: (b -> a -> m b) -> b -> arr a -> m b
foldlM' b -> a -> m b
f b
z0 = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b -> m b
go !Int
ix !b
acc1 = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          let (# a
x #) = arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
          b
acc2 <- b -> a -> m b
f b
acc1 a
x
          Int -> b -> m b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) b
acc2
        else b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
acc1
  in Int -> b -> m b
go Int
0 b
z0
{-# inline foldlM' #-}

-- | Strict left monadic fold over the elements of an array.
ifoldlM' :: (Contiguous arr, Element arr a, Monad m)
  => (b -> Int -> a -> m b) -> b -> arr a -> m b
ifoldlM' :: (b -> Int -> a -> m b) -> b -> arr a -> m b
ifoldlM' b -> Int -> a -> m b
f b
z0 = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> b -> m b
go !Int
ix !b
acc1 = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          let (# a
x #) = arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
          b
acc2 <- b -> Int -> a -> m b
f b
acc1 Int
ix a
x
          Int -> b -> m b
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) b
acc2
        else b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
acc1
  in Int -> b -> m b
go Int
0 b
z0
{-# inline ifoldlM' #-}

-- | Drop elements that do not satisfy the predicate.
filter :: (Contiguous arr, Element arr a)
  => (a -> Bool)
  -> arr a
  -> arr a
filter :: (a -> Bool) -> arr a -> arr a
filter a -> Bool
p arr a
arr = (Int -> a -> Bool) -> arr a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(Int -> a -> Bool) -> arr a -> arr a
ifilter ((a -> Bool) -> Int -> a -> Bool
forall a b. a -> b -> a
const a -> Bool
p) arr a
arr
{-# inline filter #-}

-- | Drop elements that do not satisfy the predicate which
--   is applied to values and their indices.
ifilter :: (Contiguous arr, Element arr a)
  => (Int -> a -> Bool)
  -> arr a
  -> arr a
ifilter :: (Int -> a -> Bool) -> arr a -> arr a
ifilter Int -> a -> Bool
p arr a
arr = (forall s. ST s (arr a)) -> arr a
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr a)) -> arr a)
-> (forall s. ST s (arr a)) -> arr a
forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s Word8
marr :: MutablePrimArray s Word8 <- Int -> ST s (MutablePrimArray (PrimState (ST s)) Word8)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
sz
  let go1 :: Int -> Int -> ST s Int
      go1 :: Int -> Int -> ST s Int
go1 !Int
ix !Int
numTrue = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          a
atIx <- arr a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr a
arr Int
ix
          let !keep :: Bool
keep = Int -> a -> Bool
p Int
ix a
atIx
          let !keepTag :: Int
keepTag = Int# -> Int
I# (Bool -> Int#
forall a. a -> Int#
dataToTag# Bool
keep)
          MutablePrimArray (PrimState (ST s)) Word8
-> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Word8
MutablePrimArray (PrimState (ST s)) Word8
marr Int
ix (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
keepTag)
          Int -> Int -> ST s Int
go1 (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
numTrue Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
keepTag)
        else Int -> ST s Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
numTrue
  Int
numTrue <- Int -> Int -> ST s Int
go1 Int
0 Int
0
  if Int
numTrue Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
    then arr a -> ST s (arr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure arr a
arr
    else do
      Mutable arr s a
marrTrues <- Int -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
numTrue
      let go2 :: Int -> Int -> ST s ()
go2 !Int
ixSrc !Int
ixDst = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ixDst Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
numTrue) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
            Word8
atIxKeep <- MutablePrimArray (PrimState (ST s)) Word8 -> Int -> ST s Word8
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray MutablePrimArray s Word8
MutablePrimArray (PrimState (ST s)) Word8
marr Int
ixSrc
            if Word8 -> Bool
isTrue Word8
atIxKeep
              then do
                a
atIxVal <- arr a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr a
arr Int
ixSrc
                Mutable arr (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr s a
Mutable arr (PrimState (ST s)) a
marrTrues Int
ixDst a
atIxVal
                Int -> Int -> ST s ()
go2 (Int
ixSrc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
ixDst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
              else Int -> Int -> ST s ()
go2 (Int
ixSrc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
ixDst
      Int -> Int -> ST s ()
go2 Int
0 Int
0
      Mutable arr (PrimState (ST s)) a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr s a
Mutable arr (PrimState (ST s)) a
marrTrues
  where
    !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
{-# inline ifilter #-}

-- | The 'mapMaybe' function is a version of 'map' which can throw out elements.
--   In particular, the functional arguments returns something of type @'Maybe' b@.
--   If this is 'Nothing', no element is added on to the result array. If it is
--   @'Just' b@, then @b@ is included in the result array.
mapMaybe :: forall arr1 arr2 a b.
     ( Contiguous arr1, Element arr1 a
     , Contiguous arr2, Element arr2 b
     )
  => (a -> Maybe b)
  -> arr1 a
  -> arr2 b
mapMaybe :: (a -> Maybe b) -> arr1 a -> arr2 b
mapMaybe a -> Maybe b
f arr1 a
arr = (forall s. ST s (arr2 b)) -> arr2 b
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr2 b)) -> arr2 b)
-> (forall s. ST s (arr2 b)) -> arr2 b
forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
  let go :: Int -> Int -> [b] -> ST s ([b],Int)
      go :: Int -> Int -> [b] -> ST s ([b], Int)
go !Int
ix !Int
numJusts ![b]
justs = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          a
atIx <- arr1 a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix
          case a -> Maybe b
f a
atIx of
            Maybe b
Nothing -> Int -> Int -> [b] -> ST s ([b], Int)
forall s. Int -> Int -> [b] -> ST s ([b], Int)
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
numJusts [b]
justs
            Just b
x -> Int -> Int -> [b] -> ST s ([b], Int)
forall s. Int -> Int -> [b] -> ST s ([b], Int)
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
numJustsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (b
xb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
justs)
        else ([b], Int) -> ST s ([b], Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([b]
justs,Int
numJusts)
  !([b]
bs,!Int
numJusts) <- Int -> Int -> [b] -> ST s ([b], Int)
forall s. Int -> Int -> [b] -> ST s ([b], Int)
go Int
0 Int
0 []
  !Mutable arr2 s b
marr <- Int -> [b] -> ST s (Mutable arr2 (PrimState (ST s)) b)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
numJusts [b]
bs
  Mutable arr2 (PrimState (ST s)) b -> ST s (arr2 b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 s b
Mutable arr2 (PrimState (ST s)) b
marr
{-# inline mapMaybe #-}

{-# inline isTrue #-}
isTrue :: Word8 -> Bool
isTrue :: Word8 -> Bool
isTrue Word8
0 = Bool
False
isTrue Word8
_ = Bool
True

-- | The 'catMaybes' function takes a list of 'Maybe's and returns a
--   list of all the 'Just' values.
catMaybes :: (Contiguous arr, Element arr a, Element arr (Maybe a))
  => arr (Maybe a)
  -> arr a
catMaybes :: arr (Maybe a) -> arr a
catMaybes = (Maybe a -> Maybe a) -> arr (Maybe a) -> arr a
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Element arr1 a, Contiguous arr2,
 Element arr2 b) =>
(a -> Maybe b) -> arr1 a -> arr2 b
mapMaybe Maybe a -> Maybe a
forall a. a -> a
id
{-# inline catMaybes #-}

-- | @'replicate' n x@ is an array of length @n@ with @x@ the value of every element.
replicate :: (Contiguous arr, Element arr a) => Int -> a -> arr a
replicate :: Int -> a -> arr a
replicate Int
n a
x = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int -> a -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> b -> m (Mutable arr (PrimState m) b)
replicateMut Int
n a
x)
{-# inline replicate #-}

-- | @'replicateMutM' n act@ performs the action n times, gathering the results.
replicateMutM :: (PrimMonad m, Contiguous arr, Element arr a)
  => Int
  -> m a
  -> m (Mutable arr (PrimState m) a)
replicateMutM :: Int -> m a -> m (Mutable arr (PrimState m) a)
replicateMutM Int
len m a
act = do
  Mutable arr (PrimState m) a
marr <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
x <- m a
act
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
x
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
marr
{-# inline replicateMutM #-}


-- | Create an array from a list. If the given length does
-- not match the actual length, this function has undefined
-- behavior.
unsafeFromListN :: (Contiguous arr, Element arr a)
  => Int -- ^ length of list
  -> [a] -- ^ list
  -> arr a
unsafeFromListN :: Int -> [a] -> arr a
unsafeFromListN Int
n [a]
l = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int -> [a] -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListMutableN Int
n [a]
l)
{-# inline unsafeFromListN #-}

unsafeFromListMutableN :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> [a]
  -> m (Mutable arr (PrimState m) a)
unsafeFromListMutableN :: Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListMutableN Int
n [a]
l = do
  Mutable arr (PrimState m) a
m <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
n
  let go :: Int -> [a] -> m (Mutable arr (PrimState m) a)
go !Int
_ [] = Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
m
      go !Int
ix (a
x : [a]
xs) = do
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
m Int
ix a
x
        Int -> [a] -> m (Mutable arr (PrimState m) a)
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [a]
xs
  Int -> [a] -> m (Mutable arr (PrimState m) a)
go Int
0 [a]
l
{-# inline unsafeFromListMutableN #-}

-- | Create a mutable array from a list, reversing the order of
--   the elements. If the given length does not match the actual length,
--   this function has undefined behavior.
unsafeFromListReverseMutableN :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> [a]
  -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN :: Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
n [a]
l = do
  Mutable arr (PrimState m) a
m <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
n
  let go :: Int -> [a] -> m (Mutable arr (PrimState m) a)
go !Int
_ [] = Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
m
      go !Int
ix (a
x : [a]
xs) = do
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
m Int
ix a
x
        Int -> [a] -> m (Mutable arr (PrimState m) a)
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs
  Int -> [a] -> m (Mutable arr (PrimState m) a)
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
l
{-# inline unsafeFromListReverseMutableN #-}

-- | Create an array from a list, reversing the order of the
-- elements. If the given length does not match the actual length,
-- this function has undefined behavior.
unsafeFromListReverseN :: (Contiguous arr, Element arr a)
  => Int
  -> [a]
  -> arr a
unsafeFromListReverseN :: Int -> [a] -> arr a
unsafeFromListReverseN Int
n [a]
l = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int -> [a] -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
n [a]
l)
{-# inline unsafeFromListReverseN #-}

-- | Map over a mutable array, modifying the elements in place.
mapMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => (a -> a)
  -> Mutable arr (PrimState m) a
  -> m ()
mapMutable :: (a -> a) -> Mutable arr (PrimState m) a -> m ()
mapMutable a -> a
f !Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix (a -> a
f a
a)
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
{-# inline mapMutable #-}

-- | Strictly map over a mutable array, modifying the elements in place.
mapMutable' :: (PrimMonad m, Contiguous arr, Element arr a)
  => (a -> a)
  -> Mutable arr (PrimState m) a
  -> m ()
mapMutable' :: (a -> a) -> Mutable arr (PrimState m) a -> m ()
mapMutable' a -> a
f !Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        let !b :: a
b = a -> a
f a
a
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
b
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
{-# inline mapMutable' #-}

-- | Map over a mutable array with indices, modifying the elements in place.
imapMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => (Int -> a -> a)
  -> Mutable arr (PrimState m) a
  -> m ()
imapMutable :: (Int -> a -> a) -> Mutable arr (PrimState m) a -> m ()
imapMutable Int -> a -> a
f !Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix (Int -> a -> a
f Int
ix a
a)
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
{-# inline imapMutable #-}

-- | Strictly map over a mutable array with indices, modifying the elements in place.
imapMutable' :: (PrimMonad m, Contiguous arr, Element arr a)
  => (Int -> a -> a)
  -> Mutable arr (PrimState m) a
  -> m ()
imapMutable' :: (Int -> a -> a) -> Mutable arr (PrimState m) a -> m ()
imapMutable' Int -> a -> a
f !Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        let !b :: a
b = Int -> a -> a
f Int
ix a
a
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
b
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
{-# inline imapMutable' #-}

-- | Map each element of the array to an action, evaluate these
--   actions from left to right, and collect the results in a
--   new array.
traverseP ::
     ( PrimMonad m
     , Contiguous arr1, Element arr1 a
     , Contiguous arr2, Element arr2 b
     )
  => (a -> m b)
  -> arr1 a
  -> m (arr2 b)
traverseP :: (a -> m b) -> arr1 a -> m (arr2 b)
traverseP a -> m b
f !arr1 a
arr = do
  let !sz :: Int
sz = arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
  !Mutable arr2 (PrimState m) b
marr <- Int -> m (Mutable arr2 (PrimState m) b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- arr1 a -> Int -> m a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix
        b
b <- a -> m b
f a
a
        Mutable arr2 (PrimState m) b -> Int -> b -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 (PrimState m) b
marr Int
ix b
b
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  Mutable arr2 (PrimState m) b -> m (arr2 b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 (PrimState m) b
marr
{-# inline traverseP #-}

-- | Map each element of the array to an action, evaluate these
--   actions from left to right, and collect the results in a
--   new array.
itraverseP ::
     ( PrimMonad m
     , Contiguous arr1, Element arr1 a
     , Contiguous arr2, Element arr2 b
     )
  => (Int -> a -> m b)
  -> arr1 a
  -> m (arr2 b)
itraverseP :: (Int -> a -> m b) -> arr1 a -> m (arr2 b)
itraverseP Int -> a -> m b
f !arr1 a
arr = do
  let !sz :: Int
sz = arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
  !Mutable arr2 (PrimState m) b
marr <- Int -> m (Mutable arr2 (PrimState m) b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- arr1 a -> Int -> m a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix
        b
b <- Int -> a -> m b
f Int
ix a
a
        Mutable arr2 (PrimState m) b -> Int -> b -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 (PrimState m) b
marr Int
ix b
b
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  Mutable arr2 (PrimState m) b -> m (arr2 b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 (PrimState m) b
marr
{-# inline itraverseP #-}

newtype STA v a = STA {STA v a -> forall s. Mutable v s a -> ST s (v a)
_runSTA :: forall s. Mutable v s a -> ST s (v a)}

runSTA :: (Contiguous v, Element v a) => Int -> STA v a -> v a
runSTA :: Int -> STA v a -> v a
runSTA !Int
sz (STA forall s. Mutable v s a -> ST s (v a)
m) = (forall s. ST s (v a)) -> v a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (v a)) -> v a) -> (forall s. ST s (v a)) -> v a
forall a b. (a -> b) -> a -> b
$ Int -> ST s (Mutable v (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz ST s (Mutable v s a) -> (Mutable v s a -> ST s (v a)) -> ST s (v a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable v s a -> ST s (v a)
forall s. Mutable v s a -> ST s (v a)
m
{-# inline runSTA #-}

-- | Map each element of the array to an action, evaluate these
--   actions from left to right, and collect the results.
--   For a version that ignores the results, see 'traverse_'.
traverse ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Applicative f
  )
  => (a -> f b)
  -> arr1 a
  -> f (arr2 b)
traverse :: (a -> f b) -> arr1 a -> f (arr2 b)
traverse a -> f b
f = (Int -> a -> f b) -> arr1 a -> f (arr2 b)
forall (arr1 :: * -> *) (arr2 :: * -> *) a b (f :: * -> *).
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Applicative f) =>
(Int -> a -> f b) -> arr1 a -> f (arr2 b)
itraverse ((a -> f b) -> Int -> a -> f b
forall a b. a -> b -> a
const a -> f b
f)
{-# inline traverse #-}

-- | Map each element of the array to an action, evaluate these
--   actions from left to right, and ignore the results.
--   For a version that doesn't ignore the results, see 'traverse'.
traverse_ ::
     (Contiguous arr, Element arr a, Applicative f)
  => (a -> f b)
  -> arr a
  -> f ()
traverse_ :: (a -> f b) -> arr a -> f ()
traverse_ a -> f b
f = (Int -> a -> f b) -> arr a -> f ()
forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
(Int -> a -> f b) -> arr a -> f ()
itraverse_ ((a -> f b) -> Int -> a -> f b
forall a b. a -> b -> a
const a -> f b
f)

-- | Map each element of the array and its index to an action,
--   evaluating these actions from left to right.
itraverse ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Applicative f
  )
  => (Int -> a -> f b)
  -> arr1 a
  -> f (arr2 b)
itraverse :: (Int -> a -> f b) -> arr1 a -> f (arr2 b)
itraverse Int -> a -> f b
f = \arr1 a
arr ->
  let !sz :: Int
sz = arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
      go :: Int -> f (STA arr2 b)
go !Int
ix = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
        then STA arr2 b -> f (STA arr2 b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((forall s. Mutable arr2 s b -> ST s (arr2 b)) -> STA arr2 b
forall (v :: * -> *) a.
(forall s. Mutable v s a -> ST s (v a)) -> STA v a
STA forall s. Mutable arr2 s b -> ST s (arr2 b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze)
        else case arr1 a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr1 a
arr Int
ix of
          (# a
x #) -> (b -> STA arr2 b -> STA arr2 b)
-> f b -> f (STA arr2 b) -> f (STA arr2 b)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
            (\b
b (STA forall s. Mutable arr2 s b -> ST s (arr2 b)
m) -> (forall s. Mutable arr2 s b -> ST s (arr2 b)) -> STA arr2 b
forall (v :: * -> *) a.
(forall s. Mutable v s a -> ST s (v a)) -> STA v a
STA ((forall s. Mutable arr2 s b -> ST s (arr2 b)) -> STA arr2 b)
-> (forall s. Mutable arr2 s b -> ST s (arr2 b)) -> STA arr2 b
forall a b. (a -> b) -> a -> b
$ \Mutable arr2 s b
marr -> do
              Mutable arr2 (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s b
Mutable arr2 (PrimState (ST s)) b
marr Int
ix b
b
              Mutable arr2 s b -> ST s (arr2 b)
forall s. Mutable arr2 s b -> ST s (arr2 b)
m Mutable arr2 s b
marr
            )
            (Int -> a -> f b
f Int
ix a
x)
            (Int -> f (STA arr2 b)
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
  in if Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then arr2 b -> f (arr2 b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure arr2 b
forall (arr :: * -> *) a. Contiguous arr => arr a
empty
    else Int -> STA arr2 b -> arr2 b
forall (v :: * -> *) a.
(Contiguous v, Element v a) =>
Int -> STA v a -> v a
runSTA Int
sz (STA arr2 b -> arr2 b) -> f (STA arr2 b) -> f (arr2 b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> f (STA arr2 b)
go Int
0
{-# inline itraverse #-}

-- | Map each element of the array and its index to an action,
--   evaluate these actions from left to right, and ignore the results.
--   For a version that doesn't ignore the results, see 'itraverse'.
itraverse_ ::
     (Contiguous arr, Element arr a, Applicative f)
  => (Int -> a -> f b)
  -> arr a
  -> f ()
itraverse_ :: (Int -> a -> f b) -> arr a -> f ()
itraverse_ Int -> a -> f b
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> f ()
go !Int
ix = Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (f () -> f ()) -> f () -> f ()
forall a b. (a -> b) -> a -> b
$
        Int -> a -> f b
f Int
ix (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
arr Int
ix) f b -> f () -> f ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> f ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  in Int -> f ()
go Int
0
{-# inline itraverse_ #-}

-- | 'for' is 'traverse' with its arguments flipped. For a version
--   that ignores the results see 'for_'.
for ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Applicative f
  )
  => arr1 a
  -> (a -> f b)
  -> f (arr2 b)
for :: arr1 a -> (a -> f b) -> f (arr2 b)
for = ((a -> f b) -> arr1 a -> f (arr2 b))
-> arr1 a -> (a -> f b) -> f (arr2 b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> f b) -> arr1 a -> f (arr2 b)
forall (arr1 :: * -> *) (arr2 :: * -> *) a b (f :: * -> *).
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Applicative f) =>
(a -> f b) -> arr1 a -> f (arr2 b)
traverse
{-# inline for #-}

-- | 'for_' is 'traverse_' with its arguments flipped. For a version
--   that doesn't ignore the results see 'for'.
--
--   >>> for_ (C.fromList [1..4] :: PrimArray Int) print
--   1
--   2
--   3
--   4
for_ :: (Contiguous arr, Element arr a, Applicative f)
  => arr a
  -> (a -> f b)
  -> f ()
for_ :: arr a -> (a -> f b) -> f ()
for_ = ((a -> f b) -> arr a -> f ()) -> arr a -> (a -> f b) -> f ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> f b) -> arr a -> f ()
forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
(a -> f b) -> arr a -> f ()
traverse_
{-# inline for_ #-}

-- | Monadic accumulating strict left fold over the elements on an
-- array.
mapAccumLM' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 b
  , Element arr2 c
  , Monad m
  ) => (a -> b -> m (a, c)) -> a -> arr1 b -> m (a, arr2 c)
{-# inline mapAccumLM' #-}
mapAccumLM' :: (a -> b -> m (a, c)) -> a -> arr1 b -> m (a, arr2 c)
mapAccumLM' a -> b -> m (a, c)
f a
a0 arr1 b
src = Int -> [c] -> a -> m (a, arr2 c)
go Int
0 [] a
a0 where
  !sz :: Int
sz = arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
src
  go :: Int -> [c] -> a -> m (a, arr2 c)
go !Int
ix ![c]
xs !a
acc = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
    then do
      (!a
acc',!c
x) <- a -> b -> m (a, c)
f a
acc (arr1 b -> Int -> b
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr1 b
src Int
ix)
      Int -> [c] -> a -> m (a, arr2 c)
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (c
x c -> [c] -> [c]
forall a. a -> [a] -> [a]
: [c]
xs) a
acc'
    else
      let !xs' :: arr2 c
xs' = Int -> [c] -> arr2 c
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
Int -> [a] -> arr a
unsafeFromListReverseN Int
sz [c]
xs
       in (a, arr2 c) -> m (a, arr2 c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
acc,arr2 c
xs')

mapAccum' :: forall arr1 arr2 a b c.
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 b
  , Element arr2 c
  , Monoid a
  ) => (b -> (a, c)) -> arr1 b -> (a, arr2 c)
{-# inline mapAccum' #-}
mapAccum' :: (b -> (a, c)) -> arr1 b -> (a, arr2 c)
mapAccum' b -> (a, c)
f !arr1 b
src = (forall s. ST s (a, arr2 c)) -> (a, arr2 c)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (a, arr2 c)) -> (a, arr2 c))
-> (forall s. ST s (a, arr2 c)) -> (a, arr2 c)
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr2 s c
dst <- Int -> ST s (Mutable arr2 (PrimState (ST s)) c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  a
acc <- Int -> Mutable arr2 s c -> a -> ST s a
forall s. Int -> Mutable arr2 s c -> a -> ST s a
go Int
0 Mutable arr2 s c
dst a
forall a. Monoid a => a
mempty
  arr2 c
dst' <- Mutable arr2 (PrimState (ST s)) c -> ST s (arr2 c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
dst
  (a, arr2 c) -> ST s (a, arr2 c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
acc,arr2 c
dst')
  where
  !sz :: Int
sz = arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
src
  go :: Int -> Mutable arr2 s c -> a -> ST s a
  go :: Int -> Mutable arr2 s c -> a -> ST s a
go !Int
ix !Mutable arr2 s c
dst !a
accA = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
    then do
      let (!a
accB,!c
x) = b -> (a, c)
f (arr1 b -> Int -> b
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr1 b
src Int
ix)
      Mutable arr2 (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s c
Mutable arr2 (PrimState (ST s)) c
dst Int
ix c
x
      Int -> Mutable arr2 s c -> a -> ST s a
forall s. Int -> Mutable arr2 s c -> a -> ST s a
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Mutable arr2 s c
dst (a
accA a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
accB)
    else a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
accA

-- | Map each element of a structure to a monadic action,
--   evaluate these actions from left to right, and collect
--   the results. for a version that ignores the results see
--   'mapM_'.
mapM ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Monad m
  ) => (a -> m b)
    -> arr1 a
    -> m (arr2 b)
mapM :: (a -> m b) -> arr1 a -> m (arr2 b)
mapM a -> m b
f arr1 a
arr =
  let !sz :: Int
sz = arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
  in Int -> (Int -> m b) -> m (arr2 b)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, Monad m) =>
Int -> (Int -> m a) -> m (arr a)
generateM Int
sz ((Int -> m b) -> m (arr2 b)) -> (Int -> m b) -> m (arr2 b)
forall a b. (a -> b) -> a -> b
$ \Int
ix -> arr1 a -> Int -> m a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m b
f
{-# inline mapM #-}

-- | Map each element of a structure to a monadic action,
--   evaluate these actions from left to right, and ignore
--   the results. For a version that doesn't ignore the results
--   see 'mapM'.
--
--   'mapM_' = 'traverse_'
mapM_ :: (Contiguous arr, Element arr a, Element arr b, Applicative f)
  => (a -> f b)
  -> arr a
  -> f ()
mapM_ :: (a -> f b) -> arr a -> f ()
mapM_ = (a -> f b) -> arr a -> f ()
forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
(a -> f b) -> arr a -> f ()
traverse_
{-# inline mapM_ #-}

-- | 'forM' is 'mapM' with its arguments flipped. For a version that
--   ignores its results, see 'forM_'.
forM ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Monad m
  ) => arr1 a
    -> (a -> m b)
    -> m (arr2 b)
forM :: arr1 a -> (a -> m b) -> m (arr2 b)
forM = ((a -> m b) -> arr1 a -> m (arr2 b))
-> arr1 a -> (a -> m b) -> m (arr2 b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> m b) -> arr1 a -> m (arr2 b)
forall (arr1 :: * -> *) (arr2 :: * -> *) a b (m :: * -> *).
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Monad m) =>
(a -> m b) -> arr1 a -> m (arr2 b)
mapM
{-# inline forM #-}

-- | 'forM_' is 'mapM_' with its arguments flipped. For a version that
--   doesn't ignore its results, see 'forM'.
forM_ :: (Contiguous arr, Element arr a, Element arr b, Applicative f)
  => arr a
  -> (a -> f b)
  -> f ()
forM_ :: arr a -> (a -> f b) -> f ()
forM_ = ((a -> f b) -> arr a -> f ()) -> arr a -> (a -> f b) -> f ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> f b) -> arr a -> f ()
forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
(a -> f b) -> arr a -> f ()
traverse_
{-# inline forM_ #-}

-- | Evaluate each action in the structure from left to right
--   and collect the results. For a version that ignores the
--   results see 'sequence_'.
sequence ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 (f a)
  , Element arr2 a
  , Applicative f
  ) => arr1 (f a) -> f (arr2 a)
sequence :: arr1 (f a) -> f (arr2 a)
sequence = (f a -> f a) -> arr1 (f a) -> f (arr2 a)
forall (arr1 :: * -> *) (arr2 :: * -> *) a b (f :: * -> *).
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Applicative f) =>
(a -> f b) -> arr1 a -> f (arr2 b)
traverse f a -> f a
forall a. a -> a
id
{-# inline sequence #-}

-- | Evaluate each action in the structure from left to right
--   and ignore the results. For a version that doesn't ignore
--   the results see 'sequence'.
sequence_ ::
  ( Contiguous arr
  , Element arr (f a)
  , Applicative f
  ) => arr (f a) -> f ()
sequence_ :: arr (f a) -> f ()
sequence_ = (f a -> f () -> f ()) -> f () -> arr (f a) -> f ()
forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(a -> b -> b) -> b -> arr a -> b
foldr f a -> f () -> f ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>) (() -> f ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())
{-# inline sequence_ #-}

-- | The sum of a collection of actions, generalizing 'concat'.
--
--   >>> asum (C.fromList ['Just' "Hello", 'Nothing', Just "World"] :: Array String)
--   Just "Hello"
asum ::
  ( Contiguous arr
  , Element arr (f a)
  , A.Alternative f
  ) => arr (f a) -> f a
asum :: arr (f a) -> f a
asum = (f a -> f a -> f a) -> f a -> arr (f a) -> f a
forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(a -> b -> b) -> b -> arr a -> b
foldr f a -> f a -> f a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(A.<|>) f a
forall (f :: * -> *) a. Alternative f => f a
A.empty
{-# inline asum #-}

-- | Construct an array of the given length by applying
--   the function to each index.
generate :: (Contiguous arr, Element arr a)
  => Int
  -> (Int -> a)
  -> arr a
generate :: Int -> (Int -> a) -> arr a
generate Int
len Int -> a
f = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int -> (Int -> a) -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> (Int -> a) -> m (Mutable arr (PrimState m) a)
generateMutable Int
len Int -> a
f)
{-# inline generate #-}

-- | Construct an array of the given length by applying
--   the monadic action to each index.
generateM :: (Contiguous arr, Element arr a, Monad m)
  => Int
  -> (Int -> m a)
  -> m (arr a)
generateM :: Int -> (Int -> m a) -> m (arr a)
generateM !Int
sz Int -> m a
f =
  let go :: Int -> m (STA arr a)
go !Int
ix = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then (a -> STA arr a -> STA arr a)
-> m a -> m (STA arr a) -> m (STA arr a)
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
          (\a
b (STA forall s. Mutable arr s a -> ST s (arr a)
m) -> (forall s. Mutable arr s a -> ST s (arr a)) -> STA arr a
forall (v :: * -> *) a.
(forall s. Mutable v s a -> ST s (v a)) -> STA v a
STA ((forall s. Mutable arr s a -> ST s (arr a)) -> STA arr a)
-> (forall s. Mutable arr s a -> ST s (arr a)) -> STA arr a
forall a b. (a -> b) -> a -> b
$ \Mutable arr s a
marr -> do
              Mutable arr (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr s a
Mutable arr (PrimState (ST s)) a
marr Int
ix a
b
              Mutable arr s a -> ST s (arr a)
forall s. Mutable arr s a -> ST s (arr a)
m Mutable arr s a
marr
          )
          (Int -> m a
f Int
ix)
          (Int -> m (STA arr a)
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        else STA arr a -> m (STA arr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (STA arr a -> m (STA arr a)) -> STA arr a -> m (STA arr a)
forall a b. (a -> b) -> a -> b
$ (forall s. Mutable arr s a -> ST s (arr a)) -> STA arr a
forall (v :: * -> *) a.
(forall s. Mutable v s a -> ST s (v a)) -> STA v a
STA forall s. Mutable arr s a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze
  in if Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then arr a -> m (arr a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure arr a
forall (arr :: * -> *) a. Contiguous arr => arr a
empty
    else Int -> STA arr a -> arr a
forall (v :: * -> *) a.
(Contiguous v, Element v a) =>
Int -> STA v a -> v a
runSTA Int
sz (STA arr a -> arr a) -> m (STA arr a) -> m (arr a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (STA arr a)
go Int
0

-- | Construct a mutable array of the given length by applying
--   the function to each index.
generateMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> (Int -> a)
  -> m (Mutable arr (PrimState m) a)
generateMutable :: Int -> (Int -> a) -> m (Mutable arr (PrimState m) a)
generateMutable Int
len Int -> a
f = Int -> (Int -> m a) -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> (Int -> m a) -> m (Mutable arr (PrimState m) a)
generateMutableM Int
len (a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> (Int -> a) -> Int -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
f)
{-# inline generateMutable #-}

-- | Construct a mutable array of the given length by applying
--   the monadic action to each index.
generateMutableM :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> (Int -> m a)
  -> m (Mutable arr (PrimState m) a)
generateMutableM :: Int -> (Int -> m a) -> m (Mutable arr (PrimState m) a)
generateMutableM !Int
len Int -> m a
f = do
  Mutable arr (PrimState m) a
marr <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
x <- Int -> m a
f Int
ix
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
x
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
marr
{-# inline generateMutableM #-}

-- | Apply a function @n@ times to a value and construct an array
--   where each consecutive element is the result of an additional
--   application of this function. The zeroth element is the original value.
--
--   @'iterateN' 5 ('+' 1) 0 = 'fromListN' 5 [0,1,2,3,4]@
iterateN :: (Contiguous arr, Element arr a)
  => Int
  -> (a -> a)
  -> a
  -> arr a
iterateN :: Int -> (a -> a) -> a -> arr a
iterateN Int
len a -> a
f a
z0 = (forall s. ST s (arr a)) -> arr a
forall a. (forall s. ST s a) -> a
runST (Int -> (a -> a) -> a -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> (a -> a) -> a -> m (Mutable arr (PrimState m) a)
iterateMutableN Int
len a -> a
f a
z0 ST s (Mutable arr s a)
-> (Mutable arr s a -> ST s (arr a)) -> ST s (arr a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable arr s a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze)
{-# inline iterateN #-}

-- | Apply a function @n@ times to a value and construct a mutable array
--   where each consecutive element is the result of an additional
--   application of this function. The zeroth element is the original value.
iterateMutableN :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> (a -> a)
  -> a
  -> m (Mutable arr (PrimState m) a)
iterateMutableN :: Int -> (a -> a) -> a -> m (Mutable arr (PrimState m) a)
iterateMutableN Int
len a -> a
f a
z0 = Int -> (a -> m a) -> a -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> (a -> m a) -> a -> m (Mutable arr (PrimState m) a)
iterateMutableNM Int
len (a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> (a -> a) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f) a
z0
{-# inline iterateMutableN #-}

-- | Apply a monadic function @n@ times to a value and construct a mutable array
--   where each consecutive element is the result of an additional
--   application of this function. The zeroth element is the original value.
iterateMutableNM :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> (a -> m a)
  -> a
  -> m (Mutable arr (PrimState m) a)
iterateMutableNM :: Int -> (a -> m a) -> a -> m (Mutable arr (PrimState m) a)
iterateMutableNM !Int
len a -> m a
f a
z0 = do
  Mutable arr (PrimState m) a
marr <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  -- we are strict in the accumulator because
  -- otherwise we could build up a ton of `f (f (f (f .. (f a))))`
  -- thunks for no reason.
  let go :: Int -> a -> m ()
go !Int
ix !a
acc
        | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
z0 m () -> m () -> m ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> a -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
z0
        | Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            a
a <- a -> m a
f a
acc
            Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
a
            Int -> a -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
a
  Int -> a -> m ()
go Int
0 a
z0
  Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
marr
{-# inline iterateMutableNM #-}

-- | Execute the monad action and freeze the resulting array.
create :: (Contiguous arr, Element arr a)
  => (forall s. ST s (Mutable arr s a))
  -> arr a
create :: (forall s. ST s (Mutable arr s a)) -> arr a
create forall s. ST s (Mutable arr s a)
x = (forall s. ST s (arr a)) -> arr a
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run (Mutable arr s a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze (Mutable arr s a -> ST s (arr a))
-> ST s (Mutable arr s a) -> ST s (arr a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ST s (Mutable arr s a)
forall s. ST s (Mutable arr s a)
x)
{-# inline create #-}

-- | Execute the monadic action and freeze the resulting array.
createT :: (Contiguous arr, Element arr a, Traversable f)
  => (forall s. ST s (f (Mutable arr s a)))
  -> f (arr a)
createT :: (forall s. ST s (f (Mutable arr s a))) -> f (arr a)
createT forall s. ST s (f (Mutable arr s a))
p = (forall s. ST s (f (arr a))) -> f (arr a)
forall a. (forall s. ST s a) -> a
runST ((Mutable arr s a -> ST s (arr a))
-> f (Mutable arr s a) -> ST s (f (arr a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
Prelude.mapM Mutable arr s a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze (f (Mutable arr s a) -> ST s (f (arr a)))
-> ST s (f (Mutable arr s a)) -> ST s (f (arr a))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ST s (f (Mutable arr s a))
forall s. ST s (f (Mutable arr s a))
p)
{-# inline createT #-}

-- | Construct an array by repeatedly applying a generator
--   function to a seed. The generator function yields 'Just' the
--   next element and the new seed or 'Nothing' if there are no more
--   elements.
--
-- >>> unfoldr (\n -> if n == 0 then Nothing else Just (n,n-1) 10
--     <10,9,8,7,6,5,4,3,2,1>

-- Unfortunately, because we don't know ahead of time when to stop,
-- we need to construct a list and then turn it into an array.
unfoldr :: (Contiguous arr, Element arr a)
  => (b -> Maybe (a,b))
  -> b
  -> arr a
unfoldr :: (b -> Maybe (a, b)) -> b -> arr a
unfoldr b -> Maybe (a, b)
f b
z0 = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((b -> Maybe (a, b)) -> b -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *) b.
(Contiguous arr, Element arr a, PrimMonad m) =>
(b -> Maybe (a, b)) -> b -> m (Mutable arr (PrimState m) a)
unfoldrMutable b -> Maybe (a, b)
f b
z0)
{-# inline unfoldr #-}

-- | Construct a mutable array by repeatedly applying a generator
--   function to a seed. The generator function yields 'Just' the
--   next element and the new seed or 'Nothing' if there are no more
--   elements.
--
-- >>> unfoldrMutable (\n -> if n == 0 then Nothing else Just (n,n-1) 10
--     <10,9,8,7,6,5,4,3,2,1>

-- Unfortunately, because we don't know ahead of time when to stop,
-- we need to construct a list and then turn it into an array.
unfoldrMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => (b -> Maybe (a,b))
  -> b
  -> m (Mutable arr (PrimState m) a)
unfoldrMutable :: (b -> Maybe (a, b)) -> b -> m (Mutable arr (PrimState m) a)
unfoldrMutable b -> Maybe (a, b)
f b
z0 = do
  let go :: Int -> b -> [a] -> m (Int, [a])
go !Int
sz b
s ![a]
xs = case b -> Maybe (a, b)
f b
s of
        Maybe (a, b)
Nothing -> (Int, [a]) -> m (Int, [a])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
sz,[a]
xs)
        Just (a
x,b
s') -> Int -> b -> [a] -> m (Int, [a])
go (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) b
s' (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs)
  (Int
sz,[a]
xs) <- Int -> b -> [a] -> m (Int, [a])
go Int
0 b
z0 []
  Int -> [a] -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
sz [a]
xs
{-# inline unfoldrMutable #-}

-- | Construct an array with at most n elements by repeatedly
--   applying the generator function to a seed. The generator function
--   yields 'Just' the next element and the new seed or 'Nothing' if
--   there are no more elements.
unfoldrN :: (Contiguous arr, Element arr a)
  => Int
  -> (b -> Maybe (a, b))
  -> b
  -> arr a
unfoldrN :: Int -> (b -> Maybe (a, b)) -> b -> arr a
unfoldrN Int
maxSz b -> Maybe (a, b)
f b
z0 = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int
-> (b -> Maybe (a, b))
-> b
-> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *) b.
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> (b -> Maybe (a, b)) -> b -> m (Mutable arr (PrimState m) a)
unfoldrMutableN Int
maxSz b -> Maybe (a, b)
f b
z0)
{-# inline unfoldrN #-}

-- | Construct a mutable array with at most n elements by repeatedly
--   applying the generator function to a seed. The generator function
--   yields 'Just' the next element and the new seed or 'Nothing' if
--   there are no more elements.
unfoldrMutableN :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> (b -> Maybe (a, b))
  -> b
  -> m (Mutable arr (PrimState m) a)
unfoldrMutableN :: Int -> (b -> Maybe (a, b)) -> b -> m (Mutable arr (PrimState m) a)
unfoldrMutableN !Int
maxSz b -> Maybe (a, b)
f b
z0 = do
  Mutable arr (PrimState m) a
m <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
maxSz
  let go :: Int -> b -> m Int
go !Int
ix b
s = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxSz
        then case b -> Maybe (a, b)
f b
s of
          Maybe (a, b)
Nothing -> Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
          Just (a
x,b
s') -> do
            Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
m Int
ix a
x
            Int -> b -> m Int
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) b
s'
        else Int -> m Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
  Int
sz <- Int -> b -> m Int
go Int
0 b
z0
  Mutable arr (PrimState m) a
-> Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) a.
(Contiguous arr, PrimMonad m, Element arr a) =>
Mutable arr (PrimState m) a
-> Int -> m (Mutable arr (PrimState m) a)
shrink Mutable arr (PrimState m) a
m Int
sz
{-# inline unfoldrMutableN #-}

-- | Convert an array to a list.
toList :: (Contiguous arr, Element arr a)
  => arr a
  -> [a]
toList :: arr a -> [a]
toList arr a
arr = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> (a -> b -> b) -> b -> arr a -> b
forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(a -> b -> b) -> b -> arr a -> b
foldr a -> b -> b
c b
n arr a
arr)
{-# inline toList #-}

-- | Convert a mutable array to a list.

-- I don't think this can be expressed in terms of foldr/build,
-- so we just loop through the array.
toListMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => Mutable arr (PrimState m) a
  -> m [a]
toListMutable :: Mutable arr (PrimState m) a -> m [a]
toListMutable Mutable arr (PrimState m) a
marr = do
  Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> [a] -> m [a]
go !Int
ix ![a]
acc = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
        then do
          a
x <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
          Int -> [a] -> m [a]
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc)
        else [a] -> m [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
acc
  Int -> [a] -> m [a]
go (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) []
{-# inline toListMutable #-}

-- | Given an 'Int' that is representative of the length of
--   the list, convert the list into a mutable array of the
--   given length.
--
--   /Note/: calls 'error' if the given length is incorrect.
fromListMutableN :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> [a]
  -> m (Mutable arr (PrimState m) a)
fromListMutableN :: Int -> [a] -> m (Mutable arr (PrimState m) a)
fromListMutableN Int
len [a]
vs = do
  Mutable arr (PrimState m) a
marr <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  let go :: [a] -> Int -> m ()
go [] !Int
ix = if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
        then () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        else [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Primitive.Contiguous.fromListN: list length less than specified size."
      go (a
a:[a]
as) !Int
ix = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
        then do
          Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
a
          [a] -> Int -> m ()
go [a]
as (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        else [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Data.Primitive.Contiguous.fromListN: list length greater than specified size."
  [a] -> Int -> m ()
go [a]
vs Int
0
  Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
marr
{-# inline fromListMutableN #-}

-- | Convert a list into a mutable array of the given length.
fromListMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => [a]
  -> m (Mutable arr (PrimState m) a)
fromListMutable :: [a] -> m (Mutable arr (PrimState m) a)
fromListMutable [a]
xs = Int -> [a] -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
fromListMutableN ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) [a]
xs
{-# inline fromListMutable #-}

-- | Given an 'Int' that is representative of the length of
--   the list, convert the list into a mutable array of the
--   given length.
--
--   /Note/: calls 'error' if the given length is incorrect.
fromListN :: (Contiguous arr, Element arr a)
  => Int
  -> [a]
  -> arr a
fromListN :: Int -> [a] -> arr a
fromListN Int
len [a]
vs = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int -> [a] -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
fromListMutableN Int
len [a]
vs)
{-# inline fromListN #-}

-- | Convert a list into an array.
fromList :: (Contiguous arr, Element arr a)
  => [a]
  -> arr a
fromList :: [a] -> arr a
fromList [a]
vs = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ([a] -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
[a] -> m (Mutable arr (PrimState m) a)
fromListMutable [a]
vs)
{-# inline fromList #-}

-- | Modify the elements of a mutable array in-place.
modify :: (Contiguous arr, Element arr a, PrimMonad m)
  => (a -> a)
  -> Mutable arr (PrimState m) a
  -> m ()
modify :: (a -> a) -> Mutable arr (PrimState m) a -> m ()
modify a -> a
f Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
x <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix (a -> a
f a
x)
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
{-# inline modify #-}

-- | Strictly modify the elements of a mutable array in-place.
modify' :: (Contiguous arr, Element arr a, PrimMonad m)
  => (a -> a)
  -> Mutable arr (PrimState m) a
  -> m ()
modify' :: (a -> a) -> Mutable arr (PrimState m) a -> m ()
modify' a -> a
f Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        a
x <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        let !y :: a
y = a -> a
f a
x
        Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix a
y
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
{-# inline modify' #-}

-- | Yield an array of the given length containing the values
--   @x, 'succ' x, 'succ' ('succ' x)@ etc.
enumFromN :: (Contiguous arr, Element arr a, Enum a)
  => a
  -> Int
  -> arr a
enumFromN :: a -> Int -> arr a
enumFromN a
z0 Int
sz = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (a -> Int -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m, Enum a) =>
a -> Int -> m (Mutable arr (PrimState m) a)
enumFromMutableN a
z0 Int
sz)
{-# inline enumFromN #-}

-- | Yield a mutable array of the given length containing the values
--   @x, 'succ' x, 'succ' ('succ' x)@ etc.
enumFromMutableN :: (Contiguous arr, Element arr a, PrimMonad m, Enum a)
  => a
  -> Int
  -> m (Mutable arr (PrimState m) a)
enumFromMutableN :: a -> Int -> m (Mutable arr (PrimState m) a)
enumFromMutableN a
z0 !Int
sz = do
  Mutable arr (PrimState m) a
m <- Int -> m (Mutable arr (PrimState m) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let go :: Int -> a -> m (Mutable arr (PrimState m) a)
go !Int
ix a
z = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
m Int
ix a
z
          Int -> a -> m (Mutable arr (PrimState m) a)
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (a -> a
forall a. Enum a => a -> a
succ a
z)
        else Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
m
  Int -> a -> m (Mutable arr (PrimState m) a)
go Int
0 a
z0
{-# inline enumFromMutableN #-}

-- | Lift an accumulating hash function over the elements of the array,
--   returning the final accumulated hash.
liftHashWithSalt :: (Contiguous arr, Element arr a)
  => (Int -> a -> Int)
  -> Int
  -> arr a
  -> Int
liftHashWithSalt :: (Int -> a -> Int) -> Int -> arr a -> Int
liftHashWithSalt Int -> a -> Int
f Int
s0 arr a
arr = Int -> Int -> Int
go Int
0 Int
s0 where
  sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
  go :: Int -> Int -> Int
go !Int
ix !Int
s = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
    then
      let !(# a
x #) = arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
       in Int -> Int -> Int
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> a -> Int
f Int
s a
x)
    else Int -> Int -> Int
hashIntWithSalt Int
s Int
ix
{-# inline liftHashWithSalt #-}

-- | Reverse the elements of an array.
reverse :: (Contiguous arr, Element arr a)
  => arr a
  -> arr a
reverse :: arr a -> arr a
reverse arr a
arr = (forall s. ST s (arr a)) -> arr a
forall (arr :: * -> *) a.
Contiguous arr =>
(forall s. ST s (arr a)) -> arr a
run ((forall s. ST s (arr a)) -> arr a)
-> (forall s. ST s (arr a)) -> arr a
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr s a
marr <- Int -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr)
  Mutable arr (PrimState (ST s)) a -> Int -> Sliced arr a -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> Sliced arr b -> m ()
copy Mutable arr s a
Mutable arr (PrimState (ST s)) a
marr Int
0 (arr a -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Sliced arr a
toSlice arr a
arr)
  Mutable arr (PrimState (ST s)) a -> ST s ()
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Mutable arr (PrimState m) a -> m ()
reverseMutable Mutable arr s a
Mutable arr (PrimState (ST s)) a
marr
  Mutable arr (PrimState (ST s)) a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze Mutable arr s a
Mutable arr (PrimState (ST s)) a
marr
{-# inline reverse #-}

-- | Reverse the elements of a mutable array, in-place.
reverseMutable :: (Contiguous arr, Element arr a, PrimMonad m)
  => Mutable arr (PrimState m) a
  -> m ()
reverseMutable :: Mutable arr (PrimState m) a -> m ()
reverseMutable Mutable arr (PrimState m) a
marr = do
  !Int
sz <- Mutable arr (PrimState m) a -> m Int
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  Mutable arr (PrimState m) a -> Int -> Int -> m ()
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Mutable arr (PrimState m) a -> Int -> Int -> m ()
reverseSlice Mutable arr (PrimState m) a
marr Int
0 (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
{-# inline reverseMutable #-}

-- | Reverse the elements of a slice of a mutable array, in-place.
reverseSlice :: (Contiguous arr, Element arr a, PrimMonad m)
  => Mutable arr (PrimState m) a
  -> Int -- ^ start index
  -> Int -- ^ end index
  -> m ()
reverseSlice :: Mutable arr (PrimState m) a -> Int -> Int -> m ()
reverseSlice !Mutable arr (PrimState m) a
marr !Int
start !Int
end = do
  let go :: Int -> Int -> m ()
go !Int
s !Int
e = if Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
e
        then () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        else do
          a
tmp <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
s
          Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
s (a -> m ()) -> m a -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
e
          Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
e a
tmp
          Int -> Int -> m ()
go (Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
eInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  Int -> Int -> m ()
go Int
start Int
end
{-# inline reverseSlice #-}

-- | This function does not behave deterministically. Optimization level and
-- inlining can affect its results. However, the one thing that can be counted
-- on is that if it returns 'True', the two immutable arrays are definitely the
-- same. This is useful as shortcut for equality tests. However, keep in mind
-- that a result of 'False' tells us nothing about the arguments.
same :: ContiguousU arr => arr a -> arr a -> Bool
same :: arr a -> arr a -> Bool
same arr a
a arr a
b = Int# -> Bool
isTrue# (MutableArrayArray# Any -> MutableArrayArray# Any -> Int#
forall d. MutableArrayArray# d -> MutableArrayArray# d -> Int#
sameMutableArrayArray#
  (Unlifted arr a -> MutableArrayArray# s
unsafeCoerce# (arr a -> Unlifted arr a
forall (arr :: * -> *) b.
ContiguousU arr =>
arr b -> Unlifted arr b
unlift arr a
a) :: MutableArrayArray# s)
  (Unlifted arr a -> MutableArrayArray# s
unsafeCoerce# (arr a -> Unlifted arr a
forall (arr :: * -> *) b.
ContiguousU arr =>
arr b -> Unlifted arr b
unlift arr a
b) :: MutableArrayArray# s))

hashIntWithSalt :: Int -> Int -> Int
hashIntWithSalt :: Int -> Int -> Int
hashIntWithSalt Int
salt Int
x = Int
salt Int -> Int -> Int
`combine` Int
x
{-# inline hashIntWithSalt #-}

combine :: Int -> Int -> Int
combine :: Int -> Int -> Int
combine Int
h1 Int
h2 = (Int
h1 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
16777619) Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
h2
{-# inline combine #-}

-- | Does the element occur in the structure?
elem :: (Contiguous arr, Element arr a, Eq a) => a -> arr a -> Bool
elem :: a -> arr a -> Bool
elem a
a !arr a
arr =
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> Bool
go !Int
ix
        | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz = case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
            !(# a
x #) -> if a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
x
              then Bool
True
              else Int -> Bool
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise = Bool
False
  in Int -> Bool
go Int
0
{-# inline elem #-}

-- | The largest element of a structure.
maximum :: (Contiguous arr, Element arr a, Ord a) => arr a -> Maybe a
maximum :: arr a -> Maybe a
maximum = (a -> a -> Ordering) -> arr a -> Maybe a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(a -> a -> Ordering) -> arr a -> Maybe a
maximumBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# inline maximum #-}

-- | The least element of a structure.
minimum :: (Contiguous arr, Element arr a, Ord a) => arr a -> Maybe a
minimum :: arr a -> Maybe a
minimum = (a -> a -> Ordering) -> arr a -> Maybe a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(a -> a -> Ordering) -> arr a -> Maybe a
minimumBy a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# inline minimum #-}

-- | The largest element of a structure with respect to the
--   given comparison function.
maximumBy :: (Contiguous arr, Element arr a)
  => (a -> a -> Ordering)
  -> arr a
  -> Maybe a
maximumBy :: (a -> a -> Ordering) -> arr a -> Maybe a
maximumBy a -> a -> Ordering
f arr a
arr =
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> a -> a
go !Int
ix a
o = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          !(# a
x #) -> Int -> a -> a
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (case a -> a -> Ordering
f a
x a
o of { Ordering
GT -> a
x; Ordering
_ -> a
o; })
        else a
o
  in if Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then Maybe a
forall a. Maybe a
Nothing
    else a -> Maybe a
forall a. a -> Maybe a
Just (Int -> a -> a
go Int
0 (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
arr Int
0))
{-# inline maximumBy #-}

-- | The least element of a structure with respect to the
--   given comparison function.
minimumBy :: (Contiguous arr, Element arr a)
  => (a -> a -> Ordering)
  -> arr a
  -> Maybe a
minimumBy :: (a -> a -> Ordering) -> arr a -> Maybe a
minimumBy a -> a -> Ordering
f arr a
arr =
  let !sz :: Int
sz = arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr
      go :: Int -> a -> a
go !Int
ix a
o = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then case arr a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix of
          !(# a
x #) -> Int -> a -> a
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (case a -> a -> Ordering
f a
x a
o of { Ordering
GT -> a
o; Ordering
_ -> a
x; })
        else a
o
  in if Int
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then Maybe a
forall a. Maybe a
Nothing
    else a -> Maybe a
forall a. a -> Maybe a
Just (Int -> a -> a
go Int
0 (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
arr Int
0))
{-# inline minimumBy #-}

-- | 'find' takes a predicate and an array, and returns the leftmost
--   element of the array matching the prediate, or 'Nothing' if there
--   is no such element.
find :: (Contiguous arr, Element arr a)
  => (a -> Bool)
  -> arr a
  -> Maybe a
find :: (a -> Bool) -> arr a -> Maybe a
find a -> Bool
p = Maybe (First a) -> Maybe a
coerce (Maybe (First a) -> Maybe a)
-> (arr a -> Maybe (First a)) -> arr a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a -> Maybe (First a)) -> arr a -> Maybe (First a)
forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Monoid m) =>
(a -> m) -> arr a -> m
foldMap (\a
x -> if a -> Bool
p a
x then First a -> Maybe (First a)
forall a. a -> Maybe a
Just (a -> First a
forall a. a -> First a
First a
x) else Maybe (First a)
forall a. Maybe a
Nothing))
{-# inline find #-}

-- | 'findIndex' takes a predicate and an array, and returns the index of
--   the leftmost element of the array matching the prediate, or 'Nothing'
--   if there is no such element.
findIndex :: (Contiguous arr, Element arr a)
  => (a -> Bool)
  -> arr a
  -> Maybe Int
findIndex :: (a -> Bool) -> arr a -> Maybe Int
findIndex a -> Bool
p arr a
xs = Int -> Maybe Int
loop Int
0
  where
  loop :: Int -> Maybe Int
loop Int
i
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< arr a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
xs = if a -> Bool
p (arr a -> Int -> a
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
xs Int
i) then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i else Int -> Maybe Int
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    | Bool
otherwise = Maybe Int
forall a. Maybe a
Nothing
{-# inline findIndex #-}

-- | Swap the elements of the mutable array at the given indices.
swap :: (Contiguous arr, Element arr a, PrimMonad m)
  => Mutable arr (PrimState m) a
  -> Int
  -> Int
  -> m ()
swap :: Mutable arr (PrimState m) a -> Int -> Int -> m ()
swap !Mutable arr (PrimState m) a
marr !Int
ix1 !Int
ix2 = do
  a
atIx1 <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix1
  a
atIx2 <- Mutable arr (PrimState m) a -> Int -> m a
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix2
  Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix1 a
atIx2
  Mutable arr (PrimState m) a -> Int -> a -> m ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr (PrimState m) a
marr Int
ix2 a
atIx1
{-# inline swap #-}

-- | Extracts from an array of 'Either' all the 'Left' elements.
-- All the 'Left' elements are extracted in order.
lefts :: forall arr a b.
  ( Contiguous arr
  , Element arr a
  , Element arr (Either a b)
  ) => arr (Either a b)
    -> arr a
lefts :: arr (Either a b) -> arr a
lefts !arr (Either a b)
arr = (forall s. ST s (Mutable arr s a)) -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr s a)) -> arr a)
-> (forall s. ST s (Mutable arr s a)) -> arr a
forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = arr (Either a b) -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr (Either a b)
arr
      go :: Int -> [a] -> Int -> ST s (Int, [a])
      go :: Int -> [a] -> Int -> ST s (Int, [a])
go !Int
ix ![a]
as !Int
acc = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          arr (Either a b) -> Int -> ST s (Either a b)
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr (Either a b)
arr Int
ix ST s (Either a b)
-> (Either a b -> ST s (Int, [a])) -> ST s (Int, [a])
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Left a
a -> Int -> [a] -> Int -> ST s (Int, [a])
forall s. Int -> [a] -> Int -> ST s (Int, [a])
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as) (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Right b
_ -> Int -> [a] -> Int -> ST s (Int, [a])
forall s. Int -> [a] -> Int -> ST s (Int, [a])
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
as Int
acc
        else (Int, [a]) -> ST s (Int, [a])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
acc, [a]
as)
  (Int
len, [a]
as) <- Int -> [a] -> Int -> ST s (Int, [a])
forall s. Int -> [a] -> Int -> ST s (Int, [a])
go Int
0 [] Int
0
  Int -> [a] -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
len [a]
as
{-# inline lefts #-}

-- | Extracts from an array of 'Either' all the 'Right' elements.
-- All the 'Right' elements are extracted in order.
rights :: forall arr a b.
  ( Contiguous arr
  , Element arr b
  , Element arr (Either a b)
  ) => arr (Either a b)
    -> arr b
rights :: arr (Either a b) -> arr b
rights !arr (Either a b)
arr = (forall s. ST s (Mutable arr s b)) -> arr b
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr s b)) -> arr b)
-> (forall s. ST s (Mutable arr s b)) -> arr b
forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = arr (Either a b) -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr (Either a b)
arr
      go :: Int -> [b] -> Int -> ST s (Int, [b])
      go :: Int -> [b] -> Int -> ST s (Int, [b])
go !Int
ix ![b]
bs !Int
acc = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          arr (Either a b) -> Int -> ST s (Either a b)
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr (Either a b)
arr Int
ix ST s (Either a b)
-> (Either a b -> ST s (Int, [b])) -> ST s (Int, [b])
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Left a
_ -> Int -> [b] -> Int -> ST s (Int, [b])
forall s. Int -> [b] -> Int -> ST s (Int, [b])
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [b]
bs Int
acc
            Right b
b -> Int -> [b] -> Int -> ST s (Int, [b])
forall s. Int -> [b] -> Int -> ST s (Int, [b])
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b
bb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
bs) (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        else (Int, [b]) -> ST s (Int, [b])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
acc, [b]
bs)
  (Int
len, [b]
bs) <- Int -> [b] -> Int -> ST s (Int, [b])
forall s. Int -> [b] -> Int -> ST s (Int, [b])
go Int
0 [] Int
0
  Int -> [b] -> ST s (Mutable arr (PrimState (ST s)) b)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
len [b]
bs
{-# inline rights #-}

-- | Partitions an array of 'Either' into two arrays.
-- All the 'Left' elements are extracted, in order, to the first
-- component of the output. Similarly the 'Right' elements are extracted
-- to the second component of the output.
partitionEithers :: forall arr a b.
  ( Contiguous arr
  , Element arr a
  , Element arr b
  , Element arr (Either a b)
  ) => arr (Either a b)
    -> (arr a, arr b)
partitionEithers :: arr (Either a b) -> (arr a, arr b)
partitionEithers !arr (Either a b)
arr = (forall s. ST s (arr a, arr b)) -> (arr a, arr b)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (arr a, arr b)) -> (arr a, arr b))
-> (forall s. ST s (arr a, arr b)) -> (arr a, arr b)
forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = arr (Either a b) -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr (Either a b)
arr
      go :: Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
      go :: Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
go !Int
ix ![a]
as ![b]
bs !Int
accA !Int
accB = if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
        then do
          arr (Either a b) -> Int -> ST s (Either a b)
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr (Either a b)
arr Int
ix ST s (Either a b)
-> (Either a b -> ST s (Int, Int, [a], [b]))
-> ST s (Int, Int, [a], [b])
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Left a
a -> Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
forall s.
Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as) [b]
bs (Int
accA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
accB
            Right b
b -> Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
forall s.
Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
as (b
bb -> [b] -> [b]
forall a. a -> [a] -> [a]
:[b]
bs) Int
accA (Int
accB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          else (Int, Int, [a], [b]) -> ST s (Int, Int, [a], [b])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
accA, Int
accB, [a]
as, [b]
bs)
  (Int
lenA, Int
lenB, [a]
as, [b]
bs) <- Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
forall s.
Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
go Int
0 [] [] Int
0 Int
0
  arr a
arrA <- Mutable arr s a -> ST s (arr a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze (Mutable arr s a -> ST s (arr a))
-> ST s (Mutable arr s a) -> ST s (arr a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> [a] -> ST s (Mutable arr (PrimState (ST s)) a)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
lenA [a]
as
  arr b
arrB <- Mutable arr s b -> ST s (arr b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m (arr b)
unsafeFreeze (Mutable arr s b -> ST s (arr b))
-> ST s (Mutable arr s b) -> ST s (arr b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> [b] -> ST s (Mutable arr (PrimState (ST s)) b)
forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
lenB [b]
bs
  (arr a, arr b) -> ST s (arr a, arr b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (arr a
arrA, arr b
arrB)
{-# inline partitionEithers #-}

-- | 'scanl' is similar to 'foldl', but returns an array of
--   successive reduced values from the left:
--
--   > scanl f z [x1, x2, ...] = [z, f z x1, f (f z x1) x2, ...]
--
--   Note that
--
--   > last (toList (scanl f z xs)) == foldl f z xs.
scanl ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
scanl :: (b -> a -> b) -> b -> arr1 a -> arr2 b
scanl b -> a -> b
f = (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iscanl ((b -> a -> b) -> Int -> b -> a -> b
forall a b. a -> b -> a
const b -> a -> b
f)
{-# inline scanl #-}

-- | A variant of 'scanl' whose function argument takes the current
--   index as an argument.
iscanl ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (Int -> b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
iscanl :: (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iscanl Int -> b -> a -> b
f b
q arr1 a
as = Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
internalScanl (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
as Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> b -> a -> b
f b
q arr1 a
as
{-# inline iscanl #-}

-- | A strictly accumulating version of 'scanl'.
scanl' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
scanl' :: (b -> a -> b) -> b -> arr1 a -> arr2 b
scanl' b -> a -> b
f = (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iscanl' ((b -> a -> b) -> Int -> b -> a -> b
forall a b. a -> b -> a
const b -> a -> b
f)
{-# inline scanl' #-}

-- | A strictly accumulating version of 'iscanl'.
iscanl' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (Int -> b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
iscanl' :: (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iscanl' Int -> b -> a -> b
f !b
q arr1 a
as = Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
internalScanl' (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
as Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> b -> a -> b
f b
q arr1 a
as
{-# inline iscanl' #-}

-- Internal only. The first argument is the size of the array
-- argument. This function helps prevent duplication.
internalScanl ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => Int
    -> (Int -> b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
internalScanl :: Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
internalScanl !Int
sz Int -> b -> a -> b
f !b
q arr1 a
as = (forall s. ST s (Mutable arr2 s b)) -> arr2 b
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr2 s b)) -> arr2 b)
-> (forall s. ST s (Mutable arr2 s b)) -> arr2 b
forall a b. (a -> b) -> a -> b
$ do
  !Mutable arr2 s b
marr <- Int -> ST s (Mutable arr2 (PrimState (ST s)) b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let go :: Int -> b -> ST s ()
go !Int
ix b
acc = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        Mutable arr2 (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s b
Mutable arr2 (PrimState (ST s)) b
marr Int
ix b
acc
        a
x <- arr1 a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
as Int
ix
        Int -> b -> ST s ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> b -> a -> b
f Int
ix b
acc a
x)
  Int -> b -> ST s ()
go Int
0 b
q
  Mutable arr2 s b -> ST s (Mutable arr2 s b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr2 s b
marr
{-# inline internalScanl #-}

-- Internal only. The first argument is the size of the array
-- argument. This function helps prevent duplication.
internalScanl' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => Int
    -> (Int -> b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
internalScanl' :: Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
internalScanl' !Int
sz Int -> b -> a -> b
f !b
q arr1 a
as = (forall s. ST s (Mutable arr2 s b)) -> arr2 b
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr2 s b)) -> arr2 b)
-> (forall s. ST s (Mutable arr2 s b)) -> arr2 b
forall a b. (a -> b) -> a -> b
$ do
  !Mutable arr2 s b
marr <- Int -> ST s (Mutable arr2 (PrimState (ST s)) b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let go :: Int -> b -> ST s ()
go !Int
ix !b
acc = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        Mutable arr2 (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s b
Mutable arr2 (PrimState (ST s)) b
marr Int
ix b
acc
        a
x <- arr1 a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
as Int
ix
        Int -> b -> ST s ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> b -> a -> b
f Int
ix b
acc a
x)
  Int -> b -> ST s ()
go Int
0 b
q
  Mutable arr2 s b -> ST s (Mutable arr2 s b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr2 s b
marr
{-# inline internalScanl' #-}

-- | A prescan.
--
--   @prescanl f z = init . scanl f z@
--
--   Example: @prescanl (+) 0 \<1,2,3,4\> = \<0,1,3,6\>@
prescanl ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
prescanl :: (b -> a -> b) -> b -> arr1 a -> arr2 b
prescanl b -> a -> b
f = (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iprescanl ((b -> a -> b) -> Int -> b -> a -> b
forall a b. a -> b -> a
const b -> a -> b
f)
{-# inline prescanl #-}

-- | A variant of 'prescanl' where the function argument takes
--   the current index of the array as an additional argument.
iprescanl ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (Int -> b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
iprescanl :: (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iprescanl Int -> b -> a -> b
f b
q arr1 a
as = Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
internalScanl (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
as) Int -> b -> a -> b
f b
q arr1 a
as
{-# inline iprescanl #-}

-- | Like 'prescanl', but with a strict accumulator.
prescanl' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
prescanl' :: (b -> a -> b) -> b -> arr1 a -> arr2 b
prescanl' b -> a -> b
f = (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iprescanl ((b -> a -> b) -> Int -> b -> a -> b
forall a b. a -> b -> a
const b -> a -> b
f)
{-# inline prescanl' #-}

-- | Like 'iprescanl', but with a strict accumulator.
iprescanl' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (Int -> b -> a -> b)
    -> b
    -> arr1 a
    -> arr2 b
iprescanl' :: (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
iprescanl' Int -> b -> a -> b
f !b
q arr1 a
as = Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
Int -> (Int -> b -> a -> b) -> b -> arr1 a -> arr2 b
internalScanl' (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
as) Int -> b -> a -> b
f b
q arr1 a
as
{-# inline iprescanl' #-}

-- | 'zipWith' generalises 'zip' by zipping with the function
--   given as the first argument, instead of a tupling function.
--   For example, 'zipWith' (+) is applied to two arrays to produce
--   an array of the corresponding sums.
zipWith ::
  ( Contiguous arr1
  , Contiguous arr2
  , Contiguous arr3
  , Element arr1 a
  , Element arr2 b
  , Element arr3 c
  ) => (a -> b -> c)
    -> arr1 a
    -> arr2 b
    -> arr3 c
zipWith :: (a -> b -> c) -> arr1 a -> arr2 b -> arr3 c
zipWith a -> b -> c
f = (Int -> a -> b -> c) -> arr1 a -> arr2 b -> arr3 c
forall (arr1 :: * -> *) (arr2 :: * -> *) (arr3 :: * -> *) a b c.
(Contiguous arr1, Contiguous arr2, Contiguous arr3, Element arr1 a,
 Element arr2 b, Element arr3 c) =>
(Int -> a -> b -> c) -> arr1 a -> arr2 b -> arr3 c
izipWith (\Int
_ a
a b
b -> a -> b -> c
f a
a b
b)
{-# inline zipWith #-}

-- | Variant of 'zipWith' that provides the index of each pair of elements.
izipWith ::
  ( Contiguous arr1
  , Contiguous arr2
  , Contiguous arr3
  , Element arr1 a
  , Element arr2 b
  , Element arr3 c
  ) => (Int -> a -> b -> c)
    -> arr1 a
    -> arr2 b
    -> arr3 c
izipWith :: (Int -> a -> b -> c) -> arr1 a -> arr2 b -> arr3 c
izipWith Int -> a -> b -> c
f arr1 a
as arr2 b
bs = (forall s. ST s (Mutable arr3 s c)) -> arr3 c
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr3 s c)) -> arr3 c)
-> (forall s. ST s (Mutable arr3 s c)) -> arr3 c
forall a b. (a -> b) -> a -> b
$ do
  let !sz :: Int
sz = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
as) (arr2 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr2 b
bs)
  !Mutable arr3 s c
marr <- Int -> ST s (Mutable arr3 (PrimState (ST s)) c)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let go :: Int -> ST s ()
go !Int
ix = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        a
a <- arr1 a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
as Int
ix
        b
b <- arr2 b -> Int -> ST s b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr2 b
bs Int
ix
        let !g :: c
g = Int -> a -> b -> c
f Int
ix a
a b
b
        Mutable arr3 (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr3 s c
Mutable arr3 (PrimState (ST s)) c
marr Int
ix c
g
        Int -> ST s ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
  Mutable arr3 s c -> ST s (Mutable arr3 s c)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr3 s c
marr
{-# inline izipWith #-}

-- | Variant of 'zipWith' that accepts an accumulator, performing a lazy
-- right fold over both arrays.
foldrZipWith ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (a -> b -> c -> c)
    -> c
    -> arr1 a
    -> arr2 b
    -> c
foldrZipWith :: (a -> b -> c -> c) -> c -> arr1 a -> arr2 b -> c
foldrZipWith a -> b -> c -> c
f = (Int -> a -> b -> c -> c) -> c -> arr1 a -> arr2 b -> c
forall (arr1 :: * -> *) (arr2 :: * -> *) a b c.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> a -> b -> c -> c) -> c -> arr1 a -> arr2 b -> c
ifoldrZipWith (\Int
_ a
x b
y c
c -> a -> b -> c -> c
f a
x b
y c
c)
{-# inline foldrZipWith #-}

-- | Variant of 'zipWith' that accepts an accumulator, performing a strict
-- left monadic fold over both arrays.
foldlZipWithM' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Monad m
  ) => (c -> a -> b -> m c)
    -> c
    -> arr1 a
    -> arr2 b
    -> m c
foldlZipWithM' :: (c -> a -> b -> m c) -> c -> arr1 a -> arr2 b -> m c
foldlZipWithM' c -> a -> b -> m c
f = (Int -> c -> a -> b -> m c) -> c -> arr1 a -> arr2 b -> m c
forall (arr1 :: * -> *) (arr2 :: * -> *) a b (m :: * -> *) c.
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Monad m) =>
(Int -> c -> a -> b -> m c) -> c -> arr1 a -> arr2 b -> m c
ifoldlZipWithM' (\Int
_ c
x a
y b
c -> c -> a -> b -> m c
f c
x a
y b
c)
{-# inline foldlZipWithM' #-}

-- | Variant of 'foldrZipWith' that provides the index of each pair of elements.
ifoldrZipWith ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) => (Int -> a -> b -> c -> c)
    -> c
    -> arr1 a
    -> arr2 b
    -> c
ifoldrZipWith :: (Int -> a -> b -> c -> c) -> c -> arr1 a -> arr2 b -> c
ifoldrZipWith Int -> a -> b -> c -> c
f c
z = \arr1 a
arr1 arr2 b
arr2 ->
  let !sz :: Int
sz = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr1) (arr2 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr2 b
arr2)
      go :: Int -> c
go !Int
ix = if Int
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
ix
        then case arr1 a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr1 a
arr1 Int
ix of
          (# a
x #) -> case arr2 b -> Int -> (# b #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr2 b
arr2 Int
ix of
            (# b
y #) -> Int -> a -> b -> c -> c
f Int
ix a
x b
y (Int -> c
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        else c
z
  in Int -> c
go Int
0
{-# inline ifoldrZipWith #-}

-- | Variant of 'foldlZipWithM\'' that provides the index of each pair of elements.
ifoldlZipWithM' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Monad m
  ) => (Int -> c -> a -> b -> m c)
    -> c
    -> arr1 a
    -> arr2 b
    -> m c
ifoldlZipWithM' :: (Int -> c -> a -> b -> m c) -> c -> arr1 a -> arr2 b -> m c
ifoldlZipWithM' Int -> c -> a -> b -> m c
f c
z = \arr1 a
arr1 arr2 b
arr2 ->
  let !sz :: Int
sz = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (arr1 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr1) (arr2 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr2 b
arr2)
      go :: Int -> c -> m c
go !Int
ix !c
acc = if Int
sz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
ix
        then case arr1 a -> Int -> (# a #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr1 a
arr1 Int
ix of
          (# a
x #) -> case arr2 b -> Int -> (# b #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr2 b
arr2 Int
ix of
            (# b
y #) -> do
              c
acc' <- Int -> c -> a -> b -> m c
f Int
ix c
acc a
x b
y
              Int -> c -> m c
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) c
acc'
        else c -> m c
forall (f :: * -> *) a. Applicative f => a -> f a
pure c
acc
  in Int -> c -> m c
go Int
0 c
z
{-# inline ifoldlZipWithM' #-}

-- | 'zip' takes two arrays and returns an array of
--   corresponding pairs.
--
--   > zip [1, 2] ['a', 'b'] = [(1, 'a'), (2, 'b')]
--
--   If one input array is shorter than the other, excess
--   elements of the longer array are discarded:
--
--   > zip [1] ['a', 'b'] = [(1, 'a')]
--   > zip [1, 2] ['a'] = [(1, 'a')]
--
zip ::
  ( Contiguous arr1
  , Contiguous arr2
  , Contiguous arr3
  , Element arr1 a
  , Element arr2 b
  , Element arr3 (a, b)
  ) => arr1 a
    -> arr2 b
    -> arr3 (a, b)
zip :: arr1 a -> arr2 b -> arr3 (a, b)
zip = (a -> b -> (a, b)) -> arr1 a -> arr2 b -> arr3 (a, b)
forall (arr1 :: * -> *) (arr2 :: * -> *) (arr3 :: * -> *) a b c.
(Contiguous arr1, Contiguous arr2, Contiguous arr3, Element arr1 a,
 Element arr2 b, Element arr3 c) =>
(a -> b -> c) -> arr1 a -> arr2 b -> arr3 c
zipWith (,)
{-# inline zip #-}

-- | Replace all locations in the input with the same value.
--
--   Equivalent to Data.Functor.'Data.Functor.<$'.
(<$) ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 b
  , Element arr2 a
  ) => a -> arr1 b -> arr2 a
a
a <$ :: a -> arr1 b -> arr2 a
<$ arr1 b
barr = (forall s. ST s (Mutable arr2 s a)) -> arr2 a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create (Int -> a -> ST s (Mutable arr2 (PrimState (ST s)) a)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> b -> m (Mutable arr (PrimState m) b)
replicateMut (arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
barr) a
a)
{-# inline (<$) #-}

-- | Sequential application.
--
--   Equivalent to Control.Applicative.'Control.Applicative.<*>'.
ap ::
  ( Contiguous arr1
  , Contiguous arr2
  , Contiguous arr3
  , Element arr1 (a -> b)
  , Element arr2 a
  , Element arr3 b
  ) => arr1 (a -> b) -> arr2 a -> arr3 b
ap :: arr1 (a -> b) -> arr2 a -> arr3 b
ap arr1 (a -> b)
fs arr2 a
xs = (forall s. ST s (Mutable arr3 s b)) -> arr3 b
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(forall s. ST s (Mutable arr s a)) -> arr a
create ((forall s. ST s (Mutable arr3 s b)) -> arr3 b)
-> (forall s. ST s (Mutable arr3 s b)) -> arr3 b
forall a b. (a -> b) -> a -> b
$ do
  Mutable arr3 s b
marr <- Int -> ST s (Mutable arr3 (PrimState (ST s)) b)
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (Int
szfs Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
szxs)
  let go1 :: Int -> ST s ()
go1 !Int
ix = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
szfs) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        a -> b
f <- arr1 (a -> b) -> Int -> ST s (a -> b)
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 (a -> b)
fs Int
ix
        Int -> (a -> b) -> Int -> ST s ()
go2 (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
szxs) a -> b
f Int
0
        Int -> ST s ()
go1 (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      go2 :: Int -> (a -> b) -> Int -> ST s ()
go2 !Int
off a -> b
f !Int
j = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
szxs) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
        a
x <- arr2 a -> Int -> ST s a
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr2 a
xs Int
j
        Mutable arr3 (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (arr :: * -> *) (m :: * -> *) b.
(Contiguous arr, PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> b -> m ()
write Mutable arr3 s b
Mutable arr3 (PrimState (ST s)) b
marr (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (a -> b
f a
x)
        Int -> (a -> b) -> Int -> ST s ()
go2 Int
off a -> b
f (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go1 Int
0
  Mutable arr3 s b -> ST s (Mutable arr3 s b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr3 s b
marr
  where
    !szfs :: Int
szfs = arr1 (a -> b) -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 (a -> b)
fs
    !szxs :: Int
szxs = arr2 a -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr2 a
xs
{-# inline ap #-}

all :: (Contiguous arr, Element arr a) => (a -> Bool) -> arr a -> Bool
all :: (a -> Bool) -> arr a -> Bool
all a -> Bool
f = (a -> Bool -> Bool) -> Bool -> arr a -> Bool
forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(a -> b -> b) -> b -> arr a -> b
foldr (\a
x Bool
acc -> a -> Bool
f a
x Bool -> Bool -> Bool
&& Bool
acc) Bool
True
{-# inline all #-}

any :: (Contiguous arr, Element arr a) => (a -> Bool) -> arr a -> Bool
any :: (a -> Bool) -> arr a -> Bool
any a -> Bool
f = (a -> Bool -> Bool) -> Bool -> arr a -> Bool
forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(a -> b -> b) -> b -> arr a -> b
foldr (\a
x Bool
acc -> a -> Bool
f a
x Bool -> Bool -> Bool
|| Bool
acc) Bool
False
{-# inline any #-}