implementation module int

import StdEnv

import Data.Either
import Data.Functor
import Data.Maybe
import Data.List
import Control.Applicative
import Control.Monad
import Control.Monad.State
import Control.Monad.Trans

import ast

int :: AST -> Either [String] Value
int (AST fs) = evalStateT (mapM evalFun fs >>| getStart fs >>= eval >>= printer) preamble

err :: String -> Eval a
err e = liftT (Left [e])

getStart :: [Function] -> Eval Expression
getStart [] = err "No start rule defined"
getStart [(Function ['start'] _ e):_] = pure e
getStart [_:fs] = getStart fs

:: Eval a :== StateT EvalState (Either [String]) a
:: EvalState :== [([Char], Value)]
preamble =
	[(['if'], Func 3 [] (Builtin ['if']))
	,(['eq'], Func 2 [] (Builtin ['eq']))
	,(['mul'], Func 2 [] (Builtin ['mul']))
	,(['div'], Func 2 [] (Builtin ['div']))
	,(['add'], Func 2 [] (Builtin ['add']))
	,(['sub'], Func 2 [] (Builtin ['sub']))
	]

putEnv :: [Char] Value -> Eval ()
putEnv i v = modify (\vs->[(i,v):vs])

getEnv :: [Char] -> Eval Value
getEnv v = gets (lookup v)
	>>= maybe (err ("Variable " +++ toString v +++ " not found")) pure

evalFun :: Function -> Eval ()
evalFun (Function v a b) = putEnv v (Func (length a) [] (\es->fun a es b))
where
	fun [] [] body = body
	fun [a:as] [e:es] body = fun as es (sub a e body)

printer :: Value -> Eval Value
printer t=:(Func 0 args body) = eval (body args) >>= printer
printer a = pure a

sub :: [Char] Expression Expression -> Expression
sub ident subst (Var v)
	| ident == v = subst
sub ident subst (App e1 e2) = App (sub ident subst e1) (sub ident subst e2)
sub ident subst (Lambda v b)
	| ident <> v = Lambda v (sub ident b subst)
sub _ _ x = x

eval :: Expression -> Eval Value
eval (Lit v) = pure v
eval (Var v) = getEnv v
eval (App e1 e2) = eval e1 >>= \v->case v of
	(Func 0 a b) = err "Saturated function"
	(Func n as b) = pure (Func (n-1) (as ++ [e2]) b)
	_ = err ("Cannot apply " +++ toString e2 +++ " to " +++ toString e1)
eval (Lambda a b) = pure (Func 1 [] (\[arg]->sub a arg b))
eval (Builtin i as) = case (i, as) of
	(['if'], [p,t,e]) = eval p >>= printer >>= \v->case v of
		Bool v = eval (if v t e)
		_ = err ("first argument of if must be a bool but is " +++ toString v)
	(['add'], [l,r]) = eval l >>= printer >>= \l->eval r >>= printer >>= \r->case (l,r) of
		(Int a, Int b) = pure (Int (a + b))
		_ = err "add only defined for integers"
	(['sub'], [l,r]) = eval l >>= printer >>= \l->eval r >>= printer >>= \r->case (l,r) of
		(Int a, Int b) = pure (Int (a - b))
		_ = err "sub only defined for integers"
	(['mul'], [l,r]) = eval l >>= printer >>= \l->eval r >>= printer >>= \r->case (l,r) of
		(Int a, Int b) = pure (Int (a * b))
		_ = err "mul only defined for integers"
	(['div'], [l,r]) = eval l >>= \l->eval r >>= \r->case (l,r) of
		(Int a, Int b) = pure (Int (a / b))
		_ = err "div only defined for integers"
	(['eq'], [l,r]) = eval l >>= printer >>= \l->eval r >>= printer >>= \r->case (l,r) of
		(Int a, Int b) = pure (Bool (a == b))
		_ = err "eq only defined for integers"