Template.hs

{-# LANGUAGE TemplateHaskell, QuasiQuotes, ParallelListComp #-}

-- | Template
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

-- | QuasiQuoter for producing a static stencil defintion.
--
--   A definition like
--
--   @
--     [stencil2|  0 1 0
--                 1 0 1
--                 0 1 0 |]
--   @
--
--   Is converted to:
--
--   @
--     makeStencil2 (Z:.3:.3)
--        (\\ix -> case ix of
--                  Z :. -1 :.  0  -> Just 1
--                  Z :.  0 :. -1  -> Just 1
--                  Z :.  0 :.  1  -> Just 1
--                  Z :.  1 :.  0  -> Just 1
--                  _              -> Nothing)
--   @
--
stencil2 :: QuasiQuoter
stencil2 = QuasiQuoter
		{ quoteExp	= parseStencil2
		, quotePat	= undefined
		, quoteType	= undefined
		, quoteDec	= undefined }


-- | Parse a stencil definition.
--   TODO: make this more robust.
parseStencil2 :: String -> Q Exp
parseStencil2 str
 = let
	-- Determine the extent of the stencil based on the layout.
	-- TODO: make this more robust. In particular, handle blank
	--       lines at the start of the definition.
	line1 : _	= lines str
	sizeX		= fromIntegral $ length $ lines str
	sizeY		= fromIntegral $ length $ words line1

	-- TODO: this probably doesn't work for stencils who's extents are even.
	minX		= negate (sizeX `div` 2)
	minY		= negate (sizeY `div` 2)
	maxX		= sizeX `div` 2
	maxY		= sizeY `div` 2

	-- List of coefficients for the stencil.
	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 .. maxX]
		, x	<- [minY, minY + 1 .. 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") (InlineSpec True False Nothing))
		, ValD 	  (VarP    coeffs')          (NormalB fnCoeffs) [] ]
		(VarE (mkName "coeffs"))