module Arith
where

--
-- Environment
--

data Env a = Env [(String,a)] deriving Show

lookupEnv :: Env a -> String -> Maybe a
lookupEnv (Env env) name = lookup name env

emptyEnv :: Env a
emptyEnv = Env []

addToEnv :: Env a -> String -> a -> Env a
addToEnv (Env env) name val = Env ((name,val):env)

addToEnvN :: Env a -> [(String,a)] -> Env a
addToEnvN (Env env) newdefs = Env (newdefs++env)


--
-- Expression
--

data Exp a = Const a
           | Var String
           | Binary BinOp (Exp a) (Exp a)
           | Unary  UnOp  (Exp a)
           | Let [(String,Exp a)] (Exp a)
           | Lam String (Exp a)
           | App (Exp a) (Exp a)
           deriving Show

data BinOp = Add | Sub | Mul | Div deriving Show

data UnOp = Neg | Recip deriving Show


--
-- Evaluation result
--

data Result a err = Result a
                  | Fail err
                  deriving Show

mapResult :: (a -> Result b err) -> Result a err -> Result b err
mapResult f (Result x) = f x
mapResult f (Fail   e) = Fail e

zipResult :: (a -> b -> Result c err)
          -> Result a err -> Result b err -> Result c err
zipResult f (Result x) (Result y) = f x y
zipResult _ (Fail   e) _          = Fail e
zipResult _ _          (Fail e)   = Fail e

liftResult :: [Result a err] -> Result [a] err
liftResult [] = Result []
liftResult (Result x : rest) =
    case liftResult rest of
        Result xs -> Result (x:xs)
        Fail e    -> Fail e
liftResult (Fail e : _) = Fail e


--
-- Evaluation error
--

data EvalError = DivByZero
               | Overflow
               | TypeError
               | UndefVarError
               | ApplicationError
               deriving Show

--
-- Repesentation of typed values
--

data TypedVal = TVString String
              | TVDouble Double
              | TVFun    (TypedVal -> Result TypedVal EvalError)

instance Show TypedVal where
   show (TVString s) = show s
   show (TVDouble d) = show d
   show (TVFun _)    = "<<function>>"

--
-- Evaluation for (dynamically) typed values
--

eval :: Env TypedVal
     -> Exp TypedVal
     -> Result TypedVal EvalError
eval _ (Const x)  = Result x
eval env (Var name) =
    case lookupEnv env name of
        Just x  -> Result x
        Nothing -> Fail UndefVarError
eval env (Binary op exp1 exp2) =
    zipResult (tvBinOp op) (eval env exp1) (eval env exp2)
eval env (Unary op exp) = mapResult (tvUnOp op) (eval env exp)
eval env (Let defs exp) =
    case liftResult (map (eval env) exps) of
        Result vals -> eval newEnv exp
            where
            newEnv = addToEnvN env (zip names vals)
        Fail e      -> Fail e
    where
    (names,exps) = unzip defs
eval env (Lam x exp) =
    Result $ TVFun (\val -> eval (addToEnv env x val) exp)
eval env (App exp1 exp2) = 
    case (eval env exp1, eval env exp2) of
        (Result (TVFun f), Result arg) -> f arg
        (Result (TVFun _), Fail   err) -> Fail err
        (Result _, _) -> Fail ApplicationError
        (Fail err, _) -> Fail err


tvBinOp :: BinOp -> TypedVal -> TypedVal -> Result TypedVal EvalError
tvBinOp Add (TVString s) (TVString t) = Result (TVString (s++t))
tvBinOp Add (TVString s) (TVDouble y) = Result (TVString (s++show y))
tvBinOp Add (TVDouble x) (TVString t) = Result (TVString (show x++t))
tvBinOp Add (TVDouble x) (TVDouble y) = Result (TVDouble (x+y))
tvBinOp Sub (TVDouble x) (TVDouble y) = Result (TVDouble (x-y))
tvBinOp Mul (TVDouble x) (TVDouble y) = Result (TVDouble (x*y))
tvBinOp Div (TVDouble x) (TVDouble 0) = Fail DivByZero
tvBinOp Div (TVDouble x) (TVDouble y) = Result (TVDouble (x/y))
tvBinOp _ _ _ = Fail TypeError

tvUnOp :: UnOp -> TypedVal -> Result TypedVal EvalError
tvUnOp Neg   (TVDouble x) = Result (TVDouble (negate x))
tvUnOp Recip (TVDouble 0) = Fail DivByZero
tvUnOp Recip (TVDouble x) = Result (TVDouble (recip x))
tvUnOp _ _ = Fail TypeError


--
-- Main evaluation function
--

evalExp :: Exp TypedVal -> Result TypedVal EvalError
evalExp exp = eval emptyEnv exp