{-# LANGUAGE CPP #-}
module KIMisc(
        KVar, KSubst, apKSu,
        KI, run, err, newKVar, getKSubst,
        unifyType, unifyFunc, unifyDefArgs, unifyDefAlias, unifyDefStar,
        groundK) where

import Data.List(union)
import Data.Maybe(fromMaybe)
import Control.Monad(when, ap, liftM)
import CVPrint
import PFPrint
import Error(internalError, EMsg, ErrMsg(..))
import CType(baseKVar, isKVar)
import Id(Id, getIdString)
import Debug.Trace
import Util(tracep)
import IOUtil(progArgs)
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS

infixr 4 @@

doTraceKI :: Bool
doTraceKI = "-trace-kind-inference" `elem` progArgs

type KVar = Int

data KSubst  = S (IM.IntMap Kind) (IM.IntMap IS.IntSet)

instance Eq KSubst where
  (S s1 _) == (S s2 _) = s1 == s2

instance PPrint KSubst where
    pPrint d p (S s b) = pPrint d p (IM.toList s)

nullKSubst  :: KSubst
nullKSubst   = S IM.empty IM.empty

(+->)      :: KVar -> Kind -> KSubst
u +-> t     = S (IM.singleton u t)
                (back_set u t)

back_set :: KVar -> Kind -> IM.IntMap IS.IntSet
back_set u t = (IM.fromList (map (\v -> (v, IS.singleton u)) (tv t)))

{-
-- Unused
mkKSubst :: [(KVar, Kind)] -> KSubst
mkKSubst vts = S vts
-}

apKSu :: KSubst -> Kind -> Kind
apKSu (S s _) v@(KVar u) =
    case IM.lookup u s of
    Just k  -> k
    Nothing -> v
apKSu s (Kfun l r) = Kfun (apKSu s l) (apKSu s r)
apKSu s KStar = KStar
apKSu s KNum = KNum
apKSu s KStr = KStr

tv :: Kind -> [KVar]
tv (KVar u)   = [u]
tv (Kfun l r) = tv l `union` tv r
tv KStar      = []
tv KNum       = []
tv KStr       = []

(@@) :: KSubst -> KSubst -> KSubst
(S s1 b1) @@ S s2 b2 | IS.null (IS.intersection (IM.keysSet b2)
                                                (IM.keysSet s1)) =
  S (IM.union s1 s2) (IM.unionWith IS.union b1 b2)

-- general case when s1 needs to be applied to the rhs of s2
ss1@(S s1 b1) @@ S s2 b2 = S (IM.union s1 s2') (IM.unionWith IS.union b1 b2')
  where s2' = (IM.map (apKSu ss1) s2)
        b2' = IM.unionWith IS.union b2_base b2_overlap
        b2_base = IM.difference b2 s1
        b2_overlap = (IM.fromListWith IS.union overlap_list_map)
        overlap_vars = IS.toList (IS.intersection (IM.keysSet b2) (IM.keysSet s1))
        get_mtv u = fmap tv (IM.lookup u s1)
        overlap_pairs = [ (u, fromMaybe [] (get_mtv u)) | u <- overlap_vars ]
        overlap_list_map = [ (u, IS.fromList vs) | (u, vs) <- overlap_pairs,
                                                   not (null vs) ]

{-
-- Unused
-- (was the last use removed when fixing fundep handling?)
merge :: KSubst -> KSubst -> Maybe KSubst
merge ss1@(S s1) ss2@(S s2) = if agree then Just (S s) else Nothing
 where dom s = map fst s
       s     = s1 ++ s2
       agree = all (\v -> apKSu ss1 (KVar v) ==
                          apKSu ss2 (KVar v))
                   (dom s1 `intersect` dom s2)
-}

--------------------

mgu  :: Kind -> Kind -> Maybe KSubst
mgu (Kfun l r) (Kfun l' r') = do s1 <- mgu l l'
                                 s2 <- mgu (apKSu s1 r)
                                         (apKSu s1 r')
                                 Just (s2 @@ s1)
mgu (KVar u) t        = varBind u t
mgu t (KVar u)        = varBind u t
mgu KStar KStar       = Just nullKSubst
mgu KNum KNum         = Just nullKSubst
mgu KStr KStr         = Just nullKSubst
mgu t1 t2             = Nothing


varBind :: KVar -> Kind -> Maybe KSubst
varBind u t | t == KVar u      = Just nullKSubst
            | u `elem` tv t    = Nothing
            | otherwise        = Just (u +-> t)

--------------------

data KState = KState KSubst Int

newtype KI a = M (KState -> Either EMsg (KState, a))

instance Monad KI where
    return = pure
    M a >>= f = M $ \ s ->
        case a s of
        Left e -> Left e
        Right (s', b) ->
            let M f' = f b
            in  f' s'

instance Functor KI where
  fmap = liftM

instance Applicative KI where
  pure a = M $ \ s -> Right (s, a)
  (<*>) = ap

run :: KI a -> Either EMsg a
run (M m) =
    case m (KState nullKSubst baseKVar) of
    Left e -> Left e
    Right (KState ksubst _, a) ->
        tracep doTraceKI ("KIMisc.run: " ++ ppReadable ksubst) $ Right a

err :: EMsg -> KI a
err e = M (\s->Left e)

getKSubst   :: KI KSubst
getKSubst    = M $
    \ ts@(KState s _) -> Right (ts, s)

extKSubst   :: KSubst -> KI ()
extKSubst s' = M $
    \ (KState s n) -> Right (KState (s'@@s) n, ())

newKVar' :: KI Kind
newKVar' = M $
    \ (KState s n) ->
    Right (KState s (n+1), KVar n)

newKVar :: Maybe Id -> KI Kind
newKVar Nothing = newKVar'
newKVar (Just i) = do
  v <- newKVar'
  when doTraceKI $ traceM ("newKVar: " ++ show (i, v))
  return v

--------------------

-- Unification of type declaration is split into two parts:
-- (1) unify the arguments, and get back a set of assumptions and the
--     result kind (this is "unifyDefArgs")
-- (2) unify the result, after inference has been performed with the
--     assumptions from part 1 ("unifyDefAlias" or "unifyDefStar")

unifyDefArgs :: Id -> Kind -> [Id] -> KI ([(Id, Kind)], Maybe Kind)
unifyDefArgs i con_k vs =
    -- no subst needed
    case (extractFuncArgKinds con_k) of
      [k] | (length vs > 0) && (isKVar k)
          -> -- no kind signature was given (only possible in Classic)
             do vks <- mapM (newKVar . Just) vs
                return (zip vs vks, Nothing)
      ks  -> -- we can assume that a complete kind signature has been given,
             -- but check that the number of argument kinds matches the
             -- number of argument variables (only possible in Classic)
             do let vks = init ks
                    res = last ks
                -- only check for too few; too many will be handled at the
                -- next unification step
                -- XXX we could create a new error tag for this situation
                if (length vs > length vks)
                  then err (getPosition i, ETypeTooManyArgs (pfpString i))
                  else if (length vs == length vks)
                  then return (zip vs vks, Just res)
                  else let (vks1, vks2) = splitAt (length vs) vks
                       in  return (zip vs vks1, Just (mkKFun vks2 res))

-- Unify the result of a type alias declaration
-- This is handled separately, since the result type of the alias does not
-- have to be Star, and we want to report mismatch on the result expression
-- (only matters in Classic, when a kind signature is given).
unifyDefAlias :: Id -> Kind -> [(Id, Kind)] -> Maybe Kind -> Type -> Kind ->
                 KI ()
unifyDefAlias i ck as mk t k_inferred = do
    let aks = map snd as
    case (mk) of
      Nothing -> -- unify the whole kind
                 unifyDef i (mkKFun aks k_inferred) ck
      Just k  -> -- unify the inferred type for the alias expression ("t")
                 -- with the result kind given by the user
                 unifyDefType t k_inferred k

-- Unify the result of a data type declaration (kind Star)
-- This will only fail in Classic, if the user has given a kind signature.
unifyDefStar :: Id -> Kind -> [(Id, Kind)] -> Maybe Kind -> KI ()
unifyDefStar i ck as mk = do
    let aks = map snd as
    case (mk) of
      Nothing -> -- unify the whole kind
                 unifyDef i (mkKFun aks KStar) ck
      Just k  -> -- unify just the result kind
                 unifyDef i KStar k

extractFuncArgKinds :: Kind -> [Kind]
extractFuncArgKinds (Kfun a b) = (a : extractFuncArgKinds b)
extractFuncArgKinds k = [k]

mkKFun :: [Kind] -> Kind -> Kind
mkKFun []     k = k
mkKFun (a:as) k = Kfun a (mkKFun as k)

kindErr :: Type -> Kind -> Kind -> ErrMsg
kindErr t KStar KNum  = EKindNumForStar (pfpString t)
kindErr t KStar KStr  = EKindStrForStar (pfpString t)
kindErr t KNum  KStar = EKindStarForNum (pfpString t)
kindErr t KNum  KStr  = EKindStrForNum (pfpString t)
kindErr t KStr  KStar = EKindStarForStr (pfpString t)
kindErr t KStr  KNum  = EKindNumForStr (pfpString t)
-- XXX replace Readable with String
kindErr t exp   inf   = EUnifyKind (pfpReadable t) (ppReadable inf) (ppReadable exp)


-- The following two functions are like "unifyType", except that they do
-- not report errors about being applied to too many or too few arguments.
-- In BSV, that situation cannot occur; in Classic, it only occurs if the
-- user has given a bad kind signature and there we want to report the usual
-- kind-mismatch error.
-- XXX We could have a special error for type def kind signature mismatch,
-- XXX instead of using EUnifyKind.

-- This one does not report special errors about numeric vs string vs value.
-- The only situation where that should occur is for type aliases, and
-- that's handled with "unifyDefType".
unifyDef :: Id -> Kind -> Kind -> KI ()
unifyDef i k_inferred k_expected = do
    s <- getKSubst
    let k_inferred' = apKSu s k_inferred
        k_expected' = apKSu s k_expected
    case mgu k_inferred' k_expected' of
     Just u  -> extKSubst u
     Nothing ->
         -- XXX replace Readable with String
         err (getPosition i, EUnifyKind (pfpReadable i)
                                (ppReadable k_inferred')
                                (ppReadable k_expected'))

-- This reports the error on the alias expression, and reports numeric,
-- string and value kinds specially.
unifyDefType :: Type -> Kind -> Kind -> KI ()
unifyDefType t k_inferred k_expected = do
    s <- getKSubst
    let k_inferred' = apKSu s k_inferred
        k_expected' = apKSu s k_expected
    case mgu k_inferred' k_expected' of
     Just u  -> extKSubst u
     Nothing -> err (getPosition t, kindErr t k_expected' k_inferred')


-- Takes the function and its kind and the argument to which the function
-- is applied and the arguments kind, and returns the newKVar
unifyFunc :: Type -> Kind -> Type -> Kind -> KI Kind
unifyFunc f fk a ak = do
    s <- getKSubst
    let sfk = apKSu s fk
        sak = apKSu s ak
    -- trace "unifyFunc:" $ return ()
    case sfk of
        Kfun ak' v' -> do
                         -- The function has inferred kind "ak' -> ..."
                         -- so use "ak'" as the expected kind for the argument
                         unifyType a sak ak'
                         -- Return the inferred result
                         return v'
        _ -> do v <- newKVar Nothing
                -- The function is expected to be "sak -> ..."
                unifyType f sfk (Kfun sak v)
                return v


unifyType :: Type -> Kind -> Kind -> KI ()
unifyType t k_inferred k_expected = do
    s <- getKSubst
    let k_inferred' = apKSu s k_inferred
        k_expected' = apKSu s k_expected
    case mgu k_inferred' k_expected' of
     Just u  -> extKSubst u
     Nothing ->
       let
           pos = getPosition t

           -- if a type constructor is being partially applied, get the core
           unapType = removeTypeAps t
           -- the name of the constructor, for error messages
           unapTypeName = showUnapType t
       in
         -- trace ("trace: unifyType: " ++ traceUnapType t) $
           if (isGroundNonFuncK k_inferred') && (isFuncK k_expected')
           then
             -- handle case of too many arguments
             case unapType of
                 TCon (TyNum i _) -> err (pos, ENumKindArg)
                 _ -> if (unapType == t)
                      then -- type is not partially applied
                        err (pos, ETypeNoArg unapTypeName)
                      else
                          err (pos, ETypeTooManyArgs unapTypeName)
           else
           if (isGroundNonFuncK k_expected') && (isFuncK k_inferred')
           then
             -- handle case of not enough arguments
             err (pos, ETypeTooFewArgs unapTypeName)
           else
           err (getPosition t, kindErr t k_expected' k_inferred')

--------------------

groundK :: Kind -> Kind
groundK KStar = KStar
groundK KNum = KNum
groundK KStr = KStr
groundK (Kfun l r) = Kfun (groundK l) (groundK r)
groundK (KVar k) = tracep doTraceKI ("groundK: " ++ show k) $ KStar

--------------------

-- This is like "leftCon", except that it returns the inner type
-- even when it's not a "TyCon"
removeTypeAps :: Type -> Type
removeTypeAps (TAp f _) = removeTypeAps f
removeTypeAps t = t

-- The unapplied type as a String, for error messages
showUnapType :: Type -> String
showUnapType t =
    case (removeTypeAps t) of
      TCon (TyCon i _ _) -> getIdString i
      TCon (TyNum i _) -> show i
      TCon (TyStr s _) -> show s
      TVar (TyVar i _ _) -> getIdString i
      unapType@(TGen _ _) -> show unapType
      -- These shouldn't happen
      TAp _ _ -> internalError "KIMisc.showUnapType (TAp)"
      TDefMonad _ -> internalError "KIMisc.showUnapType (TDefMonad)"

{-
-- The unapplied type as a String, for debug trace
traceUnapType :: Type -> String
traceUnapType t =
    case (removeTypeAps t) of
      TCon (TyCon i _ _) -> "TCon (TyCon) " ++ getIdString i
      TCon (TyNum i _) -> "TCon (TyNum) " ++ show i
      TVar (TyVar i _ _) -> "TVar " ++ getIdString i
      TGen _ _ -> "TGen"
      -- These shouldn't happen
      TAp _ _ -> "TAp"
      TDefMonad _ -> "TDefMonad"
-}

-- Returns whether a kind is grounded (has no variables) and is not
-- a function; i.e. a value, num or string kind.
isGroundNonFuncK :: Kind -> Bool
isGroundNonFuncK KStar = True
isGroundNonFuncK KNum  = True
isGroundNonFuncK KStr  = True
isGroundNonFuncK _     = False

-- Returns whether a kind is a function
isFuncK :: Kind -> Bool
isFuncK (Kfun _ _) = True
isFuncK _          = False

--------------------