Reification of AST using a RSW monad.

reify2-rws.hs

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

import Data.List (intercalate,foldl1)
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.Map (Map)
import qualified Data.Map as M
import Control.Monad.RWS
import Control.Monad.State.Class
import Data.Default (Default(..))

type Ident = Text

-- 2D matrix
data Mat2D r = Mat Ident | MatConst [[r]]

class Scalar a
instance Scalar Int
instance Scalar Integer
instance Scalar Double

type W a = RWST () Text CompilerState IO a

data CompilerState = CompilerState {
	variableCount :: Int
}

instance Default CompilerState where
	def = CompilerState 0

data Exp r where
	Const :: (Reifiable r) => r -> Exp r
	Add :: Exp r -> Exp r -> Exp r
	Sub :: Exp r -> Exp r -> Exp r
	Mul :: Exp r -> Exp r -> Exp r
	Abs :: (Scalar r) => Exp r -> Exp r
	Signum :: (Scalar r) => Exp r -> Exp r
	Index2D :: (Num r) => Exp Int -> Exp Int -> Exp (Mat2D r) -> Exp r

instance (Scalar r) => Num (Exp (Mat2D r)) where
	a + b = Add a b
	a - b = Sub a b
	a * b = Mul a b
	-- Caution! abs, signum, fromInteger are not defined.

instance Num (Exp Int) where
	a + b = Add a b
	a - b = Sub a b
	a * b = Mul a b
	fromInteger a = Const (fromIntegral a)
	-- Caution! abs and signum are not defined.

instance Num (Exp Double) where
	a + b = Add a b
	a - b = Sub a b
	a * b = Mul a b
	fromInteger a = Const (fromIntegral a)
	-- Caution! abs and signum are not defined.

instance Num (Exp Integer) where
	a + b = Add a b
	a - b = Sub a b
	a * b = Mul a b
	fromInteger a = Const (fromIntegral a)
	-- Caution! abs and signum are not defined.

class Reifiable a where
	reify :: a -> W Text

instance Reifiable (Exp r) where
	reify (Const exp) = reify exp
	reify (Add e1 e2) = do
		c1 <- reify e1
		c2 <- reify e2
		return $ T.concat [c1,"+",c2]
	reify (Index2D i j mat) = do
		ci <- reify i
		cj <- reify j
		m <- reify mat
		return $ T.concat [m,"(",ci,",",cj,")"]

instance (Show r) => Reifiable (Mat2D r) where
	reify (Mat ident) = return ident
	reify (MatConst xss) = do
		let code =  T.pack $ "[" ++ intercalate ";" (map (intercalate " " . map show) xss) ++ "]"
		var <- newName
		tell $ T.concat [var," = ",code,";\n"]
		return var

instance Reifiable Int where
	reify n = return $ T.pack $ show n

instance Reifiable Integer where
	reify n = return $ T.pack $ show n

instance Reifiable Double where
	reify n = return $ T.pack $ show n

newName :: W Text
newName = do
	CompilerState count <- get
	put $ CompilerState (count + 1)
	return $ T.pack $ "v" ++ show (count+1)

main = do
	putStrLn "\nm1"
	gen m1
	putStrLn "\ntest2"
	gen test2
	putStrLn "\ntest3"
	gen test3
	putStrLn "\ntest4"
	gen test4
	putStrLn "\ntest5"
	gen test5

gen :: Exp r -> IO ()
gen e = do
	(txt,_,assigned) <- runRWST (reify e) () def
	T.putStrLn $ T.append assigned txt

m1 :: Exp (Mat2D Int)
m1 = Const (MatConst [[1,2,3],[4,5,6]])

m2 :: Exp (Mat2D Int)
m2 = Const (MatConst [[10,5,3],[1,4,2]])

test4 :: Exp (Mat2D Int)
test4 = m1 + m1

test5 :: Exp (Mat2D Int)
test5 = foldl1 (+) (replicate 5 m1)

test2 :: Exp Int
test2 = Index2D 1 2 m1

test3 :: Exp Int
test3 = Index2D (Index2D 1 2 m1) 1 m2

{- Output of main
*Main> main

m1
v1 = [1 2 3;4 5 6];
v1

test2
v1 = [1 2 3;4 5 6];
v1(1,2)

test3
v1 = [1 2 3;4 5 6];
v2 = [10 5 3;1 4 2];
v2(v1(1,2),1)

test4
v1 = [1 2 3;4 5 6];
v2 = [1 2 3;4 5 6];
v1+v2

test5
v1 = [1 2 3;4 5 6];
v2 = [1 2 3;4 5 6];
v3 = [1 2 3;4 5 6];
v4 = [1 2 3;4 5 6];
v5 = [1 2 3;4 5 6];
v1+v2+v3+v4+v5
-}