module sets

// Laurens Kuijper, s4467299
// Camil Staps, s4498062

import StdBool
import StdEnum
from StdFunc import const, id, o
import StdInt
import StdList
import StdOverloaded
import StdString

import Data.Error
from Data.Functor import class Functor(fmap)
import Data.List
import qualified Data.Map as M
import Data.Maybe
import Data.Tuple
from Text import class Text(concat), instance Text String

:: Ident :== String

:: Expr = Expr
:: Upd  = Upd
:: Stmt = Stmt

:: In a b = (In) infix 0 a b
:: Do a b = (Do) infix 0 a b

class isExpr a :: a -> Bool
instance isExpr Expr where isExpr _ = True
instance isExpr Upd  where isExpr _ = True

class Type t | toString, TC t
where
	type :: t -> String

instance Type ()     where type _ = "()"
instance Type Int    where type _ = "Int"
instance Type Char   where type _ = "Char"
instance Type Bool   where type _ = "Bool"

class DSL v | AExpr, BExpr, Compare, SetExpr, Var, Stmt v

class AExpr v
where
	lit :: a -> v a Expr | Type a
	(+.) infixl 6 :: (v t p) (v t q) -> v t Expr | Type, + t
	(-.) infixl 6 :: (v t p) (v t q) -> v t Expr | Type, - t
	(*.) infixl 7 :: (v t p) (v t q) -> v t Expr | Type, * t
	(!~) :: (v t p) -> v t Expr | Type, ~ t

class BExpr v
where
	(&&.) infixr 3 :: (v Bool p) (v Bool q) -> v Bool Expr
	(||.) infixr 2 :: (v Bool p) (v Bool q) -> v Bool Expr

class Compare v
where
	(==.) infix 4 :: (v t p) (v t q) -> v Bool Expr | Type, == t
	(<.)  infix 4 :: (v t p) (v t q) -> v Bool Expr | Type, <  t

class SetExpr v
where
	set :: [a] -> v [a] Expr | Type a
	(:+) infix 5 :: (v t p) (v [t] q) -> v [t] Expr | Type, Eq t
	(+:) infix 5 :: (v [t] p) (v t q) -> v [t] Expr | Type, Eq t
	(-:) infix 5 :: (v [t] p) (v t q) -> v [t] Expr | Type, Eq t
	(:+:) infixl 6 :: (v [t] p) (v [t] q) -> v [t] Expr | Type, Eq t
	(:-:) infixl 6 :: (v [t] p) (v [t] q) -> v [t] Expr | Type, Eq t
	(:*:) infixl 7 :: (v [t] p) (v [t] q) -> v [t] Expr | Type, Eq t

class Var v
where
	var :: ((v t Upd) -> In t (v a p)) -> v a p | Type t
	(=.) infixr 2 :: (v t Upd) (v t p) -> v t p | Type t & isExpr p

class Stmt v
where
	If :: (v Bool p) (v t q) (v u r) -> v () Stmt
	For :: ((v t Upd) -> Do (v [t] Expr) (v u a)) -> v () Stmt | Type t
	(:.) infixl 1 :: (v t p) (v u q) -> v u Stmt

:: Show t s = Show (ShowState -> ShowState)
:: ShowState =
	{ output :: [String]
	, fresh  :: Int
	}

print :: (Show s t) -> String
print (Show s) = concat (s {output=[], fresh=1}).output

(<<) infixr 9 //:: (Show a p) (Show b q) -> Show c r
(<<) (Show f) (Show g) :== Show (f o g)
(-<) infixr 9; (-<) s f :== show s << f

show :: a -> Show b c | toString a
show x = Show \st -> {st & output=[toString x:st.output]}

between :: a b (Show t u) -> Show p q | toString a & toString b
between x y s = x -< s << show y

par :: (Show a p) -> Show b q
par s = between "(" ")" s

brace :: (Show a p) -> Show b q
brace s = between "{ " " }" s

showsperse :: a [Show b p] -> Show c q | toString a
showsperse _ []       = Show id
showsperse _ [Show x] = Show x
showsperse g [x:xs]   = x << g -< showsperse g xs

instance AExpr Show
where
	lit x = show x
	(+.) x y = par (x << " +. " -< y)
	(-.) x y = par (x << " -. " -< y)
	(*.) x y = par (x << " *. " -< y)
	(!~) x = "!~" -< par x

instance BExpr Show
where
	(&&.) x y = par (x << " &&. " -< y)
	(||.) x y = par (x << " ||. " -< y)

instance Compare Show
where
	(==.) x y = par (x << " ==. " -< y)
	(<.)  x y = par (x << " <. "  -< y)

instance SetExpr Show
where
	set xs = between "[" "]" (showsperse "," (map show xs))
	(:+)  xs x  = par (xs << " :+ "  -< x)
	(+:)  x  xs = par (x  << " +: "  -< xs)
	(-:)  xs x  = par (xs << " -: "  -< x)
	(:+:) xs ys = par (xs << " :+: " -< ys)
	(:-:) xs ys = par (xs << " :-: " -< ys)
	(:*:) xs ys = par (xs << " :*: " -< ys)

