{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Fold.Internal
-- Copyright   : (c) Alexey Kuleshevich 2018-2021
-- 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
  , ifoldlIO
  , ifoldrIO
  -- , splitReduce
  , any
  , anySu
  , anyPu
  ) where

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


-- | /O(n)/ - Unstructured fold of an array.
--
-- @since 0.3.0
fold ::
     (Monoid e, Source r ix e)
  => Array r ix e -- ^ Source array
  -> e
fold :: Array r ix e -> e
fold = (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
forall r ix e a b.
Source r ix e =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal e -> e -> e
forall a. Monoid a => a -> a -> a
mappend e
forall a. Monoid a => a
mempty e -> e -> e
forall a. Monoid a => a -> a -> a
mappend e
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 ::
     (Source r ix e, Monoid m)
  => (e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`.
  -> Array r ix e -- ^ Source array
  -> m
foldMono :: (e -> m) -> Array r ix e -> m
foldMono e -> m
f = (m -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
forall r ix e a b.
Source r ix e =>
(a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal (\m
a e
e -> m
a m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` e -> m
f e
e) m
forall a. Monoid a => a
mempty m -> m -> m
forall a. Monoid a => a -> a -> a
mappend m
forall a. Monoid a => a
mempty
{-# INLINE foldMono #-}


-- | /O(n)/ - Monadic left fold.
--
-- @since 0.1.0
foldlM :: (Source r ix e, Monad m) => (a -> e -> m a) -> a -> Array r ix e -> m a
foldlM :: (a -> e -> m a) -> a -> Array r ix e -> m a
foldlM a -> e -> m a
f = (a -> ix -> e -> m a) -> a -> Array r ix e -> m a
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM (\ a
a ix
_ e
b -> a -> e -> m a
f a
a e
b)
{-# INLINE foldlM #-}


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


-- | /O(n)/ - Monadic left fold with an index aware function.
--
-- @since 0.1.0
ifoldlM :: (Source r ix e, Monad m) => (a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM :: (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 =
  ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM ix
forall ix. Index ix => ix
zeroIndex (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)) (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
acc ((ix -> a -> m a) -> m a) -> (ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !a
a -> a -> ix -> e -> m a
f a
a ix
ix (Array r ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
ix)
{-# INLINE ifoldlM #-}


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


-- | /O(n)/ - Monadic right fold, that discards the result.
--
-- @since 0.1.0
foldrM_ :: (Source r ix e, Monad m) => (e -> a -> m a) -> a -> Array r ix e -> m ()
foldrM_ :: (e -> a -> m a) -> a -> Array r ix e -> m ()
foldrM_ e -> a -> m a
f = (ix -> e -> a -> m a) -> a -> Array r ix e -> m ()
forall r ix e (m :: * -> *) a.
(Source r ix 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 :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM :: (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 =
  ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
forall ix (m :: * -> *) a.
(Index ix, Monad m) =>
ix
-> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> m a) -> m a
iterM ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr))) ix
forall ix. Index ix => ix
zeroIndex (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex (-Int
1)) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(>=) a
acc ((ix -> a -> m a) -> m a) -> (ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !ix
ix !a
acc0 ->
    ix -> e -> a -> m a
f ix
ix (Array r ix e -> ix -> e
forall r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
ix) a
acc0
{-# INLINE ifoldrM #-}


-- | /O(n)/ - Monadic right fold with an index aware function, that discards the result.
--
-- @since 0.1.0
ifoldrM_ :: (Source r ix e, Monad m) => (ix -> e -> a -> m a) -> a -> Array r ix e -> m ()
ifoldrM_ :: (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 = m a -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$ (ix -> e -> a -> m a) -> a -> Array r ix e -> m a
forall r ix e (m :: * -> *) a.
(Source r ix 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 :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a
lazyFoldlS :: (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 = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)
    go :: a -> Int -> a
go a
acc Int
k
      | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = a -> Int -> a
go (a -> e -> a
f a
acc (Array r ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
k)) (Int
k Int -> Int -> Int
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 :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a
lazyFoldrS :: (e -> a -> a) -> a -> Array r ix e -> a
lazyFoldrS = (e -> a -> a) -> a -> Array r ix e -> a
forall r ix e b.
Source r ix e =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrFB
{-# INLINE lazyFoldrS #-}


-- | /O(n)/ - Left fold, computed sequentially.
--
-- @since 0.1.0
foldlS :: Source r ix e => (a -> e -> a) -> a -> Array r ix e -> a
foldlS :: (a -> e -> a) -> a -> Array r ix e -> a
foldlS a -> e -> a
f = (a -> ix -> e -> a) -> a -> Array r ix e -> a
forall r ix e a.
Source r ix e =>
(a -> ix -> e -> a) -> a -> Array r ix e -> a
ifoldlS (\ a
a ix
_ e
e -> a -> e -> a
f a
a e
e)
{-# INLINE foldlS #-}


-- | /O(n)/ - Left fold with an index aware function, computed sequentially.
--
-- @since 0.1.0
ifoldlS :: Source r ix e
        => (a -> ix -> e -> a) -> a -> Array r ix e -> a
ifoldlS :: (a -> ix -> e -> a) -> a -> Array r ix e -> a
ifoldlS a -> ix -> e -> a
f a
acc = Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a)
-> (Array r ix e -> Identity a) -> Array r ix e -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ix -> e -> Identity a) -> a -> Array r ix e -> Identity a
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
(a -> ix -> e -> m a) -> a -> Array r ix e -> m a
ifoldlM (\ a
a ix
ix e
e -> a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Identity a) -> a -> Identity a
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 :: Source r ix e => (e -> a -> a) -> a -> Array r ix e -> a
foldrS :: (e -> a -> a) -> a -> Array r ix e -> a
foldrS e -> a -> a
f = (ix -> e -> a -> a) -> a -> Array r ix e -> a
forall r ix e a.
Source r ix e =>
(ix -> e -> a -> a) -> a -> Array r ix e -> a
ifoldrS (\ix
_ e
e a
a -> e -> a -> a
f e
e a
a)
{-# INLINE foldrS #-}


-- | /O(n)/ - Right fold with an index aware function, computed sequentially.
--
-- @since 0.1.0
ifoldrS :: Source r ix e => (ix -> e -> a -> a) -> a -> Array r ix e -> a
ifoldrS :: (ix -> e -> a -> a) -> a -> Array r ix e -> a
ifoldrS ix -> e -> a -> a
f a
acc = Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a)
-> (Array r ix e -> Identity a) -> Array r ix e -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix -> e -> a -> Identity a) -> a -> Array r ix e -> Identity a
forall r ix e (m :: * -> *) a.
(Source r ix e, Monad m) =>
(ix -> e -> a -> m a) -> a -> Array r ix e -> m a
ifoldrM (\ ix
ix e
e a
a -> a -> Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Identity a) -> a -> Identity a
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 :: Source r ix e => (e -> b -> b) -> b -> Array r ix e -> b
foldrFB :: (e -> b -> b) -> b -> Array r ix e -> b
foldrFB e -> b -> b
c b
n Array r ix e
arr = Int -> b
go Int
0
  where
    !k :: Int
k = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)
    go :: Int -> b
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
k = b
n
      | Bool
otherwise = let !v :: e
v = Array r ix e -> Int -> e
forall r ix e. Source r ix e => 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 Int -> Int -> Int
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, Source r ix 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 :: (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 = IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> (Array r ix e -> IO b) -> Array r ix e -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> IO b
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix e) =>
(a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP (\ a
x ix
_ -> a -> e -> a
f a
x) a
fAcc b -> a -> b
g 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, Source r ix e) =>
           (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> m b
ifoldlP :: (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 =
  IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> (Array r ix e -> IO b) -> Array r ix e -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ix -> e -> IO a)
-> a -> (b -> a -> IO b) -> b -> Array r ix e -> IO b
forall (m :: * -> *) r ix e a b.
(MonadUnliftIO m, Source r ix e) =>
(a -> ix -> e -> m a)
-> a -> (b -> a -> m b) -> b -> Array r ix e -> m b
ifoldlIO (\a
acc ix
ix -> a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> (e -> a) -> e -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ix -> e -> a
f a
acc ix
ix) a
fAcc (\b
acc -> b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> IO b) -> (a -> b) -> a -> IO b
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, Source r ix e) =>
          (e -> a -> a) -> a -> (a -> b -> b) -> b -> Array r ix e -> m b
foldrP :: (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 = IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> (Array r ix e -> IO b) -> Array r ix e -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix -> e -> a -> a)
-> a -> (a -> b -> b) -> b -> Array r ix e -> IO b
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix e) =>
(ix -> e -> a -> a)
-> a -> (a -> b -> b) -> b -> Array r ix e -> m b
ifoldrP ((e -> a -> a) -> ix -> e -> a -> a
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, Source r ix e)
  => (ix -> e -> a -> a)
  -> a
  -> (a -> b -> b)
  -> b
  -> Array r ix e
  -> m b
ifoldrP :: (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 = IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> (Array r ix e -> IO b) -> Array r ix e -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix -> e -> a -> IO a)
-> a -> (a -> b -> IO b) -> b -> Array r ix e -> IO b
forall (m :: * -> *) r ix e a b.
(MonadUnliftIO m, Source r ix e) =>
(ix -> e -> a -> m a)
-> a -> (a -> b -> m b) -> b -> Array r ix e -> m b
ifoldrIO (\ix
ix e
e -> a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> IO a) -> (a -> a) -> a -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e -> a -> a
f ix
ix e
e) a
fAcc (\a
e -> b -> IO b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> IO b) -> (b -> b) -> b -> IO b
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 :: Source r ix e => (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
foldlInternal :: (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 = IO b -> b
forall a. IO a -> a
unsafePerformIO (IO b -> b) -> (Array r ix e -> IO b) -> Array r ix e -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> IO b
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix 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 :: Source r ix e => (a -> ix -> e -> a) -> a -> (b -> a -> b) -> b -> Array r ix e -> b
ifoldlInternal :: (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 = IO b -> b
forall a. IO a -> a
unsafePerformIO (IO b -> b) -> (Array r ix e -> IO b) -> Array r ix e -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> ix -> e -> a)
-> a -> (b -> a -> b) -> b -> Array r ix e -> IO b
forall (m :: * -> *) r ix e a b.
(MonadIO m, Source r ix 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 `ifoldlP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
ifoldlIO ::
     (MonadUnliftIO m, Source r ix 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 :: (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
  | Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr Comp -> Comp -> Bool
forall a. Eq a => a -> a -> Bool
== Comp
Seq = (a -> ix -> e -> m a) -> a -> Array r ix e -> m a
forall r ix e (m :: * -> *) a.
(Source r ix 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 m a -> (a -> m b) -> m b
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 = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
          !totalLength :: Int
totalLength = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
      [a]
results <-
        Comp -> (Scheduler m a -> m ()) -> m [a]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m [a]
withScheduler (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler m a -> m ()) -> m [a])
-> (Scheduler m a -> m ()) -> m [a]
forall a b. (a -> b) -> a -> b
$ \Scheduler m a
scheduler ->
          Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler m a -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler m a
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
            Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
              Scheduler m a -> m a -> m ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler m a
scheduler (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$
              Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
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 (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
initAcc ((Int -> ix -> a -> m a) -> m a) -> (Int -> ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
ix !a
acc ->
                a -> ix -> e -> m a
f a
acc ix
ix (Array r ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
            Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
              Scheduler m a -> m a -> m ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler m a
scheduler (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$
              Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
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
slackStart Int
totalLength Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<) a
initAcc ((Int -> ix -> a -> m a) -> m a) -> (Int -> ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
ix !a
acc ->
                a -> ix -> e -> m a
f a
acc ix
ix (Array r ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i)
      (b -> a -> m b) -> b -> [a] -> m b
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 #-}

-- -- | Split 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 0.6.0
-- splitReduce ::
--      (MonadUnliftIO m, Source r ix e)
--   => (Scheduler m a -> BatchId -> Array r Ix1 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 f g !tAcc !arr = do
--   let !sz = size arr
--       !totalLength = totalElem sz
--   results <-
--     withScheduler (getComp arr) $ \scheduler -> do
--       batchId <- getCurrentBatchId scheduler
--       splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
--         loopM_ 0 (< slackStart) (+ chunkLength) $ \ !start ->
--           scheduleWork scheduler $ f scheduler batchId $
--             unsafeLinearSlice start (SafeSz chunkLength) arr
--         when (slackStart < totalLength) $
--           scheduleWork scheduler $ f scheduler batchId $
--             unsafeLinearSlice slackStart (SafeSz (totalLength - slackStart)) arr
--   F.foldlM g tAcc results
-- {-# INLINE splitReduce #-}



-- | Similar to `ifoldrP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
ifoldrIO :: (MonadUnliftIO m, Source r ix e) =>
           (ix -> e -> a -> m a) -> a -> (a -> b -> m b) -> b -> Array r ix e -> m b
ifoldrIO :: (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
  | Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr Comp -> Comp -> Bool
forall a. Eq a => a -> a -> Bool
== Comp
Seq = (ix -> e -> a -> m a) -> a -> Array r ix e -> m a
forall r ix e (m :: * -> *) a.
(Source r ix 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 m a -> (a -> m b) -> m b
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 = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
        !totalLength :: Int
totalLength = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
    [a]
results <-
      Comp -> (Scheduler m a -> m ()) -> m [a]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m [a]
withScheduler (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler m a -> m ()) -> m [a])
-> (Scheduler m a -> m ()) -> m [a]
forall a b. (a -> b) -> a -> b
$ \ Scheduler m a
scheduler ->
        Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler m a -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler m a
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ Int
chunkLength Int
slackStart -> do
          Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
            Scheduler m a -> m a -> m ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler m a
scheduler (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$
            Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
slackStart (-Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(>=) a
initAcc ((Int -> ix -> a -> m a) -> m a) -> (Int -> ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
ix !a
acc ->
              ix -> e -> a -> m a
f ix
ix (Array r ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i) a
acc
          Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
slackStart (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
chunkLength) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
            Scheduler m a -> m a -> m ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler m a
scheduler (m a -> m ()) -> m a -> m ()
forall a b. (a -> b) -> a -> b
$
              Sz ix
-> Int
-> Int
-> Int
-> (Int -> Int -> Bool)
-> a
-> (Int -> ix -> a -> m a)
-> m a
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
chunkLength) (-Int
1) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(>=) a
initAcc ((Int -> ix -> a -> m a) -> m a) -> (Int -> ix -> a -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
ix !a
acc ->
                ix -> e -> a -> m a
f ix
ix (Array r ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix e
arr Int
i) a
acc
    (b -> a -> m b) -> b -> [a] -> m b
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM ((a -> b -> m b) -> b -> a -> m b
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 :: Source r ix a => (a -> Bool) -> Array r ix a -> Bool
anySu :: (a -> Bool) -> Array r ix a -> Bool
anySu a -> Bool
f Array r ix a
arr = Int -> Bool
go Int
0
  where
    !k :: Int
k = Array r ix a -> Int
forall r ix e. Load r ix e => Array r ix e -> Int
elemsCount Array r ix a
arr
    !k4 :: Int
k4 = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
k Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
4)
    go :: Int -> Bool
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k4 =
        a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr Int
i      ) Bool -> Bool -> Bool
||
        a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Bool -> Bool -> Bool
||
        a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)) Bool -> Bool -> Bool
||
        a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3)) Bool -> Bool -> Bool
||
        Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k = a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr Int
i) Bool -> Bool -> Bool
|| Int -> Bool
go (Int
i Int -> Int -> Int
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 ::
     Source r ix a
  => Batch IO Bool
  -> Ix1
  -> Sz1
  -> (a -> Bool)
  -> Array r ix a
  -> IO Bool
anySliceSuM :: Batch IO Bool
-> Int -> Sz1 -> (a -> Bool) -> Array r ix a -> IO Bool
anySliceSuM Batch IO Bool
batch Int
ix0 (Sz Int
k) a -> Bool
f Array r ix a
arr = Int -> IO Bool
go Int
ix0
  where
    !k' :: Int
k' = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
ix0
    !k4 :: Int
k4 = Int
ix0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
k' Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
k' Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
4))
    go :: Int -> IO Bool
go !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k4 = do
        let r :: Bool
r =
              a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr Int
i) Bool -> Bool -> Bool
||
              a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) Bool -> Bool -> Bool
||
              a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)) Bool -> Bool -> Bool
||
              a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3))
         in if Bool
