{-# LANGUAGE CPP, DeriveDataTypeable #-}
module Language.Java.Paragon.TypeCheck.Policy 
    (
     bottom, top, join, meet, recmeet, isTop, isBottom,
     TcPolicy(..), TcPolicyRec(..), TcClause(..), TcAtom(..), TcActor(..), TcMetaVar(..),
     {- toPolicyLit, -} toRecPolicy,
     lockToAtom, 
     zonkPolicy, substPolicy,
     flowAtomString
    ) where

import Language.Java.Paragon.Syntax

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

import Data.IORef
import Data.List ( (\\), groupBy, nub )
import Data.Maybe (isJust, fromJust, catMaybes)

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

flowAtomString :: String
flowAtomString = "$FLOW$"

---------------------------------------------------------------
-- Basic policies

data TcPolicy = TcPolicy [TcClause TcActor]
              | TcPolicyVar TcMetaVar
              | TcRigidVar Ident
              | TcJoin TcPolicy TcPolicy
              | TcMeet TcPolicy TcPolicy
  deriving (Eq,Ord,Show,Data,Typeable)
  
data TcPolicyRec = TcPolicyRec [TcClause TcAtom]
                | TcJoinRec TcPolicyRec TcPolicyRec
                | TcMeetRec TcPolicyRec TcPolicyRec
  deriving (Eq, Ord, Show, Data, Typeable)

data TcMetaVar = TcMetaVar Int (IORef (Maybe TcPolicy))
--  deriving (Data, Typeable)

instance Data TcMetaVar
instance Typeable TcMetaVar

instance Eq TcMetaVar where
  (TcMetaVar i1 _) == (TcMetaVar i2 _) = i1 == i2

instance Ord TcMetaVar where
  compare (TcMetaVar i1 _) (TcMetaVar i2 _) = compare i1 i2

instance Show TcMetaVar where
  show (TcMetaVar i _) = "$$" ++ show i

data TcClause a = TcClause a [TcAtom]
  deriving (Eq,Ord,Show,Data,Typeable)

data TcAtom = TcAtom Name [TcActor]
  deriving (Eq,Ord,Show,Data,Typeable)

data TcActor = TcActor ActorId
             | TcVar Ident
  deriving (Eq,Ord,Show,Data,Typeable)

--------------------------------------
-- Converting policies to recursive --
-- ones, i.e. with the head flow    --
--------------------------------------

toRecPolicy :: TcPolicy -> TcPolicyRec
toRecPolicy (TcPolicy cs)  = TcPolicyRec $ map toRecClause cs
  where toRecClause (TcClause act ats) = TcClause (TcAtom (Name [Ident flowAtomString]) [act]) ats
toRecPolicy (TcJoin p1 p2) = TcJoinRec (toRecPolicy p1) (toRecPolicy p2)
toRecPolicy (TcMeet p1 p2) = TcMeetRec (toRecPolicy p1) (toRecPolicy p2)

{-----------------------------------
-- Turning type-level policies   --
-- back to their source reps     --
-----------------------------------

toActor :: TcActor -> Actor
toActor (TcVar i) = Var i
toActor (TcActor aid) = Actor $ ActorName $ Name $ [Ident $ "a" ++ show (getId aid)] -- reprName aid

toActorName :: ActorId -> ActorName
toActorName (Fresh k) = ActorName $ Name $ [Ident $ "a" ++ show k]
toActorName (Alias k) = ActorName $ Name $ [Ident $ "a" ++ show k]

toAtom :: TcAtom -> Atom
toAtom (TcAtom n acts) = Atom n $ map toActor acts

toClause :: TcClause TcActor -> Clause ActorName
toClause (TcClause h ats) = Clause (toActor h) $ map toAtom ats

toPolicyLit :: TcPolicy -> PolicyExp
toPolicyLit (TcPolicy cs) = PolicyLit $ map toClause cs

----------------------------------}
-- Basic policies

bottom, top :: TcPolicy
bottom = TcPolicy [TcClause (TcVar $ Ident "x") []]
top = TcPolicy []

isTop, isBottom :: TcPolicy -> Bool
isTop = (== top)
isBottom (TcPolicy cs) = any isBottomC cs
  where isBottomC (TcClause (TcVar _) []) = True
        isBottomC _ = False


----------------------------------
-- meet and join

lub, join :: TcPolicy -> TcPolicy -> TcPolicy
lub (TcPolicy cs1) (TcPolicy cs2) =
      let sameFixedCs = [ TcClause (TcActor aid) (as ++ substAll senv bs) 
                              | TcClause (TcActor aid) as <- cs1,
                                TcClause (TcActor aid') bs <- cs2,
                                aid == aid',
                                let senv = zip (fuv bs) (map TcVar $ allBinders \\ fuv as)]
          bothVarCs   = [ TcClause a1 (as ++ substAll senv bs) 
                              | TcClause a1@(TcVar _) as <- cs1,
                                TcClause (TcVar i2) bs <- cs2,
                                let senv = (i2,a1): zip (fuv bs) (map TcVar $ allBinders \\ fuv as) ]
          fstVarCs    = [ TcClause a1 (as ++ substAll senv bs)
                              | TcClause a1@(TcActor _) as <- cs1,
                                TcClause (TcVar i2) bs <- cs2,
                                let senv = (i2,a1): zip (fuv bs) (map TcVar $ allBinders \\ fuv as) ]
          sndVarCs    = [ TcClause a2 (bs ++ substAll senv as)
                              | TcClause (TcVar i1) as <- cs1,
                                TcClause a2@(TcActor _) bs <- cs2,
                                let senv = (i1,a2): zip (fuv as) (map TcVar $ allBinders \\ fuv bs) ]
      in {- return $ -} TcPolicy $ sameFixedCs ++ bothVarCs ++ fstVarCs ++ sndVarCs

lub p1 p2 = TcJoin p1 p2

join = lub

type Subst = [(Ident, TcActor)]

subst :: Ident -> TcActor -> [TcAtom] -> [TcAtom]
subst i1 i2 = substAll [(i1,i2)]

substAll :: Subst -> [TcAtom] -> [TcAtom]
substAll env = map (substAtom env)
    where substAtom env (TcAtom n acts) = TcAtom n (map (substA env) acts)

substA :: Subst -> TcActor -> TcActor
substA env (TcVar i) = case lookup i env of
                         Nothing -> TcVar i
                         Just a -> a
substA env act = act

fuv :: [TcAtom] -> [Ident]
fuv = concatMap fuv'
    where fuv' (TcAtom _ as) = concatMap fuvA as

fuvA :: TcActor -> [Ident]
fuvA (TcActor _) = []
fuvA (TcVar i) = [i]

allBinders :: [Ident]
allBinders = map Ident $ 
              [ show c | c <- ['a' .. 'z']] ++
               [ c : show (i :: Int) | c <- ['a' .. 'z'], i <- [0..]]


glb, meet :: TcPolicy -> TcPolicy -> TcPolicy
-- This one could be smartified
glb (TcPolicy as) (TcPolicy bs) = TcPolicy $ as ++ bs
glb p1 p2 = TcMeet p1 p2
meet = glb

recmeet :: TcPolicyRec -> TcPolicyRec -> TcPolicyRec
recmeet (TcPolicyRec cs1) (TcPolicyRec cs2) = TcPolicyRec $ cs1 ++ cs2
recmeet p1                p2                = TcMeetRec p1 p2

----------------------------------------------
-- Specialisation

specialise :: [TcLock] -> TcPolicy -> TcPolicy
specialise ls (TcPolicy cs) = TcPolicy $ concatMap (specClause ls) cs

specClause :: [TcLock] -> TcClause TcActor -> [TcClause TcActor]
specClause ls (TcClause h atoms) = 
    -- Step 1: generate possible substitutions from each atom
    let substss = map (genSubst ls) atoms -- :: [[Subst]]
        substs  = joinSubsts $ map ([]:) substss
    in remSubsumed [] $
           [ TcClause (substA subst h) (substAll subst atoms \\ map lockToAtom ls) 
             | subst <- substs ]

  where genSubst :: [TcLock] -> TcAtom -> [Subst]
        genSubst ls (TcAtom n acts) =
            [ subst' | (TcLock ln aids) <- ls,
                       n == ln,
                       let mS = fmap concat $ sequence $ map (uncurry matchesActor) (zip acts aids),
                       isJust mS,
                       let subst = fromJust mS
                           mSubst = toConsistent subst,
                       isJust mSubst, -- is this redundant?
                       let Just subst' = mSubst ]

        joinSubsts :: [[Subst]] -> [Subst]
        joinSubsts sss = 
            let allS = map concat $ cartesian sss -- [Subst]
                conS = catMaybes $ map toConsistent allS      -- [Subst]
            in conS

        remSubsumed :: [TcClause TcActor] -> [TcClause TcActor] -> [TcClause TcActor]
        remSubsumed acc [] = acc
        remSubsumed acc (x:xs) = if any (subsumes x) (acc ++ xs)
                                  then remSubsumed acc xs
                                  else remSubsumed (x:acc) xs

        subsumes :: TcClause TcActor -> TcClause TcActor -> Bool
        subsumes (TcClause h1 as1) (TcClause h2 as2) =
            h1 == h2 && null (as2 \\ as1)

        cartesian :: [[Subst]] -> [[Subst]]
        cartesian [] = [[]]
        cartesian (xs:xss) = [ y:ys | y <- xs, ys <- cartesian xss ]

matchesActor :: TcActor -> ActorId -> Maybe [(Ident, TcActor)]
matchesActor (TcActor aid1) aid2 | aid1 == aid2 = Just []
matchesActor (TcVar i) aid = Just [(i, TcActor aid)]
matchesActor _ _ = Nothing

lockToAtom :: TcLock -> TcAtom
lockToAtom (TcLock n aids) = TcAtom n $ map TcActor aids
lockToAtom (TcLockVar i) = TcAtom (Name [i]) []

toConsistent :: Subst -> Maybe Subst
toConsistent subst = 
    let bndrs = groupBy (\a b -> fst a == fst b) subst
    in  sequence $ map checkConsistent bndrs
        where checkConsistent :: Subst -> Maybe (Ident, TcActor)
              checkConsistent s = case nub s of
                                    [s] -> Just s
                                    _   -> Nothing

---------------------------
-- Zonking

zonkPolicy :: TcPolicy -> IO TcPolicy
zonkPolicy p@(TcPolicyVar (TcMetaVar i ref)) = do
  mPol <- readIORef ref
  case mPol of
    Just pol -> do pol' <- zonkPolicy pol
                   writeIORef ref (Just pol') -- short-circuit
                   return pol'
    Nothing -> return p

zonkPolicy (TcJoin p1 p2) = do
  p1' <- zonkPolicy p1
  p2' <- zonkPolicy p2
  return $ p1' `join` p2'

zonkPolicy (TcMeet p1 p2) = do
  p1' <- zonkPolicy p1
  p2' <- zonkPolicy p2
  return $ p1' `meet` p2'

zonkPolicy p = return p

-------------------------------
-- Policy substitution

-- Invariant: policies are always closed by the env
substPolicy :: [(Ident, TcPolicy)] -> TcPolicy -> TcPolicy
substPolicy env (TcRigidVar i) = fromJust $ lookup i env
substPolicy env (TcJoin p1 p2) =
    let pol1 = substPolicy env p1
        pol2 = substPolicy env p2
    in pol1 `join` pol2
substPolicy env (TcMeet p1 p2) =
    let pol1 = substPolicy env p1
        pol2 = substPolicy env p2
    in pol1 `meet` pol2
substPolicy env p = p

mkPolicySubst :: [TcPolicy] -> [TcPolicy] -> [(Ident, TcPolicy)]
mkPolicySubst pPars pArgs = catMaybes $ map aux $ zip pPars pArgs
  where aux (TcRigidVar i, p) = Just (i,p)
        aux _ = Nothing