{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- |
-- Module      : Data.Massiv.Array.Ops.Fold
-- 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
  (
  -- ** 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, foldl, foldr, map, maximum, minimum, or, product, sum, elem)


-- | /O(n)/ - Monoidal fold over an array with an index aware function. Also known as reduce.
--
-- @since 0.2.4
ifoldMono ::
     (Source r ix e, Monoid m)
  => (ix -> e -> m) -- ^ Convert each element of an array to an appropriate `Monoid`.
  -> Array r ix e -- ^ Source array
  -> m
ifoldMono :: (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 r ix e a b.
Source r ix 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 ::
     (Source r ix 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 :: (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 r ix e a b.
Source r ix 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 ::
     (Source r ix 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 :: (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 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. 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 ix e) =>
  Dimension n -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin :: 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.
(Index (Lower ix), Source r ix e) =>
Dim
-> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' (Dimension n -> Dim
forall (n :: Nat). 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 ix e) =>
  Dimension n -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin :: 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 :: Nat) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r ix 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 ix e) =>
  Dimension n -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin :: 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.
(Index (Lower ix), Source r ix e) =>
Dim
-> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' (Dimension n -> Dim
forall (n :: Nat). 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 ix e) =>
  Dimension n -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin :: 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 :: Nat) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r ix 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' :: (Index (Lower ix), Source r ix e) =>
  Dim -> (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlWithin' :: 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.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Load r 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. Index ix => Lower ix -> Dim -> Int -> ix
insertDim' Lower ix
ixl Dim
dim Int
0)
      (Lower ix -> Dim -> Int -> ix
forall ix. 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 r ix e. Source r ix e => 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. Load r 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. 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' :: (Index (Lower ix), Source r ix e) =>
  Dim -> (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlWithin' :: 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.
(Index (Lower ix), Source r ix 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' :: (Index (Lower ix), Source r ix e) =>
  Dim -> (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrWithin' :: 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.
Construct r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array r ix e -> Comp
forall r ix e. Load r 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. 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. 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 r ix e. Source r ix e => 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. Load r 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. 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' :: (Index (Lower ix), Source r ix e) =>
  Dim -> (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrWithin' :: 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.
(Index (Lower ix), Source r ix 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), Source r ix e) =>
  (ix -> a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldlInner :: (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.
(Index (Lower ix), Source r ix 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), Source r ix e) =>
  (a -> e -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldlInner :: (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.
(Index (Lower ix), Source r ix 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), Source r ix e) =>
  (ix -> e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
ifoldrInner :: (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.
(Index (Lower ix), Source r ix 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), Source r ix e) =>
  (e -> a -> a) -> a -> Array r ix e -> Array D (Lower ix) a
foldrInner :: (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.
(Index (Lower ix), Source r ix 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), Source r ix e) => Array r ix e -> Array D (Lower ix) e
foldInner :: 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), Source r ix 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 ix a, Monoid a, Index (Lower ix), IsIndexDimension ix n)
  => Dimension n
  -> Array r ix a
  -> Array D (Lower ix) a
foldWithin :: 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 :: Nat) r e a.
(Index (Lower ix), IsIndexDimension ix n, Source r ix 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' ::
     (Source r ix a, Monoid a, Index (Lower ix))
  => Dim
  -> Array r ix a
  -> Array D (Lower ix) a
foldWithin' :: 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.
(Index (Lower ix), Source r ix 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 :: (OuterSlice r ix e, Monoid m) => (Elt r ix e -> m) -> Array r ix e -> m
foldOuterSlice :: (Elt r ix e -> m) -> Array r ix e -> m
foldOuterSlice Elt r ix e -> m
f = (Int -> Elt r ix e -> m) -> Array r ix e -> m
forall r ix e m.
(OuterSlice r ix e, Monoid m) =>
(Int -> Elt r ix e -> m) -> Array r ix e -> m
ifoldOuterSlice ((Elt r ix e -> m) -> Int -> Elt r ix e -> m
forall a b. a -> b -> a
const Elt r 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 :: (OuterSlice r ix e, Monoid m) => (Ix1 -> Elt r ix e -> m) -> Array r ix e -> m
ifoldOuterSlice :: (Int -> Elt r ix e -> m) -> Array r ix e -> m
ifoldOuterSlice Int -> Elt r ix e -> m
f Array r ix e
arr = (Int -> m) -> Array D Int Int -> m
forall r ix e m.
(Source r ix 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. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) Int
0 (ix -> Int
forall ix. Index ix => ix -> Int
headDim (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)))
  where
    g :: Int -> m
g Int
i = Int -> Elt r ix e -> m
f Int
i (Array r ix e -> Int -> Elt r ix e
forall r ix e.
OuterSlice r ix e =>
Array r ix e -> Int -> Elt r ix e
unsafeOuterSlice Array r ix e
arr 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 :: (InnerSlice r ix e, Monoid m) => (Elt r ix e -> m) -> Array r ix e -> m
foldInnerSlice :: (Elt r ix e -> m) -> Array r ix e -> m
foldInnerSlice Elt r ix e -> m
f = (Int -> Elt r ix e -> m) -> Array r ix e -> m
forall r ix e m.
(InnerSlice r ix e, Monoid m) =>
(Int -> Elt r ix e -> m) -> Array r ix e -> m
ifoldInnerSlice ((Elt r ix e -> m) -> Int -> Elt r ix e -> m
forall a b. a -> b -> a
const Elt r 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 :: (InnerSlice r ix e, Monoid m) => (Ix1 -> Elt r ix e -> m) -> Array r ix e -> m
ifoldInnerSlice :: (Int -> Elt r ix e -> m) -> Array r ix e -> m
ifoldInnerSlice Int -> Elt r ix e -> m
f Array r ix e
arr = (Int -> m) -> Array D Int Int -> m
forall r ix e m.
(Source r ix 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. Load r 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
    szs :: (Sz (Lower ix), Sz Int)
szs@(Sz (Lower ix)
_, !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. Load r ix e => Array r ix e -> Sz ix
size Array r ix e
arr)
    g :: Int -> m
g Int
i = Int -> Elt r ix e -> m
f Int
i (Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e
forall r ix e.
InnerSlice r ix e =>
Array r ix e -> (Sz (Lower ix), Sz Int) -> Int -> Elt r ix e
unsafeInnerSlice Array r ix e
arr (Sz (Lower ix), Sz Int)
szs Int
i)
    {-# INLINE g #-}
{-# INLINE ifoldInnerSlice #-}

-- | /O(n)/ - Compute maximum of all elements.
--
-- @since 0.3.0
maximumM :: (MonadThrow m, Source r ix e, Ord e) => Array r ix e -> m e
maximumM :: Array r ix e -> m e
maximumM Array r ix e
arr =
    if Array r ix e -> Bool
forall r ix e. Load r ix e => Array r ix e -> Bool
isEmpty Array r ix e
arr
      then SizeException -> m e
forall (m :: * -> *) e a. (MonadThrow m, 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. Load r 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 r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
forall ix. Index ix => ix
zeroIndex
            in e -> m e
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 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. 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' :: (Source r ix e, Ord e) => Array r ix e -> e
maximum' :: Array r ix e -> e
maximum' = (SomeException -> e) -> (e -> e) -> Either SomeException e -> e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> e
forall a e. Exception e => e -> a
throw e -> e
forall a. a -> a
id (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, Source r ix 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, Source r ix e, Ord e) => Array r ix e -> m e
minimumM :: Array r ix e -> m e
minimumM Array r ix e
arr =
    if Array r ix e -> Bool
forall r ix e. Load r ix e => Array r ix e -> Bool
isEmpty Array r ix e
arr
      then SizeException -> m e
forall (m :: * -> *) e a. (MonadThrow m, 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. Load r 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 r ix e. Source r ix e => Array r ix e -> ix -> e
unsafeIndex Array r ix e
arr ix
forall ix. Index ix => ix
zeroIndex
            in e -> m e
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 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. 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' :: (Source r ix e, Ord e) => Array r ix e -> e
minimum' :: Array r ix e -> e
minimum' = (SomeException -> e) -> (e -> e) -> Either SomeException e -> e
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SomeException -> e
forall a e. Exception e => e -> a
throw e -> e
forall a. a -> a
id (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, Source r ix 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. (Source r ix 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 :: (Source r ix e, Num e) => Array r ix e -> e
sum :: Array r ix e -> e
sum = (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. 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 :: (Source r ix e, Num e) => Array r ix e -> e
product :: Array r ix e -> e
product = (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. 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 :: Source r ix Bool => Array r ix Bool -> Bool
and :: Array r ix Bool -> Bool
and = (Bool -> Bool) -> Array r ix Bool -> Bool
forall r ix e. Source r ix 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 :: Source r ix Bool => Array r ix Bool -> Bool
or :: Array r ix Bool -> Bool
or = (Bool -> Bool) -> Array r ix Bool -> Bool
forall r ix e. Source r ix 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 :: Source r ix e => (e -> Bool) -> Array r ix e -> Bool
all :: (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 r ix e. Source r ix 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, Source r ix e) => e -> Array r ix e -> Bool
elem :: e -> Array r ix e -> Bool
elem e
e = (e -> Bool) -> Array r ix e -> Bool
forall r ix e. Source r ix 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.

-}