module Language.Java.Paragon.TypeCheck.TcState where

import Language.Java.Paragon.Syntax

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
exnPC s = foldl join bottom $ map epWrite $ Map.elems $ exnS s