{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Delayed.Pull
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Delayed.Pull (
  D (..),
  Array (..),
  delay,
  eqArrays,
  compareArrays,
  imap,
  liftArray2',
  liftArray2M,
  unsafeExtract,
  unsafeSlice,
  unsafeInnerSlice,
  zipWithInternal,
) where

import Control.Applicative
import qualified Data.Foldable as F
import Data.Massiv.Array.Ops.Fold.Internal as A
import Data.Massiv.Core.Common as A
import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec)
import Data.Massiv.Core.Operations
import qualified Data.Massiv.Vector.Stream as S
import GHC.Base (build)
import Prelude hiding (zipWith)

#include "massiv.h"

-- | Delayed representation.
data D
  = D
  deriving (Int -> D -> ShowS
[D] -> ShowS
D -> String
(Int -> D -> ShowS) -> (D -> String) -> ([D] -> ShowS) -> Show D
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> D -> ShowS
showsPrec :: Int -> D -> ShowS
$cshow :: D -> String
show :: D -> String
$cshowList :: [D] -> ShowS
showList :: [D] -> ShowS
Show)

data instance Array D ix e = DArray
  { forall ix e. Array D ix e -> Comp
dComp :: !Comp
  , forall ix e. Array D ix e -> Sz ix
dSize :: !(Sz ix)
  , forall ix e. Array D ix e -> PrefIndex ix e
dPrefIndex :: !(PrefIndex ix e)
  }

instance (Ragged L ix e, Show e) => Show (Array D ix e) where
  showsPrec :: Int -> Array D ix e -> ShowS
