{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
module Data.Massiv.Array.Stencil.Unsafe
(
makeUnsafeStencil
, unsafeTransformStencil
, unsafeMapStencil
, mapStencilUnsafe
, forStencilUnsafe
) where
import Data.Massiv.Array.Delayed.Windowed (Array(..), DW, Window(..),
insertWindow)
import Data.Massiv.Array.Stencil.Internal
import Data.Massiv.Core.Common
import GHC.Exts (inline)
forStencilUnsafe ::
(Source r ix e, Manifest r ix e)
=> Array r ix e
-> Sz ix
-> ix
-> ((ix -> Maybe e) -> a)
-> Array DW ix a
forStencilUnsafe :: Array r ix e
-> Sz ix -> ix -> ((ix -> Maybe e) -> a) -> Array DW ix a
forStencilUnsafe !Array r ix e
arr !Sz ix
sSz !ix
sCenter (ix -> Maybe e) -> a
relStencil =
Array D ix a -> Window ix a -> Array DW ix a
forall ix e.
Source D ix e =>
Array D ix e -> Window ix e -> Array DW ix e
insertWindow (Comp -> Sz ix -> (ix -> a) -> Array D ix a
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
sz ((ix -> Maybe e) -> ix -> a
stencil (Array r ix e -> ix -> Maybe e
forall r ix e. Manifest r ix e => Array r ix e -> ix -> Maybe e
index Array r ix e
arr))) Window ix a
window
where
!window :: Window ix a
window =
Window :: forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
Window
{ windowStart :: ix
windowStart = ix
sCenter
, windowSize :: Sz ix
windowSize = Sz ix
windowSz
, windowIndex :: ix -> a
windowIndex = (ix -> Maybe e) -> ix -> a
stencil (e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> (ix -> e) -> ix -> Maybe e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)
, windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = Sz Int -> Int
forall ix. Sz ix -> ix
unSz (Sz Int -> Int)
-> ((Sz Int, Sz (Lower ix)) -> Sz Int)
-> (Sz Int, Sz (Lower ix))
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Sz Int, Sz (Lower ix)) -> Sz Int
forall a b. (a, b) -> a
fst ((Sz Int, Sz (Lower ix)) -> Int)
-> Maybe (Sz Int, Sz (Lower ix)) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> Dim -> Maybe (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
windowSz Dim
2
}
!sz :: Sz 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
!windowSz :: Sz ix
windowSz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sSz)))
stencil :: (ix -> Maybe e) -> ix -> a
stencil ix -> Maybe e
getVal !ix
ix = ((ix -> Maybe e) -> a) -> (ix -> Maybe e) -> a
forall a. a -> a
inline (ix -> Maybe e) -> a
relStencil ((ix -> Maybe e) -> a) -> (ix -> Maybe e) -> a
forall a b. (a -> b) -> a -> b
$ \ !ix
ixD -> ix -> Maybe e
getVal ((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
ix ix
ixD)
{-# INLINE stencil #-}
{-# INLINE forStencilUnsafe #-}
{-# DEPRECATED forStencilUnsafe "In favor of `unsafeMapStencil`" #-}
mapStencilUnsafe ::
Manifest r ix e
=> Border e
-> Sz ix
-> ix
-> ((ix -> e) -> a)
-> Array r ix e
-> Array DW ix a
mapStencilUnsafe :: Border e
-> Sz ix -> ix -> ((ix -> e) -> a) -> Array r ix e -> Array DW ix a
mapStencilUnsafe Border e
b Sz ix
sz ix
ix (ix -> e) -> a
f = Border e
-> Sz ix
-> ix
-> (ix -> (ix -> e) -> a)
-> Array r ix e
-> Array DW ix a
forall r ix e a.
Manifest r ix e =>
Border e
-> Sz ix
-> ix
-> (ix -> (ix -> e) -> a)
-> Array r ix e
-> Array DW ix a
unsafeMapStencil Border e
b Sz ix
sz ix
ix (((ix -> e) -> a) -> ix -> (ix -> e) -> a
forall a b. a -> b -> a
const (ix -> e) -> a
f)
{-# INLINE mapStencilUnsafe #-}
{-# DEPRECATED mapStencilUnsafe "In favor of `unsafeMapStencil`" #-}
unsafeMapStencil ::
Manifest r ix e
=> Border e
-> Sz ix
-> ix
-> (ix -> (ix -> e) -> a)
-> Array r ix e
-> Array DW ix a
unsafeMapStencil :: Border e
-> Sz ix
-> ix
-> (ix -> (ix -> e) -> a)
-> Array r ix e
-> Array DW ix a
unsafeMapStencil Border e
b Sz ix
sSz ix
sCenter ix -> (ix -> e) -> a
stencilF !Array r ix e
arr = Array D ix a -> Window ix a -> Array DW ix a
forall ix e.
Source D ix e =>
Array D ix e -> Window ix e -> Array DW ix e
insertWindow Array D ix a
warr Window ix a
window
where
!warr :: Array D ix a
warr = Comp -> Sz ix -> (ix -> a) -> Array D ix a
forall ix e. Comp -> Sz ix -> (ix -> e) -> Array D ix e
DArray (Array r ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array r ix e
arr) Sz ix
sz ((ix -> e) -> ix -> a
stencil (Border e -> Array r ix e -> ix -> e
forall r ix e.
Manifest r ix e =>
Border e -> Array r ix e -> ix -> e
borderIndex Border e
b Array r ix e
arr))
!window :: Window ix a
window =
Window :: forall ix e. ix -> Sz ix -> (ix -> e) -> Maybe Int -> Window ix e
Window
{ windowStart :: ix
windowStart = ix
sCenter
, windowSize :: Sz ix
windowSize = Sz ix
windowSz
, windowIndex :: ix -> a
windowIndex = (ix -> e) -> ix -> a
stencil (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)
, windowUnrollIx2 :: Maybe Int
windowUnrollIx2 = Sz Int -> Int
forall ix. Sz ix -> ix
unSz (Sz Int -> Int)
-> ((Sz Int, Sz (Lower ix)) -> Sz Int)
-> (Sz Int, Sz (Lower ix))
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Sz Int, Sz (Lower ix)) -> Sz Int
forall a b. (a, b) -> a
fst ((Sz Int, Sz (Lower ix)) -> Int)
-> Maybe (Sz Int, Sz (Lower ix)) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sz ix -> Dim -> Maybe (Sz Int, Sz (Lower ix))
forall (m :: * -> *) ix.
(MonadThrow m, Index ix) =>
Sz ix -> Dim -> m (Sz Int, Sz (Lower ix))
pullOutSzM Sz ix
sSz Dim
2
}
!sz :: Sz 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
!windowSz :: Sz ix
windowSz = ix -> Sz ix
forall ix. Index ix => ix -> Sz ix
Sz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 (-) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sz) ((Int -> Int) -> ix -> ix
forall ix. Index ix => (Int -> Int) -> ix -> ix
liftIndex (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz Sz ix
sSz)))
stencil :: (ix -> e) -> ix -> a
stencil ix -> e
getVal !ix
ix = ((ix -> e) -> a) -> (ix -> e) -> a
forall a. a -> a
inline (ix -> (ix -> e) -> a
stencilF ix
ix) ((ix -> e) -> a) -> (ix -> e) -> a
forall a b. (a -> b) -> a -> b
$ \ !ix
ixD -> ix -> e
getVal ((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
ix ix
ixD)
{-# INLINE stencil #-}
{-# INLINE unsafeMapStencil #-}
makeUnsafeStencil
:: Index ix
=> Sz ix
-> ix
-> (ix -> (ix -> e) -> a)
-> Stencil ix e a
makeUnsafeStencil :: Sz ix -> ix -> (ix -> (ix -> e) -> a) -> Stencil ix e a
makeUnsafeStencil !Sz ix
sSz !ix
sCenter ix -> (ix -> e) -> a
relStencil = Sz ix -> ix -> ((ix -> Value e) -> ix -> Value a) -> Stencil ix e a
forall ix e a.
Sz ix -> ix -> ((ix -> Value e) -> ix -> Value a) -> Stencil ix e a
Stencil Sz ix
sSz ix
sCenter (ix -> Value e) -> ix -> Value a
stencil
where
stencil :: (ix -> Value e) -> ix -> Value a
stencil ix -> Value e
getVal !ix
ix =
a -> Value a
forall e. e -> Value e
Value (a -> Value a) -> a -> Value a
forall a b. (a -> b) -> a -> b
$ a -> a
forall a. a -> a
inline (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ ix -> (ix -> e) -> a
relStencil ix
ix (Value e -> e
forall e. Value e -> e
unValue (Value e -> e) -> (ix -> Value e) -> ix -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> Value e
getVal (ix -> Value e) -> (ix -> ix) -> ix -> Value 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
ix)
{-# INLINE stencil #-}
{-# INLINE makeUnsafeStencil #-}
unsafeTransformStencil ::
(Sz ix' -> Sz ix)
-> (ix' -> ix)
-> (((ix' -> Value e) -> ix' -> Value a) -> (ix -> Value e) -> ix -> Value a)
-> Stencil ix' e a
-> Stencil ix e a
unsafeTransformStencil :: (Sz ix' -> Sz ix)
-> (ix' -> ix)
-> (((ix' -> Value e) -> ix' -> Value a)
-> (ix -> Value e) -> ix -> Value a)
-> Stencil ix' e a
-> Stencil ix e a
unsafeTransformStencil Sz ix' -> Sz ix
transformSize ix' -> ix
transformIndex ((ix' -> Value e) -> ix' -> Value a)
-> (ix -> Value e) -> ix -> Value a
transformFunc Stencil {ix'
Sz ix'
(ix' -> Value e) -> ix' -> Value a
stencilFunc :: forall ix e a. Stencil ix e a -> (ix -> Value e) -> ix -> Value a
stencilCenter :: forall ix e a. Stencil ix e a -> ix
stencilSize :: forall ix e a. Stencil ix e a -> Sz ix
stencilFunc :: (ix' -> Value e) -> ix' -> Value a
stencilCenter :: ix'
stencilSize :: Sz ix'
..} =
Stencil :: forall ix e a.
Sz ix -> ix -> ((ix -> Value e) -> ix -> Value a) -> Stencil ix e a
Stencil
{ stencilSize :: Sz ix
stencilSize = Sz ix' -> Sz ix
transformSize Sz ix'
stencilSize
, stencilCenter :: ix
stencilCenter = ix' -> ix
transformIndex ix'
stencilCenter
, stencilFunc :: (ix -> Value e) -> ix -> Value a
stencilFunc = ((ix' -> Value e) -> ix' -> Value a)
-> (ix -> Value e) -> ix -> Value a
transformFunc (ix' -> Value e) -> ix' -> Value a
stencilFunc
}
{-# INLINE unsafeTransformStencil #-}