instance Var Show
where
	var x = Show \st -> let v = show ("v" +++ toString st.fresh) in case x v of
		In x f -> case "var " -< v << " = " -< x -< " :.\n" -< f of
			Show f -> f {st & fresh=st.fresh+1}
	(=.) v x = v << " =. " -< x

instance Stmt Show
where
	If b t e = "If " -< par b << " " -< brace t << " " -< brace e
	For x = Show \st -> let v = show ("v" +++ toString st.fresh) in case x v of
		Do s f -> case "For " -< v << " in " -< s << " do " -< brace f of
			Show f -> f {st & fresh=st.fresh+1}
	(:.) s t = s << " :.\n" -< t

:: Eval t s = Eval ((RW t) EvalState -> MaybeError String (t, EvalState))
:: RW t = R | W t
:: EvalState =
	{ vars :: 'M'.Map Int Dynamic
	, next :: Int
	}

eval :: (Eval t s) -> MaybeError String (t, EvalState)
eval (Eval x) = x R {vars='M'.newMap, next=1}

(<$>) infixl 4 // :: (a -> b) (v t a) -> v u b
(<$>) f (Eval x) = Eval \_ st -> fmap (appFst f) (x R st)

pure :: a -> Eval a t
pure x = Eval \_ st -> Ok (x,st)

liftA2 :: (a b -> c) (Eval a t) (Eval b u) -> Eval c v
liftA2 f (Eval x) (Eval y) = Eval \_ st -> case x R st of
	Ok (x0,st) -> case y R st of
		Ok (y,st) -> Ok (f x0 y, st)
		Error e -> Error e
	Error e -> Error e

instance AExpr Eval
where
	lit x = pure x
	(+.) x y = liftA2 (+) x y
	(-.) x y = liftA2 (-) x y
	(*.) x y = liftA2 (*) x y
	(!~) x = (~) <$> x

instance BExpr Eval
where
	(&&.) x y = liftA2 (&&) x y
	(||.) x y = liftA2 (||) x y

instance Compare Eval
where
	(==.) x y = liftA2 (==) x y
	(<.)  x y = liftA2 (<)  x y

addElem :: a [a] -> [a] | Eq a; addElem x xs = if (isMember x xs) xs [x:xs]
delElem :== removeMember

instance SetExpr Eval
where
	set xs = pure xs
	(:+)  x  xs = liftA2 addElem x xs
	(+:)  xs x  = liftA2 addElem x xs
	(-:)  xs x  = liftA2 delElem x xs
	(:+:) xs ys = liftA2 union xs ys
	(:*:) xs ys = liftA2 intersect xs ys
	(:-:) xs ys = liftA2 difference xs ys

rwvar :: Int (RW a) (EvalState) -> MaybeError String (a, EvalState) | TC a
rwvar i R st = case 'M'.get i st.vars of
	Just (v :: a^) -> Ok (v,st)
	Just _         -> Error "Wrong type"
	Nothing        -> Error "Undefined variable"
rwvar i (W x) st = Ok (x, {st & vars='M'.put i (dynamic x) st.vars})

instance Var Eval
where
	(=.) (Eval v) (Eval e) = Eval \_ st -> case e R st of
		Ok (e,st) -> v (W e) st
		Error e -> Error e
	var f = Eval \_ st -> let (x In (Eval r)) = f (Eval (rwvar st.next)) in
		r R {st & next=st.next+1, vars='M'.put st.next (dynamic x) st.vars}

instance Stmt Eval
where
	If (Eval b) (Eval t) (Eval e) = Eval \_ st -> case b R st of
		Ok (True, st) -> fmap (appFst (const ())) (t R st)
		Ok (False,st) -> fmap (appFst (const ())) (e R st)
		Error e       -> Error e
	For f = Eval \_ st -> let (xs Do b) = f (Eval (rwvar st.next)) in
		let (Eval e) = for st.next xs b in e R {st & next=st.next+1}
	where
		for :: Int (Eval [a] Expr) (Eval b t) -> Eval () Stmt | TC a
		for i (Eval s) (Eval b) = Eval \_ st -> case s R st of
			Ok (xs,st) -> loop xs st
			Error e    -> Error e
		where
			loop :: [a] EvalState -> MaybeError String ((), EvalState) | TC a
			loop []     st = Ok ((), st)
			loop [x:xs] st = case b R {st & vars='M'.put i (dynamic x) st.vars} of
				Ok (_,st) -> loop xs st
				Error e   -> Error e
	(:.) (Eval s) (Eval t) = Eval \_ st -> case s R st of
		Ok (s,st) -> t R st
		Error e   -> Error e

Start = (print stmt, '\n', eval stmt)
where
	stmt :: v Int Stmt | DSL v
	stmt =
		var \y = 1 In (
			For \x . set [1..10] Do
				y =. y *. x
		):.
		y