showsPrec = (Array D ix e -> Array D ix e) -> Int -> Array D ix e -> ShowS
forall r r' ix e.
(Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) =>
(Array r ix e -> Array r' ix e) -> Int -> Array r ix e -> ShowS
showsArrayPrec Array D ix e -> Array D ix e
forall a. a -> a
id
  showList :: [Array D ix e] -> ShowS
showList = [Array D ix e] -> ShowS
forall arr. Show arr => [arr] -> ShowS
showArrayList

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

instance Size D where
  size :: forall ix e. Array D ix e -> Sz ix
size = Array D ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
dSize
  {-# INLINE size #-}
  unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array D ix e -> Array D ix' e
unsafeResize !Sz ix'
sz !Array D ix e
arr =
    Comp -> Sz ix' -> (Int -> e) -> Array D ix' e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear (Array D ix e -> Comp
forall ix e. Array D ix e -> Comp
dComp Array D ix e
arr) Sz ix'
sz (Array D ix e -> ix -> e
forall ix. Index ix => Array D ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D ix e
arr (ix -> e) -> (Int -> ix) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex (Array D ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
size Array D ix e
arr))
  {-# INLINE unsafeResize #-}

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

instance Source D e where
  unsafeIndex :: forall ix. Index ix => Array D ix e -> ix -> e
unsafeIndex Array D ix e
arr =
    case Array D ix e -> PrefIndex ix e
forall ix e. Array D ix e -> PrefIndex ix e
dPrefIndex Array D ix e
arr of
      PrefIndex ix -> e
f -> ix -> e
f
      PrefIndexLinear Int -> e
f -> Int -> e
f (Int -> e) -> (ix -> Int) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex (Array D ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
size Array D ix e
arr)
  {-# INLINE unsafeIndex #-}
  unsafeLinearIndex :: forall ix. Index ix => Array D ix e -> Int -> e
unsafeLinearIndex Array D ix e
arr =
    case Array D ix e -> PrefIndex ix e
forall ix e. Array D ix e -> PrefIndex ix e
dPrefIndex Array D ix e
arr of
      PrefIndex ix -> e
f -> ix -> e
f (ix -> e) -> (Int -> ix) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> Int -> ix
forall ix. Index ix => Sz ix -> Int -> ix
fromLinearIndex (Array D ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
size Array D ix e
arr)
      PrefIndexLinear Int -> e
f -> Int -> e
f
  {-# INLINE unsafeLinearIndex #-}
  unsafePrefIndex :: forall ix. Index ix => Array D ix e -> PrefIndex ix e
unsafePrefIndex = Array D ix e -> PrefIndex ix e
forall ix e. Array D ix e -> PrefIndex ix e
dPrefIndex
  {-# INLINE unsafePrefIndex #-}
  unsafeOuterSlice :: forall ix.
(Index ix, Index (Lower ix)) =>
Array D ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeOuterSlice !Array D ix e
arr !Sz (Lower ix)
szL !Int
i =
    Comp -> Sz (Lower ix) -> (Lower ix -> e) -> Array D (Lower ix) e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray (Array D ix e -> Comp
forall ix e. Array D ix e -> Comp
dComp Array D ix e
arr) Sz (Lower ix)
szL (Array D ix e -> ix -> e
forall ix. Index ix => Array D ix e -> ix -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array D ix e
arr (ix -> e) -> (Lower ix -> ix) -> Lower ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Lower ix -> ix
forall ix. Index ix => Int -> Lower ix -> ix
consDim Int
i)
  {-# INLINE unsafeOuterSlice #-}
  unsafeLinearSlice :: forall ix. Index ix => Int -> Sz1 -> Array D ix e -> Array D Int e
unsafeLinearSlice !Int
o !Sz1
sz Array D ix e
arr =
    Comp -> Sz1 -> (Int -> e) -> Array D Int e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear (Array D ix e -> Comp
forall ix e. Array D ix e -> Comp
dComp Array D ix e
arr) Sz1
sz ((Int -> e) -> Array D Int e) -> (Int -> e) -> Array D Int e
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Array D ix e -> Int -> e
forall ix. Index ix => Array D ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array D ix e
arr (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
o)
  {-# INLINE unsafeLinearSlice #-}

-- | /O(1)/ - Extract a portion of an array. Staring index and new size are
-- not validated.
unsafeExtract :: (Source r e, Index ix) => ix -> Sz ix -> Array r ix e -> Array D ix e
unsafeExtract :: forall r e ix.
(Source r e, Index ix) =>
ix -> Sz ix -> Array r ix e -> Array D ix e
unsafeExtract !ix
sIx !Sz ix
newSz !Array r ix e
arr =
  Comp -> Sz ix -> (ix -> e) -> Array D ix e
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) Sz ix
newSz (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 -> e) -> (ix -> ix) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) ix
sIx)
{-# INLINE unsafeExtract #-}

-- | /O(1)/ - Take a slice out of an array from within
unsafeSlice
  :: (Source r e, Index ix, Index (Lower ix), MonadThrow m)
  => Array r ix e
  -> ix
  -> Sz ix
  -> Dim
  -> m (Array D (Lower ix) e)
unsafeSlice :: forall r e ix (m :: * -> *).
(Source r e, Index ix, Index (Lower ix), MonadThrow m) =>
Array r ix e -> ix -> Sz ix -> Dim -> m (Array D (Lower ix) e)
unsafeSlice Array r ix e
arr ix
start cut :: Sz ix
cut@(SafeSz ix
cutSz) Dim
dim = do
  Lower ix
newSz <- ix -> Dim -> m (Lower ix)
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
ix -> Dim -> m (Lower ix)
dropDimM ix
cutSz Dim
dim
  Array D (Lower ix) e -> m (Array D (Lower ix) e)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array D (Lower ix) e -> m (Array D (Lower ix) e))
-> Array D (Lower ix) e -> m (Array D (Lower ix) e)
forall a b. (a -> b) -> a -> b
$ Sz (Lower ix) -> Array D ix e -> Array D (Lower ix) e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array D ix e -> Array D ix' e
unsafeResize (Lower ix -> Sz (Lower ix)
forall ix. ix -> Sz ix
SafeSz Lower ix
newSz) (ix -> Sz ix -> Array r ix e -> Array D ix e
forall r e ix.
(Source r e, Index ix) =>
ix -> Sz ix -> Array r ix e -> Array D ix e
unsafeExtract ix
start Sz ix
cut Array r ix e
arr)
{-# INLINE unsafeSlice #-}

-- | /O(1)/ - Take a slice out of an array from the inside
unsafeInnerSlice
  :: (Source r e, Index ix) => Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeInnerSlice :: 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 =
  Comp
-> Sz (Lower ix) -> PrefIndex (Lower ix) e -> Array D (Lower ix) e
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray (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) Sz (Lower ix)
szL (PrefIndex (Lower ix) e -> Array D (Lower ix) e)
-> PrefIndex (Lower ix) e -> Array D (Lower ix) e
forall a b. (a -> b) -> a -> b
$ (Lower ix -> e) -> PrefIndex (Lower ix) e
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (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 -> e) -> (Lower ix -> ix) -> Lower ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Lower ix -> Int -> ix
forall ix. Index ix => Lower ix -> Int -> ix
`snocDim` Int
i))
{-# INLINE unsafeInnerSlice #-}

instance (Eq e, Index ix) => Eq (Array D ix e) where
  == :: Array D ix e -> Array D ix e -> Bool
(==) = (e -> e -> Bool) -> Array D ix e -> Array D ix e -> Bool
forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eqArrays e -> e -> Bool
forall a. Eq a => a -> a -> Bool
(==)
  {-# INLINE (==) #-}

instance (Ord e, Index ix) => Ord (Array D ix e) where
  compare :: Array D ix e -> Array D ix e -> Ordering
compare = (e -> e -> Ordering) -> Array D ix e -> Array D ix e -> Ordering
forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
compareArrays e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
  {-# INLINE compare #-}

instance Functor (Array D ix) where
  fmap :: forall a b. (a -> b) -> Array D ix a -> Array D ix b
fmap a -> b
f (DArray Comp
c Sz ix
sz PrefIndex ix a
g) = Comp -> Sz ix -> PrefIndex ix b -> Array D ix b
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray Comp
c Sz ix
sz ((a -> b) -> PrefIndex ix a -> PrefIndex ix b
forall a b. (a -> b) -> PrefIndex ix a -> PrefIndex ix b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f PrefIndex ix a
g)
  {-# INLINE fmap #-}
  <$ :: forall a b. a -> Array D ix b -> Array D ix a
(<$) a
e (DArray Comp
c Sz ix
sz PrefIndex ix b
g) = Comp -> Sz ix -> PrefIndex ix a -> Array D ix a
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray Comp
c Sz ix
sz (a
e a -> PrefIndex ix b -> PrefIndex ix a
forall a b. a -> PrefIndex ix b -> PrefIndex ix a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ PrefIndex ix b
g)
  {-# INLINE (<$) #-}

instance Index ix => Applicative (Array D ix) where
  pure :: forall a. a -> Array D ix a
pure = a -> Array D ix a
forall r ix e. Load r ix e => e -> Array r ix e
singleton
  {-# INLINE pure #-}
  <*> :: forall a b. Array D ix (a -> b) -> Array D ix a -> Array D ix b
(<*>) = ((a -> b) -> a -> b)
-> Array D ix (a -> b) -> Array D ix a -> Array D ix b
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' (a -> b) -> a -> b
forall a. a -> a
id
  {-# INLINE (<*>) #-}
#if MIN_VERSION_base(4,10,0)
  liftA2 :: forall a b c.
(a -> b -> c) -> Array D ix a -> Array D ix b -> Array D ix c
liftA2 = (a -> b -> c) -> Array D ix a -> Array D ix b -> Array D ix c
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2'
  {-# INLINE liftA2 #-}
#endif

-- | Row-major sequential folding over a Delayed array.
instance Index ix => Foldable (Array D ix) where
  fold :: forall m. Monoid m => Array D ix m -> m
fold = Array D ix m -> m
forall e ix r.
(Monoid e, Index ix, Source r e) =>
Array r ix e -> e
A.fold
  {-# INLINE fold #-}
  foldMap :: forall m a. Monoid m => (a -> m) -> Array D ix a -> m
foldMap = (a -> m) -> Array D ix a -> m
forall ix r e m.
(Index ix, Source r e, Monoid m) =>
(e -> m) -> Array r ix e -> m
A.foldMono
  {-# INLINE foldMap #-}
  foldl :: forall b a. (b -> a -> b) -> b -> Array D ix a -> b
foldl = (b -> a -> b) -> b -> Array D ix a -> b
forall ix r e a.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> a
lazyFoldlS
  {-# INLINE foldl #-}
  foldl' :: forall b a. (b -> a -> b) -> b -> Array D ix a -> b
foldl' = (b -> a -> b) -> b -> Array D ix a -> b
forall ix r e a.
(Index ix, Source r e) =>
(a -> e -> a) -> a -> Array r ix e -> a
foldlS
  {-# INLINE foldl' #-}
  foldr :: forall a b. (a -> b -> b) -> b -> Array D ix a -> b
foldr = (a -> b -> b) -> b -> Array D ix a -> b
forall ix r e b.
(Index ix, Source r e) =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrFB
  {-# INLINE foldr #-}
  foldr' :: forall a b. (a -> b -> b) -> b -> Array D ix a -> b
foldr' = (a -> b -> b) -> b -> Array D ix a -> b
forall ix r e b.
(Index ix, Source r e) =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrS
  {-# INLINE foldr' #-}
  null :: forall a. Array D ix a -> Bool
null (DArray Comp
_ Sz ix
sz PrefIndex ix a
_) = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
  {-# INLINE null #-}
  length :: forall a. Array D ix a -> Int
length = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Sz ix -> Int) -> (Array D ix a -> Sz ix) -> Array D ix a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array D ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
size
  {-# INLINE length #-}
  elem :: forall a. Eq a => a -> Array D ix a -> Bool
elem a
e = (a -> Bool) -> Array D ix a -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
A.any (a
e a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==)
  {-# INLINE elem #-}
  toList :: forall a. Array D ix a -> [a]
toList Array D ix a
arr = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> (a -> b -> b) -> b -> Array D ix a -> b
forall ix r e b.
(Index ix, Source r e) =>
(e -> b -> b) -> b -> Array r ix e -> b
foldrFB a -> b -> b
c b
n Array D ix a
arr)
  {-# INLINE toList #-}

instance Index ix => Load D ix e where
  makeArray :: Comp -> Sz ix -> (ix -> e) -> Array D ix e
makeArray Comp
comp Sz ix
sz = Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray Comp
comp Sz ix
sz (PrefIndex ix e -> Array D ix e)
-> ((ix -> e) -> PrefIndex ix e) -> (ix -> e) -> Array D ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix -> e) -> PrefIndex ix e
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex
  {-# INLINE makeArray #-}
  makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array D ix e
makeArrayLinear Comp
comp Sz ix
sz = Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray Comp
comp Sz ix
sz (PrefIndex ix e -> Array D ix e)
-> ((Int -> e) -> PrefIndex ix e) -> (Int -> e) -> Array D ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> e) -> PrefIndex ix e
forall ix e. (Int -> e) -> PrefIndex ix e
PrefIndexLinear
  {-# INLINE makeArrayLinear #-}
  iterArrayLinearST_ :: forall s.
Scheduler s () -> Array D ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ !Scheduler s ()
scheduler DArray{Comp
Sz ix
PrefIndex ix e
dComp :: forall ix e. Array D ix e -> Comp
dSize :: forall ix e. Array D ix e -> Sz ix
dPrefIndex :: forall ix e. Array D ix e -> PrefIndex ix e
dComp :: Comp
dSize :: Sz ix
dPrefIndex :: PrefIndex ix e
..} Int -> e -> ST s ()
uWrite =
    case PrefIndex ix e
dPrefIndex of
      PrefIndex ix -> e
f ->
        RowMajor
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
forall ix s.
Index ix =>
RowMajor
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetFullST_ RowMajor
defRowMajor Scheduler s ()
scheduler Int
0 Sz ix
dSize ((Int -> ix -> ST s ()) -> ST s ())
-> (Int -> ix -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i -> Int -> e -> ST s ()
uWrite Int
i (e -> ST s ()) -> (ix -> e) -> ix -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
f
      PrefIndexLinear Int -> e
f ->
        RowMajorLinear
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
forall ix s.
Index ix =>
RowMajorLinear
-> Scheduler s ()
-> Int
-> Sz ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetFullST_ RowMajorLinear
defRowMajorLinear Scheduler s ()
scheduler Int
0 Sz ix
dSize ((Int -> ix -> ST s ()) -> ST s ())
-> (Int -> ix -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \ !Int
i ix
_ -> Int -> e -> ST s ()
uWrite Int
i (Int -> e
f Int
i)
  {-# INLINE iterArrayLinearST_ #-}

instance Index ix => StrideLoad D ix e where
  iterArrayLinearWithStrideST_ :: forall s.
Scheduler s ()
-> Stride ix
-> Sz ix
-> Array D ix e
-> (Int -> e -> ST s ())
-> ST s ()
iterArrayLinearWithStrideST_ !Scheduler s ()
scheduler !Stride ix
stride Sz ix
sz DArray{Comp
Sz ix
PrefIndex ix e
dComp :: forall ix e. Array D ix e -> Comp
dSize :: forall ix e. Array D ix e -> Sz ix
dPrefIndex :: forall ix e. Array D ix e -> PrefIndex ix e
dComp :: Comp
dSize :: Sz ix
dPrefIndex :: PrefIndex ix e
..} Int -> e -> ST s ()
uWrite =
    case PrefIndex ix e
dPrefIndex of
      PrefIndex ix -> e
f ->
        RowMajor
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> (Int -> ix -> ST s ())
-> ST s ()
forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetFullWithStrideST_ RowMajor
defRowMajor Scheduler s ()
scheduler Int
0 Sz ix
sz Stride ix
stride ((Int -> ix -> ST s ()) -> ST s ())
-> (Int -> ix -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
          Int -> e -> ST s ()
uWrite Int
i (e -> ST s ()) -> (ix -> e) -> ix -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
f
      PrefIndexLinear Int -> e
f -> do
        RowMajor
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> (Int -> ix -> ST s ())
-> ST s ()
forall it ix s.
(Iterator it, Index ix) =>
it
-> Scheduler s ()
-> Int
-> Sz ix
-> Stride ix
-> (Int -> ix -> ST s ())
-> ST s ()
iterTargetFullWithStrideST_ RowMajor
defRowMajor Scheduler s ()
scheduler Int
0 Sz ix
sz Stride ix
stride ((Int -> ix -> ST s ()) -> ST s ())
-> (Int -> ix -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i ->
          Int -> e -> ST s ()
uWrite Int
i (e -> ST s ()) -> (ix -> e) -> ix -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
f (Int -> e) -> (ix -> Int) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
dSize
  {-# INLINE iterArrayLinearWithStrideST_ #-}

instance Index ix => Stream D ix e where
  toStream :: Array D ix e -> Steps Id e
toStream = Array D ix e -> Steps Id e
forall r ix e (m :: * -> *).
(Monad m, Index ix, Source r e) =>
Array r ix e -> Steps m e
S.steps
  {-# INLINE toStream #-}
  toStreamIx :: Array D ix e -> Steps Id (ix, e)
toStreamIx = Array D ix (ix, e) -> Steps Id (ix, e)
forall r ix e (m :: * -> *).
(Monad m, Index ix, Source r e) =>
Array r ix e -> Steps m e
S.steps (Array D ix (ix, e) -> Steps Id (ix, e))
-> (Array D ix e -> Array D ix (ix, e))
-> Array D ix e
-> Steps Id (ix, e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ix -> e -> (ix, e)) -> Array D ix e -> Array D ix (ix, e)
forall r ix e a.
(Index ix, Source r e) =>
(ix -> e -> a) -> Array r ix e -> Array D ix a
imap (,)
  {-# INLINE toStreamIx #-}

-- | Map an index aware function over an array
--
-- @since 0.1.0
imap
  :: forall r ix e a
   . (Index ix, Source r e)
  => (ix -> e -> a)
  -> Array r ix e
  -> Array D ix a
imap :: forall r ix e a.
(Index ix, Source r e) =>
(ix -> e -> a) -> Array r ix e -> Array D ix a
imap ix -> e -> a
f !Array r ix e
arr =
  case Array r ix e -> PrefIndex ix e
forall ix. Index ix => Array r ix e -> PrefIndex ix e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
    PrefIndex ix -> e
gix -> Comp -> Sz ix -> PrefIndex ix a -> Array D ix a
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray (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) Sz ix
sz (PrefIndex ix a -> Array D ix a) -> PrefIndex ix a -> Array D ix a
forall a b. (a -> b) -> a -> b
$ (ix -> a) -> PrefIndex ix a
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (\ !ix
ix -> ix -> e -> a
f ix
ix (ix -> e
gix ix
ix))
    PrefIndexLinear Int -> e
gi ->
      Comp -> Sz ix -> PrefIndex ix a -> Array D ix a
forall ix e. Comp -> Sz ix -> PrefIndex ix e -> Array D ix e
DArray (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) Sz ix
sz (PrefIndex ix a -> Array D ix a) -> PrefIndex ix a -> Array D ix a
forall a b. (a -> b) -> a -> b
$ (ix -> a) -> PrefIndex ix a
forall ix e. (ix -> e) -> PrefIndex ix e
PrefIndex (\ !ix
ix -> ix -> e -> a
f ix
ix (Int -> e
gi (Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)))
  where
    !sz :: Sz 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
{-# INLINE imap #-}

instance Num e => FoldNumeric D e where
  unsafeDotProduct :: forall ix. Index ix => Array D ix e -> Array D ix e -> e
unsafeDotProduct = Array D ix e -> Array D ix e -> e
forall e ix r.
(Num e, Index ix, Source r e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray :: forall ix. Index ix => Array D ix e -> Int -> e
powerSumArray = Array D ix e -> Int -> e
forall ix r e.
(Index ix, Source r e, Num e) =>
Array r ix e -> Int -> e
defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray :: forall ix. Index ix => (e -> e -> e) -> e -> Array D ix e -> e
foldArray = (e -> e -> e) -> e -> Array D ix e -> e
forall ix r e.
(Index ix, Source r e) =>
(e -> e -> e) -> e -> Array r ix e -> e
defaultFoldArray
  {-# INLINE foldArray #-}

instance Num e => Numeric D e where
  unsafeLiftArray :: forall ix. Index ix => (e -> e) -> Array D ix e -> Array D ix e
unsafeLiftArray e -> e
f Array D ix e
arr = Array D ix e
arr{dPrefIndex = f <$> dPrefIndex arr}
  {-# INLINE unsafeLiftArray #-}
  unsafeLiftArray2 :: forall ix.
Index ix =>
(e -> e -> e) -> Array D ix e -> Array D ix e -> Array D ix e
unsafeLiftArray2 e -> e -> e
f Array D ix e
a1 Array D ix e
a2 = Sz ix
-> (e -> e -> e) -> Array D ix e -> Array D ix e -> Array D ix e
forall ix r1 e1 r2 e2 e3.
(Index ix, Source r1 e1, Source r2 e2) =>
Sz ix
-> (e1 -> e2 -> e3)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array D ix e3
zipWithInternal (Array D ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array D ix e -> Sz ix
size Array D ix e
a1) e -> e -> e
f Array D ix e
a1 Array D ix e
a2
  {-# INLINE unsafeLiftArray2 #-}

instance Floating e => NumericFloat D e

-- | /O(1)/ Conversion from a source array to `D` representation.
delay :: (Index ix, Source r e) => Array r ix e -> Array D ix e
delay :: forall ix r e.
(Index ix, Source r e) =>
Array r ix e -> Array D ix e
delay Array r ix e
arr =
  case Array r ix e -> PrefIndex ix e
forall ix. Index ix => Array r ix e -> PrefIndex ix e
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r ix e
arr of
    PrefIndex ix -> e
gix -> Comp -> Sz ix -> (ix -> e) -> Array D ix e
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) (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) ix -> e
gix
    PrefIndexLinear Int -> e
gi -> Comp -> Sz ix -> (Int -> e) -> Array D ix e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear (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) (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 -> e
gi
{-# INLINE [1] delay #-}

{-# RULES
"delay" [~1] forall (arr :: Array D ix e). delay arr = arr
  #-}

-- | Compute array equality by applying a comparing function to each
-- element. Empty arrays are always equal, regardless of their size.
--
-- @since 0.5.7
eqArrays
  :: (Index ix, Source r1 e1, Source r2 e2)
  => (e1 -> e2 -> Bool)
  -> Array r1 ix e1
  -> Array r2 ix e2
  -> Bool
eqArrays :: forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eqArrays e1 -> e2 -> Bool
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
  let sz1 :: Sz ix
sz1 = Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix e1
arr1
      sz2 :: Sz ix
sz2 = Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 ix e -> Sz ix
size Array r2 ix e2
arr2
   in ( Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2
          Bool -> Bool -> Bool
&& Bool -> Bool
not
            ( (Bool -> Bool) -> Array D ix Bool -> Bool
forall ix r e.
(Index ix, Source r e) =>
(e -> Bool) -> Array r ix e -> Bool
A.any
                Bool -> Bool
not
                ( forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray @D (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r1 ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r2 ix e -> Comp
getComp Array r2 ix e2
arr2) (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix e1
arr1) ((ix -> Bool) -> Array D ix Bool)
-> (ix -> Bool) -> Array D ix Bool
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
                    e1 -> e2 -> Bool
f (Array r1 ix e1 -> ix -> e1
forall ix. Index ix => Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall ix. Index ix => Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix)
                )
            )
      )
        Bool -> Bool -> Bool
|| (Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz1 Bool -> Bool -> Bool
&& Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz2)
{-# INLINE eqArrays #-}

-- | Compute array ordering by applying a comparing function to each element.
-- The exact ordering is unspecified so this is only intended for use in maps and the like where
-- you need an ordering but do not care about which one is used.
--
-- @since 0.5.7
compareArrays
  :: (Index ix, Source r1 e1, Source r2 e2)
  => (e1 -> e2 -> Ordering)
  -> Array r1 ix e1
  -> Array r2 ix e2
  -> Ordering
compareArrays :: forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
compareArrays e1 -> e2 -> Ordering
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
  Sz ix -> Sz ix -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix e1
arr1) (Array r2 ix e2 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 ix e -> Sz ix
size Array r2 ix e2
arr2)
    Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Array D ix Ordering -> Ordering
forall e ix r.
(Monoid e, Index ix, Source r e) =>
Array r ix e -> e
A.fold
      ( forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray @D (Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r1 ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r2 ix e -> Comp
getComp Array r2 ix e2
arr2) (Array r1 ix e1 -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix e1
arr1) ((ix -> Ordering) -> Array D ix Ordering)
-> (ix -> Ordering) -> Array D ix Ordering
forall a b. (a -> b) -> a -> b
$ \ix
ix ->
          e1 -> e2 -> Ordering
f (Array r1 ix e1 -> ix -> e1
forall ix. Index ix => Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall ix. Index ix => Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix)
      )
{-# INLINE compareArrays #-}

-- | Same as `liftArray2M`, but throws an imprecise exception on mismatched
-- sizes.
--
-- @since 1.0.0
liftArray2'
  :: (HasCallStack, Index ix, Source r1 a, Source r2 b)
  => (a -> b -> e)
  -> Array r1 ix a
  -> Array r2 ix b
  -> Array D ix e
liftArray2' :: forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' a -> b -> e
f Array r1 ix a
arr1 Array r2 ix b
arr2 = Either SomeException (Array D ix e) -> Array D ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array D ix e) -> Array D ix e)
-> Either SomeException (Array D ix e) -> Array D ix e
forall a b. (a -> b) -> a -> b
$ (a -> b -> e)
-> Array r1 ix a
-> Array r2 ix b
-> Either SomeException (Array D ix e)
forall ix r1 a r2 b (m :: * -> *) e.
(Index ix, Source r1 a, Source r2 b, MonadThrow m) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> m (Array D ix e)
liftArray2M a -> b -> e
f Array r1 ix a
arr1 Array r2 ix b
arr2
{-# INLINE liftArray2' #-}

-- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays
-- have to be the same, otherwise it throws `SizeMismatchException`.
--
-- @since 1.0.0
liftArray2M
  :: (Index ix, Source r1 a, Source r2 b, MonadThrow m)
  => (a -> b -> e)
  -> Array r1 ix a
  -> Array r2 ix b
  -> m (Array D ix e)
liftArray2M :: forall ix r1 a r2 b (m :: * -> *) e.
(Index ix, Source r1 a, Source r2 b, MonadThrow m) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> m (Array D ix e)
liftArray2M a -> b -> e
f !Array r1 ix a
arr1 !Array r2 ix b
arr2
  | Sz ix
sz1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Sz ix
sz2 = Array D ix e -> m (Array D ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array D ix e -> m (Array D ix e))
-> Array D ix e -> m (Array D ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix
-> (a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
forall ix r1 e1 r2 e2 e3.
(Index ix, Source r1 e1, Source r2 e2) =>
Sz ix
-> (e1 -> e2 -> e3)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array D ix e3
zipWithInternal Sz ix
sz1 a -> b -> e
f Array r1 ix a
arr1 Array r2 ix b
arr2
  | Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz1 Bool -> Bool -> Bool
&& Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz2 = Array D ix e -> m (Array D ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array D ix e
forall r ix e. Load r ix e => Array r ix e
A.empty
  | Bool
otherwise = SizeException -> m (Array D ix e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Array D ix e))
-> SizeException -> m (Array D ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Array r1 ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix a
arr1) (Array r2 ix b -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 ix e -> Sz ix
size Array r2 ix b
arr2)
  where
    sz1 :: Sz ix
sz1 = Array r1 ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r1 ix e -> Sz ix
size Array r1 ix a
arr1
    sz2 :: Sz ix
sz2 = Array r2 ix b -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r2 ix e -> Sz ix
size Array r2 ix b
arr2
{-# INLINE liftArray2M #-}

zipWithInternal
  :: (Index ix, Source r1 e1, Source r2 e2)
  => Sz ix
  -> (e1 -> e2 -> e3)
  -> Array r1 ix e1
  -> Array r2 ix e2
  -> Array D ix e3
zipWithInternal :: forall ix r1 e1 r2 e2 e3.
(Index ix, Source r1 e1, Source r2 e2) =>
Sz ix
-> (e1 -> e2 -> e3)
-> Array r1 ix e1
-> Array r2 ix e2
-> Array D ix e3
zipWithInternal Sz ix
sz e1 -> e2 -> e3
f Array r1 ix e1
arr1 Array r2 ix e2
arr2 =
  case Array r1 ix e1 -> PrefIndex ix e1
forall ix. Index ix => Array r1 ix e1 -> PrefIndex ix e1
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r1 ix e1
arr1 of
    PrefIndexLinear Int -> e1
gi1
      | PrefIndexLinear Int -> e2
gi2 <- Array r2 ix e2 -> PrefIndex ix e2
forall ix. Index ix => Array r2 ix e2 -> PrefIndex ix e2
forall r e ix.
(Source r e, Index ix) =>
Array r ix e -> PrefIndex ix e
unsafePrefIndex Array r2 ix e2
arr2 ->
          Comp -> Sz ix -> (Int -> e3) -> Array D ix e3
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear Comp
comp Sz ix
sz (\ !Int
i -> e1 -> e2 -> e3
f (Int -> e1
gi1 Int
i) (Int -> e2
gi2 Int
i))
    PrefIndex ix e1
_ -> Comp -> Sz ix -> (ix -> e3) -> Array D ix e3
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz (\ !ix
ix -> e1 -> e2 -> e3
f (Array r1 ix e1 -> ix -> e1
forall ix. Index ix => Array r1 ix e1 -> ix -> e1
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r1 ix e1
arr1 ix
ix) (Array r2 ix e2 -> ix -> e2
forall ix. Index ix => Array r2 ix e2 -> ix -> e2
forall r e ix. (Source r e, Index ix) => Array r ix e -> ix -> e
unsafeIndex Array r2 ix e2
arr2 ix
ix))
  where
    comp :: Comp
comp = Array r1 ix e1 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r1 ix e -> Comp
getComp Array r1 ix e1
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r2 ix e2 -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r2 ix e -> Comp
getComp Array r2 ix e2
arr2
{-# INLINE zipWithInternal #-}