contiguous-0.5.2: src/Data/Primitive/Contiguous.hs

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

-- | The contiguous typeclass parameterises over a contiguous array type.
--   This allows us to have a common API to a number of contiguous
--   array types and their mutable counterparts.
module Data.Primitive.Contiguous
  (
    -- * Accessors
    -- ** Length Information
    size
  , sizeMutable
  , null
    -- ** Indexing
  , index
  , index#
  , read
    -- ** Monadic indexing
  , indexM

    -- * Construction
    -- ** Initialisation
  , empty
  , new
  , singleton
  , doubleton
  , tripleton
  , quadrupleton
  , replicate
  , replicateMutable
  , generate
  , generateM
  , generateMutable
  , iterateN
  , iterateMutableN
  , write
    -- ** Running
  , run
    -- ** Monadic initialisation
  , replicateMutableM
  , generateMutableM
  , iterateMutableNM
  , create
  , createT
    -- ** Unfolding
  , unfoldr
  , unfoldrN
  , unfoldrMutable
    -- ** Enumeration
  , enumFromN
  , enumFromMutableN
    -- ** Concatenation
  , append
    -- ** Splitting and Splicing
  , insertAt
  , insertSlicing
    -- * Modifying arrays
  , replaceAt
  , modifyAt
  , modifyAt'
  , modifyAtF
  , modifyAtF'
    -- ** Permutations
  , reverse
  , reverseMutable
  , reverseSlice

    -- ** Resizing
  , resize

    -- * 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
  , equalsMutable
  , same

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

    -- * Traversals
  , traverse
  , traverse_
  , itraverse
  , itraverse_
  , traverseP
  , 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
  , unlift
    -- ** Between mutable and immutable variants
  , clone
  , cloneMutable
  , copy
  , copyMutable
  , freeze
  , thaw
  , unsafeFreeze

    -- * Hashing
  , liftHashWithSalt

    -- * Forcing an array and its contents
  , rnf

    -- * Classes
  , Contiguous(Mutable,Element)
  , Always

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

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

import Control.Applicative (liftA2)
import Control.DeepSeq (NFData)
import Control.Monad (when)
import Control.Monad.ST (runST,ST)
import Control.Monad.ST.Run (runPrimArrayST,runSmallArrayST,runUnliftedArrayST,runArrayST)
import Data.Bits (xor)
import Data.Coerce (coerce)
import Data.Kind (Type)
import Data.Primitive.Unlifted.Class (PrimUnlifted)
import Data.Semigroup (First(..))
import Data.Word (Word8)
import GHC.Base (build)
import GHC.Exts (MutableArrayArray#,ArrayArray#,Constraint,sizeofByteArray#,sizeofArray#,sizeofArrayArray#,unsafeCoerce#,sameMutableArrayArray#,isTrue#,dataToTag#,Int(..))

import qualified Control.Applicative as A
import qualified Control.DeepSeq as DS
import qualified Prelude

-- | A typeclass that is satisfied by all types. This is used
-- used to provide a fake constraint for 'Array' and 'SmallArray'.
class Always a
instance Always a

-- | The 'Contiguous' typeclass as an interface to a multitude of
--   contiguous structures.
class Contiguous (arr :: Type -> Type) where
  -- | The Mutable counterpart to the array.
  type family Mutable arr = (r :: Type -> Type -> Type) | r -> arr
  -- | The constraint needed to store elements in the array.
  type family Element arr :: Type -> Constraint
  -- | The empty array.
  empty :: arr a
  -- | Test whether the array is empty.
  null :: arr b -> Bool
  -- | Allocate a new mutable array of the given size.
  new :: (PrimMonad m, Element arr b) => Int -> m (Mutable arr (PrimState m) b)
  -- | @'replicateMutable' n x@ is a mutable array of length @n@ with @x@ the value of every element.
  replicateMutable :: (PrimMonad m, Element arr b) => Int -> b -> m (Mutable arr (PrimState m) b)
  -- | Index into an array at the given index.
  index :: Element arr b => arr b -> Int -> b
  -- | Index into an array at the given index, yielding an unboxed one-tuple of the element.
  index# :: Element arr b => arr b -> Int -> (# b #)
  -- | Indexing in a monad.
  --
  --   The monad allows operations to be strict in the array
  --   when necessary. Suppose array copying is implemented like this:
  --
  --   > copy mv v = ... write mv i (v ! i) ...
  --
  --   For lazy arrays, @v ! i@ would not be not be evaluated,
  --   which means that @mv@ would unnecessarily retain a reference
  --   to @v@ in each element written.
  --
  --   With 'indexM', copying can be implemented like this instead:
  --
  --   > copy mv v = ... do
  --   >   x <- indexM v i
  --   >   write mv i x
  --
  --   Here, no references to @v@ are retained because indexing
  --   (but /not/ the elements) is evaluated eagerly.
  indexM :: (Element arr b, Monad m) => arr b -> Int -> m b
  -- | Read a mutable array at the given index.
  read :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> m b
  -- | Write to a mutable array at the given index.
  write :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> b -> m ()
  -- | Resize an array into one with the given size.
  resize :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> Int -> m (Mutable arr (PrimState m) b)
  -- | The size of the array
  size :: Element arr b => arr b -> Int
  -- | The size of the mutable array
  sizeMutable :: (PrimMonad m, Element arr b) => Mutable arr (PrimState m) b -> m Int
  -- | Turn a mutable array into an immutable one without copying.
  --   The mutable array should not be used after this conversion.
  unsafeFreeze :: PrimMonad m => Mutable arr (PrimState m) b -> m (arr b)
  -- | Turn a mutable array into an immutable one with copying, using a slice of the mutable array.
  freeze :: (PrimMonad m, Element arr b)
    => Mutable arr (PrimState m) b
    -> Int -- ^ offset into the array
    -> Int -- ^ length of the slice
    -> m (arr b)
  -- | Copy a slice of an immutable array into a new mutable array.
  thaw :: (PrimMonad m, Element arr b)
    => arr b
    -> Int -- ^ offset into the array
    -> Int -- ^ length of the slice
    -> m (Mutable arr (PrimState m) b)
  -- | Copy a slice of an array into a mutable array.
  copy :: (PrimMonad m, Element arr b)
    => Mutable arr (PrimState m) b -- ^ destination array
    -> Int -- ^ offset into destination array
    -> arr b -- ^ source array
    -> Int -- ^ offset into source array
    -> Int -- ^ number of elements to copy
    -> m ()
  -- | Copy a slice of a mutable array into another mutable array.
  --   In the case that the destination and source arrays are the
  --   same, the regions may overlap.
  copyMutable :: (PrimMonad m, Element arr b)
    => Mutable arr (PrimState m) b -- ^ destination array
    -> Int -- ^ offset into destination array
    -> Mutable arr (PrimState m) b -- ^ source array
    -> Int -- ^ offset into source array
    -> Int -- ^ number of elements to copy
    -> m ()
  -- | Clone a slice of an array.
  clone :: Element arr b
    => arr b -- ^ Array to copy a slice of
    -> Int -- ^ Offset into the array
    -> Int -- ^ Length of the slice
    -> arr b
  -- | Clone a slice of a mutable array.
  cloneMutable :: (PrimMonad m, Element arr b)
    => Mutable arr (PrimState m) b -- ^ Array to copy a slice of
    -> Int -- ^ Offset into the array
    -> Int -- ^ Length of the slice
    -> m (Mutable arr (PrimState m) b)
  -- | Copy a slice of an array an then insert an element into that array.
  --
  -- The default implementation performs a memset which would be unnecessary
  -- except that the garbage collector might trace the uninitialized array.
  insertSlicing :: Element arr b
    => arr b -- ^ array to copy a slice from
    -> Int -- ^ offset into source array
    -> Int -- ^ length of the slice
    -> Int -- ^ index in the output array to insert at
    -> b -- ^ element to insert
    -> arr b
  insertSlicing src off len0 i x = run $ do
    dst <- replicateMutable (len0 + 1) x
    copy dst 0 src off i
    copy dst (i + 1) src (off + i) (len0 - i)
    unsafeFreeze dst
  {-# inline insertSlicing #-}
  -- | Test the two arrays for equality.
  equals :: (Element arr b, Eq b) => arr b -> arr b -> Bool
  -- | Test the two mutable arrays for pointer equality.
  --   Does not check equality of elements.
  equalsMutable :: Mutable arr s a -> Mutable arr s a -> Bool
  -- | Unlift an array into an 'ArrayArray#'.
  unlift :: arr b -> ArrayArray#
  -- | Lift an 'ArrayArray#' into an array.
  lift :: ArrayArray# -> arr b
  -- | Create a singleton array.
  singleton :: Element arr a => a -> arr a
  -- | Create a doubleton array.
  doubleton :: Element arr a => a -> a -> arr a
  -- | Create a tripleton array.
  tripleton :: Element arr a => a -> a -> a -> arr a
  -- | Create a quadrupleton array.
  quadrupleton :: Element arr a => a -> a -> a -> a -> arr a
  -- | Reduce the array and all of its elements to WHNF.
  rnf :: (NFData a, Element arr a) => arr a -> ()
  -- | Run an effectful computation that produces an array.
  run :: (forall s. ST s (arr a)) -> arr a

instance Contiguous SmallArray where
  type Mutable SmallArray = SmallMutableArray
  type Element SmallArray = Always
  empty = mempty
  new n = newSmallArray n errorThunk
  index = indexSmallArray
  indexM = indexSmallArrayM
  index# = indexSmallArray##
  read = readSmallArray
  write = writeSmallArray
  null a = case sizeofSmallArray a of
    0 -> True
    _ -> False
  freeze = freezeSmallArray
  size = sizeofSmallArray
  sizeMutable = (\x -> pure $! sizeofSmallMutableArray x)
  unsafeFreeze = unsafeFreezeSmallArray
  thaw = thawSmallArray
  equals = (==)
  equalsMutable = (==)
  singleton a = runST $ do
    marr <- newSmallArray 1 errorThunk
    writeSmallArray marr 0 a
    unsafeFreezeSmallArray marr
  doubleton a b = runST $ do
    m <- newSmallArray 2 errorThunk
    writeSmallArray m 0 a
    writeSmallArray m 1 b
    unsafeFreezeSmallArray m
  tripleton a b c = runST $ do
    m <- newSmallArray 3 errorThunk
    writeSmallArray m 0 a
    writeSmallArray m 1 b
    writeSmallArray m 2 c
    unsafeFreezeSmallArray m
  quadrupleton a b c d = runST $ do
    m <- newSmallArray 4 errorThunk
    writeSmallArray m 0 a
    writeSmallArray m 1 b
    writeSmallArray m 2 c
    writeSmallArray m 3 d
    unsafeFreezeSmallArray m
  rnf !ary =
    let !sz = sizeofSmallArray ary
        go !ix = if ix < sz
          then
            let !(# x #) = indexSmallArray## ary ix
             in DS.rnf x `seq` go (ix + 1)
          else ()
     in go 0
  clone = cloneSmallArray
  cloneMutable = cloneSmallMutableArray
  lift x = SmallArray (unsafeCoerce# x)
  unlift (SmallArray x) = unsafeCoerce# x
  copy = copySmallArray
  copyMutable = copySmallMutableArray
  replicateMutable = replicateSmallMutableArray
  resize = resizeSmallArray
  run = runSmallArrayST
  {-# inline empty #-}
  {-# inline null #-}
  {-# inline new #-}
  {-# inline replicateMutable #-}
  {-# inline index #-}
  {-# inline index# #-}
  {-# inline indexM #-}
  {-# inline read #-}
  {-# inline write #-}
  {-# inline resize #-}
  {-# inline size #-}
  {-# inline sizeMutable #-}
  {-# inline unsafeFreeze #-}
  {-# inline freeze #-}
  {-# inline thaw #-}
  {-# inline copy #-}
  {-# inline copyMutable #-}
  {-# inline clone #-}
  {-# inline cloneMutable #-}
  {-# inline equals #-}
  {-# inline equalsMutable #-}
  {-# inline unlift #-}
  {-# inline lift #-}
  {-# inline singleton #-}
  {-# inline doubleton #-}
  {-# inline tripleton #-}
  {-# inline quadrupleton #-}
  {-# inline rnf #-}
  {-# inline run #-}

instance Contiguous PrimArray where
  type Mutable PrimArray = MutablePrimArray
  type Element PrimArray = Prim
  empty = mempty
  new = newPrimArray
  replicateMutable = replicateMutablePrimArray
  index = indexPrimArray
  index# arr ix = (# indexPrimArray arr ix #)
  indexM arr ix = pure (indexPrimArray arr ix)
  read = readPrimArray
  write = writePrimArray
  resize = resizeMutablePrimArray
  size = sizeofPrimArray
  sizeMutable = getSizeofMutablePrimArray
  freeze = freezePrimArrayShim
  unsafeFreeze = unsafeFreezePrimArray
  thaw = thawPrimArray
  copy = copyPrimArray
  copyMutable = copyMutablePrimArray
  clone = clonePrimArrayShim
  cloneMutable = cloneMutablePrimArrayShim
  equals = (==)
  unlift (PrimArray x) = unsafeCoerce# x
  lift x = PrimArray (unsafeCoerce# x)
  null (PrimArray a) = case sizeofByteArray# a of
    0# -> True
    _ -> False
  equalsMutable = sameMutablePrimArray
  rnf (PrimArray !_) = ()
  singleton a = runPrimArrayST $ do
    marr <- newPrimArray 1
    writePrimArray marr 0 a
    unsafeFreezePrimArray marr
  doubleton a b = runPrimArrayST $ do
    m <- newPrimArray 2
    writePrimArray m 0 a
    writePrimArray m 1 b
    unsafeFreezePrimArray m
  tripleton a b c = runPrimArrayST $ do
    m <- newPrimArray 3
    writePrimArray m 0 a
    writePrimArray m 1 b
    writePrimArray m 2 c
    unsafeFreezePrimArray m
  quadrupleton a b c d = runPrimArrayST $ do
    m <- newPrimArray 4
    writePrimArray m 0 a
    writePrimArray m 1 b
    writePrimArray m 2 c
    writePrimArray m 3 d
    unsafeFreezePrimArray m
  insertSlicing src off len0 i x = runPrimArrayST $ do
    dst <- new (len0 + 1)
    copy dst 0 src off i
    write dst i x
    copy dst (i + 1) src (off + i) (len0 - i)
    unsafeFreeze dst
  run = runPrimArrayST
  {-# inline empty #-}
  {-# inline null #-}
  {-# inline new #-}
  {-# inline replicateMutable #-}
  {-# inline index #-}
  {-# inline index# #-}
  {-# inline indexM #-}
  {-# inline read #-}
  {-# inline write #-}
  {-# inline resize #-}
  {-# inline size #-}
  {-# inline sizeMutable #-}
  {-# inline unsafeFreeze #-}
  {-# inline freeze #-}
  {-# inline thaw #-}
  {-# inline copy #-}
  {-# inline copyMutable #-}
  {-# inline clone #-}
  {-# inline cloneMutable #-}
  {-# inline insertSlicing #-}
  {-# inline equals #-}
  {-# inline equalsMutable #-}
  {-# inline unlift #-}
  {-# inline lift #-}
  {-# inline singleton #-}
  {-# inline doubleton #-}
  {-# inline tripleton #-}
  {-# inline quadrupleton #-}
  {-# inline rnf #-}
  {-# inline run #-}

instance Contiguous Array where
  type Mutable Array = MutableArray
  type Element Array = Always
  empty = mempty
  new n = newArray n errorThunk
  replicateMutable = newArray
  index = indexArray
  index# = indexArray##
  indexM = indexArrayM
  read = readArray
  write = writeArray
  resize = resizeArray
  size = sizeofArray
  sizeMutable = (\x -> pure $! sizeofMutableArray x)
  freeze = freezeArray
  unsafeFreeze = unsafeFreezeArray
  thaw = thawArray
  copy = copyArray
  copyMutable = copyMutableArray
  clone = cloneArray
  cloneMutable = cloneMutableArray
  equals = (==)
  unlift (Array x) = unsafeCoerce# x
  lift x = Array (unsafeCoerce# x)
  null (Array a) = case sizeofArray# a of
    0# -> True
    _ -> False
  equalsMutable = sameMutableArray
  rnf !ary =
    let !sz = sizeofArray ary
        go !i
          | i == sz = ()
          | otherwise =
              let !(# x #) = indexArray## ary i
               in DS.rnf x `seq` go (i+1)
     in go 0
  singleton a = runArrayST (newArray 1 a >>= unsafeFreezeArray)
  doubleton a b = runArrayST $ do
    m <- newArray 2 a
    writeArray m 1 b
    unsafeFreezeArray m
  tripleton a b c = runArrayST $ do
    m <- newArray 3 a
    writeArray m 1 b
    writeArray m 2 c
    unsafeFreezeArray m
  quadrupleton a b c d = runArrayST $ do
    m <- newArray 4 a
    writeArray m 1 b
    writeArray m 2 c
    writeArray m 3 d
    unsafeFreezeArray m
  run = runArrayST
  {-# inline empty #-}
  {-# inline null #-}
  {-# inline new #-}
  {-# inline replicateMutable #-}
  {-# inline index #-}
  {-# inline index# #-}
  {-# inline indexM #-}
  {-# inline read #-}
  {-# inline write #-}
  {-# inline resize #-}
  {-# inline size #-}
  {-# inline sizeMutable #-}
  {-# inline unsafeFreeze #-}
  {-# inline freeze #-}
  {-# inline thaw #-}
  {-# inline copy #-}
  {-# inline copyMutable #-}
  {-# inline clone #-}
  {-# inline cloneMutable #-}
  {-# inline equals #-}
  {-# inline equalsMutable #-}
  {-# inline unlift #-}
  {-# inline lift #-}
  {-# inline singleton #-}
  {-# inline doubleton #-}
  {-# inline tripleton #-}
  {-# inline quadrupleton #-}
  {-# inline rnf #-}
  {-# inline run #-}

instance Contiguous UnliftedArray where
  type Mutable UnliftedArray = MutableUnliftedArray
  type Element UnliftedArray = PrimUnlifted
  empty = emptyUnliftedArray
  new = unsafeNewUnliftedArray
  replicateMutable = newUnliftedArray
  index = indexUnliftedArray
  index# arr ix = (# indexUnliftedArray arr ix #)
  indexM arr ix = pure (indexUnliftedArray arr ix)
  read = readUnliftedArray
  write = writeUnliftedArray
  resize = resizeUnliftedArray
  size = sizeofUnliftedArray
  sizeMutable = pure . sizeofMutableUnliftedArray
  freeze = freezeUnliftedArray
  unsafeFreeze = unsafeFreezeUnliftedArray
  thaw = thawUnliftedArray
  copy = copyUnliftedArray
  copyMutable = copyMutableUnliftedArray
  clone = cloneUnliftedArray
  cloneMutable = cloneMutableUnliftedArray
  equals = (==)
  unlift (UnliftedArray x) = x
  lift x = UnliftedArray x
  null (UnliftedArray a) = case sizeofArrayArray# a of
    0# -> True
    _ -> False
  equalsMutable = sameMutableUnliftedArray
  rnf !ary =
    let !sz = sizeofUnliftedArray ary
        go !i
          | i == sz = ()
          | otherwise =
              let x = indexUnliftedArray ary i
               in DS.rnf x `seq` go (i+1)
     in go 0
  singleton a = runUnliftedArrayST (newUnliftedArray 1 a >>= unsafeFreezeUnliftedArray)
  doubleton a b = runUnliftedArrayST $ do
    m <- newUnliftedArray 2 a
    writeUnliftedArray m 1 b
    unsafeFreezeUnliftedArray m
  tripleton a b c = runUnliftedArrayST $ do
    m <- newUnliftedArray 3 a
    writeUnliftedArray m 1 b
    writeUnliftedArray m 2 c
    unsafeFreezeUnliftedArray m
  quadrupleton a b c d = runUnliftedArrayST $ do
    m <- newUnliftedArray 4 a
    writeUnliftedArray m 1 b
    writeUnliftedArray m 2 c
    writeUnliftedArray m 3 d
    unsafeFreezeUnliftedArray m
  run = runUnliftedArrayST
  {-# inline empty #-}
  {-# inline null #-}
  {-# inline new #-}
  {-# inline replicateMutable #-}
  {-# inline index #-}
  {-# inline index# #-}
  {-# inline indexM #-}
  {-# inline read #-}
  {-# inline write #-}
  {-# inline resize #-}
  {-# inline size #-}
  {-# inline sizeMutable #-}
  {-# inline unsafeFreeze #-}
  {-# inline freeze #-}
  {-# inline thaw #-}
  {-# inline copy #-}
  {-# inline copyMutable #-}
  {-# inline clone #-}
  {-# inline cloneMutable #-}
  {-# inline equals #-}
  {-# inline equalsMutable #-}
  {-# inline unlift #-}
  {-# inline lift #-}
  {-# inline singleton #-}
  {-# inline doubleton #-}
  {-# inline tripleton #-}
  {-# inline quadrupleton #-}
  {-# inline rnf #-}
  {-# inline run #-}

errorThunk :: a
errorThunk = error "Contiguous typeclass: unitialized element"
{-# noinline errorThunk #-}

freezePrimArrayShim :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a)
freezePrimArrayShim !src !off !len = do
  dst <- newPrimArray len
  copyMutablePrimArray dst 0 src off len
  unsafeFreezePrimArray dst
{-# inline freezePrimArrayShim #-}

resizeArray :: PrimMonad m => MutableArray (PrimState m) a -> Int -> m (MutableArray (PrimState m) a)
resizeArray !src !sz = do
  dst <- newArray sz errorThunk
  copyMutableArray dst 0 src 0 (min sz (sizeofMutableArray src))
  pure dst
{-# inline resizeArray #-}

resizeSmallArray :: PrimMonad m => SmallMutableArray (PrimState m) a -> Int -> m (SmallMutableArray (PrimState m) a)
resizeSmallArray !src !sz = do
  dst <- newSmallArray sz errorThunk
  copySmallMutableArray dst 0 src 0 (min sz (sizeofSmallMutableArray src))
  pure dst
{-# inline resizeSmallArray #-}

resizeUnliftedArray :: (PrimMonad m, PrimUnlifted a) => MutableUnliftedArray (PrimState m) a -> Int -> m (MutableUnliftedArray (PrimState m) a)
resizeUnliftedArray !src !sz = do
  dst <- unsafeNewUnliftedArray sz
  copyMutableUnliftedArray dst 0 src 0 (min sz (sizeofMutableUnliftedArray src))
  pure dst
{-# inline resizeUnliftedArray #-}

-- | Append two arrays.
append :: (Contiguous arr, Element arr a) => arr a -> arr a -> arr a
append !a !b = run $ do
  let !szA = size a
  let !szB = size b
  m <- new (szA + szB)
  copy m 0 a 0 szA
  copy m szA b 0 szB
  unsafeFreeze m
{-# inline append #-}

-- | Insert an element into an array at the given index.
insertAt :: (Contiguous arr, Element arr a) => arr a -> Int -> a -> arr a
insertAt src i x = insertSlicing src 0 (size src) i x

-- | 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 src i x = create $ do
  dst <- thaw src 0 (size src)
  write dst i x
  pure dst
{-# inline replaceAt #-}

modifyAt :: (Contiguous arr, Element arr a)
  => (a -> a) -> arr a -> Int -> arr a
modifyAt f src i = replaceAt src i $ f (index src 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' f src i = replaceAt src i $! f (index src i)
{-# inline modifyAt' #-}

modifyAtF :: (Contiguous arr, Element arr a, Functor f)
  => (a -> f a) -> arr a -> Int -> f (arr a)
modifyAtF f src i = replaceAt src i <$> f (index src 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' f src i = do
  !r <- f (index src i)
  let !dst = replaceAt src i 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 f a = run $ do
  mb <- new (size a)
  let go !i
        | i == size a = pure ()
        | otherwise = do
            x <- indexM a 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' f a = run $ do
  mb <- new (size a)
  let go !i
        | i == size a = pure ()
        | otherwise = do
            x <- indexM a i
            let !b = f i 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 f a = run $ do
  mb <- new (size a)
  let go !i
        | i == size a = pure ()
        | otherwise = do
            x <- indexM a 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' f a = run $ do
  mb <- new (size a)
  let go !i
        | i == size a = pure ()
        | otherwise = do
            x <- indexM a i
            let !b = f 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 a = map id 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 f z = \arr ->
  let !sz = size arr
      go !ix = if sz > ix
        then case index# arr ix of
          (# x #) -> f x (go (ix + 1))
        else z
  in go 0

-- | Strict right fold over the elements of an array.
foldr' :: (Contiguous arr, Element arr a) => (a -> b -> b) -> b -> arr a -> b
foldr' f !z = \arr ->
  let go !ix !acc = if ix == -1
        then acc
        else case index# arr ix of
          (# x #) -> go (ix - 1) (f x acc)
  in go (size arr - 1) 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 f z = \arr ->
  let !sz = size arr
      go !ix acc = if ix == sz
        then acc
        else case index# arr ix of
          (# x #) -> go (ix + 1) (f acc x)
  in go 0 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' f !z = \arr ->
  let !sz = size arr
      go !ix !acc = if ix == sz
        then acc
        else case index# arr ix of
          (# x #) -> go (ix + 1) (f acc x)
  in go 0 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' f !z = \arr ->
  let !sz = size arr
      go !ix !acc = if ix == sz
        then acc
        else case index# arr ix of
          (# x #) -> go (ix + 1) (f acc ix x)
  in go 0 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' f !z = \arr ->
  let !sz = size arr
      go !ix !acc = if ix == (-1)
        then acc
        else case index# arr ix of
          (# x #) -> go (ix - 1) (f ix x acc)
  in go (sz - 1) 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 f = \arr ->
  let !sz = size arr
      go !ix = if sz > ix
        then case index# arr ix of
          (# x #) -> mappend (f x) (go (ix + 1))
        else mempty
  in go 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' f = \arr ->
  let !sz = size arr
      go !ix !acc = if ix == sz
        then acc
        else case index# arr ix
          of (# x #) -> go (ix + 1) (mappend acc (f x))
  in go 0 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' = 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' f = \arr ->
  let !sz = size arr
      go !ix !acc = if ix == sz
        then acc
        else case index# arr ix of
          (# x #) -> go (ix + 1) (mappend acc (f ix x))
  in go 0 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' f = \arr ->
  let !sz = size arr
      go !ix !acc = if ix == sz
        then acc
        else case index# arr ix of
          (# x #) -> go (ix + 1) (acc <> f ix x)
      !(# e0 #) = index# arr 0
  in go 1 (f 0 e0)
{-# inline ifoldlMap1' #-}

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

-- | Drop elements that do not satisfy the predicate.
filter :: (Contiguous arr, Element arr a)
  => (a -> Bool)
  -> arr a
  -> arr a
filter p arr = ifilter (const p) 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 p arr = run $ do
  marr :: MutablePrimArray s Word8 <- newPrimArray sz
  let go1 :: Int -> Int -> ST s Int
      go1 !ix !numTrue = if ix < sz
        then do
          atIx <- indexM arr ix
          let !keep = p ix atIx
          let !keepTag = I# (dataToTag# keep)
          writePrimArray marr ix (fromIntegral keepTag)
          go1 (ix + 1) (numTrue + keepTag)
        else pure numTrue
  numTrue <- go1 0 0
  if numTrue == sz
    then pure arr
    else do
      marrTrues <- new numTrue
      let go2 !ixSrc !ixDst = when (ixDst < numTrue) $ do
            atIxKeep <- readPrimArray marr 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 = size 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 f arr = run $ do
  let !sz = size arr
  let go :: Int -> Int -> [b] -> ST s ([b],Int)
      go !ix !numJusts !justs = if ix < sz
        then do
          atIx <- indexM arr ix
          case f atIx of
            Nothing -> go (ix+1) numJusts justs
            Just x -> go (ix+1) (numJusts+1) (x:justs)
        else pure (justs,numJusts)
  !(bs,!numJusts) <- go 0 0 []
  !marr <- unsafeFromListReverseMutableN numJusts bs
  unsafeFreeze marr
{-# inline mapMaybe #-}

{-# inline isTrue #-}
isTrue :: Word8 -> Bool
isTrue 0 = False
isTrue _ = 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 = mapMaybe id
{-# inline catMaybes #-}

clonePrimArrayShim :: Prim a => PrimArray a -> Int -> Int -> PrimArray a
clonePrimArrayShim !arr !off !len = runPrimArrayST $ do
  marr <- newPrimArray len
  copyPrimArray marr 0 arr off len
  unsafeFreezePrimArray marr
{-# inline clonePrimArrayShim #-}

cloneMutablePrimArrayShim :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArrayShim !arr !off !len = do
  marr <- newPrimArray len
  copyMutablePrimArray marr 0 arr off len
  pure marr
{-# inline cloneMutablePrimArrayShim #-}

-- | @'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 n x = create (replicateMutable n x)
{-# inline replicate #-}

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

replicateMutablePrimArray :: (PrimMonad m, Prim a)
  => Int -- ^ length
  -> a -- ^ element
  -> m (MutablePrimArray (PrimState m) a)
replicateMutablePrimArray len a = do
  marr <- newPrimArray len
  setPrimArray marr 0 len a
  pure marr
{-# inline replicateMutablePrimArray #-}

replicateSmallMutableArray :: (PrimMonad m)
  => Int
  -> a
  -> m (SmallMutableArray (PrimState m) a)
replicateSmallMutableArray len a = do
  marr <- newSmallArray len errorThunk
  let go !ix = when (ix < len) $ do
        writeSmallArray marr ix a
        go (ix + 1)
  go 0
  pure marr
{-# inline replicateSmallMutableArray #-}

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

unsafeFromListMutableN :: (Contiguous arr, Element arr a, PrimMonad m)
  => Int
  -> [a]
  -> m (Mutable arr (PrimState m) a)
unsafeFromListMutableN n l = do
  m <- new n
  let go !_ [] = pure m
      go !ix (x : xs) = do
        write m ix x
        go (ix+1) 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 n l = do
  m <- new n
  let go !_ [] = pure m
      go !ix (x : xs) = do
        write m ix x
        go (ix-1) 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 n l = create (unsafeFromListReverseMutableN n 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 f !marr = do
  !sz <- sizeMutable marr
  let go !ix = when (ix < sz) $ do
        a <- read marr 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' f !marr = do
  !sz <- sizeMutable marr
  let go !ix = when (ix < sz) $ do
        a <- read marr ix
        let !b = f 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 f !marr = do
  !sz <- sizeMutable marr
  let go !ix = when (ix < sz) $ do
        a <- read marr 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' f !marr = do
  !sz <- sizeMutable marr
  let go !ix = when (ix < sz) $ do
        a <- read marr ix
        let !b = f ix 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, Contiguous arr2, Element arr1 a, Element arr2 b)
  => (a -> m b)
  -> arr1 a
  -> m (arr2 b)
traverseP f !arr = do
  let !sz = size arr
  !marr <- new sz
  let go !ix = when (ix < sz) $ do
        a <- indexM arr ix
        b <- f a
        write marr ix b
        go (ix + 1)
  go 0
  unsafeFreeze marr
{-# inline traverseP #-}

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

runSTA :: (Contiguous v, Element v a) => Int -> STA v a -> v a
runSTA !sz (STA m) = runST $ new sz >>= 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 f = itraverse (const 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_ f = itraverse_ (const 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 f = \arr ->
  let !sz = size arr
      go !ix = if ix == sz
        then pure (STA unsafeFreeze)
        else case index# arr ix of
          (# x #) -> liftA2
            (\b (STA m) -> STA $ \marr -> do
              write marr ix b
              m marr
            )
            (f ix x)
            (go (ix + 1))
  in if sz == 0
    then pure empty
    else runSTA sz <$> go 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_ f = \arr ->
  let !sz = size arr
      go !ix = when (ix < sz) $
        f ix (index arr ix) *> go (ix + 1)
  in go 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 = flip 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_ = flip 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' f a0 src = go 0 [] a0 where
  !sz = size src
  go !ix !xs !acc = if ix < sz
    then do
      (!acc',!x) <- f acc (index src ix)
      go (ix + 1) (x : xs) acc'
    else
      let !xs' = unsafeFromListReverseN sz xs
       in pure (acc,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' f !src = runST $ do
  dst <- new sz
  acc <- go 0 dst mempty
  dst' <- unsafeFreeze dst
  pure (acc,dst')
  where
  !sz = size src
  go :: Int -> Mutable arr2 s c -> a -> ST s a
  go !ix !dst !accA = if ix < sz
    then do
      let (!accB,!x) = f (index src ix)
      write dst ix x
      go (ix + 1) dst (accA <> accB)
    else pure 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 f arr =
  let !sz = size arr
  in generateM sz $ \ix -> indexM arr ix >>= 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_ = 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 = flip 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_ = flip 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 = traverse 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_ = foldr (*>) (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 = foldr (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 len f = create (generateMutable len f)
{-# inline generate #-}

-- | Construct an array of the given length by applying
--   the monadic action to each index.
generateM :: (Contiguous arr, Element arr a, Monad m)
  => Int
  -> (Int -> m a)
  -> m (arr a)
generateM !sz f =
  let go !ix = if ix < sz
        then liftA2
          (\b (STA m) -> STA $ \marr -> do
              write marr ix b
              m marr
          )
          (f ix)
          (go (ix + 1))
        else pure $ STA unsafeFreeze
  in if sz == 0
    then pure empty
    else runSTA sz <$> go 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 len f = generateMutableM len (pure . 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 !len f = do
  marr <- new len
  let go !ix = when (ix < len) $ do
        x <- f 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 len f z0 = runST (iterateMutableN len f z0 >>= 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 len f z0 = iterateMutableNM len (pure . f) 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 !len f z0 = do
  marr <- new 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 !ix !acc
        | ix <= 0 = write marr ix z0 >> go (ix + 1) z0
        | ix == len = pure ()
        | otherwise = do
            a <- f 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 x = run (unsafeFreeze =<< 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 p = runST (Prelude.mapM unsafeFreeze =<< 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 f z0 = create (unfoldrMutable f 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 f z0 = do
  let go !sz s !xs = case f s of
        Nothing -> pure (sz,xs)
        Just (x,s') -> go (sz + 1) s' (x : xs)
  (sz,xs) <- go 0 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 maxSz f z0 = create (unfoldrMutableN maxSz f 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 !maxSz f z0 = do
  m <- new maxSz
  let go !ix s = if ix < maxSz
        then case f s of
          Nothing -> pure ix
          Just (x,s') -> do
            write m ix x
            go (ix + 1) s'
        else pure ix
  sz <- go 0 z0
  case compare maxSz sz of
    EQ -> pure m
    GT -> resize m sz
    LT -> error "Data.Primitive.Contiguous.unfoldrMutableN: internal error"
{-# inline unfoldrMutableN #-}

-- | Convert an array to a list.
toList :: (Contiguous arr, Element arr a)
  => arr a
  -> [a]
toList arr = build (\c n -> foldr c n 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 marr = do
  sz <- sizeMutable marr
  let go !ix !acc = if ix >= 0
        then do
          x <- read marr ix
          go (ix - 1) (x : acc)
        else pure 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 len vs = do
  marr <- new len
  let go [] !ix = if ix == len
        then pure ()
        else error "Data.Primitive.Contiguous.fromListN: list length less than specified size."
      go (a:as) !ix = if ix < len
        then do
          write marr ix a
          go as (ix + 1)
        else error "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 xs = fromListMutableN (length xs) 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 len vs = create (fromListMutableN len vs)
{-# inline fromListN #-}

-- | Convert a list into an array.
fromList :: (Contiguous arr, Element arr a)
  => [a]
  -> arr a
fromList vs = create (fromListMutable 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 f marr = do
  !sz <- sizeMutable marr
  let go !ix = when (ix < sz) $ do
        x <- read marr 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' f marr = do
  !sz <- sizeMutable marr
  let go !ix = when (ix < sz) $ do
        x <- read marr ix
        let !y = f 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 z0 sz = create (enumFromMutableN z0 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 z0 !sz = do
  m <- new sz
  let go !ix z = if ix < sz
        then do
          write m ix z
          go (ix + 1) (succ z)
        else pure 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 f s0 arr = go 0 s0 where
  sz = size arr
  go !ix !s = if ix < sz
    then
      let !(# x #) = index# arr ix
       in go (ix + 1) (f s x)
    else hashIntWithSalt s ix
{-# inline liftHashWithSalt #-}

-- | Reverse the elements of an array.
reverse :: (Contiguous arr, Element arr a)
  => arr a
  -> arr a
reverse arr = run $ do
  marr <- new sz
  copy marr 0 arr 0 sz
  reverseMutable marr
  unsafeFreeze marr
  where
    !sz = size arr
{-# 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 marr = do
  !sz <- sizeMutable 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
  -> Int -- ^ start index
  -> Int -- ^ end index
  -> m ()
reverseSlice !marr !start !end = do
  let go !s !e = if s >= e
        then pure ()
        else do
          tmp <- read marr s
          write marr s =<< read marr e
          write marr e tmp
          go (s+1) (e-1)
  go start 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 :: Contiguous arr => arr a -> arr a -> Bool
same a b = isTrue# (sameMutableArrayArray# (unsafeCoerce# (unlift a) :: MutableArrayArray# s) (unsafeCoerce# (unlift b) :: MutableArrayArray# s))

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

combine :: Int -> Int -> Int
combine h1 h2 = (h1 * 16777619) `xor` h2
{-# inline combine #-}

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

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

-- | The least element of a structure.
minimum :: (Contiguous arr, Element arr a, Ord a) => arr a -> Maybe a
minimum = minimumBy 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 f arr =
  let !sz = size arr
      go !ix o = if ix < sz
        then case index# arr ix of
          !(# x #) -> go (ix + 1) (case f x o of { GT -> x; _ -> o; })
        else o
  in if sz == 0
    then Nothing
    else Just (go 0 (index arr 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 f arr =
  let !sz = size arr
      go !ix o = if ix < sz
        then case index# arr ix of
          !(# x #) -> go (ix + 1) (case f x o of { GT -> o; _ -> x; })
        else o
  in if sz == 0
    then Nothing
    else Just (go 0 (index arr 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 p = coerce . (foldMap (\x -> if p x then Just (First x) else 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 p xs = loop 0
  where
  loop i
    | i < size xs = if p (index xs i) then Just i else loop (i + 1)
    | otherwise = 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 !marr !ix1 !ix2 = do
  atIx1 <- read marr 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 !arr = create $ do
  let !sz = size arr
      go :: Int -> [a] -> Int -> ST s (Int, [a])
      go !ix !as !acc = if ix < sz
        then do
          indexM arr ix >>= \case
            Left a -> go (ix + 1) (a:as) (acc + 1)
            Right _ -> go (ix + 1) as acc
        else pure (acc, as)
  (len, as) <- go 0 [] 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 !arr = create $ do
  let !sz = size arr
      go :: Int -> [b] -> Int -> ST s (Int, [b])
      go !ix !bs !acc = if ix < sz
        then do
          indexM arr ix >>= \case
            Left _ -> go (ix + 1) bs acc
            Right b -> go (ix + 1) (b:bs) (acc + 1)
        else pure (acc, bs)
  (len, bs) <- go 0 [] 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 !arr = runST $ do
  let !sz = size arr
      go :: Int -> [a] -> [b] -> Int -> Int -> ST s (Int, Int, [a], [b])
      go !ix !as !bs !accA !accB = if ix < sz
        then do
          indexM arr ix >>= \case
            Left a -> go (ix + 1) (a:as) bs (accA + 1) accB
            Right b -> go (ix + 1) as (b:bs) accA (accB + 1)
          else pure (accA, accB, as, bs)
  (lenA, lenB, as, bs) <- go 0 [] [] 0 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 f = iscanl (const 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 f q as = internalScanl (size as + 1) f q 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' f = iscanl' (const 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' f !q as = internalScanl' (size as + 1) f q 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 !sz f !q as = create $ do
  !marr <- new sz
  let go !ix acc = when (ix < sz) $ do
        write marr ix acc
        x <- indexM as 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' !sz f !q as = create $ do
  !marr <- new sz
  let go !ix !acc = when (ix < sz) $ do
        write marr ix acc
        x <- indexM as 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 f = iprescanl (const 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 f q as = internalScanl (size as) f q 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' f = iprescanl (const 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' f !q as = internalScanl' (size as) f q 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 f = izipWith (\_ a b -> f a 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 f as bs = create $ do
  let !sz = min (size as) (size bs)
  !marr <- new sz
  let go !ix = when (ix < sz) $ do
        a <- indexM as ix
        b <- indexM bs ix
        let !g = f ix a 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 f = ifoldrZipWith (\_ x y c -> f x y 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' f = ifoldlZipWithM' (\_ x y c -> f x y 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 f z = \arr1 arr2 ->
  let !sz = min (size arr1) (size arr2)
      go !ix = if sz > ix
        then case index# arr1 ix of
          (# x #) -> case index# arr2 ix of
            (# y #) -> f ix x y (go (ix + 1))
        else z
  in go 0
{-# inline ifoldrZipWith #-}

-- | Variant of 'foldlZipWithM\'' that provides the index of each pair of elements.
ifoldlZipWithM' ::
  ( Contiguous arr1
  , Contiguous arr2
  , Element arr1 a
  , Element arr2 b
  , Monad m
  ) => (Int -> c -> a -> b -> m c)
    -> c
    -> arr1 a
    -> arr2 b
    -> m c
ifoldlZipWithM' f z = \arr1 arr2 ->
  let !sz = min (size arr1) (size arr2)
      go !ix !acc = if sz > ix
        then case index# arr1 ix of
          (# x #) -> case index# arr2 ix of
            (# y #) -> do
              acc' <- f ix acc x y
              go (ix + 1) acc'
        else pure acc
  in go 0 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 = 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 <$ barr = create (replicateMutable (size barr) 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 fs xs = create $ do
  marr <- new (szfs * szxs)
  let go1 !ix = when (ix < szfs) $ do
        f <- indexM fs ix
        go2 (ix * szxs) f 0
        go1 (ix + 1)
      go2 !off f !j = when (j < szxs) $ do
        x <- indexM xs j
        write marr (off + j) (f x)
        go2 off f (j + 1)
  go1 0
  pure marr
  where
    !szfs = size fs
    !szxs = size xs
{-# inline ap #-}

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

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