{-# OPTIONS -fglasgow-exts #-}
-- Module      : TypeInference
-- Authors     : Sean Seefried (http://www.cse.unsw.edu.au/~sseefried)
-- Copyright   : (c) 2006
-- License     : BSD3
-- Created     : 17 Feb 2006
--
-- In this file I implement constructor classes as described in:
--
-- Mark P. Jones' paper "A system of constructor classes: overloading
-- and implicit higher order polymorphism". In FPCA'93: Proceedings of
-- the conference on Functional programming language and computer
-- architecture. pages 52--61, New York, NY, USA. 1993. ACM Press
--
-- When a term is type checked a trace of the derivation tree
-- corresponding to the syntax direction type inference algorithm W
-- (Figure 3, p60) is printed to the screen.  
--
-- I do not print all of the information that would appear in a
-- derivation tree; this would take up too much space. However, I do
-- show the value of all the important variables in each rule.
--
-- Figure 3 is interesting. It is presented as relations but, in fact,
-- it defines a function. One can think of the type assignment as the
-- input and the predicates, substituion and type as the outputs.  To
-- clarify I will now highlight which values are the inputs and
-- outputs for each rule.  (Don't bother reading this unless you have
-- a copy of the paper handy.)
--
-- Rule (var)^W 
-- Bottom line: 
--   Input: A
--   Output: [b_i/a_i]P, (empty subtitution), tau
--
-- Rule (->E)^W
-- Top line (1st relation)
--   Input: A
--   Output: P, T, tau
-- Top line (2nd relation)
--   Input: TA 
--   Output: Q, T', tau'
--   Note: The input is TA; this means 1st relation must be "executed"
--         before this one.)
-- Top line (Unification)
--   Input:  T'tau, tau' -> alpha
--   Output: U
-- Bottom line:
--   Input:  A
--   Output: U(T'P union Q), UT'T, U alpha
--
-- Rule (->I)^W
-- Top line
--   Input:  A_x, x: alpha
--   Output: P, T, tau
--
-- Bottom line
--   Input:  A
--   Output: P, T, tau
--
-- Rule (let)^W
--   Top line (1st relation)
--     Input:  A
--     Output: P, T, tau
--   Top line (2nd relation)
--     Input: TA_x, x: Gen (TA, P => tau)
--     Ouput: P', T', tau'
--   Bottom line
--     Input:  A
--     Output: P', T'T, tau' 

module TypeInference where

import Data.Maybe
import Control.Monad.State
import Data.List

data TyVar = TyVar String Kind deriving (Eq) 
           -- By convention begins with lower case letter.
data Con = Con String Kind deriving (Eq)
           -- By convention first letter is capitalised. 

data Kind = Star | Kind :->: Kind deriving (Eq)

data Pred = Pred Name [Type] deriving (Eq)

data Type = ConT Con        -- type constructor
	  | VarT TyVar        -- type variables
	  | AppT Type Type  -- type application
           deriving (Eq)


-- Creates function types
to :: Type -> Type -> Type
to t1 t2 = AppT (AppT arrowCon t1) t2
  where arrowCon = ConT (Con "->" (Star :->: (Star :->: Star)))

-- fixity declaration for @to.
infixr 5 `to`  -- right associative

data Scheme = Scheme [TyVar] [Pred] Type  -- universally quantified
                                         -- type variables
              deriving (Eq)

type Var = String

data Exp = VarE Var          -- term variables
	 | AppE Exp Exp      -- application
	 | LamE Var Exp      -- lambda abstraction
	 | LetE Var Exp Exp  -- (non-recursive) let binding

--
-- Normally one would represent a substitution as a function (Type ->
-- Type) but I want to print these out.  This complicates the
-- composition of substitutions which would normally be done using
-- normal function composition (.).
--
type Subst = [ (TyVar, Type) ]

--
-- Substitution composition.  Using definition from Benjamin Pierce,
-- Types and Programming Languages, p318.
--
(.:) :: Subst -> Subst -> Subst
sub2 .: sub1 = mapSnd (appSubst sub2) sub1 ++ newSub2
  where mapSnd f  = map (\(x,y) -> (x, f y))
        newSub2 = filter (keyOf sub1) sub2
        keyOf sub (key, _) = 
           case lookup key sub of
             Just _  -> False
             Nothing -> True

kind :: Type -> Kind
kind (VarT (TyVar _ k)) = k
kind (ConT (Con _ k)) = k 
kind (AppT t t') 
 | (k'' :->: k) <- kt, k'' == k' = k
 | otherwise = error $ "Error in: " ++ show (AppT t t') ++ ". " ++ 
                       show t ++ " has kind " ++ show kt ++ " but " ++ 
                       show t' ++ " has kind " ++ show k'
 where k'  = kind t'
       kt  = kind t  




instance Show Type where
  show (VarT (TyVar v _)) = v
  show (AppT (AppT (ConT (Con "->" _ )) t1) t2) = show_t1 ++ " -> " ++ show t2
    where show_t1 
            | AppT (AppT (ConT (Con "->" _)) _) _  <- t1 = "(" ++ show t1 ++ ")"
            | otherwise      = show t1
  show (ConT (Con name _ )) = name
  show (AppT t1 t2) = show t1 ++ " " ++ show_t2
    where show_t2 
            | AppT _ _ <- t2 = "(" ++ show t2 ++ ")"
            | otherwise      = show t2
instance Show Pred where
  show (Pred nm typs) = concat $ intersperse " " (nm : map show typs)

instance Show Exp where
  show (VarE v) = v
  show (AppE e1 e2) = "(" ++ show e1 ++ " " ++ show e2 ++ ")"
  show (LamE v body) = "\\" ++ show v ++ " -> " ++ show body
  show (LetE x e1 e2) = "let " ++ show x ++ " = " ++ show e1 ++
                        " in " ++ show e2

instance Show Scheme where
  show (Scheme vs ps typ) = showVars ++ showPreds ++ show typ
    where showVars  
             | null vs = ""
             | otherwise = let vars = concat $ intersperse " "  (map show vs)
                           in  "forall " ++ vars ++ "."
          showPreds 
            | null ps = "" 
            | otherwise = "(" ++ (concat $ intersperse ", " (map show ps))  ++ 
                          ") => " 

instance Show TyVar where
  show (TyVar nm k) = nm

instance Show Kind where
  show Star = "*"
  show (k :->: k') = show_k k ++ " -> " ++ show_k k'
    where 
      show_k  k
        | Star <- k = show k
        | otherwise = "(" ++ show k  ++ ")"

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

-- Define a subtitution using a single mapping.
(|->) :: TyVar -> Type -> Subst
v |-> t = [ (v,t) ]

-- Applies a substitution to a type
appSubst :: Subst -> Type -> Type
appSubst sub typ@(VarT v) = 
  case lookup v sub of
    Just typ' -> typ'
    Nothing -> typ
appSubst _ con@(ConT _) = con
appSubst sub (AppT t1 t2) = AppT (appSubst sub t1) (appSubst sub t2)
    
showSubst sub = "{" ++ (concat $ intersperse ", " (map showOneSub sub)) ++ "}"
  where showOneSub (a,b) = show b ++ "/" ++ show a

printGamma gamma = do
  let lines = map' (++",\n") (map showAssign gamma)
  tPutStr "  A = {"
  lines <- if null lines 
             then return lines
             else do
               liftIO $ putStr (head lines)
               return (tail lines)
  sequence (map (\s -> tPutStr ("       " ++ s)) lines)
  liftIO $ putStrLn "}"
  where 
    showAssign (a,b) = show a ++ ": " ++ show b
    map' _ []  = []
    map' f [x] = [x] -- not mapped to last thing
    map' f (x:xs) = f x : (map' f xs)

-- Free variables in ...
-- ...types
fv :: Type -> [TyVar]
fv (ConT _) = []
fv (VarT var) = [var]
fv (AppT t t') = fv t ++ fv t'

-- ... predicates
fv_pred :: Pred -> [TyVar]
fv_pred (Pred _ typs)= nub (concat $ map fv typs)

fv_preds ps = nub (concat (map fv_pred ps))

fv_scheme :: Scheme -> [TyVar]
fv_scheme (Scheme vs ps typ) = nub (fv_preds ps ++ fv typ) \\ vs

-- ... environments
fv_gamma :: Gamma -> [TyVar]
fv_gamma gamma = nub (concatMap (fv_scheme . snd) gamma)


-- Kind preserving unification.
unify :: Type -> Type -> Maybe Subst
unify (VarT v) (VarT v') 
  | v == v' = Just [] -- identity substitution
  | otherwise = Nothing

unify (ConT (Con c k)) (ConT (Con c' k')) 
  | c == c' && k == k' = Just []
  | otherwise = Nothing

unify (VarT v) t 
  | not (elem v (fv t)) = Just (v |-> t)
  | otherwise = Nothing
  
unify t (VarT v) 
  | not (elem v (fv t)) = Just (v |-> t)
  | otherwise = Nothing


unify (AppT t1 t2) (AppT t1' t2')
 | isJust mbU && isAppKind k1 && k1 == k1' 
   && lhs k1 == k2  && lhs k1 == k2' = 
    let u = fromJust mbU
    in  case unify (appSubst u t2) (appSubst u t2') of
         Just u' -> Just (u' .: u)
         Nothing -> Nothing
 | otherwise = Nothing
 where 
   mbU = unify t1 t1'
   k1 = kind t1
   k2 = kind t2
   k1' = kind t1'
   k2' = kind t2'
   isAppKind (a :->: b) = True
   isAppKind _          = False
   lhs (a :->: b) = a 

unify (ConT c) t = Nothing
unify t (ConT c) = Nothing


type Name       = String
type NameSupply = [Name]

type Gamma = [(String, Scheme)]

-- Type inference function type
type TI a = Gamma -> a -> TIM ([Pred], Subst, Type)

--
-- The driver of the type inference algorithm. 
--
typeInf :: Gamma -> Exp -> IO () 
typeInf gamma exp = do
  let io = runTIM (typeInf' gamma exp)
  (ps, _, typ) <- io
  putStrLn (show (ps, typ))

-- | typeInf' dispatches to other type inference functions based on 
--   structure of expression
typeInf' :: TI Exp
typeInf' gamma (VarE v)       = subTIM $ typeInfVar gamma v
typeInf' gamma (AppE e1 e2)   = subTIM $ typeInfApp gamma (e1,e2)
typeInf' gamma (LamE v e)     = subTIM $ typeInfLam gamma (v, e)
typeInf' gamma (LetE v e1 e2) = subTIM $ typeInfLet gamma (v, e1, e2)

-- | Type inference for a variable
typeInfVar :: TI Var
typeInfVar gamma v = do
  case lookup v gamma of 
    Just (Scheme vs ps typ) -> do
      let createSub var@(TyVar nm k) = do
            t <- newTyVarWithKind k
            return (var |-> t)
      subs <- sequence (map createSub vs)
      let sub = foldr (.:) [] subs

      tPutStrLn $ "Rule (var) invoked on: " ++ show v
      printGamma gamma
      tPutStrLn $ "  P         = " ++ show ps 
      tPutStrLn $ "  tau       = " ++ show typ
      tPutStrLn $ "  {b_i/a_i} = " ++ showSubst sub  

      return (subPreds sub ps, [], appSubst sub typ)
    Nothing -> fail (show v ++ " not found in environment")

-- | Type inference for an application
typeInfApp :: TI (Exp, Exp)
typeInfApp gamma (e, f) = do

  (ps, sub, tau) <- typeInf' gamma e
  (qs, sub', tau') <- typeInf' (sub `subGamma` gamma) f
  alpha <- newTyVar
  let u = case unify (appSubst sub' tau) (tau' `to` alpha) of
            Just u' -> u'
            Nothing -> fail ("Cannot unify " ++ show (appSubst sub tau) ++
                             " and " ++ show (tau' `to` alpha))

      newPreds = subPreds u (subPreds sub' ps ++ qs)
      newSub   = u .: sub' .: sub 
      newType  = appSubst u alpha

  tPutStrLn $ "Rule (->E) invoked on: " ++ show e ++ " " ++ show f
  printGamma gamma
  tPutStrLn $ "  alpha = " ++ show alpha
  tPutStrLn $ "  P   = " ++ show ps
  tPutStrLn $ "  T   = " ++ showSubst sub
  tPutStrLn $ "  tau = " ++ show tau

  tPutStrLn $ "  Q   = " ++ show qs
  tPutStrLn $ "  T'  = " ++ showSubst sub'
  tPutStrLn $ "  tau'= " ++ show tau'
  tPutStrLn $ "  Unifying " ++ show (appSubst sub' tau) ++ " and " ++ 
            show (tau' `to` alpha)
  tPutStrLn $ "  U   = " ++ showSubst u

  return ( newPreds, newSub, newType )
	   

-- Removes a variable from the environment
gammaRestrict :: Var -> Gamma -> Gamma
gammaRestrict v gamma = filter (\(x, _) -> x /= v) gamma

-- | Type inference for a lambda abstraction
typeInfLam :: TI (Var, Exp)
typeInfLam gamma (v, e) = do

  alpha <- newTyVar
  let gamma' = (v, Scheme [] [] alpha): gammaRestrict v gamma
  (ps, sub, etyp) <- typeInf' gamma' e

  tPutStrLn $ "Rule (->I) invoked on: \\" ++ show v ++ " -> " ++ show e
  printGamma gamma
  tPutStrLn $ "  alpha = " ++ show alpha

  tPutStrLn $ "  P   = " ++ show ps
  tPutStrLn $ "  T   = " ++ showSubst sub
  tPutStrLn $ "  tau = " ++ show etyp

  return ( ps, sub, (appSubst sub alpha) `to` etyp )

-- | Type inference for a let expression
typeInfLet :: TI (Var, Exp, Exp)
typeInfLet gamma (x, e, f) = do
  (ps, sub, tau) <- typeInf' gamma e
  let sub_gamma = sub `subGamma` gamma
      gen = generalise (sub_gamma, ps, tau)
      gamma' = (x, gen) : gammaRestrict x sub_gamma
  (ps', sub', tau') <- typeInf' gamma' f

  tPutStrLn $ "Rule let invoked on: let " ++ show x ++ " = " ++
            show e ++ " in " ++ show f
  printGamma gamma
  tPutStrLn $ "P   = " ++ show ps
  tPutStrLn $ "T   = " ++ showSubst sub
  tPutStrLn $ "tau = " ++ show tau
  tPutStrLn $ "P'  = " ++ show ps'
  tPutStrLn $ "T'  = " ++ showSubst sub'
  tPutStrLn $ "tau' = " ++ show tau'
  tPutStrLn $ "Gen(TA, P => tau) = " ++ show gen

  return (ps', sub' .: sub, tau')


generalise :: (Gamma, [Pred], Type) -> Scheme 
generalise (gamma, ps, tau) = 
   let vars = nub (fv_preds ps ++ fv tau) \\ fv_gamma gamma
   in Scheme vars ps tau
  
---------


subGamma :: Subst -> Gamma -> Gamma
subGamma sub gamma = mapSnd (subScheme sub) gamma
  where
    mapSnd f = map (\(x,y) -> (x, f y))

subPred :: Subst -> Pred -> Pred
subPred sub (Pred name typs) = Pred name (map (appSubst sub) typs)

subPreds :: Subst -> [Pred] -> [Pred]
subPreds sub ps = map (subPred sub) ps

subScheme :: Subst -> Scheme -> Scheme
subScheme sub (Scheme vs ps typ) = Scheme vs (map (subPred sub) ps) 
                                             (appSubst sub typ)

----------

--
-- Type inference monad
--

type TIMState = (NameSupply, Int {-depth-})

initialTIMState = (newNS, 0)
  where newNS = zipWith (++) (repeat "t") (map show [1..])

type TIM a = StateT TIMState IO a

newName :: TIM Name
newName = do
  (n:ns, depth) <- get
  put (ns, depth)
  return n

-- These new variables have kind of Star.
newTyVar :: TIM Type
newTyVar = fmap (\x -> VarT (TyVar x Star)) newName

newTyVarWithKind :: Kind -> TIM Type
newTyVarWithKind k = fmap (\x -> VarT (TyVar x k)) newName

newTyVars :: Int -> TIM [Type]
newTyVars n = sequence [ newTyVar | x <- [1..n]]

timSpaces :: Int -> String
timSpaces n = take (n*2) (repeat ' ')

tPutStr :: String -> TIM ()
tPutStr str   = do
  (_, depth) <- get
  liftIO (putStr $ timSpaces depth ++ str)

tPutStrLn :: String -> TIM ()
tPutStrLn str = tPutStr (str ++ "\n")


-- Calls something in the type inference monad. First increases depth
-- and then sets it back after the call returns.
subTIM :: TIM a -> TIM a
subTIM tim = do
  (ns, depth) <- get
  put (ns, depth+1)
  result <- tim
  (ns, depth) <- get
  put (ns, depth-1)
  return result

runTIM :: TIM a -> IO a
runTIM tim = do 
  (a, _) <- runStateT tim initialTIMState
  return a