{-# OPTIONS -fglasgow-exts #-}

module LISS.Checker where

import LISS.Syntax

import Control.Monad.State
import Control.Monad.Error
import Control.Monad

import Data.Map hiding (map)

import Data.Generics

import Prelude hiding (lookup)

type Env = Map String Type

type ErrorS = Either String

myError :: (Show a, Show b, Monad m) => a -> b -> m c
myError a b = fail $ "Type error: cannot unify '" ++ show a ++ "' with '" ++ show b ++ "'."

gatherInfo :: Liss -> ErrorS Env
gatherInfo (Liss _ (Body decs _)) = mapM readDecl decs >>= return . foldr union empty

readDecl :: Declaration -> ErrorS Env
readDecl (VarDec (Arg var ty)) = return $ singleton var ty
readDecl (SubprogDec specs pn args ty body rs) = 
  let equality = case (getTypeRS rs) of
		    (Array _) -> eqArrayType
		    _ -> (==)
  in do
    readArgs <- mapM readDecl (map VarDec args) >>= return . foldr union empty
{-
    if (equality (getTypeRS rs) ty) 
      then gatherInfo (Liss "" body) >>= return . (readArgs `union` singleton pn ty `union`) 
        else myError rs ty
-}
    gatherInfo (Liss "" body) >>= return . (readArgs `union` singleton pn ty `union`) 

eqArrayType (Array _) (Array _) = True
eqArrayType _ _ = False

getTypeA (IntAssign _ _) = Integer
getTypeA (SeqAssign _ _) = Sequence
getTypeA (SetAssign _ _) = Set
getTypeA (BoolAssign _ _) = Boolean
getTypeA (ArrayPosAssign _ _ _) = Array 0

getTypeRS (ReturnInt _) = Integer
getTypeRS (ReturnBool _) = Boolean
getTypeRS (ReturnSet _) = Set
getTypeRS (ReturnSeq _) = Sequence
getTypeRS (ReturnArray _) = Array 0

check :: Liss -> ErrorS Liss
check p = do 
	    env <- gatherInfo p 
	    p' <- everywhereM (mkM (rewriteD env)) p
	    everywhereM (mkM (rewriteS env)) p'
  where
    rewriteS :: Env -> Statement -> ErrorS Statement
    rewriteS env (IntAssign iv (IntFunction fn vars)) = f env iv fn vars
    rewriteS env (IntAssign v1 (IntVarExp v2)) = h env v1 v2
    rewriteS env e@(IntAssign iv ie) = g env e iv ie
    rewriteS env (SetAssign sv (SetFunction fn vars)) = f env sv fn vars
    rewriteS env (SetAssign v1 (SetVarExp v2)) = h env v1 v2
    rewriteS env e@(SetAssign sv se) = g env e sv se
    rewriteS env (SeqAssign sv (SeqFunction fn vars)) = f env sv fn vars
    rewriteS env (SeqAssign v1 (SeqVarExp v2)) = h env v1 v2
    rewriteS env e@(SeqAssign sv se) = g env e sv se
    rewriteS env (BoolAssign bv (BoolFunction fn vars)) = f env bv fn vars
    rewriteS env (BoolAssign v1 (BoolVarExp v2)) = h env v1 v2
    rewriteS env e@(BoolAssign bv be) = g env e bv be
--    rewriteS env (ArrayPosAssign av pos (IntFunction fn vars)) = undefined -- Array bounds checked elsewhere
    rewriteS env e@(ArrayPosAssign av pos ie) =
      case (lookup av env) of
	Just x -> if (eqArrayType x (getTypeA e)) then (return e) else (myError av ie)
  	Nothing -> myError av ie
    rewriteS _ x = return x

    -- Correct return statements
    rewriteD :: Env -> Declaration -> ErrorS Declaration
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnInt (IntVarExp v))) = 
      let newRs = case ty of
                    Integer  -> ReturnInt (IntVarExp v)
                    Sequence -> ReturnSeq (SeqVarExp v)
                    Set      -> ReturnSet (SetVarExp v)
                    Boolean  -> ReturnBool (BoolVarExp v)
                    Array n  -> ReturnArray (ArrayVarExp v)
      in return $ SubprogDec specs pn args ty body newRs
    rewriteD env e@(SubprogDec specs pn args Integer body (ReturnInt _)) = return e
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnInt x)) = myError ty x

    rewriteD env e@(SubprogDec specs pn args ty body (ReturnBool (BoolVarExp v))) = 
      let newRs = case ty of
                    Integer  -> ReturnInt (IntVarExp v)
                    Sequence -> ReturnSeq (SeqVarExp v)
                    Set      -> ReturnSet (SetVarExp v)
                    Boolean  -> ReturnBool (BoolVarExp v)
                    Array n  -> ReturnArray (ArrayVarExp v)
      in return $ SubprogDec specs pn args ty body newRs
    rewriteD env e@(SubprogDec specs pn args Boolean body (ReturnBool _)) = return e
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnBool x)) = myError ty x
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnSeq (SeqVarExp v))) = 
      let newRs = case ty of
                    Integer  -> ReturnInt (IntVarExp v)
                    Sequence -> ReturnSeq (SeqVarExp v)
                    Set      -> ReturnSet (SetVarExp v)
                    Boolean  -> ReturnBool (BoolVarExp v)
                    Array n  -> ReturnArray (ArrayVarExp v)
      in return $ SubprogDec specs pn args ty body newRs
    rewriteD env e@(SubprogDec specs pn args Sequence body (ReturnSeq _)) = return e
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnSeq x)) = myError ty x
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnSet (SetVarExp v))) = 
      let newRs = case ty of
                    Integer  -> ReturnInt (IntVarExp v)
                    Sequence -> ReturnSeq (SeqVarExp v)
                    Set      -> ReturnSet (SetVarExp v)
                    Boolean  -> ReturnBool (BoolVarExp v)
                    Array n  -> ReturnArray (ArrayVarExp v)
      in return $ SubprogDec specs pn args ty body newRs
    rewriteD env e@(SubprogDec specs pn args Set body (ReturnSet _)) = return e
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnSet x)) = myError ty x
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnArray (ArrayVarExp v))) = 
      let newRs = case ty of
                    Integer  -> ReturnInt (IntVarExp v)
                    Sequence -> ReturnSeq (SeqVarExp v)
                    Set      -> ReturnSet (SetVarExp v)
                    Boolean  -> ReturnBool (BoolVarExp v)
                    Array n  -> case lookup v env of
				  Just (Array m) -> 
				    if n==m 
				      then ReturnArray (ArrayVarExp v)
					else error $ "Cannot unify " ++ show (Array n) ++ " with " ++ show (Array m)
				  _		 -> error $ "undeclared variable " ++ show v ++ "?" 
      in return $ SubprogDec specs pn args ty body newRs
    rewriteD env e@(SubprogDec specs pn args (Array n) body (ReturnArray _)) = return e
    rewriteD env e@(SubprogDec specs pn args ty body (ReturnArray x)) = myError ty x
    rewriteD _ x = return x

    -- shortcuts, to avoid repetition
    f env v fn vars = 
      case ((lookup v env), (lookup fn env)) of
	(Just ty1, Just ty2) -> 
	  if (ty1 == ty2) 
	    then (case ty1 of
		Integer -> return $ IntAssign v (IntFunction fn vars)
		Boolean -> return $ BoolAssign v (BoolFunction fn vars)
		Sequence -> return $ SeqAssign v (SeqFunction fn vars)
		Set -> return $ SetAssign v (SetFunction fn vars)
		Array n -> error "Internal error #1093"
	      )
	      else myError ty1 ty2
	(a,b) -> myError a b
    h env v1 v2 = 
      case ((lookup v1 env), (lookup v2 env)) of
	(Just ty1, Just ty2) -> 
	  if (ty1 == ty2) 
	    then (case ty1 of
		Integer -> return $ IntAssign v1 (IntVarExp v2)
		Boolean -> return $ BoolAssign v1 (BoolVarExp v2)
		Sequence -> return $ SeqAssign v1 (SeqVarExp v2)
		Set -> return $ SetAssign v1 (SetVarExp v2)
		Array n -> error "Internal error #1093"
	      )
	      else myError ty1 ty2
	(a,b) -> myError a b
    g env e v e' = if (lookup v env == Just (getTypeA e)) then (return e) else (myError v e')
