module Type_Checking where

data Type = TVar String
          | TFun Type Type
  deriving Eq

instance Show Type where
  show (TVar name) = "'" ++ name
  show (TFun a r) =
    case a of
      TVar name -> show a ++ " -> " ++ show r
      TFun _ _ -> "(" ++ show a ++ ") -> " ++ show r

data Term = Var String
          | App Term Term
          | Abs String Type Term
  deriving Eq

instance Show Term where
  show (Var name) = name
  show (App fun arg) = f ++ " " ++ a
    where
      f = case fun of
        Abs _ _ _ -> "(" ++ show fun ++ ")"
        _ -> show fun
      a = case arg of
        App _ _ -> "(" ++ show arg ++ ")"
        _ -> show arg
  show (Abs arg ty t) = "\\" ++ arg ++ "::" ++ show ty ++ ". " ++ show t 

{-- Typed variant of subst and beta reduction --}

remove x xs = filter (\y -> y /= x) xs

free_vars :: Term -> [String]
free_vars (Var s) = [s]
free_vars (App t1 t2) = free_vars t1 ++ free_vars t2
free_vars (Abs n ty t) = remove n (free_vars t)

new_name :: String -> [String] -> String
new_name n xs =
  if n `elem` xs
    then new_name (n ++ "'") xs
    else n

subst :: String -> Term -> Term -> Term
subst n s (Var n') = if n == n' then s else Var n'
subst n s t'@(Abs n' ty t) = Abs n'' ty (subst n s (subst n' (Var n'') t))
  where
    n'' = new n' (frees s ++ remove n' (frees t) ++ [n])
subst n s (App t1 t2) = App (subst n s t1) (subst n s t2)

beta_redux :: Term -> Term
beta_redux t@(Var _) = t
beta_redux t@(Abs _ _ _) = t
beta_redux (App t@(App _ _) t') = App (beta_redux t) t'
beta_redux (App (Abs n ty t) t') =
  subst n t' t

{-- Representing certain data elements in lambda calculus --}

id_t :: Type -> Term
id_t t = Abs "x" t (Var "x")

true_t :: Type -> Type -> Term
true_t t e = Abs "t" t (Abs "e" e (Var "t"))

false_t :: Type -> Type -> Term
false_t t e = Abs "t" t (Abs "e" e (Var "e"))

pair_ty :: Type -> Type -> Type -> Type
pair_ty f s x = TFun (TFun f (TFun s x)) x

pair_t :: Ctxt -> Term -> Term -> Type -> Term
pair_t ctxt f s x_ty =
  Abs "f" (TFun f_ty (TFun s_ty x_ty)) (App (App (Var "f") f) s)
  where
    Just f_ty = typing ctxt f
    Just s_ty = typing ctxt s


{-- Type checking --}

type Ctxt = [(String, Type)]

empty_ctxt :: Ctxt
empty_ctxt = []

update_ctxt :: String -> Type -> Ctxt -> Ctxt
update_ctxt name ty ctxt = (name, ty) : ctxt

lookup_var :: Ctxt -> String -> Maybe Type
lookup_var [] name = Nothing
lookup_var ((name', ty):ctxt) name = 
  if name == name'
    then Just ty
    else lookup_var ctxt name

typing :: Ctxt -> Term -> Maybe Type
typing ctxt (Var name) = lookup_var ctxt name
typing ctxt (App t1 t2) =
    case typing ctxt t1 of
      Just (TFun arg res) ->
        if Just arg == typing ctxt t2
          then Just res
          else Nothing
      _ -> Nothing
typing ctxt (Abs var arg t) =
  case typing (update_ctxt var arg ctxt) t of
    Just ty -> Just (TFun arg ty)
    Nothing -> Nothing