1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
|
{-# 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 :: 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"))
|
