{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Ops.Fold.Internal
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Ops.Fold.Internal (
  foldlS,
  foldrS,
  ifoldlS,
  ifoldrS,
  -- Monadic
  foldlM,
  foldrM,
  foldlM_,
  foldrM_,
  ifoldlM,
  ifoldrM,
  ifoldlM_,
  ifoldrM_,
  -- Special folds
  fold,
  foldMono,
  foldlInternal,
  ifoldlInternal,
  foldrFB,
  lazyFoldlS,
  lazyFoldrS,
  -- Parallel folds
  foldlP,
  foldrP,
  ifoldlP,
  ifoldrP,
  foldlIO,
  ifoldlIO,
  ifoldrIO,
  splitReduce,
  any,
  anySu,
  anyPu,
) where

import Control.Monad (void, when)
import Control.Monad.Primitive
import Control.Scheduler
import qualified Data.Foldable as F
import Data.Functor.Identity (runIdentity)
import Data.Massiv.Core.Common
import System.IO.Unsafe (unsafePerformIO)
import Prelude hiding (any, foldl, foldr)

-- | /O(n)/ - Unstructured fold of an array.
--
-- @since 0.3.0
fold
  :: (Monoid e, Index ix, Source r e)
  => Array r ix e
  -- ^ Source array
  -> e
fold :: forall e ix r.
(Monoid e, Index ix, Source r e) =>
Array r ix e -> e
fold = forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal forall a. Monoid a => a -> a -> a
mappend forall a. Monoid a => a
mempty forall a. Monoid a => a -> a -> a
mappend forall a. Monoid a => a
mempty
{-# INLINE fold #-}

-- | /O(n)/ - This is exactly like `Data.Foldable.foldMap`, but for arrays. Fold over an array,
-- while converting each element into a `Monoid`. Also known as map-reduce. If elements of the array
-- are already a `Monoid` you can use `fold` instead.
--
-- @since 0.1.4
foldMono
  :: (Index ix, Source r e, Monoid m)
  => (e -> m)
  -- ^ Convert each element of an array to an appropriate `Monoid`.
  -> Array r ix e
  -- ^ Source array
  -> m
foldMono :: forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono e -> m
f = forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal (\m
a e
e -> m
a forall a. Monoid a => a -> a -> a
`mappend` e -> m
f e
e) forall a. Monoid a => a
mempty forall a. Monoid a => a -> a -> a
mappend forall a. Monoid a => a
mempty
{-# INLINE foldMono #-}

-- | /O(n)/ - Monadic left fold.
--
-- @since 0.1.0
foldlM :: (Index ix, Source r e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a
foldlM :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> e -> m a) -> a -> Array r ix e -> m a
foldlM a -> e -> m a
f a
acc Array r ix e
arr =
  case forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
    PrefIndex ix -> e
gix ->
      forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM forall ix. Index ix => ix
zeroIndex (forall ix. Sz ix -> ix
unSz Sz ix
sz) (forall ix. Index ix => Int -> ix
pureIndex Int
1) forall a. Ord a => a -> a -> Bool
(<) a
acc forall a b. (a -> b) -> a -> b
$ \ !ix
ix !a
a -> a -> e -> m a
f a
a (ix -> e
gix ix
ix)
    PrefIndexLinear Int -> e
gi ->
      forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM Int
0 (forall a. Ord a => a -> a -> Bool
< forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) (forall a. Num a => a -> a -> a
+ Int
1) a
acc forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
a -> a -> e -> m a
f a
a (Int -> e
gi Int
i)
  where
    sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
{-# INLINE foldlM #-}

-- | /O(n)/ - Monadic left fold, that discards the result.
--
-- @since 0.1.0
foldlM_ :: (Index ix, Source r e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m ()
foldlM_ :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> e -> m a) -> a -> Array r ix e -> m ()
foldlM_ a -> e -> m a
f a
acc = forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> e -> m a) -> a -> Array r ix e -> m a
foldlM a -> e -> m a
f a
acc
{-# INLINE foldlM_ #-}

-- | /O(n)/ - Monadic left fold with an index aware function.
--
-- @since 0.1.0
ifoldlM :: (Index ix, Source r e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM a -> ix -> e -> m a
f !a
acc !Array r ix e
arr =
  case forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
    PrefIndex ix -> e
gix ->
      forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM forall ix. Index ix => ix
zeroIndex (forall ix. Sz ix -> ix
unSz (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)) (forall ix. Index ix => Int -> ix
pureIndex Int
1) forall a. Ord a => a -> a -> Bool
(<) a
acc forall a b. (a -> b) -> a -> b
$ \ !ix
ix !a
a -> a -> ix -> e -> m a
f a
a ix
ix (ix -> e
gix ix
ix)
    PrefIndexLinear Int -> e
gi ->
      forall it ix (m :: * -> *) a.
(Iterator it, Index ix, Monad m) =>
it
-> Int
-> Sz ix
-> ix
-> Stride ix
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterTargetM RowMajor
defRowMajor Int
0 (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr) forall ix. Index ix => ix
zeroIndex forall ix. Index ix => Stride ix
oneStride a
acc forall a b. (a -> b) -> a -> b
$ \Int
i ix
ix !a
a -> a -> ix -> e -> m a
f a
a ix
ix (Int -> e
gi Int
i)
{-# INLINE ifoldlM #-}

-- | /O(n)/ - Monadic left fold with an index aware function, that discards the result.
--
-- @since 0.1.0
ifoldlM_ :: (Index ix, Source r e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m ()
ifoldlM_ :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m ()
ifoldlM_ a -> ix -> e -> m a
f a
acc = forall (f :: * -> *) a. Functor f => f a -> f ()
void forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM a -> ix -> e -> m a
f a
acc
{-# INLINE ifoldlM_ #-}

-- | /O(n)/ - Monadic right fold.
--
-- @since 0.1.0
foldrM :: (Index ix, Source r e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m a
foldrM :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(e -> a -> m a) -> a -> Array r ix e -> m a
foldrM e -> a -> m a
f a
acc Array r ix e
arr =
  case forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
    PrefIndex ix -> e
gix ->
      forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (forall a. Num a => a -> a -> a
subtract Int
1) (forall ix. Sz ix -> ix
unSz Sz ix
sz)) forall ix. Index ix => ix
zeroIndex (forall ix. Index ix => Int -> ix
pureIndex (-Int
1)) forall a. Ord a => a -> a -> Bool
(>=) a
acc (e -> a -> m a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
gix)
    PrefIndexLinear Int -> e
gi ->
      forall (m :: * -> *) a.
Monad m =>
Int
-> (Int -> Bool) -> (Int -> Int) -> a -> (Int -> a -> m a) -> m a
loopM (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz forall a. Num a => a -> a -> a
- Int
1) (forall a. Ord a => a -> a -> Bool
>= Int
0) (forall a. Num a => a -> a -> a
subtract Int
1) a
acc (e -> a -> m a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
gi)
  where
    !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
{-# INLINE foldrM #-}

-- | /O(n)/ - Monadic right fold, that discards the result.
--
-- @since 0.1.0
foldrM_ :: (Index ix, Source r e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m ()
foldrM_ :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(e -> a -> m a) -> a -> Array r ix e -> m ()
foldrM_ e -> a -> m a
f = forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m ()
ifoldrM_ (\ix
_ e
e a
a -> e -> a -> m a
f e
e a
a)
{-# INLINE foldrM_ #-}

-- | /O(n)/ - Monadic right fold with an index aware function.
--
-- @since 0.1.0
ifoldrM :: (Index ix, Source r e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM ix -> e -> a -> m a
f !a
acc !Array r ix e
arr =
  forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM (forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (forall a. Num a => a -> a -> a
subtract Int
1) (forall ix. Sz ix -> ix
unSz (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr))) forall ix. Index ix => ix
zeroIndex (forall ix. Index ix => Int -> ix
pureIndex (-Int
1)) forall a. Ord a => a -> a -> Bool
(>=) a
acc forall a b. (a -> b) -> a -> b
$ \ !ix
ix ->
    ix -> e -> a -> m a
f ix
ix (forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
ix)
{-# INLINE ifoldrM #-}

-- | /O(n)/ - Monadic right fold with an index aware function, that discards the result.
--
-- @since 0.1.0
ifoldrM_ :: (Index ix, Source r e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m ()
ifoldrM_ :: forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m ()
ifoldrM_ ix -> e -> a -> m a
f !a
acc !Array r ix e
arr = forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM ix -> e -> a -> m a
f a
acc Array r ix e
arr
{-# INLINE ifoldrM_ #-}

-- | /O(n)/ - Left fold, computed sequentially with lazy accumulator.
--
-- @since 0.1.0
lazyFoldlS :: (Index ix, Source r e) => (a -> e -> a) -> a -> Array r ix e -> a
lazyFoldlS :: forall ix r e a.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> a
lazyFoldlS a -> e -> a
f a
initAcc Array r ix e
arr = a -> Int -> a
go a
initAcc Int
0
  where
    len :: Int
len = forall ix. Index ix => Sz ix -> Int
totalElem (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
    go :: a -> Int -> a
go a
acc !Int
k
      | Int
k forall a. Ord a => a -> a -> Bool
< Int
len = a -> Int -> a
go (a -> e -> a
f a
acc (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
k)) (Int
k forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = a
acc
{-# INLINE lazyFoldlS #-}

-- | /O(n)/ - Right fold, computed sequentially with lazy accumulator.
--
-- @since 0.1.0
lazyFoldrS :: (Index ix, Source r e) => (e -> a -> a) -> a -> Array r ix e -> a
lazyFoldrS :: forall ix r e a.
(Index ix, Source r e) =>
(e -> a -> a) -> a -> Array r ix e -> a
lazyFoldrS = forall ix r e a.
(Index ix, Source r e) =>
(e -> a -> a) -> a -> Array r ix e -> a
foldrFB
{-# INLINE lazyFoldrS #-}

-- | /O(n)/ - Left fold, computed sequentially.
--
-- @since 0.1.0
foldlS :: (Index ix, Source r e) => (a -> e -> a) -> a -> Array r ix e -> a
foldlS :: forall ix r e a.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> a
foldlS a -> e -> a
f a
acc = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> e -> m a) -> a -> Array r ix e -> m a
foldlM (\a
a e
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! a -> e -> a
f a
a e
e) a
acc
{-# INLINE foldlS #-}

-- | /O(n)/ - Left fold with an index aware function, computed sequentially.
--
-- @since 0.1.0
ifoldlS
  :: (Index ix, Source r e)
  => (a -> ix -> e -> a)
  -> a
  -> Array r ix e
  -> a
ifoldlS :: forall ix r e a.
(Index ix, Source r e) =>
(a -> ix -> e -> a) -> a -> Array r ix e -> a
ifoldlS a -> ix -> e -> a
f a
acc = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM (\a
a ix
ix e
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! a -> ix -> e -> a
f a
a ix
ix e
e) a
acc
{-# INLINE ifoldlS #-}

-- | /O(n)/ - Right fold, computed sequentially.
--
-- @since 0.1.0
foldrS :: (Index ix, Source r e) => (e -> a -> a) -> a -> Array r ix e -> a
foldrS :: forall ix r e a.
(Index ix, Source r e) =>
(e -> a -> a) -> a -> Array r ix e -> a
foldrS e -> a -> a
f a
acc = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(e -> a -> m a) -> a -> Array r ix e -> m a
foldrM (\e
e a
a -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! e -> a -> a
f e
e a
a) a
acc
{-# INLINE foldrS #-}

-- | /O(n)/ - Right fold with an index aware function, computed sequentially.
--
-- @since 0.1.0
ifoldrS :: (Index ix, Source r e) => (ix -> e -> a -> a) -> a -> Array r ix e -> a
ifoldrS :: forall ix r e a.
(Index ix, Source r e) =>
(ix -> e -> a -> a) -> a -> Array r ix e -> a
ifoldrS ix -> e -> a -> a
f a
acc = forall a. Identity a -> a
runIdentity forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM (\ix
ix e
e a
a -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! ix -> e -> a -> a
f ix
ix e
e a
a) a
acc
{-# INLINE ifoldrS #-}

-- | Version of foldr that supports @foldr/build@ list fusion implemented by GHC.
--
-- @since 0.1.0
foldrFB :: (Index ix, Source r e) => (e -> b -> b) -> b -> Array r ix e -> b
foldrFB :: forall ix r e a.
(Index ix, Source r e) =>
(e -> a -> a) -> a -> Array r ix e -> a
foldrFB e -> b -> b
c b
n Array r ix e
arr = Int -> b
go Int
0
  where
    !k :: Int
k = forall ix. Index ix => Sz ix -> Int
totalElem (forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr)
    go :: Int -> b
go !Int
i
      | Int
i forall a. Eq a => a -> a -> Bool
== Int
k = b
n
      | Bool
otherwise = let v :: e
v = forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i in e
v e -> b -> b
`c` Int -> b
go (Int
i forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE [0] foldrFB #-}

-- | /O(n)/ - Left fold, computed with respect of array's computation strategy. Because we do
-- potentially split the folding among many threads, we also need a combining function and an
-- accumulator for the results. Depending on the number of threads being used, results can be
-- different, hence is the `MonadIO` constraint.
--
-- ===__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> foldlP (flip (:)) [] (flip (:)) [] $ makeArrayR D Seq (Sz1 6) id
-- [[5,4,3,2,1,0]]
-- >>> foldlP (flip (:)) [] (++) [] $ makeArrayR D Seq (Sz1 6) id
-- [5,4,3,2,1,0]
-- >>> foldlP (flip (:)) [] (flip (:)) [] $ makeArrayR D (ParN 3) (Sz1 6) id
-- [[5,4],[3,2],[1,0]]
-- >>> foldlP (flip (:)) [] (++) [] $ makeArrayR D (ParN 3) (Sz1 6) id
-- [1,0,3,2,5,4]
--
-- @since 0.1.0
foldlP
  :: (MonadIO m, Index ix, Source r e)
  => (a -> e -> a)
  -- ^ Folding function @g@.
  -> a
  -- ^ Accumulator. Will be applied to @g@ multiple times, thus must be neutral.
  -> (b -> a -> b)
  -- ^ Chunk results folding function @f@.
  -> b
  -- ^ Accumulator for results of chunks folding.
  -> Array r ix e
  -> m b
foldlP :: forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> m b
foldlP a -> e -> a
f a
fAcc b -> a -> b
g b
gAcc =
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(a -> e -> m a) -> a -> (b -> a -> m b) -> b -> Array r ix e -> m b
foldlIO (\a
acc -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> e -> a
f a
acc) a
fAcc (\b
acc -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> a -> b
g b
acc) b
gAcc
{-# INLINE foldlP #-}

-- | /O(n)/ - Left fold with an index aware function, computed in parallel. Just
-- like `foldlP`, except that folding function will receive an index of an
-- element it is being applied to.
--
-- @since 0.1.0
ifoldlP
  :: (MonadIO m, Index ix, Source r e)
  => (a -> ix -> e -> a)
  -> a
  -> (b -> a -> b)
  -> b
  -> Array r ix e
  -> m b
ifoldlP :: forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP a -> ix -> e -> a
f a
fAcc b -> a -> b
g b
gAcc =
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(a -> ix -> e -> m a)
-> a -> (b -> a -> m b) -> b -> Array r ix e -> m b
ifoldlIO (\a
acc ix
ix -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ix -> e -> a
f a
acc ix
ix) a
fAcc (\b
acc -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> a -> b
g b
acc) b
gAcc
{-# INLINE ifoldlP #-}

-- | /O(n)/ - Right fold, computed with respect to computation strategy. Same as `foldlP`, except
-- directed from the last element in the array towards beginning.
--
-- ==== __Examples__
--
-- >>> import Data.Massiv.Array
-- >>> foldrP (:) [] (++) [] $ makeArrayR D (ParN 2) (Sz2 2 3) fromIx2
-- [(0,0),(0,1),(0,2),(1,0),(1,1),(1,2)]
-- >>> foldrP (:) [] (:) [] $ makeArrayR D Seq (Sz1 6) id
-- [[0,1,2,3,4,5]]
-- >>> foldrP (:) [] (:) [] $ makeArrayR D (ParN 3) (Sz1 6) id
-- [[0,1],[2,3],[4,5]]
--
-- @since 0.1.0
foldrP
  :: (MonadIO m, Index ix, Source r e)
  => (e -> a -> a)
  -> a
  -> (a -> b -> b)
  -> b
  -> Array r ix e
  -> m b
foldrP :: forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> m b
foldrP e -> a -> a
f a
fAcc a -> b -> b
g b
gAcc = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(ix -> e -> a -> a)
-> a -> (a -> b -> b) -> b -> Array r ix e -> m b
ifoldrP (forall a b. a -> b -> a
const e -> a -> a
f) a
fAcc a -> b -> b
g b
gAcc
{-# INLINE foldrP #-}

-- | /O(n)/ - Right fold with an index aware function, while respecting the computation strategy.
-- Same as `ifoldlP`, except directed from the last element in the array towards
-- beginning, but also row-major.
--
-- @since 0.1.0
ifoldrP
  :: (MonadIO m, Index ix, Source r e)
  => (ix -> e -> a -> a)
  -> a
  -> (a -> b -> b)
  -> b
  -> Array r ix e
  -> m b
ifoldrP :: forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(ix -> e -> a -> a)
-> a -> (a -> b -> b) -> b -> Array r ix e -> m b
ifoldrP ix -> e -> a -> a
f a
fAcc a -> b -> b
g b
gAcc = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(ix -> e -> a -> m a)
-> a -> (a -> b -> m b) -> b -> Array r ix e -> m b
ifoldrIO (\ix
ix e
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e -> a -> a
f ix
ix e
e) a
fAcc (\a
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b -> b
g a
e) b
gAcc
{-# INLINE ifoldrP #-}

-- | This folding function breaks referential transparency on some functions
-- @f@, therefore it is kept here for internal use only.
foldlInternal
  :: (Index ix, Source r e) => (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal :: forall ix r e a b.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal a -> e -> a
g a
initAcc b -> a -> b
f b
resAcc = forall a. IO a -> a
unsafePerformIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> m b
foldlP a -> e -> a
g a
initAcc b -> a -> b
f b
resAcc
{-# INLINE foldlInternal #-}

ifoldlInternal
  :: (Index ix, Source r e) => (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
ifoldlInternal :: forall ix r e a b.
(Index ix, Source r e) =>
(a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
ifoldlInternal a -> ix -> e -> a
g a
initAcc b -> a -> b
f b
resAcc = forall a. IO a -> a
unsafePerformIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) ix r e a b.
(MonadIO m, Index ix, Source r e) =>
(a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP a -> ix -> e -> a
g a
initAcc b -> a -> b
f b
resAcc
{-# INLINE ifoldlInternal #-}

-- | Similar to `foldlP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
foldlIO
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (a -> e -> m a)
  -- ^ Index aware folding IO action
  -> a
  -- ^ Accumulator
  -> (b -> a -> m b)
  -- ^ Folding action that is applied to the results of a parallel fold
  -> b
  -- ^ Accumulator for chunks folding
  -> Array r ix e
  -> m b
foldlIO :: forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(a -> e -> m a) -> a -> (b -> a -> m b) -> b -> Array r ix e -> m b
foldlIO a -> e -> m a
f !a
initAcc b -> a -> m b
g !b
tAcc !Array r ix e
arr
  | forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr forall a. Eq a => a -> a -> Bool
== Comp
Seq = forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> e -> m a) -> a -> Array r ix e -> m a
foldlM a -> e -> m a
f a
initAcc Array r ix e
arr forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> a -> m b
g b
tAcc
  | Bool
otherwise = do
      let splitAcc :: a -> ST RealWorld (a, a)
splitAcc a
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
initAcc, a
initAcc)
          !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
      [a]
results <-
        forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld a
scheduler ->
          forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
            forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim forall a b. (a -> b) -> a -> b
$
              case forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
                PrefIndex ix -> e
gix ->
                  forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterFullAccST RowMajor
defRowMajor Scheduler RealWorld a
scheduler forall ix. Index ix => ix
zeroIndex Sz ix
sz a
initAcc a -> ST RealWorld (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$ \ !ix
ix !a
acc ->
                    forall (m :: * -> *) a.
(PrimMonad m, PrimState m ~ RealWorld) =>
IO a -> m a
ioToPrim (forall a. m a -> IO a
run (a -> e -> m a
f a
acc (ix -> e
gix ix
ix)))
                PrefIndexLinear Int -> e
gi ->
                  forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterFullAccST RowMajor
defRowMajor Scheduler RealWorld a
scheduler Int
0 (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz) a
initAcc a -> ST RealWorld (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$ \ !Int
i !a
acc ->
                    forall (m :: * -> *) a.
(PrimMonad m, PrimState m ~ RealWorld) =>
IO a -> m a
ioToPrim (forall a. m a -> IO a
run (a -> e -> m a
f a
acc (Int -> e
gi Int
i)))
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM b -> a -> m b
g b
tAcc [a]
results
{-# INLINE foldlIO #-}

-- | Similar to `ifoldlP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
ifoldlIO
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (a -> ix -> e -> m a)
  -- ^ Index aware folding IO action
  -> a
  -- ^ Accumulator
  -> (b -> a -> m b)
  -- ^ Folding action that is applied to the results of a parallel fold
  -> b
  -- ^ Accumulator for chunks folding
  -> Array r ix e
  -> m b
ifoldlIO :: forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(a -> ix -> e -> m a)
-> a -> (b -> a -> m b) -> b -> Array r ix e -> m b
ifoldlIO a -> ix -> e -> m a
f !a
initAcc b -> a -> m b
g !b
tAcc !Array r ix e
arr
  | forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr forall a. Eq a => a -> a -> Bool
== Comp
Seq = forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM a -> ix -> e -> m a
f a
initAcc Array r ix e
arr forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> a -> m b
g b
tAcc
  | Bool
otherwise = do
      let !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
          splitAcc :: a -> ST RealWorld (a, a)
splitAcc a
_ = forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
initAcc, a
initAcc)
      [a]
results <-
        forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld a
scheduler ->
          forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
            forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim forall a b. (a -> b) -> a -> b
$
              case forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
                PrefIndex ix -> e
gix ->
                  forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> ix
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (ix -> a -> ST s a)
-> ST s a
iterFullAccST RowMajor
defRowMajor Scheduler RealWorld a
scheduler forall ix. Index ix => ix
zeroIndex Sz ix
sz a
initAcc a -> ST RealWorld (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$ \ !ix
ix !a
acc ->
                    forall (m :: * -> *) a.
(PrimMonad m, PrimState m ~ RealWorld) =>
IO a -> m a
ioToPrim (forall a. m a -> IO a
run (a -> ix -> e -> m a
f a
acc ix
ix (ix -> e
gix ix
ix)))
                PrefIndexLinear Int -> e
gi ->
                  forall it ix s a.
(Iterator it, Index ix) =>
it
-> Scheduler s a
-> Int
-> Sz ix
-> a
-> (a -> ST s (a, a))
-> (Int -> ix -> a -> ST s a)
-> ST s a
iterTargetFullAccST RowMajor
defRowMajor Scheduler RealWorld a
scheduler Int
0 Sz ix
sz a
initAcc a -> ST RealWorld (a, a)
splitAcc forall a b. (a -> b) -> a -> b
$ \ !Int
i !ix
ix !a
acc ->
                    forall (m :: * -> *) a.
(PrimMonad m, PrimState m ~ RealWorld) =>
IO a -> m a
ioToPrim (forall a. m a -> IO a
run (a -> ix -> e -> m a
f a
acc ix
ix (Int -> e
gi Int
i)))
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM b -> a -> m b
g b
tAcc [a]
results
{-# INLINE ifoldlIO #-}

-- | Slice an array into linear row-major vector chunks and apply an action to each of
-- them. Number of chunks will depend on the computation strategy. Results of each action
-- will be combined with a folding function.
--
-- @since 1.0.0
splitReduce
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (Scheduler RealWorld a -> Vector r e -> m a)
  -> (b -> a -> m b)
  -- ^ Folding action that is applied to the results of a parallel fold
  -> b
  -- ^ Accumulator for chunks folding
  -> Array r ix e
  -> m b
splitReduce :: forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(Scheduler RealWorld a -> Vector r e -> m a)
-> (b -> a -> m b) -> b -> Array r ix e -> m b
splitReduce Scheduler RealWorld a -> Vector r e -> m a
f b -> a -> m b
g !b
tAcc !Array r ix e
arr = do
  let !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
      !totalLength :: Int
totalLength = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
  [a]
results <-
    forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld a
scheduler -> do
      forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> do
        forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler RealWorld a
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
          forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
            forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld a
scheduler forall a b. (a -> b) -> a -> b
$
              forall a. m a -> IO a
run forall a b. (a -> b) -> a -> b
$
                Scheduler RealWorld a -> Vector r e -> m a
f Scheduler RealWorld a
scheduler forall a b. (a -> b) -> a -> b
$
                  forall r e ix.
(Source r e, Index ix) =>
Int -> Sz1 -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
start (forall ix. ix -> Sz ix
SafeSz Int
chunkLength) Array r ix e
arr
          forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
            forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld a
scheduler forall a b. (a -> b) -> a -> b
$
              forall a. m a -> IO a
run forall a b. (a -> b) -> a -> b
$
                Scheduler RealWorld a -> Vector r e -> m a
f Scheduler RealWorld a
scheduler forall a b. (a -> b) -> a -> b
$
                  forall r e ix.
(Source r e, Index ix) =>
Int -> Sz1 -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
slackStart (forall ix. ix -> Sz ix
SafeSz (Int
totalLength forall a. Num a => a -> a -> a
- Int
slackStart)) Array r ix e
arr
  forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM b -> a -> m b
g b
tAcc [a]
results
{-# INLINE splitReduce #-}

-- | Similar to `ifoldrP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
ifoldrIO
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (ix -> e -> a -> m a)
  -> a
  -> (a -> b -> m b)
  -> b
  -> Array r ix e
  -> m b
ifoldrIO :: forall (m :: * -> *) ix r e a b.
(MonadUnliftIO m, Index ix, Source r e) =>
(ix -> e -> a -> m a)
-> a -> (a -> b -> m b) -> b -> Array r ix e -> m b
ifoldrIO ix -> e -> a -> m a
f !a
initAcc a -> b -> m b
g !b
tAcc !Array r ix e
arr
  | forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr forall a. Eq a => a -> a -> Bool
== Comp
Seq = forall ix r e (m :: * -> *) a.
(Index ix, Source r e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM ix -> e -> a -> m a
f a
initAcc Array r ix e
arr forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> b -> m b
`g` b
tAcc)
  | Bool
otherwise = do
      let !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
          !totalLength :: Int
totalLength = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
      [a]
results <-
        forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> do
          forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld a
scheduler ->
            forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler RealWorld a
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
              forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
                forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld a
scheduler forall a b. (a -> b) -> a -> b
$
                  forall a. m a -> IO a
run forall a b. (a -> b) -> a -> b
$
                    forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterLinearM Sz ix
sz (Int
totalLength forall a. Num a => a -> a -> a
- Int
1) Int
slackStart (-Int
1) forall a. Ord a => a -> a -> Bool
(>=) a
initAcc forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
ix ->
                      ix -> e -> a -> m a
f ix
ix (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
              forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
slackStart (forall a. Ord a => a -> a -> Bool
> Int
0) (forall a. Num a => a -> a -> a
subtract Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
                forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld a
scheduler forall a b. (a -> b) -> a -> b
$
                  forall a. m a -> IO a
run forall a b. (a -> b) -> a -> b
$
                    forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
iterLinearM Sz ix
sz (Int
start forall a. Num a => a -> a -> a
- Int
1) (Int
start forall a. Num a => a -> a -> a
- Int
chunkLength) (-Int
1) forall a. Ord a => a -> a -> Bool
(>=) a
initAcc forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
ix ->
                      ix -> e -> a -> m a
f ix
ix (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
      forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM (forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> m b
g) b
tAcc [a]
results
{-# INLINE ifoldrIO #-}

-- | Sequential implementation of `any` with unrolling
anySu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool
anySu :: forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
anySu e -> Bool
f Array r ix e
arr = Int -> Bool
go Int
0
  where
    !k :: Int
k = forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount Array r ix e
arr
    !k4 :: Int
k4 = Int
k forall a. Num a => a -> a -> a
- (Int
k forall a. Integral a => a -> a -> a
`rem` Int
4)
    go :: Int -> Bool
go !Int
i
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
k4 =
          e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
            Bool -> Bool -> Bool
|| e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int
i forall a. Num a => a -> a -> a
+ Int
1))
            Bool -> Bool -> Bool
|| e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int
i forall a. Num a => a -> a -> a
+ Int
2))
            Bool -> Bool -> Bool
|| e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int
i forall a. Num a => a -> a -> a
+ Int
3))
            Bool -> Bool -> Bool
|| Int -> Bool
go (Int
i forall a. Num a => a -> a -> a
+ Int
4)
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
k = e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i) Bool -> Bool -> Bool
|| Int -> Bool
go (Int
i forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = Bool
False
{-# INLINE anySu #-}

-- | Implementaton of `any` on a slice of an array with short-circuiting using batch cancellation.
anySliceSuM
  :: (Index ix, Source r e)
  => Batch RealWorld Bool
  -> Ix1
  -> Sz1
  -> (e -> Bool)
  -> Array r ix e
  -> IO Bool
anySliceSuM :: forall ix r e.
(Index ix, Source r e) =>
Batch RealWorld Bool
-> Int -> Sz1 -> (e -> Bool) -> Array r ix e -> IO Bool
anySliceSuM Batch RealWorld Bool
batch Int
ix0 (Sz1 Int
k) e -> Bool
f Array r ix e
arr = Int -> IO Bool
go Int
ix0
  where
    !k' :: Int
k' = Int
k forall a. Num a => a -> a -> a
- Int
ix0
    !k4 :: Int
k4 = Int
ix0 forall a. Num a => a -> a -> a
+ (Int
k' forall a. Num a => a -> a -> a
- (Int
k' forall a. Integral a => a -> a -> a
`rem` Int
4))
    go :: Int -> IO Bool
go !Int
i
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
k4 = do
          let r :: Bool
r =
                e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
                  Bool -> Bool -> Bool
|| e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int
i forall a. Num a => a -> a -> a
+ Int
1))
                  Bool -> Bool -> Bool
|| e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int
i forall a. Num a => a -> a -> a
+ Int
2))
                  Bool -> Bool -> Bool
|| e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr (Int
i forall a. Num a => a -> a -> a
+ Int
3))
           in if Bool
r
                then forall s (m :: * -> *) a. MonadPrim s m => Batch s a -> a -> m Bool
cancelBatchWith Batch RealWorld Bool
batch Bool
True
                else do
                  Bool
done <- forall s (m :: * -> *) a. MonadPrim s m => Batch s a -> m Bool
hasBatchFinished Batch RealWorld Bool
batch
                  if Bool
done
                    then forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
                    else Int -> IO Bool
go (Int
i forall a. Num a => a -> a -> a
+ Int
4)
      | Int
i forall a. Ord a => a -> a -> Bool
< Int
k =
          if e -> Bool
f (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
            then forall s (m :: * -> *) a. MonadPrim s m => Batch s a -> a -> m Bool
cancelBatchWith Batch RealWorld Bool
batch Bool
True
            else Int -> IO Bool
go (Int
i forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
{-# INLINE anySliceSuM #-}

-- | Parallelizable implementation of `any` with unrolling
anyPu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> IO Bool
-- TODO: switch to splitReduce
-- anyPu f arr =
--   splitReduce anySu (\r acc -> pure (r || acc)) False
anyPu :: forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> IO Bool
anyPu e -> Bool
f Array r ix e
arr = do
  let !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
      !totalLength :: Int
totalLength = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
  [Bool]
results <-
    forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld Bool
scheduler -> do
      Batch RealWorld Bool
batch <- forall s (m :: * -> *) a.
MonadPrim s m =>
Scheduler s a -> m (Batch s a)
getCurrentBatch Scheduler RealWorld Bool
scheduler
      forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (forall s a. Scheduler s a -> Int
numWorkers Scheduler RealWorld Bool
scheduler) Int
totalLength forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
        forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (forall a. Ord a => a -> a -> Bool
< Int
slackStart) (forall a. Num a => a -> a -> a
+ Int
chunkLength) forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
          forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld Bool
scheduler forall a b. (a -> b) -> a -> b
$ forall ix r e.
(Index ix, Source r e) =>
Batch RealWorld Bool
-> Int -> Sz1 -> (e -> Bool) -> Array r ix e -> IO Bool
anySliceSuM Batch RealWorld Bool
batch Int
start (forall ix. Index ix => ix -> Sz ix
Sz (Int
start forall a. Num a => a -> a -> a
+ Int
chunkLength)) e -> Bool
f Array r ix e
arr
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart forall a. Ord a => a -> a -> Bool
< Int
totalLength) forall a b. (a -> b) -> a -> b
$
          forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld Bool
scheduler forall a b. (a -> b) -> a -> b
$
            forall ix r e.
(Index ix, Source r e) =>
Batch RealWorld Bool
-> Int -> Sz1 -> (e -> Bool) -> Array r ix e -> IO Bool
anySliceSuM Batch RealWorld Bool
batch Int
slackStart (forall ix. Index ix => ix -> Sz ix
Sz Int
totalLength) e -> Bool
f Array r ix e
arr
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' Bool -> Bool -> Bool
(||) Bool
False [Bool]
results
{-# INLINE anyPu #-}

-- | /O(n)/ - Determines whether any element of the array satisfies a predicate.
--
-- @since 0.1.0
any :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool
any :: forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any e -> Bool
f Array r ix e
arr =
  case forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr of
    Comp
Seq -> forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
anySu e -> Bool
f Array r ix e
arr
    Comp
_ -> forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> IO Bool
anyPu e -> Bool
f Array r ix e
arr
{-# INLINE any #-}