module Language.Java.Paragon.TypeCheck.TcState where

import Language.Java.Paragon.Syntax
import Language.Java.Paragon.Pretty
import Text.PrettyPrint (text)

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.TypeCheck.Uniq

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

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

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

mergeStates :: Uniq -> TcState -> TcState -> IO TcState
mergeStates u s1 s2 = do
  newActors <- mergeActors u (actorSt s1) (actorSt s2)
  newExns   <- mergeExns   u (exnS   s1) (exnS   s2)
  return $ TcState { 
               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 (Name ())
  deriving (Eq, Show)

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

scramble :: Uniq -> Stability -> TcState -> IO TcState
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 aid 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 :: TcState, epWrite :: TcPolicy }
  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)

-- This should probably be pre-computed each time the map is updated instead
exnPC :: TcState -> [(TcPolicy, 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