r
              then Batch IO Bool -> Bool -> IO Bool
forall (m :: * -> *) a. Batch m a -> a -> m Bool
cancelBatchWith Batch IO Bool
batch Bool
True
              else do
                Bool
done <- Batch IO Bool -> IO Bool
forall (m :: * -> *) a. Functor m => Batch m a -> m Bool
hasBatchFinished Batch IO Bool
batch
                if Bool
done
                  then Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
                  else Int -> IO Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
4)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
k =
        if a -> Bool
f (Array r ix a -> Int -> a
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array r ix a
arr Int
i)
          then Batch IO Bool -> Bool -> IO Bool
forall (m :: * -> *) a. Batch m a -> a -> m Bool
cancelBatchWith Batch IO Bool
batch Bool
True
          else Int -> IO Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      | Bool
otherwise = Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
{-# INLINE anySliceSuM #-}



-- | Parallelizable implementation of `any` with unrolling
anyPu :: Source r ix e => (e -> Bool) -> Array r ix e -> IO Bool
anyPu :: (e -> Bool) -> Array r ix e -> IO Bool
anyPu e -> Bool
f Array r ix e
arr = do
  let !sz :: Sz ix
sz = Array r ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr
      !totalLength :: Int
totalLength = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
  [Bool]
results <-
    Comp -> (Scheduler IO Bool -> IO ()) -> IO [Bool]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler m a -> m b) -> m [a]
