{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DeriveFunctor         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuasiQuotes           #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TemplateHaskell       #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Dense.Stencil
-- Copyright   :  (c) Christopher Chalmers
-- License     :  BSD3
--
-- Maintainer  :  Christopher Chalmers
-- Stability   :  provisional
-- Portability :  non-portable
--
-- Stencils can be used to sum (or any fold) over neighbouring sites to
-- the current position on a 'Focused'.
-----------------------------------------------------------------------------
module Data.Dense.Stencil
  ( -- * The Stencil type
    Stencil (..)
  , mkStencil
  , mkStencilUnboxed

    -- ** Using stencils
  , stencilSum

  ) where

import           Control.Lens
import           Data.Dense.Base
import           Data.Dense.Generic   (Boundary (..), peekRelativeB)
import           Data.Dense.Index
import qualified Data.Foldable        as F
import           Data.Functor.Classes
import qualified Data.Vector.Unboxed  as U
import           Text.Show

-- Types ---------------------------------------------------------------

-- | Stencils are used to fold over neighbouring array sites. To
--   construct a stencil use 'mkStencil', 'mkStencilUnboxed'. For
--   static sized stencils you can use the quasiquoter
--   'Data.Dense.TH.stencil'.
--
--   To use a stencil you can use 'stencilSum' or use the 'Foldable' and
--   'FoldableWithIndex' instances.
newtype Stencil f a = Stencil (forall b. (f Int -> a -> b -> b) -> b -> b)

instance (Show1 f, Show a) => Show (Stencil f a) where
  showsPrec :: Int -> Stencil f a -> ShowS
showsPrec Int
_ Stencil f a
s = ((f Int, a) -> ShowS) -> [(f Int, a)] -> ShowS
forall a. (a -> ShowS) -> [a] -> ShowS
showListWith (f Int, a) -> ShowS
forall (f :: * -> *) a a.
(Show1 f, Show a, Show a) =>
(f a, a) -> ShowS
g (Stencil f a -> [(f Int, a)]
forall i (f :: * -> *) a. FoldableWithIndex i f => f a -> [(i, a)]
itoList Stencil f a
s) where
    g :: (f a, a) -> ShowS
g (f a
i,a
x) = Char -> ShowS
showChar Char
'(' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> f a -> ShowS
forall (f :: * -> *) a. (Show1 f, Show a) => Int -> f a -> ShowS
showsPrec1 Int
0 f a
i ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
',' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
0 a
x ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
')'

instance F.Foldable (Stencil f) where
  foldr :: (a -> b -> b) -> b -> Stencil f a -> b
foldr a -> b -> b
f b
z (Stencil forall b. (f Int -> a -> b -> b) -> b -> b
s) = (f Int -> a -> b -> b) -> b -> b
forall b. (f Int -> a -> b -> b) -> b -> b
s (\f Int
_ a
a b
b -> a -> b -> b
f a
a b
b) b
z
  {-# INLINE foldr #-}

instance FoldableWithIndex (f Int) (Stencil f) where
  ifoldr :: (f Int -> a -> b -> b) -> b -> Stencil f a -> b
ifoldr f Int -> a -> b -> b
f b
b (Stencil forall b. (f Int -> a -> b -> b) -> b -> b
s) = (f Int -> a -> b -> b) -> b -> b
forall b. (f Int -> a -> b -> b) -> b -> b
s f Int -> a -> b -> b
f b
b
  {-# INLINE ifoldr #-}
  ifoldMap :: (f Int -> a -> m) -> Stencil f a -> m
ifoldMap = IndexedGetting (f Int) m (Stencil f a) a
-> (f Int -> a -> m) -> Stencil f a -> m
forall i m s a. IndexedGetting i m s a -> (i -> a -> m) -> s -> m
ifoldMapOf (((f Int -> a -> Const m a -> Const m a)
 -> Const m a -> Stencil f a -> Const m a)
-> IndexedGetting (f Int) m (Stencil f a) a
forall i (p :: * -> * -> *) (f :: * -> *) a s t b.
(Indexable i p, Contravariant f, Applicative f) =>
((i -> a -> f a -> f a) -> f a -> s -> f a) -> Over p f s t a b
ifoldring (f Int -> a -> Const m a -> Const m a)
-> Const m a -> Stencil f a -> Const m a
forall i (f :: * -> *) a b.
FoldableWithIndex i f =>
(i -> a -> b -> b) -> b -> f a -> b
ifoldr)
  {-# INLINE ifoldMap #-}

instance Functor (Stencil f) where
  fmap :: (a -> b) -> Stencil f a -> Stencil f b
fmap a -> b
f (Stencil forall b. (f Int -> a -> b -> b) -> b -> b
s) = (forall b. (f Int -> b -> b -> b) -> b -> b) -> Stencil f b
forall (f :: * -> *) a.
(forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
Stencil ((forall b. (f Int -> b -> b -> b) -> b -> b) -> Stencil f b)
-> (forall b. (f Int -> b -> b -> b) -> b -> b) -> Stencil f b
forall a b. (a -> b) -> a -> b
$ \f Int -> b -> b -> b
g b
z -> (f Int -> a -> b -> b) -> b -> b
forall b. (f Int -> a -> b -> b) -> b -> b
s (\f Int
x a
a b
b -> f Int -> b -> b -> b
g f Int
x (a -> b
f a
a) b
b) b
z
  {-# INLINE [0] fmap #-}

-- | Make a stencil folding over a list.
--
--   If the list is staticlly known this should expand at compile time
--   via rewrite rules, similar to 'Data.Dense.TH.makeStencilTH' but less reliable. If
--   that does not happen the resulting could be slow. If the list is
--   not know at compile time, 'mkStencilUnboxed' can be signifcantly
--   faster (but isn't subject expending via rewrite rules).
mkStencil :: [(f Int, a)] -> Stencil f a
mkStencil :: [(f Int, a)] -> Stencil f a
mkStencil [(f Int, a)]
l = (forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
forall (f :: * -> *) a.
(forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
Stencil ((forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a)
-> (forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
forall a b. (a -> b) -> a -> b
$ \f Int -> a -> b -> b
g b
z -> ((f Int, a) -> b -> b) -> b -> [(f Int, a)] -> b
forall a b. (a -> b -> b) -> b -> [a] -> b
myfoldr (\(f Int
i,a
a) b
b -> f Int -> a -> b -> b
g f Int
i a
a b
b) b
z [(f Int, a)]
l
{-# INLINE mkStencil #-}

-- Version of foldr that recursivly expands the list via rewrite rules.
myfoldr :: (a -> b -> b) -> b -> [a] -> b
myfoldr :: (a -> b -> b) -> b -> [a] -> b
myfoldr a -> b -> b
f b
b = [a] -> b
go where
  go :: [a] -> b
go []     = b
b
  go (a
a:[a]
as) = a -> b -> b
f a
a ([a] -> b
go [a]
as)
{-# INLINE [0] myfoldr #-}

{-# RULES
"mkStencil/cons" forall f b a as.
 myfoldr f b (a:as) = f a (myfoldr f b as)
 #-}

-- | Make a stencil folding over an unboxed vector from the list.
mkStencilUnboxed :: (U.Unbox (f Int), U.Unbox a) => [(f Int, a)] -> Stencil f a
mkStencilUnboxed :: [(f Int, a)] -> Stencil f a
mkStencilUnboxed [(f Int, a)]
l = (forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
forall (f :: * -> *) a.
(forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
Stencil ((forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a)
-> (forall b. (f Int -> a -> b -> b) -> b -> b) -> Stencil f a
forall a b. (a -> b) -> a -> b
$ \f Int -> a -> b -> b
g b
z -> ((f Int, a) -> b -> b) -> b -> Vector (f Int, a) -> b
forall a b. Unbox a => (a -> b -> b) -> b -> Vector a -> b
U.foldr (\(f Int
i,a
a) b
b -> f Int -> a -> b -> b
g f Int
i a
a b
b) b
z Vector (f Int, a)
v
  where !v :: Vector (f Int, a)
v = [(f Int, a)] -> Vector (f Int, a)
forall a. Unbox a => [a] -> Vector a
U.fromList [(f Int, a)]
l
{-# INLINE mkStencilUnboxed #-}

-- | Sum the elements around a 'Focused' using a 'Boundary' condition
--   and a 'Stencil'.
--
--   This is often used in conjunction with 'Data.Dense.extendFocus'.
stencilSum :: (Shape f, Num a) => Boundary -> Stencil f a -> Focused f a -> a
stencilSum :: Boundary -> Stencil f a -> Focused f a -> a
stencilSum Boundary
bnd Stencil f a
s = \Focused f a
w ->
  let f :: f Int -> a -> a -> a
f f Int
i a
b a
a = a
b a -> a -> a
forall a. Num a => a -> a -> a
+ a
a a -> a -> a
forall a. Num a => a -> a -> a
* Boundary -> f Int -> Focused f a -> a
forall (f :: * -> *) a.
Shape f =>
Boundary -> f Int -> Focused f a -> a
peekRelativeB Boundary
bnd f Int
i Focused f a
w
      {-# INLINE [0] f #-}
  in  (f Int -> a -> a -> a) -> a -> Stencil f a -> a
forall i (f :: * -> *) b a.
FoldableWithIndex i f =>
(i -> b -> a -> b) -> b -> f a -> b
ifoldl' f Int -> a -> a -> a
f a
0 Stencil f a
s
{-# INLINE stencilSum #-}