{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
-- |
-- Module      : Data.Massiv.Array.Delayed.Push
-- Copyright   : (c) Alexey Kuleshevich 2019-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Delayed.Push
  ( DL(..)
  , Array(..)
  , Loader
  , toLoadArray
  , makeLoadArrayS
  , makeLoadArray
  , unsafeMakeLoadArray
  , unsafeMakeLoadArrayAdjusted
  , fromStrideLoad
  , appendOuterM
  , concatOuterM
  ) where

import Control.Monad
import Control.Scheduler as S (traverse_)
import Data.Foldable as F
import Data.Massiv.Core.Common
import Prelude hiding (map, zipWith)

#include "massiv.h"

-- | Delayed load representation. Also known as Push array.
data DL = DL deriving Ix1 -> DL -> ShowS
[DL] -> ShowS
DL -> String
forall a.
(Ix1 -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DL] -> ShowS
$cshowList :: [DL] -> ShowS
show :: DL -> String
$cshow :: DL -> String
showsPrec :: Ix1 -> DL -> ShowS
$cshowsPrec :: Ix1 -> DL -> ShowS
Show

type Loader e =
  forall s. Scheduler s () -- ^ Scheduler that will be used for loading
         -> Ix1 -- ^ Start loading at this linear index
         -> (Ix1 -> e -> ST s ()) -- ^ Linear element writing action
         -> (Ix1 -> Sz1 -> e -> ST s ()) -- ^ Linear region setting action
         -> ST s ()


data instance Array DL ix e = DLArray
  { forall ix e. Array DL ix e -> Comp
dlComp    :: !Comp
  , forall ix e. Array DL ix e -> Sz ix
dlSize    :: !(Sz ix)
  , forall ix e. Array DL ix e -> Loader e
dlLoad    :: Loader e
  }

instance Strategy DL where
  getComp :: forall ix e. Array DL ix e -> Comp
getComp = forall ix e. Array DL ix e -> Comp
dlComp
  {-# INLINE getComp #-}
  setComp :: forall ix e. Comp -> Array DL ix e -> Array DL ix e
setComp Comp
c Array DL ix e
arr = Array DL ix e
arr {dlComp :: Comp
dlComp = Comp
c}
  {-# INLINE setComp #-}
  repr :: DL
repr = DL
DL


instance Index ix => Shape DL ix where
  maxLinearSize :: forall e. Array DL ix e -> Maybe Sz1
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Size r) => Array r ix e -> Ix1
elemsCount
  {-# INLINE maxLinearSize #-}


instance Size DL where
  size :: forall ix e. Array DL ix e -> Sz ix
size = forall ix e. Array DL ix e -> Sz ix
dlSize
  {-# INLINE size #-}
  unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array DL ix e -> Array DL ix' e
unsafeResize !Sz ix'
sz !Array DL ix e
arr = Array DL ix e
arr { dlSize :: Sz ix'
dlSize = Sz ix'
sz }
  {-# INLINE unsafeResize #-}

instance Semigroup (Array DL Ix1 e) where
  <> :: Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
(<>) = forall e. Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL
  {-# INLINE (<>) #-}

instance Monoid (Array DL Ix1 e) where
  mempty :: Array DL Ix1 e
mempty = DLArray {dlComp :: Comp
dlComp = forall a. Monoid a => a
mempty, dlSize :: Sz1
dlSize = forall ix. Index ix => Sz ix
zeroSz, dlLoad :: Loader e
dlLoad = \Scheduler s ()
_ Ix1
_ Ix1 -> e -> ST s ()
_ Ix1 -> Sz1 -> e -> ST s ()
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()}
  {-# INLINE mempty #-}
#if !MIN_VERSION_base(4,11,0)
  mappend = mappendDL
  {-# INLINE mappend #-}
#endif
  mconcat :: [Array DL Ix1 e] -> Array DL Ix1 e
mconcat [] = forall a. Monoid a => a
mempty
  mconcat [Array DL Ix1 e
x] = Array DL Ix1 e
x
  mconcat [Array DL Ix1 e
x, Array DL Ix1 e
y] = Array DL Ix1 e
x forall a. Semigroup a => a -> a -> a
<> Array DL Ix1 e
y
  mconcat [Array DL Ix1 e]
xs = forall e. [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL [Array DL Ix1 e]
xs
  {-# INLINE mconcat #-}

mconcatDL :: forall e . [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL :: forall e. [Array DL Ix1 e] -> Array DL Ix1 e
mconcatDL ![Array DL Ix1 e]
arrs =
  DLArray {dlComp :: Comp
dlComp = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap forall r ix e. Strategy r => Array r ix e -> Comp
getComp [Array DL Ix1 e]
arrs, dlSize :: Sz1
dlSize = forall ix. ix -> Sz ix
SafeSz Ix1
k, dlLoad :: Loader e
dlLoad = Loader e
load}
  where
    !k :: Ix1
k = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' forall a. Num a => a -> a -> a
(+) Ix1
0 (forall ix. Sz ix -> ix
unSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e. Size r => Array r ix e -> Sz ix
size forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Array DL Ix1 e]
arrs)
    load :: forall s .
      Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
    load :: Loader e
load Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet =
      let loadArr :: Ix1 -> Array DL Ix1 e -> ST s Ix1
loadArr !Ix1
startAtCur DLArray {dlSize :: forall ix e. Array DL ix e -> Sz ix
dlSize = SafeSz Ix1
kCur, Loader e
dlLoad :: Loader e
dlLoad :: forall ix e. Array DL ix e -> Loader e
dlLoad} = do
            let !endAtCur :: Ix1
endAtCur = Ix1
startAtCur forall a. Num a => a -> a -> a
+ Ix1
kCur
            forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
dlLoad Scheduler s ()
scheduler Ix1
startAtCur Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
            forall (f :: * -> *) a. Applicative f => a -> f a
pure Ix1
endAtCur
          {-# INLINE loadArr #-}
       in forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ Ix1 -> Array DL Ix1 e -> ST s Ix1
loadArr Ix1
startAt [Array DL Ix1 e]
arrs
    {-# INLINE load #-}
{-# INLINE mconcatDL #-}


mappendDL :: forall e . Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL :: forall e. Array DL Ix1 e -> Array DL Ix1 e -> Array DL Ix1 e
mappendDL (DLArray Comp
c1 Sz1
sz1 Loader e
load1) (DLArray Comp
c2 Sz1
sz2 Loader e
load2) =
  DLArray {dlComp :: Comp
dlComp = Comp
c1 forall a. Semigroup a => a -> a -> a
<> Comp
c2, dlSize :: Sz1
dlSize = forall ix. ix -> Sz ix
SafeSz (Ix1
k1 forall a. Num a => a -> a -> a
+ Ix1
k2), dlLoad :: Loader e
dlLoad = Loader e
load}
  where
    !k1 :: Ix1
k1 = forall ix. Sz ix -> ix
unSz Sz1
sz1
    !k2 :: Ix1
k2 = forall ix. Sz ix -> ix
unSz Sz1
sz2
    load :: forall s.
      Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
    load :: Loader e
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet = do
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load1 Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load2 Scheduler s ()
scheduler (Ix1
startAt forall a. Num a => a -> a -> a
+ Ix1
k1) Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
    {-# INLINE load #-}
{-# INLINE mappendDL #-}

-- | Append two arrays together along the outer most dimension. Inner dimensions must
-- agree, otherwise `SizeMismatchException`.
--
-- @since 0.4.4
appendOuterM ::
     forall ix e m. (Index ix, MonadThrow m)
  => Array DL ix e
  -> Array DL ix e
  -> m (Array DL ix e)
appendOuterM :: forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
Array DL ix e -> Array DL ix e -> m (Array DL ix e)
appendOuterM (DLArray Comp
c1 Sz ix
sz1 Loader e
load1) (DLArray Comp
c2 Sz ix
sz2 Loader e
load2) = do
  let (!Sz1
i1, !Sz (Lower ix)
szl1) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz ix
sz1
      (!Sz1
i2, !Sz (Lower ix)
szl2) = forall ix. Index ix => Sz ix -> (Sz1, Sz (Lower ix))
unconsSz Sz ix
sz2
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Sz (Lower ix)
szl1 forall a. Eq a => a -> a -> Bool
== Sz (Lower ix)
szl2) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz ix
sz1 Sz ix
sz2
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    DLArray {dlComp :: Comp
dlComp = Comp
c1 forall a. Semigroup a => a -> a -> a
<> Comp
c2, dlSize :: Sz ix
dlSize = forall ix. Index ix => Sz1 -> Sz (Lower ix) -> Sz ix
consSz (forall ix.
Index ix =>
(Ix1 -> Ix1 -> Ix1) -> Sz ix -> Sz ix -> Sz ix
liftSz2 forall a. Num a => a -> a -> a
(+) Sz1
i1 Sz1
i2) Sz (Lower ix)
szl1, dlLoad :: Loader e
dlLoad = Loader e
load}
  where
    load :: Loader e
    load :: Loader e
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet = do
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load1 Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
      forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler forall a b. (a -> b) -> a -> b
$ Loader e
load2 Scheduler s ()
scheduler (Ix1
startAt forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz1) Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet
    {-# INLINE load #-}
{-# INLINE appendOuterM #-}

-- | Concat arrays together along the outer most dimension. Inner dimensions must agree
-- for all arrays in the list, otherwise `SizeMismatchException`.
--
-- @since 0.4.4
concatOuterM ::
     forall ix e m. (Index ix, MonadThrow m)
  => [Array DL ix e]
  -> m (Array DL ix e)
concatOuterM :: forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
[Array DL ix e] -> m (Array DL ix e)
concatOuterM =
  \case
    []     -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall r ix e. Load r ix e => Array r ix e
empty
    (Array DL ix e
x:[Array DL ix e]
xs) -> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM forall ix e (m :: * -> *).
(Index ix, MonadThrow m) =>
Array DL ix e -> Array DL ix e -> m (Array DL ix e)
appendOuterM Array DL ix e
x [Array DL ix e]
xs
{-# INLINE concatOuterM #-}


-- | Describe how an array should be loaded into memory sequentially. For parallelizable
-- version see `makeLoadArray`.
--
-- @since 0.3.1
makeLoadArrayS ::
     forall ix e. Index ix
  => Sz ix
  -- ^ Size of the resulting array
  -> e
  -- ^ Default value to use for all cells that might have been ommitted by the writing function
  -> (forall m. Monad m => (ix -> e -> m Bool) -> m ())
  -- ^ Writing function that described which elements to write into the target array.
  -> Array DL ix e
makeLoadArrayS :: forall ix e.
Index ix =>
Sz ix
-> e
-> (forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArrayS Sz ix
sz e
defVal forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ()
writer = forall ix e.
Index ix =>
Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray Comp
Seq Sz ix
sz e
defVal (forall a b. a -> b -> a
const forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ()
writer)
{-# INLINE makeLoadArrayS #-}

-- | Specify how an array should be loaded into memory. Unlike `makeLoadArrayS`, loading
-- function accepts a scheduler, thus can be parallelized. If you need an unsafe version
-- of this function see `unsafeMakeLoadArray`.
--
-- @since 0.4.0
makeLoadArray ::
     forall ix e. Index ix
  => Comp
  -- ^ Computation strategy to use. Directly affects the scheduler that gets created for
  -- the loading function.
  -> Sz ix
  -- ^ Size of the resulting array
  -> e
  -- ^ Default value to use for all cells that might have been ommitted by the writing function
  -> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
  -- ^ Writing function that described which elements to write into the target array. It
  -- accepts a scheduler, that can be used for parallelization, as well as a safe element
  -- writing function.
  -> Array DL ix e
makeLoadArray :: forall ix e.
Index ix =>
Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray Comp
comp Sz ix
sz e
defVal forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ()
writer = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load
  where
    load :: forall s.
      Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
    load :: forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
uWrite Ix1 -> Sz1 -> e -> ST s ()
uSet = do
      Ix1 -> Sz1 -> e -> ST s ()
uSet Ix1
startAt (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz) e
defVal
      let safeWrite :: ix -> e -> ST s Bool
safeWrite !ix
ix !e
e
            | forall ix. Index ix => Sz ix -> ix -> Bool
isSafeIndex Sz ix
sz ix
ix = Bool
True forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Ix1 -> e -> ST s ()
uWrite (Ix1
startAt forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Ix1
toLinearIndex Sz ix
sz ix
ix) e
e
            | Bool
otherwise = forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
          {-# INLINE safeWrite #-}
      forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ()
writer Scheduler s ()
scheduler ix -> e -> ST s Bool
safeWrite
    {-# INLINE load #-}
{-# INLINE makeLoadArray #-}

-- | Specify how an array can be loaded/computed through creation of a `DL` array. Unlike
-- `makeLoadArrayS` or `makeLoadArray` this function is unsafe, since there is no
-- guarantee that all elements will be initialized and the supplied element writing
-- function does not perform any bounds checking.
--
-- @since 0.3.1
unsafeMakeLoadArray ::
     forall ix e. Index ix
  => Comp
  -- ^ Computation strategy to use. Directly affects the scheduler that gets created for
  -- the loading function.
  -> Sz ix
  -- ^ Size of the array
  -> Maybe e
  -- ^ An element to use for initialization of the mutable array that will be created in
  -- the future
  -> (forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ())
  -- ^ This function accepts:
  --
  -- * A scheduler that can be used for parallelization of loading
  --
  -- * Linear index at which this load array will start (an offset that should be added to
  --   the linear writng function)
  --
  -- * Linear element writing function
  -> Array DL ix e
unsafeMakeLoadArray :: forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall s.
    Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArray Comp
comp Sz ix
sz Maybe e
mDefVal forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
writer = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz Loader e
load
  where
    load :: Loader e
    load :: Loader e
load Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
uWrite Ix1 -> Sz1 -> e -> ST s ()
uSet = do
      forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (Ix1 -> Sz1 -> e -> ST s ()
uSet Ix1
startAt (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz)) Maybe e
mDefVal
      forall s. Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> ST s ()
writer Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
uWrite
    {-# INLINE load #-}
{-# INLINE unsafeMakeLoadArray #-}

-- | Same as `unsafeMakeLoadArray`, except will ensure that starting index is correctly
-- adjusted. Which means the writing function gets one less argument.
--
-- @since 0.5.2
unsafeMakeLoadArrayAdjusted ::
     forall ix e. Index ix
  => Comp
  -> Sz ix
  -> Maybe e
  -> (forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ())
  -> Array DL ix e
unsafeMakeLoadArrayAdjusted :: forall ix e.
Index ix =>
Comp
-> Sz ix
-> Maybe e
-> (forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ())
-> Array DL ix e
unsafeMakeLoadArrayAdjusted Comp
comp Sz ix
sz Maybe e
mDefVal forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ()
writer = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load
  where
    load :: forall s.
      Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
    load :: forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
uWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet = do
      forall (f :: * -> *) (t :: * -> *) a.
(Applicative f, Foldable t) =>
(a -> f ()) -> t a -> f ()
S.traverse_ (Ix1 -> Sz1 -> e -> ST s ()
dlSet Ix1
startAt (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
sz)) Maybe e
mDefVal
      forall s. Scheduler s () -> (Ix1 -> e -> ST s ()) -> ST s ()
writer Scheduler s ()
scheduler (\Ix1
i -> Ix1 -> e -> ST s ()
uWrite (Ix1
startAt forall a. Num a => a -> a -> a
+ Ix1
i))
    {-# INLINE load #-}
{-# INLINE unsafeMakeLoadArrayAdjusted #-}

-- | Convert any `Load`able array into `DL` representation.
--
-- @since 0.3.0
toLoadArray ::
     forall r ix e. (Size r, Load r ix e)
  => Array r ix e
  -> Array DL ix e
toLoadArray :: forall r ix e.
(Size r, Load r ix e) =>
Array r ix e -> Array DL ix e
toLoadArray Array r ix e
arr = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
sz forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load
  where
    !sz :: Sz ix
sz = forall r ix e. Size r => Array r ix e -> Sz ix
size Array r ix e
arr
    load :: forall s.
      Scheduler s () -> Ix1 -> (Ix1 -> e -> ST s ()) -> (Ix1 -> Sz1 -> e -> ST s ()) -> ST s ()
    load :: forall s.
Scheduler s ()
-> Ix1
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
dlSet =
      forall r ix e s.
Load r ix e =>
Scheduler s ()
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
scheduler Array r ix e
arr (Ix1 -> e -> ST s ()
dlWrite forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Num a => a -> a -> a
+ Ix1
startAt)) (\Ix1
offset -> Ix1 -> Sz1 -> e -> ST s ()
dlSet (Ix1
offset forall a. Num a => a -> a -> a
+ Ix1
startAt))
    {-# INLINE load #-}
{-# INLINE[1] toLoadArray #-}
{-# RULES "toLoadArray/id" toLoadArray = id #-}

-- | Convert an array that can be loaded with stride into `DL` representation.
--
-- @since 0.3.0
fromStrideLoad ::
     forall r ix e. (StrideLoad r ix e)
  => Stride ix
  -> Array r ix e
  -> Array DL ix e
fromStrideLoad :: forall r ix e.
StrideLoad r ix e =>
Stride ix -> Array r ix e -> Array DL ix e
fromStrideLoad Stride ix
stride Array r ix e
arr =
  forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray (forall r ix e. Strategy r => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
newsz Loader e
load
  where
    !newsz :: Sz ix
newsz = forall ix. Index ix => Stride ix -> Sz ix -> Sz ix
strideSize Stride ix
stride (forall r ix e. Shape r ix => Array r ix e -> Sz ix
outerSize Array r ix e
arr)
    load :: Loader e
    load :: Loader e
load Scheduler s ()
scheduler !Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
_ =
      forall r ix e s.
StrideLoad r ix e =>
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array r ix e
-> (Ix1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ Scheduler s ()
scheduler Stride ix
stride Sz ix
newsz Array r ix e
arr (\ !Ix1
i -> Ix1 -> e -> ST s ()
dlWrite (Ix1
i forall a. Num a => a -> a -> a
+ Ix1
startAt))
    {-# INLINE load #-}
{-# INLINE fromStrideLoad #-}

instance Index ix => Load DL ix e where
  makeArrayLinear :: Comp -> Sz ix -> (Ix1 -> e) -> Array DL ix e
makeArrayLinear Comp
comp Sz ix
sz Ix1 -> e
f = forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp Sz ix
sz Loader e
load
    where
      load :: Loader e
      load :: Loader e
load Scheduler s ()
scheduler Ix1
startAt Ix1 -> e -> ST s ()
dlWrite Ix1 -> Sz1 -> e -> ST s ()
_ =
        forall s (m :: * -> *) b c.
MonadPrimBase s m =>
Scheduler s ()
-> Ix1 -> Ix1 -> (Ix1 -> m b) -> (Ix1 -> b -> m c) -> m ()
splitLinearlyWithStartAtM_ Scheduler s ()
scheduler Ix1
startAt (forall ix. Index ix => Sz ix -> Ix1
totalElem Sz ix
sz) (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ix1 -> e
f) Ix1 -> e -> ST s ()
dlWrite
      {-# INLINE load #-}
  {-# INLINE makeArrayLinear #-}
  replicate :: Comp -> Sz ix -> e -> Array DL ix e
replicate Comp
comp !Sz ix
sz !e
e = forall ix e.
Index ix =>
Comp
-> Sz ix
-> e
-> (forall s. Scheduler s () -> (ix -> e -> ST s Bool) -> ST s ())
-> Array DL ix e
makeLoadArray Comp
comp Sz ix
sz e
e forall a b. (a -> b) -> a -> b
$ \Scheduler s ()
_ ix -> e -> ST s Bool
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  {-# INLINE replicate #-}
  iterArrayLinearWithSetST_ :: forall s.
Scheduler s ()
-> Array DL ix e
-> (Ix1 -> e -> ST s ())
-> (Ix1 -> Sz1 -> e -> ST s ())
-> ST s ()
iterArrayLinearWithSetST_ Scheduler s ()
scheduler DLArray {Loader e
dlLoad :: Loader e
dlLoad :: forall ix e. Array DL ix e -> Loader e
dlLoad} = Loader e
dlLoad Scheduler s ()
scheduler Ix1
0
  {-# INLINE iterArrayLinearWithSetST_ #-}

instance Index ix => Functor (Array DL ix) where
  fmap :: forall a b. (a -> b) -> Array DL ix a -> Array DL ix b
fmap a -> b
f Array DL ix a
arr = Array DL ix a
arr {dlLoad :: Loader b
dlLoad = forall ix a b s.
Array DL ix a
-> (a -> b)
-> Scheduler s ()
-> Ix1
-> (Ix1 -> b -> ST s ())
-> (Ix1 -> Sz1 -> b -> ST s ())
-> ST s ()
loadFunctor Array DL ix a
arr a -> b
f}
  {-# INLINE fmap #-}
  <$ :: forall a b. a -> Array DL ix b -> Array DL ix a
(<$) = forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a
overwriteFunctor
  {-# INLINE (<$) #-}

overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a
overwriteFunctor :: forall ix a b. Index ix => a -> Array DL ix b -> Array DL ix a
overwriteFunctor a
e Array DL ix b
arr = Array DL ix b
arr {dlLoad :: Loader a
dlLoad = Loader a
load}
  where
    load :: Loader a
    load :: Loader a
load Scheduler s ()
_ !Ix1
startAt Ix1 -> a -> ST s ()
_ Ix1 -> Sz1 -> a -> ST s ()
dlSet = Ix1 -> Sz1 -> a -> ST s ()
dlSet Ix1
startAt (forall r ix e. Shape r ix => Array r ix e -> Sz1
linearSize Array DL ix b
arr) a
e
    {-# INLINE load #-}
{-# INLINE overwriteFunctor #-}


loadFunctor ::
     Array DL ix a
  -> (a -> b)
  -> Scheduler s ()
  -> Ix1
  -> (Ix1 -> b -> ST s ())
  -> (Ix1 -> Sz1 -> b -> ST s ())
  -> ST s ()
loadFunctor :: forall ix a b s.
Array DL ix a
-> (a -> b)
-> Scheduler s ()
-> Ix1
-> (Ix1 -> b -> ST s ())
-> (Ix1 -> Sz1 -> b -> ST s ())
-> ST s ()
loadFunctor Array DL ix a
arr a -> b
f Scheduler s ()
scheduler Ix1
startAt Ix1 -> b -> ST s ()
uWrite Ix1 -> Sz1 -> b -> ST s ()
uSet =
  forall ix e. Array DL ix e -> Loader e
dlLoad Array DL ix a
arr Scheduler s ()
scheduler Ix1
startAt (\ !Ix1
i a
e -> Ix1 -> b -> ST s ()
uWrite Ix1
i (a -> b
f a
e)) (\Ix1
o Sz1
sz a
e -> Ix1 -> Sz1 -> b -> ST s ()
uSet Ix1
o Sz1
sz (a -> b
f a
e))
{-# INLINE loadFunctor #-}