withScheduler (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) ((Scheduler IO Bool -> IO ()) -> IO [Bool])
-> (Scheduler IO Bool -> IO ()) -> IO [Bool]
forall a b. (a -> b) -> a -> b
$ \Scheduler IO Bool
scheduler -> do
      Batch IO Bool
batch <- Scheduler IO Bool -> IO (Batch IO Bool)
forall (m :: * -> *) a. Monad m => Scheduler m a -> m (Batch m a)
getCurrentBatch Scheduler IO Bool
scheduler
      Int -> Int -> (Int -> Int -> IO ()) -> IO ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler IO Bool -> Int
forall (m :: * -> *) a. Scheduler m a -> Int
numWorkers Scheduler IO Bool
scheduler) Int
totalLength ((Int -> Int -> IO ()) -> IO ()) -> (Int -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> do
        Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO ()) -> IO ()
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m a) -> m ()
loopM_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
          Scheduler IO Bool -> IO Bool -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO Bool
scheduler (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Batch IO Bool
-> Int -> Sz1 -> (e -> Bool) -> Array r ix e -> IO Bool
forall r ix a.
Source r ix a =>
Batch IO Bool
-> Int -> Sz1 -> (a -> Bool) -> Array r ix a -> IO Bool
anySliceSuM Batch IO Bool
batch Int
start (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength)) e -> Bool
f Array r ix e
arr
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
          Scheduler IO Bool -> IO Bool -> IO ()
