{-# 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
  , quintupleton
  , sextupleton
  , replicate
  , replicateMut
  , generate
  , generateM
  , generateMutable
  , iterateN
  , iterateMutableN
  , write

    -- ** Fixed Length
  , construct1
  , construct2
  , construct3
  , construct4
  , construct5
  , construct6

    -- ** 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'
  , foldrM'
  , asum
  , all
  , any

    -- ** Zipping Folds
  , foldrZipWith
  , ifoldrZipWith
  , foldlZipWith'
  , ifoldlZipWith'
  , 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
  , SmallUnliftedArray
  , SmallMutableUnliftedArray
  ) where

import Control.Monad.Primitive
import Data.Primitive
import Data.Primitive.Unlifted.Array
import Data.Primitive.Unlifted.SmallArray
import Prelude hiding (Foldable (..), all, any, filter, map, mapM, mapM_, read, replicate, reverse, scanl, sequence, sequence_, traverse, zip, zipWith, (<$))

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

import qualified Control.Applicative as A
import qualified Prelude

construct1 ::
  (Contiguous arr, Element arr a) =>
  a ->
  arr a
{-# INLINE construct1 #-}
construct1 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> arr a
construct1 = a -> arr a
forall a. Element arr a => a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> arr a
singleton

construct2 ::
  (Contiguous arr, Element arr a) =>
  a ->
  a ->
  arr a
{-# INLINE construct2 #-}
construct2 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> arr a
construct2 = a -> a -> arr a
forall a. Element arr a => a -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> arr a
doubleton

construct3 ::
  (Contiguous arr, Element arr a) =>
  a ->
  a ->
  a ->
  arr a
{-# INLINE construct3 #-}
construct3 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> arr a
construct3 = a -> a -> a -> arr a
forall a. Element arr a => a -> a -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> arr a
tripleton

construct4 ::
  (Contiguous arr, Element arr a) =>
  a ->
  a ->
  a ->
  a ->
  arr a
{-# INLINE construct4 #-}
construct4 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> a -> arr a
construct4 = a -> a -> a -> a -> arr a
forall a. Element arr a => a -> a -> a -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> a -> arr a
quadrupleton

construct5 ::
  (Contiguous arr, Element arr a) =>
  a ->
  a ->
  a ->
  a ->
  a ->
  arr a
{-# INLINE construct5 #-}
construct5 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> a -> a -> arr a
construct5 = a -> a -> a -> a -> a -> arr a
forall a. Element arr a => a -> a -> a -> a -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> a -> a -> arr a
quintupleton

construct6 ::
  (Contiguous arr, Element arr a) =>
  a ->
  a ->
  a ->
  a ->
  a ->
  a ->
  arr a
{-# INLINE construct6 #-}
construct6 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> a -> a -> a -> arr a
construct6 = a -> a -> a -> a -> a -> a -> arr a
forall a. Element arr a => a -> a -> a -> a -> a -> a -> arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
a -> a -> a -> a -> a -> a -> arr a
sextupleton

-- | Append two arrays.
append :: (Contiguous arr, Element arr a) => arr a -> arr a -> arr a
append :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> arr a -> arr a
append !arr a
a !arr a
b = (forall s. ST s (arr a)) -> arr a
forall a. (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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
b)
  copy m 0 (toSlice a)
  copy m (size a) (toSlice b)
  unsafeFreeze m
{-# INLINE append #-}

-- | Delete the element at the given position.
deleteAt :: (Contiguous arr, Element arr a) => arr a -> Int -> arr a
deleteAt :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Int -> arr a
deleteAt arr a
src Int
i = (forall s. ST s (arr a)) -> arr a
forall a. (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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Sliced arr b -> m (Mutable arr (PrimState m) b)
thaw (arr a -> Int -> Int -> Sliced arr a
forall a. Element arr a => 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 b. Element arr b => arr b -> 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 -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
  copy dst i (slice src i' (size src - i'))
  unsafeFreeze 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Sliced arr b -> m (Mutable arr (PrimState m) b)
thaw (arr a -> Sliced arr a
forall a. Element arr a => arr a -> Sliced arr a
forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> Sliced arr a
toSlice arr a
src)
  write dst i x
  pure dst
{-# INLINE replaceAt #-}

modifyAt ::
  (Contiguous arr, Element arr a) =>
  (a -> a) ->
  arr a ->
  Int ->
  arr a
modifyAt :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> Int -> b
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' :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> Int -> b
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 :: forall (arr :: * -> *) a (f :: * -> *).
(Contiguous arr, Element arr a, Functor f) =>
(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 b. Element arr b => arr b -> Int -> b
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' :: forall (arr :: * -> *) a (f :: * -> *).
(Contiguous arr, Element arr a, Monad f) =>
(a -> f a) -> arr a -> Int -> f (arr a)
modifyAtF' a -> f a
f arr a
src Int
i = do
  !r <- a -> f a
f (arr a -> Int -> a
forall b. Element arr b => arr b -> Int -> b
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
src Int
i)
  let !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
  pure 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 :: forall (arr1 :: * -> *) b (arr2 :: * -> *) c.
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 c) =>
(Int -> b -> c) -> arr1 b -> arr2 c
imap Int -> b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall a. (forall s. ST s (arr2 a)) -> arr2 a
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new (arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            x <- arr1 b -> Int -> ST s b
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 b
a Int
i
            write mb i (f i x)
            go (i + 1)
  go 0
  unsafeFreeze 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' :: forall (arr1 :: * -> *) b (arr2 :: * -> *) c.
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 c) =>
(Int -> b -> c) -> arr1 b -> arr2 c
imap' Int -> b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall a. (forall s. ST s (arr2 a)) -> arr2 a
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new (arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            x <- arr1 b -> Int -> ST s b
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m 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 = Int -> b -> c
f Int
i b
x
            write mb i b
            go (i + 1)
  go 0
  unsafeFreeze 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 :: forall (arr1 :: * -> *) b (arr2 :: * -> *) c.
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 c) =>
(b -> c) -> arr1 b -> arr2 c
map b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall a. (forall s. ST s (arr2 a)) -> arr2 a
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new (arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            x <- arr1 b -> Int -> ST s b
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 b
a Int
i
            write mb i (f x)
            go (i + 1)
  go 0
  unsafeFreeze 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' :: forall (arr1 :: * -> *) b (arr2 :: * -> *) c.
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 c) =>
(b -> c) -> arr1 b -> arr2 c
map' b -> c
f arr1 b
a = (forall s. ST s (arr2 c)) -> arr2 c
forall a. (forall s. ST s (arr2 a)) -> arr2 a
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new (arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a)
  let go !Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== arr1 b -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 b
a = () -> ST s ()
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            x <- arr1 b -> Int -> ST s b
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m 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 = b -> c
f b
x
            write mb i b
            go (i + 1)
  go 0
  unsafeFreeze 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 :: forall (arr1 :: * -> *) b (arr2 :: * -> *).
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 b) =>
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 :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> Int -> (# b #)
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 b. Element arr b => arr b -> 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 :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 :: forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Monoid m) =>
(a -> m) -> arr a -> m
foldMap a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Monoid m) =>
(a -> m) -> arr a -> m
foldMap' a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Monoid m) =>
(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' :: forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Monoid m) =>
(Int -> a -> m) -> arr a -> m
ifoldlMap' Int -> a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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' :: forall (arr :: * -> *) a m.
(Contiguous arr, Element arr a, Semigroup m) =>
(Int -> a -> m) -> arr a -> m
ifoldlMap1' Int -> a -> m
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 b. Element arr b => arr b -> Int -> (# b #)
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 right monadic fold over the elements of an array.
foldrM' ::
  (Contiguous arr, Element arr a, Monad m) =>
  (a -> b -> m b) ->
  b ->
  arr a ->
  m b
foldrM' :: forall (arr :: * -> *) a (m :: * -> *) b.
(Contiguous arr, Element arr a, Monad m) =>
(a -> b -> m b) -> b -> arr a -> m b
foldrM' a -> b -> m b
f !b
z0 = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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
0
          then do
            let (# a
x #) = arr a -> Int -> (# a #)
forall b. Element arr b => arr b -> Int -> (# b #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
            acc2 <- a -> b -> m b
f a
x b
acc1
            go (ix - 1) acc2
          else b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
acc1
   in Int -> b -> m b
go (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) b
z0
{-# INLINE foldrM' #-}

-- | 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' :: forall (arr :: * -> *) a (m :: * -> *) b.
(Contiguous arr, Element arr a, Monad m) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
            acc2 <- b -> a -> m b
f b
acc1 a
x
            go (ix + 1) acc2
          else b -> m b
forall a. a -> m a
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' :: forall (arr :: * -> *) a (m :: * -> *) b.
(Contiguous arr, Element arr a, Monad m) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> (# b #)
index# arr a
arr Int
ix
            acc2 <- b -> Int -> a -> m b
f b
acc1 Int
ix a
x
            go (ix + 1) acc2
          else b -> m b
forall a. a -> m a
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(Int -> a -> Bool) -> arr a -> arr a
ifilter Int -> a -> Bool
p arr a
arr = (forall s. ST s (arr a)) -> arr a
forall a. (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
  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
ix !Int
numTrue =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
          then do
            atIx <- arr a -> Int -> ST s a
forall b (m :: * -> *).
(Element arr b, Monad m) =>
arr b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr a
arr Int
ix
            let !keep = Int -> a -> Bool
p Int
ix a
atIx
            let !keepTag = Int# -> Int
I# (Bool -> Int#
forall a. DataToTag a => a -> Int#
dataToTag# Bool
keep)
            writePrimArray marr ix (fromIntegral keepTag)
            go1 (ix + 1) (numTrue + keepTag)
          else Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
numTrue
  numTrue <- go1 0 0
  if numTrue == sz
    then pure arr
    else do
      marrTrues <- new numTrue
      let 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
            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 isTrue atIxKeep
              then do
                atIxVal <- indexM arr ixSrc
                write marrTrues ixDst atIxVal
                go2 (ixSrc + 1) (ixDst + 1)
              else go2 (ixSrc + 1) ixDst
      go2 0 0
      unsafeFreeze marrTrues
 where
  !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 :: 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
f arr1 a
arr = (forall s. ST s (arr2 b)) -> arr2 b
forall a. (forall s. ST s (arr2 a)) -> arr2 a
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 b. Element arr1 b => arr1 b -> 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 :: forall s. 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
            atIx <- arr1 a -> Int -> ST s a
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix
            case f atIx of
              Maybe b
Nothing -> Int -> Int -> [b] -> ST s ([b], Int)
forall s. Int -> Int -> [b] -> ST s ([b], Int)
go (Int
ix Int -> 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
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
numJusts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b
x b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
justs)
          else ([b], Int) -> ST s ([b], Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([b]
justs, Int
numJusts)
  !(bs, !numJusts) <- Int -> Int -> [b] -> ST s ([b], Int)
forall s. Int -> Int -> [b] -> ST s ([b], Int)
go Int
0 Int
0 []
  !marr <- unsafeFromListReverseMutableN numJusts bs
  unsafeFreeze 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a, Element arr (Maybe a)) =>
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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)
forall (m :: * -> *) b.
(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 :: forall (m :: * -> *) (arr :: * -> *) a.
(PrimMonad m, Contiguous arr, Element arr a) =>
Int -> m a -> m (Mutable arr (PrimState m) a)
replicateMutM Int
len m a
act = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  let 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
        x <- m a
act
        write marr ix x
        go (ix + 1)
  go 0
  pure 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) =>
  -- | length of list
  Int ->
  -- | list
  [a] ->
  arr a
unsafeFromListN :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListMutableN Int
n [a]
l = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
n
  let go !Int
_ [] = Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall a. a -> 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 ()
forall (m :: * -> *) b.
(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
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs
  go 0 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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
unsafeFromListReverseMutableN Int
n [a]
l = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
n
  let go !Int
_ [] = Mutable arr (PrimState m) a -> m (Mutable arr (PrimState m) a)
forall a. a -> 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 ()
forall (m :: * -> *) b.
(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
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs
  go (n - 1) 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
(a -> a) -> Mutable arr (PrimState m) a -> m ()
mapMutable a -> a
f !Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let 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 <- 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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        write marr ix (f a)
        go (ix + 1)
  go 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' :: forall (m :: * -> *) (arr :: * -> *) a.
(PrimMonad m, Contiguous arr, Element arr a) =>
(a -> a) -> Mutable arr (PrimState m) a -> m ()
mapMutable' a -> a
f !Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let 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 <- 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
forall (m :: * -> *) b.
(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 -> a
f a
a
        write marr ix b
        go (ix + 1)
  go 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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
(Int -> a -> a) -> Mutable arr (PrimState m) a -> m ()
imapMutable Int -> a -> a
f !Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let 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 <- 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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        write marr ix (f ix a)
        go (ix + 1)
  go 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' :: forall (m :: * -> *) (arr :: * -> *) a.
(PrimMonad m, Contiguous arr, Element arr a) =>
(Int -> a -> a) -> Mutable arr (PrimState m) a -> m ()
imapMutable' Int -> a -> a
f !Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let 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 <- 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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        let !b = Int -> a -> a
f Int
ix a
a
        write marr ix b
        go (ix + 1)
  go 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 :: forall (m :: * -> *) (arr1 :: * -> *) a (arr2 :: * -> *) b.
(PrimMonad m, Contiguous arr1, Element arr1 a, Contiguous arr2,
 Element arr2 b) =>
(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 b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
  !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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new Int
sz
  let 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 <- arr1 a -> Int -> m a
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix
        b <- f a
        write marr ix b
        go (ix + 1)
  go 0
  unsafeFreeze 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 :: forall (m :: * -> *) (arr1 :: * -> *) a (arr2 :: * -> *) b.
(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
f !arr1 a
arr = do
  let !sz :: Int
sz = arr1 a -> Int
forall b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr
  !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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new Int
sz
  let 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 <- arr1 a -> Int -> m a
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
arr Int
ix
        b <- f ix a
        write marr ix b
        go (ix + 1)
  go 0
  unsafeFreeze marr
{-# INLINE itraverseP #-}

newtype STA v a = STA {forall (v :: * -> *) a.
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 :: forall (v :: * -> *) a.
(Contiguous v, Element v a) =>
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)
forall (m :: * -> *) b.
(PrimMonad m, Element v b) =>
Int -> m (Mutable v (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 a b. ST s a -> (a -> ST s b) -> ST s b
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 :: 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 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_ :: forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
(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 :: 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 Int -> a -> f b
f = \arr1 a
arr ->
  let !sz :: Int
sz = arr1 a -> Int
forall b. Element arr1 b => arr1 b -> 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 a. a -> f a
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 Mutable arr2 s b -> ST s (arr2 b)
Mutable arr2 (PrimState (ST s)) b -> ST s (arr2 b)
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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Mutable arr2 (PrimState m) b -> m (arr2 b)
unsafeFreeze)
          else case arr1 a -> Int -> (# a #)
forall b. Element arr1 b => arr1 b -> Int -> (# b #)
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 a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
A.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 ()
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Mutable arr2 (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 a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure arr2 b
forall a. arr2 a
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_ :: forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
(Int -> a -> f b) -> arr a -> f ()
itraverse_ Int -> a -> f b
f = \arr a
arr ->
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> b
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr a
arr Int
ix) f b -> f () -> f ()
forall a b. f a -> f b -> f b
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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b (f :: * -> *).
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Applicative f) =>
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_ :: forall (arr :: * -> *) a (f :: * -> *) b.
(Contiguous arr, Element arr a, Applicative f) =>
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' :: forall (arr1 :: * -> *) (arr2 :: * -> *) b c (m :: * -> *) a.
(Contiguous arr1, Contiguous arr2, Element arr1 b, Element arr2 c,
 Monad m) =>
(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 b. Element arr1 b => 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
        (!acc', !x) <- a -> b -> m (a, c)
f a
acc (arr1 b -> Int -> b
forall b. Element arr1 b => arr1 b -> Int -> b
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int -> b
index arr1 b
src Int
ix)
        go (ix + 1) (x : xs) 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 a. a -> m a
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' :: 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)
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new Int
sz
  acc <- go 0 dst mempty
  dst' <- unsafeFreeze dst
  pure (acc, dst')
 where
  !sz :: Int
sz = arr1 b -> Int
forall b. Element arr1 b => 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 :: forall s. 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 b. Element arr1 b => 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 ()
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Mutable arr2 (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 a. 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 :: 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 a -> m b
f arr1 a
arr =
  let !sz :: Int
sz = arr1 a -> Int
forall b. Element arr1 b => arr1 b -> 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 b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
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 a b. 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_ :: forall (arr :: * -> *) a b (f :: * -> *).
(Contiguous arr, Element arr a, Element arr b, Applicative f) =>
(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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b (m :: * -> *).
(Contiguous arr1, Contiguous arr2, Element arr1 a, Element arr2 b,
 Monad m) =>
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_ :: forall (arr :: * -> *) a b (f :: * -> *).
(Contiguous arr, Element arr a, Element arr b, Applicative f) =>
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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) (f :: * -> *) a.
(Contiguous arr1, Contiguous arr2, Element arr1 (f a),
 Element arr2 a, Applicative f) =>
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_ :: forall (arr :: * -> *) (f :: * -> *) a.
(Contiguous arr, Element arr (f a), Applicative f) =>
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 a b. f a -> f b -> f b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>) (() -> f ()
forall a. a -> f a
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 :: forall (arr :: * -> *) (f :: * -> *) a.
(Contiguous arr, Element arr (f a), Alternative f) =>
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 a. f a -> f a -> f a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(A.<|>) f a
forall 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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)
{-# INLINE generateM #-}
generateM :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, Monad m) =>
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 a b c. (a -> b -> c) -> m a -> m b -> m c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
A.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 ()
forall (m :: * -> *) b.
(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 a. a -> m 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 Mutable arr s a -> ST s (arr a)
Mutable arr (PrimState (ST s)) a -> ST s (arr a)
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)
forall (m :: * -> *) b.
(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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure arr a
forall a. 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 :: 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 = 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 a. 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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> (Int -> m a) -> m (Mutable arr (PrimState m) a)
generateMutableM !Int
len Int -> m a
f = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  let 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
        x <- Int -> m a
f Int
ix
        write marr ix x
        go (ix + 1)
  go 0
  pure 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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 a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable arr s a -> ST s (arr a)
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)
forall (m :: * -> *) b.
(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 :: 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 = 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 a. 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 :: 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
f a
z0 = do
  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)
forall (m :: * -> *) b.
(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
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 ()
forall (m :: * -> *) b.
(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 a b. m a -> m b -> m b
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        | Bool
otherwise = do
            a <- a -> m a
f a
acc
            write marr ix a
            go (ix + 1) a
  go 0 z0
  pure 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 (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)
x = (forall s. ST s (arr a)) -> arr a
forall a. (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)
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)
forall (m :: * -> *) b.
(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 (arr :: * -> *) a (f :: * -> *).
(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))
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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> f a -> m (f b)
Prelude.mapM Mutable arr s a -> ST s (arr a)
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)
forall (m :: * -> *) b.
(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 :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
(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 :: 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 = 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 a. a -> m 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)
  (sz, xs) <- Int -> b -> [a] -> m (Int, [a])
go Int
0 b
z0 []
  unsafeFromListReverseMutableN sz 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 :: forall (arr :: * -> *) a b.
(Contiguous arr, Element arr a) =>
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 :: 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 = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
maxSz
  let 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 a. a -> m a
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 ()
forall (m :: * -> *) b.
(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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
ix
  sz <- go 0 z0
  shrink m sz
{-# INLINE unfoldrMutableN #-}

-- | Convert an array to a list.
toList ::
  (Contiguous arr, Element arr a) =>
  arr a ->
  [a]
toList :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Mutable arr (PrimState m) a -> m [a]
toListMutable Mutable arr (PrimState m) a
marr = do
  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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let go !Int
ix ![a]
acc =
        if Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
          then do
            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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
            go (ix - 1) (x : acc)
          else [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
acc
  go (sz - 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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Int -> [a] -> m (Mutable arr (PrimState m) a)
fromListMutableN Int
len [a]
vs = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
len
  let go [] !Int
ix =
        if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
          then () -> m ()
forall a. a -> m a
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 ()
forall (m :: * -> *) b.
(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."
  go vs 0
  pure 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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
[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 a. [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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
[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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
(a -> a) -> Mutable arr (PrimState m) a -> m ()
modify a -> a
f Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let 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
        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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix
        write marr ix (f x)
        go (ix + 1)
  go 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' :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
(a -> a) -> Mutable arr (PrimState m) a -> m ()
modify' a -> a
f Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  let 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
        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
forall (m :: * -> *) b.
(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 -> a
f a
x
        write marr ix y
        go (ix + 1)
  go 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a, Enum a) =>
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 :: 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 = do
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new Int
sz
  let 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 ()
forall (m :: * -> *) b.
(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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Mutable arr (PrimState m) a
m
  go 0 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
arr a -> arr a
reverse arr a
arr = (forall s. ST s (arr a)) -> arr a
forall a. (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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Int -> m (Mutable arr (PrimState m) b)
new (arr a -> Int
forall b. Element arr b => arr b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr a
arr)
  copy marr 0 (toSlice arr)
  reverseMutable marr
  unsafeFreeze 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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Mutable arr (PrimState m) a -> m ()
reverseMutable Mutable arr (PrimState m) a
marr = do
  !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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> m Int
sizeMut Mutable arr (PrimState m) a
marr
  reverseSlice marr 0 (sz - 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 ->
  -- | start index
  Int ->
  -- | end index
  Int ->
  m ()
reverseSlice :: 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
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          else do
            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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
s
            write marr s =<< read marr e
            write marr e tmp
            go (s + 1) (e - 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 :: forall (arr :: * -> *) a. ContiguousU arr => arr a -> arr a -> Bool
same arr a
a arr a
b =
  Int# -> Bool
isTrue#
    ( MutableArrayArray# (ZonkAny 0)
-> MutableArrayArray# (ZonkAny 0) -> Int#
forall s. MutableArrayArray# s -> MutableArrayArray# s -> Int#
sameMutableArrayArray#
        (Unlifted arr a -> MutableArrayArray# s
forall a b. a -> b
unsafeCoerce# (arr a -> Unlifted arr a
forall b. arr b -> Unlifted arr b
forall (arr :: * -> *) b.
ContiguousU arr =>
arr b -> Unlifted arr b
unlift arr a
a) :: MutableArrayArray# s)
        (Unlifted arr a -> MutableArrayArray# s
forall a b. a -> b
unsafeCoerce# (arr a -> Unlifted arr a
forall b. arr b -> Unlifted arr b
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a, Eq a) =>
a -> arr a -> Bool
elem a
a !arr a
arr =
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a, Ord a) =>
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a, Ord a) =>
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(a -> a -> Ordering) -> arr a -> Maybe a
maximumBy a -> a -> Ordering
f arr a
arr =
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 b. Element arr b => arr b -> Int -> b
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(a -> a -> Ordering) -> arr a -> Maybe a
minimumBy a -> a -> Ordering
f arr a
arr =
  let !sz :: Int
sz = arr a -> Int
forall b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> (# b #)
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 b. Element arr b => arr b -> Int -> b
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(a -> Bool) -> arr a -> Maybe a
find a -> Bool
p = Maybe (First a) -> Maybe a
forall a b. Coercible a b => a -> b
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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 b. Element arr b => arr b -> 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 b. Element arr b => arr b -> Int -> b
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 :: forall (arr :: * -> *) a (m :: * -> *).
(Contiguous arr, Element arr a, PrimMonad m) =>
Mutable arr (PrimState m) a -> Int -> Int -> m ()
swap !Mutable arr (PrimState m) a
marr !Int
ix1 !Int
ix2 = do
  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
forall (m :: * -> *) b.
(PrimMonad m, Element arr b) =>
Mutable arr (PrimState m) b -> Int -> m b
read Mutable arr (PrimState m) a
marr Int
ix1
  atIx2 <- read marr ix2
  write marr ix1 atIx2
  write marr ix2 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 :: 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 = (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 b. Element arr b => arr 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 :: forall s. 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 b (m :: * -> *).
(Element arr b, Monad m) =>
arr b -> Int -> m 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 a b. ST s a -> (a -> ST s b) -> ST s b
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
a a -> [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 a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
acc, [a]
as)
  (len, as) <- Int -> [a] -> Int -> ST s (Int, [a])
forall s. Int -> [a] -> Int -> ST s (Int, [a])
go Int
0 [] Int
0
  unsafeFromListReverseMutableN len 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 :: 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 = (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 b. Element arr b => arr 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 :: forall s. 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 b (m :: * -> *).
(Element arr b, Monad m) =>
arr b -> Int -> m 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 a b. ST s a -> (a -> ST s b) -> ST s 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
b b -> [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 a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
acc, [b]
bs)
  (len, bs) <- Int -> [b] -> Int -> ST s (Int, [b])
forall s. Int -> [b] -> Int -> ST s (Int, [b])
go Int
0 [] Int
0
  unsafeFromListReverseMutableN len 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 :: 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 = (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 b. Element arr b => arr 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 :: forall s.
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 b (m :: * -> *).
(Element arr b, Monad m) =>
arr b -> Int -> m 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 a b. ST s a -> (a -> ST s b) -> ST s 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
a a -> [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
b b -> [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 a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
accA, Int
accB, [a]
as, [b]
bs)
  (lenA, lenB, as, 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
  arrA <- unsafeFreeze =<< unsafeFromListReverseMutableN lenA as
  arrB <- unsafeFreeze =<< unsafeFromListReverseMutableN lenB bs
  pure (arrA, 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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(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 :: 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 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 b. Element arr1 b => arr1 b -> 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' :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(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' :: 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' 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 b. Element arr1 b => arr1 b -> 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 :: 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 !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
  !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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new Int
sz
  let 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 ()
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Mutable arr2 (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s b
Mutable arr2 (PrimState (ST s)) b
marr Int
ix b
acc
        x <- arr1 a -> Int -> ST s a
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
as Int
ix
        go (ix + 1) (f ix acc x)
  go 0 q
  pure 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' :: 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' !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
  !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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> m (Mutable arr2 (PrimState m) b)
new Int
sz
  let 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 ()
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Mutable arr2 (PrimState m) b -> Int -> b -> m ()
write Mutable arr2 s b
Mutable arr2 (PrimState (ST s)) b
marr Int
ix b
acc
        x <- arr1 a -> Int -> ST s a
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
as Int
ix
        go (ix + 1) (f ix acc x)
  go 0 q
  pure 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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(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 :: 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 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 b. Element arr1 b => arr1 b -> 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' :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(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' :: 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' 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 b. Element arr1 b => arr1 b -> 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 :: 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 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 :: 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 -> 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 b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
as) (arr2 b -> Int
forall b. Element arr2 b => arr2 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr2 b
bs)
  !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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr3 b) =>
Int -> m (Mutable arr3 (PrimState m) b)
new Int
sz
  let 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 <- arr1 a -> Int -> ST s a
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 a
as Int
ix
        b <- indexM bs ix
        let !g = Int -> a -> b -> c
f Int
ix a
a b
b
        write marr ix g
        go (ix + 1)
  go 0
  pure 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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b c.
(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
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' :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b (m :: * -> *) c.
(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
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 :: 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 -> 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 b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr1) (arr2 b -> Int
forall b. Element arr2 b => 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 b. Element arr1 b => arr1 b -> Int -> (# b #)
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 b. Element arr2 b => 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 #-}

foldlZipWith' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) =>
  (c -> a -> b -> c) ->
  c ->
  arr1 a ->
  arr2 b ->
  c
foldlZipWith' :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b c.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(c -> a -> b -> c) -> c -> arr1 a -> arr2 b -> c
foldlZipWith' c -> a -> b -> c
f = (Int -> c -> a -> b -> c) -> c -> arr1 a -> arr2 b -> c
forall (arr1 :: * -> *) (arr2 :: * -> *) a b c.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> c -> a -> b -> c) -> c -> arr1 a -> arr2 b -> c
ifoldlZipWith' (\Int
_ c
x a
y b
c -> c -> a -> b -> c
f c
x a
y b
c)
{-# INLINE foldlZipWith' #-}

ifoldlZipWith' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  ) =>
  (Int -> c -> a -> b -> c) ->
  c ->
  arr1 a ->
  arr2 b ->
  c
ifoldlZipWith' :: forall (arr1 :: * -> *) (arr2 :: * -> *) a b c.
(Contiguous arr1, Contiguous arr2, Element arr1 a,
 Element arr2 b) =>
(Int -> c -> a -> b -> c) -> c -> arr1 a -> arr2 b -> c
ifoldlZipWith' Int -> c -> a -> b -> 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 b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr1) (arr2 b -> Int
forall b. Element arr2 b => arr2 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr2 b
arr2)
      go :: Int -> c -> c
go !Int
ix !c
acc =
        if Int
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
          then c
acc
          else case arr1 a -> Int -> (# a #)
forall b. Element arr1 b => arr1 b -> Int -> (# b #)
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 b. Element arr2 b => 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 -> c -> c
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> c -> a -> b -> c
f Int
ix c
acc a
x b
y)
   in Int -> c -> c
go Int
0 c
z
{-# INLINE ifoldlZipWith' #-}

-- | 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' :: 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 -> 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 b. Element arr1 b => arr1 b -> Int
forall (arr :: * -> *) b.
(Contiguous arr, Element arr b) =>
arr b -> Int
size arr1 a
arr1) (arr2 b -> Int
forall b. Element arr2 b => 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 b. Element arr1 b => arr1 b -> Int -> (# b #)
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 b. Element arr2 b => 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
                acc' <- Int -> c -> a -> b -> m c
f Int
ix c
acc a
x b
y
                go (ix + 1) acc'
          else c -> m c
forall a. a -> m a
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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) (arr3 :: * -> *) a b.
(Contiguous arr1, Contiguous arr2, Contiguous arr3, Element arr1 a,
 Element arr2 b, Element arr3 (a, b)) =>
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 <$ :: forall (arr1 :: * -> *) (arr2 :: * -> *) b a.
(Contiguous arr1, Contiguous arr2, Element arr1 b,
 Element arr2 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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr2 b) =>
Int -> b -> m (Mutable arr2 (PrimState m) b)
replicateMut (arr1 b -> Int
forall b. Element arr1 b => 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 :: forall (arr1 :: * -> *) (arr2 :: * -> *) (arr3 :: * -> *) a b.
(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)
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
  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)
forall (m :: * -> *) b.
(PrimMonad m, Element arr3 b) =>
Int -> m (Mutable arr3 (PrimState m) b)
new (Int
szfs Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
szxs)
  let 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
        f <- arr1 (a -> b) -> Int -> ST s (a -> b)
forall b (m :: * -> *).
(Element arr1 b, Monad m) =>
arr1 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr1 (a -> b)
fs Int
ix
        go2 (ix * szxs) f 0
        go1 (ix + 1)
      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
        x <- arr2 a -> Int -> ST s a
forall b (m :: * -> *).
(Element arr2 b, Monad m) =>
arr2 b -> Int -> m b
forall (arr :: * -> *) b (m :: * -> *).
(Contiguous arr, Element arr b, Monad m) =>
arr b -> Int -> m b
indexM arr2 a
xs Int
j
        write marr (off + j) (f x)
        go2 off f (j + 1)
  go1 0
  pure marr
 where
  !szfs :: Int
szfs = arr1 (a -> b) -> Int
forall b. Element arr1 b => arr1 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 b. Element arr2 b => arr2 b -> 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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 :: forall (arr :: * -> *) a.
(Contiguous arr, Element arr a) =>
(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 #-}