module Language.Java.Paragon.TypeCheck.Monad.CodeState where

import Language.Java.Paragon.Syntax
import Language.Java.Paragon.Pretty
import Language.Java.Paragon.Interaction

import Language.Java.Paragon.TypeCheck.Policy
import Language.Java.Paragon.TypeCheck.Actors
import Language.Java.Paragon.TypeCheck.Locks
import Language.Java.Paragon.TypeCheck.Types

import Language.Java.Paragon.Monad.Uniq

import qualified Data.Map as Map
import Data.List (intersect, union)
import Data.Maybe (fromJust)

codeStateModule :: String
codeStateModule = typeCheckerBase ++ ".Monad.CodeState"

data CodeState = CodeState {
      actorSt   :: ActorMap,
      lockMods  :: LockMods,
      exnS      :: ExnsMap
    }
  deriving (Eq, Show)


------------------------------------------
-- Merging of states in parallel
------------------------------------------

mergeStates :: Uniq -> CodeState -> CodeState -> IO CodeState
mergeStates u s1 s2 = do
  newActors <- mergeActors u (actorSt s1) (actorSt s2)
  newExns   <- mergeExns   u (exnS   s1) (exnS   s2)
  return $ CodeState { 
               actorSt = newActors,
               lockMods = lockMods s1 <++> lockMods s2,
               exnS = newExns
             }

------------------------------------------
-- Actor analysis
------------------------------------------

type ActorMap = Map.Map (Name ()) ActorInfo

data ActorInfo = AI { aID :: ActorId, stability :: Stability }
  deriving (Eq, Show)

data Stability = Stable | FullV | FieldV (Maybe (Name ()))
  deriving (Eq, Show)

scrambles :: Stability -> Stability -> Bool
scrambles FullV (FieldV _) = True
scrambles x y = x == y

scramble :: Uniq -> Stability -> CodeState -> IO CodeState
scramble u stab state = do
    let acts = Map.toAscList $ actorSt state
    newActs <- mapM (\(k,v) -> scramble' u stab v >>= \v' -> return (k, v')) acts
    return state { actorSt = Map.fromAscList newActs }

scramble' :: Uniq -> Stability -> ActorInfo -> IO ActorInfo
scramble' u stab a@(AI _ stab') =
    if scrambles stab stab' 
     then do aid' <- newAlias u
             return $ AI aid' stab'
     else return a
             

mergeActors :: Uniq -> ActorMap -> ActorMap -> IO ActorMap
mergeActors u a1 a2 = do
    let newKeys = Map.keys a1 `intersect` Map.keys a2
        oldVals = map (\k -> (fromJust (Map.lookup k a1), fromJust (Map.lookup k a1))) newKeys
    newVals <- mapM (mergeInfo u) oldVals
    return $ Map.fromList $ zip newKeys newVals

mergeInfo :: Uniq -> (ActorInfo, ActorInfo) -> IO ActorInfo
mergeInfo _ (ai1,ai2) | ai1 == ai2 = return ai1
mergeInfo u ((AI _ st),_) = do
  aid <- newAlias u
  return $ AI aid st

------------------------------------------
-- Exception states
------------------------------------------

type ExnsMap  = Map.Map ExnType ExnPoint
data ExnType  = ExnType TcType | ExnContinue | ExnBreak | ExnReturn
  deriving (Eq, Ord, Show)
data ExnPoint = ExnPoint { epState :: CodeState, epWrite :: (TcPolicy TcActor) }
  deriving (Eq, Show)



mergeExns :: Uniq -> ExnsMap -> ExnsMap -> IO ExnsMap
mergeExns u em1 em2 = do
    let newKeys = Map.keys em1 `union` Map.keys em2
        oldVals = map (\k -> (Map.lookup k em1, Map.lookup k em2)) newKeys
    newVals <-mapM (uncurry (mergePoints u)) oldVals
    return $ Map.fromList $ zip newKeys newVals

-- Invariant: At most one of the two arguments can be Nothing
mergePoints :: Uniq -> Maybe ExnPoint -> Maybe ExnPoint -> IO ExnPoint
mergePoints _ Nothing (Just e) = return e
mergePoints _ (Just e) Nothing = return e
mergePoints u (Just (ExnPoint st1 w1)) (Just (ExnPoint st2 w2)) = do
  st <- mergeStates u st1 st2
  let w = w1 `join` w2
  return (ExnPoint st w)
mergePoints _ _ _ = panic (codeStateModule ++ ".mergePoints")
                    "Both ExnPoint arguments cannot be missing!"

-- This should probably be pre-computed each time the map is updated instead
exnPC :: CodeState -> [((TcPolicy TcActor), String)]
exnPC s = map (\(tyX,ptX) -> (epWrite ptX, errorSrc tyX)) $ Map.assocs $ exnS s

errorSrc :: ExnType -> String
errorSrc et = "area of influence of " ++
    case et of
      ExnBreak -> "a break statement"
      ExnContinue -> "a continue statement"
      ExnReturn -> "a return statement"
      ExnType tX -> "exception " ++ prettyPrint tX