module Data.Array.Repa.Stencil.Template
(stencil2)
where
import Data.Array.Repa.Index
import Language.Haskell.TH
import Language.Haskell.TH.Quote
import qualified Data.List as List
stencil2 :: QuasiQuoter
stencil2 = QuasiQuoter
{ quoteExp = parseStencil2
, quotePat = undefined
, quoteType = undefined
, quoteDec = undefined }
parseStencil2 :: String -> Q Exp
parseStencil2 str
= let
line1 : _ = lines str
sizeX = fromIntegral $ length $ lines str
sizeY = fromIntegral $ length $ words line1
minX = negate (sizeX `div` 2)
minY = negate (sizeY `div` 2)
maxX = sizeX `div` 2
maxY = sizeY `div` 2
coeffs = (List.map read $ words str) :: [Integer]
in makeStencil2' sizeX sizeY
$ filter (\(_, _, v) -> v /= 0)
$ [ (fromIntegral y, fromIntegral x, fromIntegral v)
| y <- [minX, minX + (1 :: Integer) .. maxX]
, x <- [minY, minY + (1 :: Integer) .. maxY]
| v <- coeffs ]
makeStencil2'
:: Integer -> Integer
-> [(Integer, Integer, Integer)]
-> Q Exp
makeStencil2' sizeX sizeY coeffs
= do ix' <- newName "ix"
z' <- [p| Z |]
coeffs' <- newName "coeffs"
let fnCoeffs
= LamE [VarP ix']
$ CaseE (VarE (mkName "ix"))
$ [ Match (InfixP (InfixP z' (mkName ":.") (LitP (IntegerL oy)))
(mkName ":.") (LitP (IntegerL ox)))
(NormalB $ ConE (mkName "Just") `AppE` LitE (IntegerL v))
[] | (oy, ox, v) <- coeffs ]
++ [Match WildP
(NormalB $ ConE (mkName "Nothing")) []]
return
$ AppE (VarE (mkName "makeStencil2")
`AppE` (LitE (IntegerL sizeX))
`AppE` (LitE (IntegerL sizeY)))
$ LetE [ PragmaD (InlineP (mkName "coeffs") Inline FunLike (BeforePhase 0))
, ValD (VarP coeffs') (NormalB fnCoeffs) [] ]
(VarE (mkName "coeffs"))