{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Data.Massiv.Array.Stencil.Unsafe
(
makeUnsafeStencil
, 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 !arr !sSz !sCenter relStencil =
insertWindow (DArray (getComp arr) sz (stencil (index arr))) window
where
!window =
Window
{ windowStart = sCenter
, windowSize = windowSz
, windowIndex = stencil (Just . unsafeIndex arr)
, windowUnrollIx2 = unSz . fst <$> pullOutSzM windowSz 2
}
!sz = size arr
!windowSz = Sz (liftIndex2 (-) (unSz sz) (liftIndex (subtract 1) (unSz sSz)))
stencil getVal !ix = inline relStencil $ \ !ixD -> getVal (liftIndex2 (+) ix ixD)
{-# INLINE stencil #-}
{-# INLINE forStencilUnsafe #-}
makeUnsafeStencil
:: Index ix
=> Sz ix
-> ix
-> (ix -> (ix -> e) -> a)
-> Stencil ix e a
makeUnsafeStencil !sSz !sCenter relStencil = Stencil sSz sCenter stencil
where
stencil getVal !ix =
Value $ inline $ relStencil ix (unValue . getVal . liftIndex2 (+) ix)
{-# INLINE stencil #-}
{-# INLINE makeUnsafeStencil #-}