Prototype of typed EDSL in Haskell for Matlab code generation.

gadt-dsl.hs

{-# LANGUAGE NoImplicitPrelude,OverloadedStrings,GADTs,MultiParamTypeClasses,FlexibleInstances #-}

-- GADTを使った型付きDSLのごく簡単なプロトタイプ。
-- 出力のロジックに、この記事のアイデア・実装を一部利用：　http://d.hatena.ne.jp/keigoi/20111206/haskell_tagless_dsl

import BasicPrelude as B
import Prelude (Show(..))

import Control.Monad.Writer
import Control.Monad.State

data CompilerState = CompilerState {
    varNames :: [(VarName,Int)],
    funcTable :: [FuncDef']
}

data FuncDef' = FuncDef' (Exp ()) | FuncForeign Foreign

emptyState :: CompilerState
emptyState = CompilerState [] []

type W a b = WriterT b (State CompilerState) a
type W' a = W a Text
type W'' = W' ()

data Size = D1 Int | D2 Int Int | D3 Int Int Int | D4 Int Int Int Int | UnknownSize deriving Show

data Mat r =
    Mat Size Text | MatC1 [r] | MatC2 [[r]] | MatC3 [[[r]]]
    -- | MatAdd (Mat r) (Mat r) | MatMul (Mat r) (Mat r)
    deriving Show --stub

data Mat' = MatI (Mat Int) | MatD (Mat Double)
data Statement = SInt (Exp Int) | SUnit (Exp ()) | SMat (Exp Mat')

data Op1' a r = Op1' (Exp a -> W' (Maybe VarName, RawCode))
data Op2' a b r = Op2' (Exp a -> Exp b -> W' (Maybe VarName, RawCode))

data Arg1 a = Arg1 (Exp a)
data Arg2 a b = Arg2 (Exp a) (Exp b)

data Index = Index1 Int | Index2 Int Int

-- As long as Exp is purely functional style, it should not be too hard to implement translation.
data Exp r where
    Op1 :: Op1' a r -> Exp a -> Exp r
    Op2 :: Op2' a b r -> Exp a -> Exp b -> Exp r
    Const :: (Show r,Num r) => r -> Exp r
    ConstMat :: (Show r,Num r) => (Mat r) -> Exp (Mat r)
    Func1 :: Text -> Args -> Arg1 a -> Exp r -- Function that does not return a value or needs assignment before use.
    Func1Exp :: Text -> Args -> Arg1 a -> Exp r
        -- Function that can be used as an expression.
    Raw :: Text -> Exp r
    Elem :: (Num r) => Index -> Exp (Mat r) -> Exp r
    SubMat :: (Num r) => Index -> Mat r -> Exp (Mat r)
    FuncDef :: Text -> Args -> Args -> Exp a -> Exp ()

data Foreign where
    Foreign :: Text -> Args -> Args -> Foreign

i :: Int -> Exp Int
i = Const

m :: (Show a,Num a) => Text -> Exp (Mat a)
m = ConstMat . Mat UnknownSize


class (Monad q) => NewName q where
    newName :: Text -> q Text
    runWM :: W' r -> q (VarName,CompilerState)
    assign :: Exp a -> q VarName

instance (Monoid w, NewName m) => NewName (WriterT w m) where
  newName = lift . newName

type VarName = Text
type RawCode = Text

instance NewName (State CompilerState) where
    newName str = do
        (CompilerState varmap fs) <- get
        let cnt = fromMaybe 0 $ lookup str varmap
        let newmap = (str,cnt+1):deleteBy (\(a,_)(b,_)->a==b) (str,0) varmap
        put $ CompilerState newmap fs
        return $ str ++ (if cnt == 0 then "" else B.show cnt)

class GenCode a where
    genCode :: a -> W' (Maybe VarName, RawCode)

defOp2' :: Text -> Exp a -> Exp b -> W' (Maybe VarName, RawCode)
defOp2' t a b = do
    (mva,codea) <- genCode a
    (mvb,codeb) <- genCode b
    let code = B.concat ["(",codea," ",t," ",codeb,")"]
    return (Nothing,code)

defFn1 :: Text -> Exp a -> W' (Maybe VarName, RawCode)
defFn1 t a = do
    (mva,codea) <- genCode a
    let code = B.concat [t,"(",codea,")"]
    return (Nothing,code)

instance (Show a,Num a) => Num (Mat a)

instance (Show a,Num a) => Num (Exp a) where
    a + b = Op2 (Op2' (defOp2' "+")) a b
    a - b = Op2 (Op2' (defOp2' "-")) a b
    a * b = Op2 (Op2' (defOp2' "*")) a b
    abs a = Op1 (Op1' (defFn1 "abs")) a
    signum a = Op1 (Op1' (defFn1 "sign")) a
    fromInteger n = Const (fromIntegral n)


raw :: Text -> Exp r
raw t = Raw t

genTopLevel exp =
    let (codeBefore,(_,code))= runW $ do
        (_,code) <- genCode exp
        return (Nothing,code)
    in
    codeBefore ++ code

runW :: W' r -> (Text,r)
runW m = (t,r) where ((r,t),_) = runState (runWriterT m) emptyState

data Args = Args [Text]
    deriving (Show,Eq)

instance GenCode Args where
    genCode as@(Args ns) = do
        return (Nothing,intercalate "," ns)

instance GenCode (Exp a) where
    genCode (ConstMat (MatC2 xss)) = do
        var <- newName "mat"
        tell $ var ++ " = [" ++ intercalate ";" (flip map xss $ \xs ->
                    intercalate " " (map B.show xs)) ++ "];\n"
        return (Just var,var)

    genCode (Const a) = return (Nothing,B.show a)

    genCode (Func1 name rs@(Args ts) (Arg1 a)) = do
        (mv,code) <- genCode a
        tell $ B.concat [mkRets rs,name,"(",code,");\n"]
        mapM_ newName ts
        return (Nothing,"") -- is this empty text okay?

    genCode (Func1Exp name rs (Arg1 a)) = do
        (mv,code) <- genCode a
        let code2 = B.concat [name,"(",code,")"]
        return (Nothing,code2)
 
    genCode (FuncDef name rs as exp) = do
        (mv,code) <- genCode exp
        v <- newName "res"
        let code2 = B.concat [mkRets (Args [v]),"function ",name,"(",mkArgs as,")\n",v," = ",code,"\nend\n"]
        return (Nothing,code2)


    genCode (Raw t) = return (Nothing,t)

    genCode (Elem ind exp) = do
        (mv,code) <- genCode exp
        let str = B.concat [code,"(",B.show ind,")"]
        return (Nothing,str)

    genCode (Op1 (Op1' w) a ) = w a

    genCode (Op2 (Op2' w) a1 a2) = w a1 a2
        
    genCode a = return (Nothing,"")

instance Show Index where
    show (Index1 i) = Prelude.show i
    show (Index2 i j) = Prelude.show i ++ "," ++ Prelude.show j

mkArgs :: Args -> Text
mkArgs (Args ts) = intercalate "," ts

mkRets :: Args -> Text
mkRets (Args ts) =
    case length ts of
        0 -> ""
        1 -> head ts ++ " = "
        2 -> intercalate "," ts ++ " = "

---------
main = do
    putStrLn . genTopLevel $ test1
    putStrLn ""
    putStrLn . genTopLevel $ test2
    putStrLn ""
    putStrLn . genTopLevel $ test3a
    putStrLn ""
    putStrLn . genTopLevel $ test3b
    putStrLn ""
    putStrLn . genTopLevel $ test4

test1 :: Exp Int
test1 = (signum (abs (-3) + 2)) * 4

test2 :: Exp Double
test2 = Elem (Index2 2 1) (ConstMat $ MatC2 [[1,2,3],[4,5,6]])

test3 :: (Show a,Real a) => Exp a
test3 = (Elem (Index2 1 2) (ConstMat $ MatC2 [[1,2,3],[4,5,6]]))
			+ (Elem (Index2 2 3) (ConstMat $ MatC2 [[10,20,15],[2,5,2]]))

test3a :: Exp Int
test3a = test3

test3b :: Exp Double
test3b = test3

test4 :: Exp ()
test4 = 
    FuncDef
        "testFunc"
        (Args ["res"])
        (Args ["Im"])
        (Func1Exp "imread" (Args []) (Arg1 (raw "'test.tiff'")))


-- 出力結果
{-
(sign((abs((0 - 3)) + 2)) * 4)

mat = [1.0 2.0 3.0;4.0 5.0 6.0];
mat(2,1)

mat = [1 2 3;4 5 6];
mat1 = [10 20 15;2 5 2];
(mat(1,2) + mat1(2,3))

mat = [1.0 2.0 3.0;4.0 5.0 6.0];
mat1 = [10.0 20.0 15.0;2.0 5.0 2.0];
(mat(1,2) + mat1(2,3))

res = function testFunc(Im)
res = imread('test.tiff')
end

-}