{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Data.Massiv.Array.Ops.Fold
-- 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 (
  -- ** Unstructured folds
  -- $unstruct_folds
  fold,
  ifoldMono,
  foldMono,
  ifoldSemi,
  foldSemi,
  foldOuterSlice,
  ifoldOuterSlice,
  foldInnerSlice,
  ifoldInnerSlice,
  minimumM,
  minimum',
  maximumM,
  maximum',
  sum,
  product,
  and,
  or,
  all,
  any,
  elem,
  eqArrays,
  compareArrays,

  -- ** Single dimension folds

  -- *** Safe inner most

  --
  -- Folding along the inner most dimension will always be faster when compared to doing the same
  -- operation along any other dimension, this is due to the fact that inner most folds follow the
  -- memory layout of data.
  ifoldlInner,
  foldlInner,
  ifoldrInner,
  foldrInner,
  foldInner,

  -- *** Type safe within
  ifoldlWithin,
  foldlWithin,
  ifoldrWithin,
  foldrWithin,
  foldWithin,

  -- *** Partial within
  ifoldlWithin',
  foldlWithin',
  ifoldrWithin',
  foldrWithin',
  foldWithin',

  -- ** Sequential folds
  -- $seq_folds
  foldlS,
  foldrS,
  ifoldlS,
  ifoldrS,

  -- *** Monadic
  foldlM,
  foldrM,
  foldlM_,
  foldrM_,
  ifoldlM,
  ifoldrM,
  ifoldlM_,
  ifoldrM_,

  -- *** Special folds
  foldrFB,
  lazyFoldlS,
  lazyFoldrS,

  -- ** Parallel folds
  -- $par_folds
  foldlP,
  foldrP,
  ifoldlP,
  ifoldrP,
  ifoldlIO,
  ifoldrIO,
  -- , splitReduce
) where

import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Ops.Construct
import Data.Massiv.Array.Ops.Fold.Internal
import Data.Massiv.Core
import Data.Massiv.Core.Common
import Prelude hiding (all, and, any, elem, foldl, foldr, map, maximum, minimum, or, product, sum)

-- | /O(n)/ - Monoidal fold over an array with an index aware function. Also known as reduce.
--
-- @since 0.2.4
ifoldMono
  :: (Index ix, Source r e, Monoid m)
  => (ix -> e -> m)
  -- ^ Convert each element of an array to an appropriate `Monoid`.
  -> Array r ix e
  -- ^ Source array
  -> m
ifoldMono :: forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(ix -> e -> m) -> Array r ix e -> m
ifoldMono ix -> e -> m
f = (m -> ix -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
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 (\m
a ix
ix e
e -> m
a m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` ix -> e -> m
f ix
ix 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 ifoldMono #-}

-- | /O(n)/ - Semigroup fold over an array with an index aware function.
--
-- @since 0.2.4
ifoldSemi
  :: (Index ix, Source r e, Semigroup m)
  => (ix -> e -> m)
  -- ^ Convert each element of an array to an appropriate `Semigroup`.
  -> m
  -- ^ Initial element that must be neutral to the (`<>`) function.
  -> Array r ix e
  -- ^ Source array
  -> m
ifoldSemi :: forall ix r e m.
(Index ix, Source r e, Semigroup m) =>
(ix -> e -> m) -> m -> Array r ix e -> m
ifoldSemi ix -> e -> m
f m
m = (m -> ix -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
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 (\m
a ix
ix e
e -> m
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> ix -> e -> m
f ix
ix e
e) m
m m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) m
m
{-# INLINE ifoldSemi #-}

-- | /O(n)/ - Semigroup fold over an array.
--
-- @since 0.1.6
foldSemi
  :: (Index ix, Source r e, Semigroup m)
  => (e -> m)
  -- ^ Convert each element of an array to an appropriate `Semigroup`.
  -> m
  -- ^ Initial element that must be neutral to the (`<>`) function.
  -> Array r ix e
  -- ^ Source array
  -> m
foldSemi :: forall ix r e m.
(Index ix, Source r e, Semigroup m) =>
(e -> m) -> m -> Array r ix e -> m
foldSemi e -> m
f m
m = (m -> e -> m) -> m -> (m -> m -> m) -> m -> Array r ix e -> m
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 m -> m -> m
forall a. Semigroup a => a -> a -> a
<> e -> m
f e
e) m
m m -> m -> m
forall a. Semigroup a => a -> a -> a
(<>) m
m
{-# INLINE foldSemi #-}

-- | Left fold along a specified dimension with an index aware function.
--
-- @since 0.2.4
ifoldlWithin
  :: (Index (Lower ix), IsIndexDimension ix n, Source r e)
  => Dimension n
  -> (ix -> a -> e -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
ifoldlWithin :: forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin Dimension n
dim = Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' (Dimension n -> Dim
forall (n :: Natural). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE ifoldlWithin #-}

-- | Left fold along a specified dimension.
--
-- ====__Example__
--
-- >>> import Data.Massiv.Array
-- >>> :set -XTypeApplications
-- >>> arr = makeArrayLinear @U Seq (Sz (2 :. 5)) id
-- >>> arr
-- Array U Seq (Sz (2 :. 5))
--   [ [ 0, 1, 2, 3, 4 ]
--   , [ 5, 6, 7, 8, 9 ]
--   ]
-- >>> foldlWithin Dim1 (flip (:)) [] arr
-- Array D Seq (Sz1 2)
--   [ [4,3,2,1,0], [9,8,7,6,5] ]
-- >>> foldlWithin Dim2 (flip (:)) [] arr
-- Array D Seq (Sz1 5)
--   [ [5,0], [6,1], [7,2], [8,3], [9,4] ]
--
-- @since 0.2.4
foldlWithin
  :: (Index (Lower ix), IsIndexDimension ix n, Source r e)
  => Dimension n
  -> (a -> e -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
foldlWithin :: forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin Dimension n
dim a -> e -> a
f = Dimension n
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin Dimension n
dim ((a -> e -> a) -> ix -> a -> e -> a
forall a b. a -> b -> a
const a -> e -> a
f)
{-# INLINE foldlWithin #-}

-- | Right fold along a specified dimension with an index aware function.
--
-- @since 0.2.4
ifoldrWithin
  :: (Index (Lower ix), IsIndexDimension ix n, Source r e)
  => Dimension n
  -> (ix -> e -> a -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
ifoldrWithin :: forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin Dimension n
dim = Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' (Dimension n -> Dim
forall (n :: Natural). KnownNat n => Dimension n -> Dim
fromDimension Dimension n
dim)
{-# INLINE ifoldrWithin #-}

-- | Right fold along a specified dimension.
--
-- @since 0.2.4
foldrWithin
  :: (Index (Lower ix), IsIndexDimension ix n, Source r e)
  => Dimension n
  -> (e -> a -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
foldrWithin :: forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin Dimension n
dim e -> a -> a
f = Dimension n
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin Dimension n
dim ((e -> a -> a) -> ix -> e -> a -> a
forall a b. a -> b -> a
const e -> a -> a
f)
{-# INLINE foldrWithin #-}

-- | Similar to `ifoldlWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
-- @since 0.2.4
ifoldlWithin'
  :: (HasCallStack, Index (Lower ix), Index ix, Source r e)
  => Dim
  -> (ix -> a -> e -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
ifoldlWithin' :: forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' Dim
dim ix -> a -> e -> a
f a
acc0 Array r ix e
arr =
  Comp -> Sz (Lower ix) -> (Lower ix -> a) -> Array D (Lower ix) a
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szl) ((Lower ix -> a) -> Array D (Lower ix) a)
-> (Lower ix -> a) -> Array D (Lower ix) a
forall a b. (a -> b) -> a -> b
$ \Lower ix
ixl ->
    ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
forall ix a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim Int
0)
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
      (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex Int
1)
      Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(<=)
      a
acc0
      (\ix
ix a
acc' -> ix -> a -> e -> a
f ix
ix a
acc' (Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
ix))
  where
    SafeSz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr
    (Int
k, Lower ix
szl) = ix -> Dim -> (Int, Lower ix)
forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
sz Dim
dim
{-# INLINE ifoldlWithin' #-}

-- | Similar to `foldlWithin`, except that dimension is specified at a value level, which means it will
-- throw an exception on an invalid dimension.
--
-- @since 0.2.4
foldlWithin'
  :: (HasCallStack, Index (Lower ix), Index ix, Source r e)
  => Dim
  -> (a -> e -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
foldlWithin' :: forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' Dim
dim a -> e -> a
f = Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' Dim
dim ((a -> e -> a) -> ix -> a -> e -> a
forall a b. a -> b -> a
const a -> e -> a
f)
{-# INLINE foldlWithin' #-}

-- | Similar to `ifoldrWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
--
-- @since 0.2.4
ifoldrWithin'
  :: (HasCallStack, Index (Lower ix), Index ix, Source r e)
  => Dim
  -> (ix -> e -> a -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
ifoldrWithin' :: forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' Dim
dim ix -> e -> a -> a
f a
acc0 Array r ix e
arr =
  Comp -> Sz (Lower ix) -> (Lower ix -> a) -> Array D (Lower ix) a
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
szl) ((Lower ix -> a) -> Array D (Lower ix) a)
-> (Lower ix -> a) -> Array D (Lower ix) a
forall a b. (a -> b) -> a -> b
$ \Lower ix
ixl ->
    ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
forall ix a.
Index ix =>
ix -> ix -> ix -> (Int -> Int -> Bool) -> a -> (ix -> a -> a) -> a
iter
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
      (Lower ix -> Dim -> Int -> ix
forall ix. (HasCallStack, Index ix) => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim Int
0)
      (Int -> ix
forall ix. Index ix => Int -> ix
pureIndex (-Int
1))
      Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
(>=)
      a
acc0
      (\ix
ix a
acc' -> ix -> e -> a -> a
f ix
ix (Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
ix) a
acc')
  where
    SafeSz ix
sz = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr
    (Int
k, Lower ix
szl) = ix -> Dim -> (Int, Lower ix)
forall ix. (HasCallStack, Index ix) => ix -> Dim -> (Int, Lower ix)
pullOutDim' ix
sz Dim
dim
{-# INLINE ifoldrWithin' #-}

-- | Similar to `foldrWithin`, except that dimension is specified at a value level, which means it
-- will throw an exception on an invalid dimension.
--
-- @since 0.2.4
foldrWithin'
  :: (HasCallStack, Index (Lower ix), Index ix, Source r e)
  => Dim
  -> (e -> a -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
foldrWithin' :: forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin' Dim
dim e -> a -> a
f = Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' Dim
dim ((e -> a -> a) -> ix -> e -> a -> a
forall a b. a -> b -> a
const e -> a -> a
f)
{-# INLINE foldrWithin' #-}

-- | Left fold over the inner most dimension with index aware function.
--
-- @since 0.2.4
ifoldlInner
  :: (Index (Lower ix), Index ix, Source r e)
  => (ix -> a -> e -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
ifoldlInner :: forall ix r e a.
(Index (Lower ix), Index ix, Source r e) =>
(ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlInner = Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' Dim
1
{-# INLINE ifoldlInner #-}

-- | Left fold over the inner most dimension.
--
-- @since 0.2.4
foldlInner
  :: (Index (Lower ix), Index ix, Source r e)
  => (a -> e -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
foldlInner :: forall ix r e a.
(Index (Lower ix), Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlInner = Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' Dim
1
{-# INLINE foldlInner #-}

-- | Right fold over the inner most dimension with index aware function.
--
-- @since 0.2.4
ifoldrInner
  :: (Index (Lower ix), Index ix, Source r e)
  => (ix -> e -> a -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
ifoldrInner :: forall ix r e a.
(Index (Lower ix), Index ix, Source r e) =>
(ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrInner = Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' Dim
1
{-# INLINE ifoldrInner #-}

-- | Right fold over the inner most dimension.
--
-- @since 0.2.4
foldrInner
  :: (Index (Lower ix), Index ix, Source r e)
  => (e -> a -> a)
  -> a
  -> Array r ix e
  -> Array D (Lower ix) a
foldrInner :: forall ix r e a.
(Index (Lower ix), Index ix, Source r e) =>
(e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrInner = Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin' Dim
1
{-# INLINE foldrInner #-}

-- | Monoidal fold over the inner most dimension.
--
-- @since 0.4.3
foldInner :: (Monoid e, Index (Lower ix), Index ix, Source r e) => Array r ix e -> Array D (Lower ix) e
foldInner :: forall e ix r.
(Monoid e, Index (Lower ix), Index ix, Source r e) =>
Array r ix e -> Array D (Lower ix) e
foldInner = (e -> e -> e) -> e -> Array r ix e -> Array D (Lower ix) e
forall ix r e a.
(Index (Lower ix), Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlInner e -> e -> e
forall a. Monoid a => a -> a -> a
mappend e
forall a. Monoid a => a
mempty
{-# INLINE foldInner #-}

-- | Monoidal fold over some internal dimension.
--
-- @since 0.4.3
foldWithin
  :: (Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n)
  => Dimension n
  -> Array r ix a
  -> Array D (Lower ix) a
foldWithin :: forall r a ix (n :: Natural).
(Source r a, Monoid a, Index (Lower ix), IsIndexDimension ix n) =>
Dimension n -> Array r ix a -> Array D (Lower ix) a
foldWithin Dimension n
dim = Dimension n
-> (a -> a -> a) -> a -> Array r ix a -> Array D (Lower ix) a
forall ix (n :: Natural) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r e) =>
Dimension n
-> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin Dimension n
dim a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
forall a. Monoid a => a
mempty
{-# INLINE foldWithin #-}

-- | Monoidal fold over some internal dimension. This is a pratial function and will
-- result in `IndexDimensionException` if supplied dimension is invalid.
--
-- @since 0.4.3
foldWithin'
  :: (HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix))
  => Dim
  -> Array r ix a
  -> Array D (Lower ix) a
foldWithin' :: forall ix r a.
(HasCallStack, Index ix, Source r a, Monoid a, Index (Lower ix)) =>
Dim -> Array r ix a -> Array D (Lower ix) a
foldWithin' Dim
dim = Dim -> (a -> a -> a) -> a -> Array r ix a -> Array D (Lower ix) a
forall ix r e a.
(HasCallStack, Index (Lower ix), Index ix, Source r e) =>
Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' Dim
dim a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
forall a. Monoid a => a
mempty
{-# INLINE foldWithin' #-}

-- | Reduce each outer slice into a monoid and mappend results together
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> import Data.Monoid (Product(..))
-- >>> arr = computeAs P $ iterateN (Sz2 2 3) (+1) (10 :: Int)
-- >>> arr
-- Array P Seq (Sz (2 :. 3))
--   [ [ 11, 12, 13 ]
--   , [ 14, 15, 16 ]
--   ]
-- >>> getProduct $ foldOuterSlice (\row -> Product (A.sum row)) arr
-- 1620
-- >>> (11 + 12 + 13) * (14 + 15 + 16) :: Int
-- 1620
--
-- @since 0.4.3
foldOuterSlice
  :: (Index ix, Index (Lower ix), Source r e, Monoid m)
  => (Array r (Lower ix) e -> m)
  -> Array r ix e
  -> m
foldOuterSlice :: forall ix r e m.
(Index ix, Index (Lower ix), Source r e, Monoid m) =>
(Array r (Lower ix) e -> m) -> Array r ix e -> m
foldOuterSlice Array r (Lower ix) e -> m
f = (Int -> Array r (Lower ix) e -> m) -> Array r ix e -> m
forall ix r e m.
(Index ix, Index (Lower ix), Source r e, Monoid m) =>
(Int -> Array r (Lower ix) e -> m) -> Array r ix e -> m
ifoldOuterSlice ((Array r (Lower ix) e -> m) -> Int -> Array r (Lower ix) e -> m
forall a b. a -> b -> a
const Array r (Lower ix) e -> m
f)
{-# INLINE foldOuterSlice #-}

-- | Reduce each outer slice into a monoid with an index aware function and mappend results
-- together
--
-- @since 0.4.3
ifoldOuterSlice
  :: (Index ix, Index (Lower ix), Source r e, Monoid m)
  => (Ix1 -> Array r (Lower ix) e -> m)
  -> Array r ix e
  -> m
ifoldOuterSlice :: forall ix r e m.
(Index ix, Index (Lower ix), Source r e, Monoid m) =>
(Int -> Array r (Lower ix) e -> m) -> Array r ix e -> m
ifoldOuterSlice Int -> Array r (Lower ix) e -> m
f Array r ix e
arr = (Int -> m) -> Array D Int Int -> m
forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono Int -> m
g (Array D Int Int -> m) -> Array D Int Int -> m
forall a b. (a -> b) -> a -> b
$ Comp -> Int -> Int -> Array D Int Int
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) Int
0 Int
k
  where
    (Sz1 Int
k, Sz (Lower ix)
szL) = Sz ix -> (Sz Int, Sz (Lower ix))
forall ix. Index ix => Sz ix -> (Sz Int, Sz (Lower ix))
unconsSz (Sz ix -> (Sz Int, Sz (Lower ix)))
-> Sz ix -> (Sz Int, Sz (Lower ix))
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr
    g :: Int -> m
g Int
i = Int -> Array r (Lower ix) e -> m
f Int
i (Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
forall ix.
(Index ix, Index (Lower ix)) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
forall r e ix.
(Source r e, Index ix, Index (Lower ix)) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array r (Lower ix) e
unsafeOuterSlice Array r ix e
arr Sz (Lower ix)
szL Int
i)
    {-# INLINE g #-}
{-# INLINE ifoldOuterSlice #-}

-- | Reduce each inner slice into a monoid and mappend results together
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array as A
-- >>> import Data.Monoid (Product(..))
-- >>> arr = computeAs P $ iterateN (Sz2 2 3) (+1) (10 :: Int)
-- >>> arr
-- Array P Seq (Sz (2 :. 3))
--   [ [ 11, 12, 13 ]
--   , [ 14, 15, 16 ]
--   ]
-- >>> getProduct $ foldInnerSlice (\column -> Product (A.sum column)) arr
-- 19575
-- >>> (11 + 14) * (12 + 15) * (13 + 16) :: Int
-- 19575
--
-- @since 0.4.3
foldInnerSlice
  :: (Source r e, Index ix, Monoid m) => (Array D (Lower ix) e -> m) -> Array r ix e -> m
foldInnerSlice :: forall r e ix m.
(Source r e, Index ix, Monoid m) =>
(Array D (Lower ix) e -> m) -> Array r ix e -> m
foldInnerSlice Array D (Lower ix) e -> m
f = (Int -> Array D (Lower ix) e -> m) -> Array r ix e -> m
forall r e ix m.
(Source r e, Index ix, Monoid m) =>
(Int -> Array D (Lower ix) e -> m) -> Array r ix e -> m
ifoldInnerSlice ((Array D (Lower ix) e -> m) -> Int -> Array D (Lower ix) e -> m
forall a b. a -> b -> a
const Array D (Lower ix) e -> m
f)
{-# INLINE foldInnerSlice #-}

-- | Reduce each inner slice into a monoid with an index aware function and mappend
-- results together
--
-- @since 0.4.3
ifoldInnerSlice
  :: (Source r e, Index ix, Monoid m) => (Ix1 -> Array D (Lower ix) e -> m) -> Array r ix e -> m
ifoldInnerSlice :: forall r e ix m.
(Source r e, Index ix, Monoid m) =>
(Int -> Array D (Lower ix) e -> m) -> Array r ix e -> m
ifoldInnerSlice Int -> Array D (Lower ix) e -> m
f Array r ix e
arr = (Int -> m) -> Array D Int Int -> m
forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(e -> m) -> Array r ix e -> m
foldMono Int -> m
g (Array D Int Int -> m) -> Array D Int Int -> m
forall a b. (a -> b) -> a -> b
$ Comp -> Int -> Int -> Array D Int Int
forall ix. Index ix => Comp -> ix -> ix -> Array D ix ix
range (Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
arr) Int
0 (Sz Int -> Int
forall ix. Sz ix -> ix
unSz Sz Int
k)
  where
    (Sz (Lower ix)
szL, !Sz Int
k) = Sz ix -> (Sz (Lower ix), Sz Int)
forall ix. Index ix => Sz ix -> (Sz (Lower ix), Sz Int)
unsnocSz (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr)
    g :: Int -> m
g Int
i = Int -> Array D (Lower ix) e -> m
f Int
i (Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeInnerSlice Array r ix e
arr Sz (Lower ix)
szL Int
i)
    {-# INLINE g #-}
{-# INLINE ifoldInnerSlice #-}

-- | /O(n)/ - Compute maximum of all elements.
--
-- @since 0.3.0
maximumM :: (MonadThrow m, Shape r ix, Source r e, Ord e) => Array r ix e -> m e
maximumM :: forall (m :: * -> *) r ix e.
(MonadThrow m, Shape r ix, Source r e, Ord e) =>
Array r ix e -> m e
maximumM Array r ix e
arr =
  if Array r ix e -> Bool
forall e. Array r ix e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNull Array r ix e
arr
    then SizeException -> m e
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr))
    else
      let !e0 :: e
e0 = Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
forall ix. Index ix => ix
zeroIndex
       in e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> e -> m e
forall a b. (a -> b) -> a -> b
$ (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
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 e -> e -> e
forall a. Ord a => a -> a -> a
max e
e0 e -> e -> e
forall a. Ord a => a -> a -> a
max e
e0 Array r ix e
arr
{-# INLINE maximumM #-}

-- | /O(n)/ - Compute maximum of all elements.
--
-- @since 0.3.0
maximum'
  :: forall r ix e
   . (HasCallStack, Shape r ix, Source r e, Ord e)
  => Array r ix e
  -> e
maximum' :: forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
maximum' = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException e -> e)
-> (Array r ix e -> Either SomeException e) -> Array r ix e -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Either SomeException e
forall (m :: * -> *) r ix e.
(MonadThrow m, Shape r ix, Source r e, Ord e) =>
Array r ix e -> m e
maximumM
{-# INLINE maximum' #-}

-- | /O(n)/ - Compute minimum of all elements.
--
-- @since 0.3.0
minimumM :: (MonadThrow m, Shape r ix, Source r e, Ord e) => Array r ix e -> m e
minimumM :: forall (m :: * -> *) r ix e.
(MonadThrow m, Shape r ix, Source r e, Ord e) =>
Array r ix e -> m e
minimumM Array r ix e
arr =
  if Array r ix e -> Bool
forall e. Array r ix e -> Bool
forall r ix e. Shape r ix => Array r ix e -> Bool
isNull Array r ix e
arr
    then SizeException -> m e
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (Sz ix -> SizeException
forall ix. Index ix => Sz ix -> SizeException
SizeEmptyException (Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
arr))
    else
      let !e0 :: e
e0 = Array r ix e -> ix -> e
forall ix. Index ix => Array r ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
forall ix. Index ix => ix
zeroIndex
       in e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> e -> m e
forall a b. (a -> b) -> a -> b
$ (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
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 e -> e -> e
forall a. Ord a => a -> a -> a
min e
e0 e -> e -> e
forall a. Ord a => a -> a -> a
min e
e0 Array r ix e
arr
{-# INLINE minimumM #-}

-- | /O(n)/ - Compute minimum of all elements.
--
-- @since 0.3.0
minimum' :: forall r ix e. (HasCallStack, Shape r ix, Source r e, Ord e) => Array r ix e -> e
minimum' :: forall r ix e.
(HasCallStack, Shape r ix, Source r e, Ord e) =>
Array r ix e -> e
minimum' = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException e -> e)
-> (Array r ix e -> Either SomeException e) -> Array r ix e -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array r ix e -> Either SomeException e
forall (m :: * -> *) r ix e.
(MonadThrow m, Shape r ix, Source r e, Ord e) =>
Array r ix e -> m e
minimumM
{-# INLINE minimum' #-}

-- -- | /O(n)/ - Compute sum of all elements.
-- --
-- -- @since 0.1.0
-- sum' ::
--      forall r ix e. (Index ix, Source r e, Numeric r e)
--   => Array r ix e
--   -> IO e
-- sum' = splitReduce (\_ -> pure . sumArray) (\x y -> pure (x + y)) 0
-- {-# INLINE sum' #-}

-- | /O(n)/ - Compute sum of all elements.
--
-- @since 0.1.0
sum :: (Index ix, Source r e, Num e) => Array r ix e -> e
sum :: forall ix r e. (Index ix, Source r e, Num e) => Array r ix e -> e
sum = (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
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 e -> e -> e
forall a. Num a => a -> a -> a
(+) e
0 e -> e -> e
forall a. Num a => a -> a -> a
(+) e
0
{-# INLINE sum #-}

-- | /O(n)/ - Compute product of all elements.
--
-- @since 0.1.0
product :: (Index ix, Source r e, Num e) => Array r ix e -> e
product :: forall ix r e. (Index ix, Source r e, Num e) => Array r ix e -> e
product = (e -> e -> e) -> e -> (e -> e -> e) -> e -> Array r ix e -> e
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 e -> e -> e
forall a. Num a => a -> a -> a
(*) e
1 e -> e -> e
forall a. Num a => a -> a -> a
(*) e
1
{-# INLINE product #-}

-- | /O(n)/ - Compute conjunction of all elements.
--
-- @since 0.1.0
and :: (Index ix, Source r Bool) => Array r ix Bool -> Bool
and :: forall ix r. (Index ix, Source r Bool) => Array r ix Bool -> Bool
and = (Bool -> Bool) -> Array r ix Bool -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
all Bool -> Bool
forall a. a -> a
id
{-# INLINE and #-}

-- | /O(n)/ - Compute disjunction of all elements.
--
-- @since 0.1.0
or :: (Index ix, Source r Bool) => Array r ix Bool -> Bool
or :: forall ix r. (Index ix, Source r Bool) => Array r ix Bool -> Bool
or = (Bool -> Bool) -> Array r ix Bool -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any Bool -> Bool
forall a. a -> a
id
{-# INLINE or #-}

-- | /O(n)/ - Determines whether all elements of the array satisfy a predicate.
--
-- @since 0.1.0
all :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> Bool
all :: forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
all e -> Bool
f = Bool -> Bool
not (Bool -> Bool) -> (Array r ix e -> Bool) -> Array r ix e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> Bool) -> Array r ix e -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any (Bool -> Bool
not (Bool -> Bool) -> (e -> Bool) -> e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> Bool
f)
{-# INLINE all #-}

-- | /O(n)/ - Determines whether an element is present in the array.
--
-- @since 0.5.5
elem :: (Eq e, Index ix, Source r e) => e -> Array r ix e -> Bool
elem :: forall e ix r.
(Eq e, Index ix, Source r e) =>
e -> Array r ix e -> Bool
elem e
e = (e -> Bool) -> Array r ix e -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
any (e
e e -> e -> Bool
forall a. Eq a => a -> a -> Bool
==)
{-# INLINE elem #-}

-- $unstruct_folds
--
-- Functions in this section will fold any `Source` array with respect to the inner
-- `Comp`utation strategy setting.

-- $seq_folds
--
-- Functions in this section will fold any `Source` array sequentially, regardless of the inner
-- `Comp`utation strategy setting.

-- $par_folds
--
-- __Note__ It is important to compile with @-threaded -with-rtsopts=-N@ flags, otherwise
-- there will be no parallelization.
--
-- Functions in this section will fold any `Source` array in parallel, regardless of the
-- inner `Comp`utation strategy setting. All of the parallel structured folds are
-- performed inside `IO` monad, because referential transparency can't generally be
-- preserved and results will depend on the number of cores/capabilities that computation
-- is being performed on.
--
-- In contrast to sequential folds, each parallel folding function accepts two functions
-- and two initial elements as arguments. This is necessary because an array is first
-- split into chunks, which folded individually on separate cores with the first function,
-- and the results of those folds are further folded with the second function.