Algw.hs

{-# Language TypeSynonymInstances, FlexibleInstances #-}

module Algw where

import Data.Maybe
import qualified Data.Map as M
import qualified Data.Set as S
import Control.Monad.Trans.State

import Text.Parsec hiding (State, token)
import Text.Parsec.Char

-----------------------------------------------------------------------------------
-----------------------------------------------------------------------------------
--                              AST
-----------------------------------------------------------------------------------
-----------------------------------------------------------------------------------


data Exp 
  = Var String
  | Int Int
  | Pair Exp Exp
  | App Exp Exp
  | Fun String Exp
  | Let String Exp Exp
  deriving (Eq, Show)

data Type
  = TVar String
  | TInt
  | TPair Type Type
  | TFun Type Type
  deriving Eq

data Scheme = Forall (S.Set String) Type
              deriving (Eq)

instance Show Type where
  show (TVar x) = x
  show TInt = "Int"
  show (TPair t1 t2) = "(" ++ show t1 ++ ", " ++ show t2 ++ ")"
  show (TFun tyf@(TFun _ _) ty) = "(" ++ show tyf ++ ")" ++ " -> " ++ show ty
  show (TFun ty1 ty2) = show ty1 ++ " → " ++ show ty2

instance Show Scheme where
  show (Forall xs m) = "∀(" ++ unwords (S.toList xs) ++ "). " ++ show m

-----------------------------------------------------------------------------------
-----------------------------------------------------------------------------------
--                              Parser
-----------------------------------------------------------------------------------
-----------------------------------------------------------------------------------

-- e = e e |  |  | (e)
-- fact = n | x | fn(x) e | let x = e in e

type Parser = Parsec String ()

keywords = ["fn", "let", "in"]

token :: Parser a -> Parser a
token p = p <* spaces

sym :: String -> Parser String
sym s = (token . try) (string s)

kw :: String -> Parser String
kw s = (token . try) (string s)

parens p = sym "(" *> p <* sym ")" 

identifier = (token . try) ((many1 letter) >>= check) where
  check s = if s `elem` keywords
              then fail (s ++ " can't be a variable")
              else return s

int = Int . read <$> token (many1 digit)


factor :: Parser Exp
factor =
  int
  <|> Var <$> identifier
  <|> let_ 
  <|> fn
  <|> parens (try pair <|> expr)

fn :: Parser Exp
fn = Fun <$> (kw "fn" *> parens identifier) <*> expr

let_ :: Parser Exp
let_ = Let <$> (kw "let" *> identifier) <*> (sym "=" *> expr) <*> (kw "in" *> expr)

pair :: Parser Exp
pair = Pair <$> expr <*> (sym "," *> expr)

expr :: Parser Exp
expr = foldApp <$> (many1 factor) where
  foldApp (e:es) = foldl App e es

main = spaces *> expr <* eof

testParser :: Parser a -> String -> a
testParser p s = case parse p "Test" s of
  Left err -> error (show err)
  Right x  -> x

-----------------------------------------------------------------------------------
-----------------------------------------------------------------------------------
--                              Type checker
-----------------------------------------------------------------------------------
-----------------------------------------------------------------------------------

type Ti = State Int

type Ctx = M.Map String Scheme
type Subst = M.Map String Type

emptySub :: Subst
emptySub = M.empty

infixr 5 |>
(|>) :: Subst -> Subst ->Subst
s1 |> s2 = M.union s1 ((s1 #) <$> s2)

infix 5 #

class Types t where
  ftv :: t -> S.Set String
  (#) :: Subst -> t -> t

instance Types Type where
  ftv (TVar x) = S.singleton x
  ftv (TPair t1 t2) = S.union (ftv t1) (ftv t2)
  ftv (TFun t1 t2) = S.union (ftv t1) (ftv t2)
  ftv _ = S.empty
  s # (TVar x) = fromMaybe (TVar x) $ M.lookup x s
  s # (TPair t1 t2) = TPair (s # t1) (s # t2)
  s # (TFun t1 t2) = TFun (s # t1) (s # t2)
  s # ty = ty

instance Types Scheme where
  ftv (Forall xs t) = S.difference (ftv t) xs
  s # (Forall xs t) = 
    let s' = M.withoutKeys s xs
    in Forall xs (s' # t)

instance Types Ctx where
  ftv ctx = mconcat (ftv <$> M.elems ctx)
  s # ctx = (s #) <$> ctx

fresh :: Ti Type
fresh = do
  i <- get
  put (i + 1)
  return $ TVar ("a" ++ show i)


instantiate :: Ctx -> Scheme -> Ti Type
instantiate ctx (Forall xs t) = do
  s <- sequenceA (M.fromSet (const fresh) xs)
  return (s # t)

generalise :: Ctx -> Type -> Scheme
generalise ctx t =
  let fvs = S.difference (ftv t) (ftv ctx)
  in Forall fvs t


infer :: Ctx -> Exp -> Ti (Subst, Type)
infer ctx (Int _) = return (emptySub, TInt)
infer ctx (Var x) = 
  case M.lookup x ctx of
    Nothing  -> error $ "Unkown variable " ++ x
    Just sch -> ((,) emptySub) <$> instantiate ctx sch

infer ctx (Pair e1 e2) = do
  (s1, t1) <- infer ctx e1
  (s2, t2) <- infer (s1 # ctx) e2
  return (s1 |> s2, TPair (s2 # t1) t2)

infer ctx (Fun x e) = do
  tv <- fresh
  (s, t) <- infer (M.insert x (Forall S.empty tv) ctx) e
  return (s, TFun (s # tv) t)

infer ctx (App e1 e2) = do
  tres <- fresh
  (s1, t1) <- infer ctx e1
  (s2, t2) <- infer (s1 # ctx) e2
  let s3 = unify (s2 # t1) (TFun t2 tres)
  return (s3 |> s2 |> s1, s3 # tres)

infer ctx (Let x e1 e2) = do
  (s1, t1) <- infer ctx e1
  let ts = generalise (s1 # ctx) t1
  (s2, t2) <- infer (M.insert x ts ctx) e2
  return (s1 |> s2, t2)


unify :: Type -> Type -> Subst
unify TInt TInt = emptySub
unify (TVar x) ty  = varBind x ty
unify ty (TVar x)  = varBind x ty
unify (TPair t1 t2) (TPair u1 u2) = 
  let s1 = unify t1 u1
      s2 = unify (s1 # t2) (s1 # u2)
  in s1 |> s2
unify (TFun t1 t2) (TFun u1 u2) = 
  let s1 = unify t1 u1
      s2 = unify (s1 # t2) (s1 # u2)
  in s1 |> s2
unify t1 t2 = error $ "Can't unify " ++ show t1 ++ " and " ++ show t2


varBind :: String -> Type -> Subst
varBind x ty
    | ty == TVar x           = emptySub
    | x `S.member` (ftv ty) = error "Occurs check failed"
    | otherwise             = M.singleton x ty

typeof :: Exp -> Type
typeof e = 
  let (s, ty) = evalState (infer M.empty e) 0
  in s # ty


typeCheck :: String -> Type
typeCheck = typeof . testParser main