Internal.hs

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-- | Similar to `foldlP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
foldlIO
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (a -> e -> m a)
  -- ^ Index aware folding IO action
  -> a
  -- ^ Accumulator
  -> (b -> a -> m b)
  -- ^ Folding action that is applied to the results of a parallel fold
  -> b
  -- ^ Accumulator for chunks folding
  -> Array r ix e
  -> m b
foldlIO f !initAcc g !tAcc !arr
  | getComp arr == Seq = foldlM f initAcc arr >>= g tAcc
  | otherwise = do
      let splitAcc _ = pure (initAcc, initAcc)
          !sz = size arr
      results <-
        withScheduler (getComp arr) $ \scheduler ->
          withRunInIO $ \run ->
            stToPrim $
              case unsafePrefIndex arr of
                PrefIndex gix ->
                  iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc ->
                    ioToPrim (run (f acc (gix ix)))
                PrefIndexLinear gi ->
                  iterFullAccST defRowMajor scheduler 0 (toLinearSz sz) initAcc splitAcc $ \ !i !acc ->
                    ioToPrim (run (f acc (gi i)))
      F.foldlM g tAcc results
{-# INLINE foldlIO #-}

-- | Similar to `ifoldlP`, except that folding functions themselves do live in IO
--
-- @since 0.1.0
ifoldlIO
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (a -> ix -> e -> m a)
  -- ^ Index aware folding IO action
  -> a
  -- ^ Accumulator
  -> (b -> a -> m b)
  -- ^ Folding action that is applied to the results of a parallel fold
  -> b
  -- ^ Accumulator for chunks folding
  -> Array r ix e
  -> m b
ifoldlIO f !initAcc g !tAcc !arr
  | getComp arr == Seq = ifoldlM f initAcc arr >>= g tAcc
  | otherwise = do
      let !sz = size arr
          splitAcc _ = pure (initAcc, initAcc)
      results <-
        withScheduler (getComp arr) $ \scheduler ->
          withRunInIO $ \run ->
            stToPrim $
              case unsafePrefIndex arr of
                PrefIndex gix ->
                  iterFullAccST defRowMajor scheduler zeroIndex sz initAcc splitAcc $ \ !ix !acc ->
                    ioToPrim (run (f acc ix (gix ix)))
                PrefIndexLinear gi ->
                  iterTargetFullAccST defRowMajor scheduler 0 sz initAcc splitAcc $ \ !i !ix !acc ->
                    ioToPrim (run (f acc ix (gi i)))
      F.foldlM g tAcc results
{-# INLINE ifoldlIO #-}

-- | Slice an array into linear row-major vector chunks and apply an action to each of
-- them. Number of chunks will depend on the computation strategy. Results of each action
-- will be combined with a folding function.
--
-- @since 1.0.0
splitReduce
  :: (MonadUnliftIO m, Index ix, Source r e)
  => (Scheduler RealWorld a -> Vector r e -> m a)
  -> (b -> a -> m b)
  -- ^ Folding action that is applied to the results of a parallel fold
  -> b
  -- ^ Accumulator for chunks folding
  -> Array r ix e
  -> m b
splitReduce f g !tAcc !arr = do
  let !sz = size arr
      !totalLength = totalElem sz
  results <-
    withScheduler (getComp arr) $ \scheduler -> do
      withRunInIO $ \run -> do
        splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
          loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start ->
            scheduleWork scheduler $
              run $
                f scheduler $
                  unsafeLinearSlice start (SafeSz chunkLength) arr
          when (slackStart < totalLength) $
            scheduleWork scheduler $
              run $
                f scheduler $
                  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, Index ix, Source r e)
  => (ix -> e -> a -> m a)
  -> a
  -> (a -> b -> m b)
  -> b
  -> Array r ix e
  -> m b
ifoldrIO f !initAcc g !tAcc !arr
  | getComp arr == Seq = ifoldrM f initAcc arr >>= (`g` tAcc)
  | otherwise = do
      let !sz = size arr
          !totalLength = totalElem sz
      results <-
        withRunInIO $ \run -> do
          withScheduler (getComp arr) $ \scheduler ->
            splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
              when (slackStart < totalLength) $
                scheduleWork scheduler $
                  run $
                    iterLinearM sz (totalLength - 1) slackStart (-1) (>=) initAcc $ \ !i ix ->
                      f ix (unsafeLinearIndex arr i)
              loopA_ slackStart (> 0) (subtract chunkLength) $ \ !start ->
                scheduleWork scheduler $
                  run $
                    iterLinearM sz (start - 1) (start - chunkLength) (-1) (>=) initAcc $ \ !i ix ->
                      f ix (unsafeLinearIndex arr i)
      F.foldlM (flip g) tAcc results
{-# INLINE ifoldrIO #-}

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

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

-- | Parallelizable implementation of `any` with unrolling
anyPu :: (Index ix, Source r e) => (e -> Bool) -> Array r ix e -> IO Bool
-- TODO: switch to splitReduce
-- anyPu f arr =
--   splitReduce anySu (\r acc -> pure (r || acc)) False
anyPu f arr = do
  let !sz = size arr
      !totalLength = totalElem sz
  results <-
    withScheduler (getComp arr) $ \scheduler -> do
      batch <- getCurrentBatch scheduler
      splitLinearly (numWorkers scheduler) totalLength $ \chunkLength slackStart -> do
        loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start ->
          scheduleWork scheduler $ anySliceSuM batch start (Sz (start + chunkLength)) f arr
        when (slackStart < totalLength) $
          scheduleWork scheduler $
            anySliceSuM batch slackStart (Sz totalLength) f arr
  pure $ F.foldl' (||) False results
{-# INLINE anyPu #-}

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