{-# 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
stencil2 = QuasiQuoter :: (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
                { quoteExp :: String -> Q Exp
quoteExp      = String -> Q Exp
parseStencil2
                , quotePat :: String -> Q Pat
quotePat      = String -> Q Pat
forall a. HasCallStack => a
undefined
                , quoteType :: String -> Q Type
quoteType     = String -> Q Type
forall a. HasCallStack => a
undefined
                , quoteDec :: String -> Q [Dec]
quoteDec      = String -> Q [Dec]
forall a. HasCallStack => a
undefined }


-- | Parse a stencil definition.
--   TODO: make this more robust.
parseStencil2 :: String -> Q Exp
parseStencil2 :: String -> Q Exp
parseStencil2 String
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.
        String
line1 : [String]
_       = String -> [String]
lines String
str
        sizeX :: Integer
sizeX           = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([String] -> Int) -> [String] -> Int
forall a b. (a -> b) -> a -> b
$ String -> [String]
lines String
str
        sizeY :: Integer
sizeY           = Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([String] -> Int) -> [String] -> Int
forall a b. (a -> b) -> a -> b
$ String -> [String]
words String
line1

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

        -- List of coefficients for the stencil.
        coeffs :: [Integer]
coeffs          = ((String -> Integer) -> [String] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
List.map String -> Integer
forall a. Read a => String -> a
read ([String] -> [Integer]) -> [String] -> [Integer]
forall a b. (a -> b) -> a -> b
$ String -> [String]
words String
str) :: [Integer]

   in   Integer -> Integer -> [(Integer, Integer, Integer)] -> Q Exp
makeStencil2' Integer
sizeX Integer
sizeY
         ([(Integer, Integer, Integer)] -> Q Exp)
-> [(Integer, Integer, Integer)] -> Q Exp
forall a b. (a -> b) -> a -> b
$ ((Integer, Integer, Integer) -> Bool)
-> [(Integer, Integer, Integer)] -> [(Integer, Integer, Integer)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Integer
_, Integer
_, Integer
v) -> Integer
v Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Integer
0)
         ([(Integer, Integer, Integer)] -> [(Integer, Integer, Integer)])
-> [(Integer, Integer, Integer)] -> [(Integer, Integer, Integer)]
forall a b. (a -> b) -> a -> b
$ [ (Integer -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
y, Integer -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
x, Integer -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
v)
                | Integer
y     <- [Integer
minX, Integer
minX Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ (Integer
1 :: Integer) .. Integer
maxX]
                , Integer
x     <- [Integer
minY, Integer
minY Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ (Integer
1 :: Integer) .. Integer
maxY]
                | Integer
v     <- [Integer]
coeffs ]


makeStencil2'
        :: Integer -> Integer
        -> [(Integer, Integer, Integer)]
        -> Q Exp

makeStencil2' :: Integer -> Integer -> [(Integer, Integer, Integer)] -> Q Exp
makeStencil2' Integer
sizeX Integer
sizeY [(Integer, Integer, Integer)]
coeffs
 = do   Name
ix'             <- String -> Q Name
newName String
"ix"
        Pat
z'              <- [p| Z |]
        Name
coeffs'         <- String -> Q Name
newName String
"coeffs"

        let fnCoeffs :: Exp
fnCoeffs
                = [Pat] -> Exp -> Exp
LamE  [Name -> Pat
VarP Name
ix']
                (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE (String -> Name
mkName String
"ix"))
                ([Match] -> Exp) -> [Match] -> Exp
forall a b. (a -> b) -> a -> b
$   [ Pat -> Body -> [Dec] -> Match
Match     (Pat -> Name -> Pat -> Pat
InfixP (Pat -> Name -> Pat -> Pat
InfixP Pat
z' (String -> Name
mkName String
":.") (Lit -> Pat
LitP (Integer -> Lit
IntegerL Integer
oy)))
                                        (String -> Name
mkName String
":.") (Lit -> Pat
LitP (Integer -> Lit
IntegerL Integer
ox)))
                                (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE (String -> Name
mkName String
"Just") Exp -> Exp -> Exp
`AppE` Lit -> Exp
LitE (Integer -> Lit
IntegerL Integer
v))
                                [] | (Integer
oy, Integer
ox, Integer
v) <- [(Integer, Integer, Integer)]
coeffs ]
                    [Match] -> [Match] -> [Match]
forall a. [a] -> [a] -> [a]
++ [Pat -> Body -> [Dec] -> Match
Match Pat
WildP
                                (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE (String -> Name
mkName String
"Nothing")) []]

        Exp -> Q Exp
forall (m :: * -> *) a. Monad m => a -> m a
return
         (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE (String -> Name
mkName String
"makeStencil2") 
                        Exp -> Exp -> Exp
`AppE` (Lit -> Exp
LitE (Integer -> Lit
IntegerL Integer
sizeX)) 
                        Exp -> Exp -> Exp
`AppE` (Lit -> Exp
LitE (Integer -> Lit
IntegerL Integer
sizeY)))
         (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ [Dec] -> Exp -> Exp
LetE [ Pragma -> Dec
PragmaD (Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP (String -> Name
mkName String
"coeffs") Inline
Inline RuleMatch
FunLike (Int -> Phases
BeforePhase Int
0))
                , Pat -> Body -> [Dec] -> Dec
ValD    (Name -> Pat
VarP    Name
coeffs')          (Exp -> Body
NormalB Exp
fnCoeffs) [] ]
                (Name -> Exp
VarE (String -> Name
mkName String
"coeffs"))