{-# LANGUAGE CPP, DeriveDataTypeable, PatternGuards, RelaxedPolyRec #-}
module Language.Java.Paragon.TypeCheck.TcEnv where

import Language.Java.Paragon.Syntax

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

import qualified Data.Map as Map
import Control.Monad ( msum )

import Data.Generics.Uniplate.Data

#ifdef BASE4
import Data.Data
#else
import Data.Generics (Data(..),Typeable(..))
#endif

type Map = Map.Map

data TypeMap = TypeMap {
      -- signatures
      fields      :: Map Ident VarFieldSig,
      methods     :: Map (Ident, [TcType]) ([TypeParam], MethodSig),
      constrs     :: Map [TcType] ([TypeParam], ConstrSig),
      lockArities :: Map Ident Int,
      -- known policy-level entities
      policies    :: Map Ident TcPolicy,
      actors      :: Map Ident ActorId,
      -- typemethod eval info
      typemethods :: Map Ident ([Ident], Block),
      -- types and packages
      types       :: Map Ident ([TypeParam], TypeSig),
      packages    :: Map Ident TypeMap
    }
  deriving (Show, Data, Typeable)

emptyTM :: TypeMap
emptyTM = TypeMap {
--                this = typ,
                fields = Map.empty,
                methods = Map.empty,
                constrs = Map.empty,
                lockArities = Map.empty,
                policies = Map.empty,
                actors = Map.empty,
                typemethods = Map.empty,
                types = Map.empty,
                packages = Map.empty
              }

clearToPkgs :: TypeMap -> TypeMap
clearToPkgs tm = emptyTM { packages = packages tm }
                   
clearToTypes :: TypeMap -> TypeMap
clearToTypes tm = emptyTM { packages = packages tm, types = types tm }

data TcEnv = TcEnv  {
      -- lockProps :: ??,
      vars      :: Map Ident VarFieldSig,
      lockstate :: [TcLock],
      returnI   :: (TcType, TcPolicy),
      exnsE     :: Map TcType (TcPolicy, TcPolicy),
      branchPCE :: (Map Entity TcPolicy, TcPolicy)
    }
  deriving (Show, Data, Typeable)

-- Env to use when typechecking expressions not inside method
-- bodies, e.g. in field initializers and policy modifiers
simpleEnv :: TcPolicy -> TcEnv
simpleEnv brPol = TcEnv {
                    vars = Map.empty,
                    lockstate = [],
                    returnI = error "No returns in simple env",
                    exnsE = Map.empty,
                    branchPCE = (Map.empty, brPol)
                  }

data Entity = VarEntity Name
            | ThisFieldEntity Ident
            | ExnEntity TcType
            | LockEntity Name
            | BreakE | ContinueE | ReturnE
  deriving (Show, Eq, Ord, Data, Typeable)

varE, lockE :: Name -> Entity
varE = VarEntity
lockE = LockEntity

exnE :: TcType -> Entity
exnE = ExnEntity

thisFE :: Ident -> Entity
thisFE = ThisFieldEntity 

breakE, continueE, returnE :: Entity
breakE = BreakE
continueE = ContinueE
returnE = ReturnE

data VarFieldSig = VSig { 
      varType :: TcType,
      varPol :: TcPolicy,
      varStatic :: Bool,
      varFinal  :: Bool
    }
  deriving (Show, Data, Typeable)

data MethodSig = MSig {
      mRetType :: TcType,
      mRetPol  :: TcPolicy,
      mPars    :: [TcPolicy],
      mWrites  :: TcPolicy,
      mExpects :: [TcLock],
      mLMods   :: ([TcLock],[TcLock]),
      mExns    :: [(TcType, ExnSig)]
    }
  deriving (Show, Data, Typeable)

data ExnSig = ExnSig { 
      exnReads :: TcPolicy, 
      exnWrites :: TcPolicy, 
      exnMods :: ([TcLock],[TcLock]) 
    }
  deriving (Show, Data, Typeable)

data ConstrSig = CSig {
      cPars    :: [TcPolicy],
      cWrites  :: TcPolicy,
      cExpects :: [TcLock],
      cLMods   :: ([TcLock],[TcLock]),
      cExns    :: [(TcType, ExnSig)]
    }
  deriving (Show, Data, Typeable)


data TypeSig = TSig {
      tIsClass    :: Bool,
      tIsFinal    :: Bool,
      tSupers     :: [TcType],
      tImpls      :: [TcType],
      tMembers    :: TypeMap
    }
  deriving (Show, Data, Typeable)

--------------------------------------
--   Type argument instantiation    --
--------------------------------------

instantiate :: Data a => [(TypeParam,TcTypeArg)] -> a -> a
instantiate pas = transformBi instT
                     . transformBi instA
                     . transformBi instP
                     . transformBi instLs
    where instT :: TcRefType -> TcRefType
          instT tv@(TcTypeVar i) =
              case lookup i typs of
                Just rt -> rt
                Nothing -> tv
          instT rt = rt

          instA :: ActorId -> ActorId
{-          instA av@(ActorTPVar i) =
              case lookup i as of
                Just a -> a
                Nothing -> av -}
          instA a = a

          instP :: TcPolicy -> TcPolicy
          instP pv@(TcRigidVar i) =
              case lookup i ps of
                Just p -> p
                Nothing -> pv
          instP p = p

          instLs :: [TcLock] -> [TcLock]
          instLs = concatMap instL

          instL :: TcLock -> [TcLock]
          instL lv@(TcLockVar i) =
              case lookup i locks of
                Just le -> le
                Nothing -> [lv]
          instL l = [l]

          typs  = [ (i, rt) | (TypeParam   i _, TcActualType  rt) <- pas ] 
          as    = [ (i, n ) | (ActorParam  i  , TcActualActor  n) <- pas ]
          ps    = [ (i, p ) | (PolicyParam i  , TcActualPolicy p) <- pas ]
          locks = [ (i, ls) | (LockStateParam i, TcActualLockState ls) <- pas ]



--------------------------------------
--    Working with the branchPC     --
--------------------------------------

branchPC :: Maybe Entity -> TcEnv -> TcPolicy
branchPC men (TcEnv { branchPCE = (bm, def) }) = 
    flip (maybe def) men $ \en ->
        maybe def id (Map.lookup en bm)

joinBranchPC :: TcPolicy -> TcEnv -> TcEnv
joinBranchPC p env = let (bm, def) = branchPCE env
                      in env { branchPCE = (Map.map (`join` p) bm, def `join` p) }

--------------------------------------
--    Working with the lookups      --
--------------------------------------

lookupNamed :: (TypeMap -> Map Ident a) -> Name -> TypeMap -> Maybe a
lookupNamed recf (Name is) tm =
   let actualTm = lookupTypeOfN (Name $ init is) tm
   in Map.lookup (last is) (recf actualTm)

lookupNamedMethod :: Name -> [TcType] -> TypeMap -> Maybe ([TypeParam],MethodSig)
lookupNamedMethod (Name is) ts tm = 
    let actualTm = lookupTypeOfN (Name $ init is) tm
    in Map.lookup (last is, ts) (methods actualTm)

pkgsAndTypes :: TypeMap -> Map Ident TypeMap
pkgsAndTypes tm = Map.union (packages tm)
                    -- disregard type parameters
                    (Map.map (tMembers . snd) $ types tm)

lookupTypeOfN :: Name -> TypeMap -> TypeMap
lookupTypeOfN (Name is) tm = aux is tm tm
    where aux :: [Ident] -> TypeMap -> TypeMap -> TypeMap
          aux [] tm _ = tm
          aux (i:is) tm baseTm =
            let newTm = 
                 case Map.lookup i (fields tm) of
                   Just (VSig typ _ _ _) -> lookupTypeOfT typ baseTm
                   Nothing -> case Map.lookup i (pkgsAndTypes tm) of
                                Just aTm -> aTm
                                Nothing -> error "lookupTypeOfN: No such name"
            in aux is newTm baseTm

lookupTypeOfT :: TcType -> TypeMap -> TypeMap
lookupTypeOfT (TcRefT (TcClsRefT (TcClassT iargs))) tm = aux iargs tm
    where aux :: [(Ident, [TcTypeArg])] -> TypeMap -> TypeMap
          aux [] tm = tm
          aux ((i, args):iargs) tm =
              let newTm =
                   case Map.lookup i (types tm) of
                     Just (pars, tsig) -> instantiate (zip pars args) (tMembers tsig)
                     Nothing -> 
                         case Map.lookup i (packages tm) of
                           Just tm -> tm
                           Nothing -> error $ "lookupTypeOfT: No such type: " ++ show i
              in aux iargs newTm
lookupTypeOfT t tm = error $ "lookupTypeOfT: Unexpected type: " ++ show t