forall (m :: * -> *) a. Scheduler m a -> m a -> m ()
scheduleWork Scheduler IO Bool
scheduler (IO Bool -> IO ()) -> IO Bool -> IO ()
forall a b. (a -> b) -> a -> b
$ Batch IO Bool
-> Int -> Sz1 -> (e -> Bool) -> Array r ix e -> IO Bool
forall r ix a.
Source r ix a =>
Batch IO Bool
-> Int -> Sz1 -> (a -> Bool) -> Array r ix a -> IO Bool
anySliceSuM Batch IO Bool
batch Int
slackStart (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz Int
totalLength) e -> Bool
f Array r ix e
arr
  Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> Bool -> Bool) -> Bool -> [Bool] -> Bool
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 :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool
any :: (e -> Bool) -> Array r ix e -> Bool
any e -> Bool
f Array r ix e
arr =
  case Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr of
    Comp
Seq -> (e -> Bool) -> Array r ix e -> Bool
forall r ix a. Source r ix a => (a -> Bool) -> Array r ix a -> Bool
anySu e -> Bool
f Array r ix e
arr
    Comp
_ -> IO Bool -> Bool
forall a. IO a -> a
unsafePerformIO (IO Bool -> Bool) -> IO Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (e -> Bool) -> Array r ix e -> IO Bool
forall r ix e.
Source r ix e =>
(e -> Bool) -> Array r ix e -> IO Bool
anyPu e -> Bool
f Array r ix e
arr
{-# INLINE any #-}