Start adding type inference

Author: emptyflash

package.yaml

@@ -25,6 +25,9 @@ dependencies:
 - text
 - megaparsec >= 7.0.4 && < 7.1.0
 - parser-combinators
+- mtl
+- containers
+- bifunctors
 
 library:
   source-dirs: src

src/Infer.hs

@@ -0,0 +1,225 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE TypeSynonymInstances #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+module Infer where
+
+import Prelude hiding (foldr)
+
+import Type
+import Syntax
+
+import Control.Monad.State
+import Control.Monad.Except
+
+import Data.Monoid
+import Data.List (nub)
+import Data.Foldable (foldr)
+import qualified Data.Map as Map
+import qualified Data.Set as Set
+
+newtype TypeEnv = TypeEnv (Map.Map Var Scheme)
+  deriving (Semigroup, Monoid, Show)
+
+
+data Unique = Unique { count :: Int }
+
+type Infer = ExceptT TypeError (State Unique)
+type Subst = Map.Map TVar Type
+
+data TypeError
+  = UnificationFail Type Type
+  | InfiniteType TVar Type
+  | UnboundVariable String
+  deriving (Show, Eq)
+
+runInfer :: Infer (Subst, Type) -> Either TypeError Scheme
+runInfer m = case evalState (runExceptT m) initUnique of
+  Left err  -> Left err
+  Right res -> Right $ closeOver res
+
+closeOver :: (Map.Map TVar Type, Type) -> Scheme
+closeOver (sub, ty) = normalize sc
+  where sc = generalize emptyTyenv (apply sub ty)
+
+initUnique :: Unique
+initUnique = Unique { count = 0 }
+
+extend :: TypeEnv -> (Var, Scheme) -> TypeEnv
+extend (TypeEnv env) (x, s) = TypeEnv $ Map.insert x s env
+
+emptyTyenv :: TypeEnv
+emptyTyenv = TypeEnv Map.empty
+
+typeof :: TypeEnv -> Var -> Maybe Type.Scheme
+typeof (TypeEnv env) name = Map.lookup name env
+
+class Substitutable a where
+  apply :: Subst -> a -> a
+  ftv   :: a -> Set.Set TVar
+
+instance Substitutable Type where
+  apply _ (TCon a)       = TCon a
+  apply s t@(TVar a)     = Map.findWithDefault t a s
+  apply s (t1 `TArr` t2) = apply s t1 `TArr` apply s t2
+
+  ftv TCon{}         = Set.empty
+  ftv (TVar a)       = Set.singleton a
+  ftv (t1 `TArr` t2) = ftv t1 `Set.union` ftv t2
+
+instance Substitutable Scheme where
+  apply s (Forall as t)   = Forall as $ apply s' t
+                            where s' = foldr Map.delete s as
+  ftv (Forall as t) = ftv t `Set.difference` Set.fromList as
+
+instance Substitutable a => Substitutable [a] where
+  apply = fmap . apply
+  ftv   = foldr (Set.union . ftv) Set.empty
+
+instance Substitutable TypeEnv where
+  apply s (TypeEnv env) =  TypeEnv $ Map.map (apply s) env
+  ftv (TypeEnv env) = ftv $ Map.elems env
+
+
+nullSubst :: Subst
+nullSubst = Map.empty
+
+compose :: Subst -> Subst -> Subst
+s1 `compose` s2 = Map.map (apply s1) s2 `Map.union` s1
+
+unify ::  Type -> Type -> Infer Subst
+unify (l `TArr` r) (l' `TArr` r')  = do
+  s1 <- unify l l'
+  s2 <- unify (apply s1 r) (apply s1 r')
+  return (s2 `compose` s1)
+
+unify (TVar a) t = bind a t
+unify t (TVar a) = bind a t
+unify (TCon a) (TCon b) | a == b = return nullSubst
+unify t1 t2 = throwError $ UnificationFail t1 t2
+
+bind ::  TVar -> Type -> Infer Subst
+bind a t
+  | t == TVar a     = return nullSubst
+  | occursCheck a t = throwError $ InfiniteType a t
+  | otherwise       = return $ Map.singleton a t
+
+occursCheck ::  Substitutable a => TVar -> a -> Bool
+occursCheck a t = a `Set.member` ftv t
+
+letters :: [String]
+letters = [1..] >>= flip replicateM ['a'..'z']
+
+fresh :: Infer Type
+fresh = do
+  s <- get
+  put s{count = count s + 1}
+  return $ TVar $ TV (letters !! count s)
+
+instantiate ::  Scheme -> Infer Type
+instantiate (Forall as t) = do
+  as' <- mapM (const fresh) as
+  let s = Map.fromList $ zip as as'
+  return $ apply s t
+
+generalize :: TypeEnv -> Type -> Scheme
+generalize env t  = Forall as t
+  where as = Set.toList $ ftv t `Set.difference` ftv env
+
+ops :: Binop -> Type
+ops Add = typeInt `TArr` typeInt `TArr` typeInt
+ops Mul = typeInt `TArr` typeInt `TArr` typeInt
+ops Sub = typeInt `TArr` typeInt `TArr` typeInt
+ops Eql = typeInt `TArr` typeInt `TArr` typeBool
+
+lookupEnv :: TypeEnv -> Var -> Infer (Subst, Type)
+lookupEnv (TypeEnv env) x =
+  case Map.lookup x env of
+    Nothing -> throwError $ UnboundVariable (show x)
+    Just s  -> do t <- instantiate s
+                  return (nullSubst, t)
+
+extendDecl :: TypeEnv -> Decl -> Infer (Subst, TypeEnv)
+extendDecl env (name, e) = do
+  (s, t) <- infer env e
+  let env' = apply s env
+      t'   = generalize env' t
+  pure $ (s, env' `extend` (name, t'))
+
+extendDecls :: TypeEnv -> [Decl] -> Infer (Subst, TypeEnv)
+extendDecls env =
+  foldM step (nullSubst, env)
+  where
+    step (s, e) decl = do
+      (s1, e2) <- extendDecl e decl
+      pure (s1 `compose` s, e2)
+
+infer :: TypeEnv -> Expr -> Infer (Subst, Type)
+infer env ex = case ex of
+
+  Var x -> lookupEnv env x
+
+  Lam x e -> do
+    tv <- fresh
+    let env' = env `extend` (x, Forall [] tv)
+    (s1, t1) <- infer env' e
+    return (s1, apply s1 tv `TArr` t1)
+
+  App e1 e2 -> do
+    tv <- fresh
+    (s1, t1) <- infer env e1
+    (s2, t2) <- infer (apply s1 env) e2
+    s3       <- unify (apply s2 t1) (TArr t2 tv)
+    return (s3 `compose` s2 `compose` s1, apply s3 tv)
+
+  Let decls e2 -> do
+    (s1, env') <- extendDecls env decls
+    (s2, t2) <- infer env' e2
+    return (s2 `compose` s1, t2)
+
+  If cond tr fl -> do
+    tv <- fresh
+    inferPrim env [cond, tr, fl] (typeBool `TArr` tv `TArr` tv `TArr` tv)
+
+  Op op e1 e2 -> do
+    inferPrim env [e1, e2] (ops op)
+
+  Lit (LInt _)  -> return (nullSubst, typeInt)
+  Lit (LBool _) -> return (nullSubst, typeBool)
+  Lit (LFloat _) -> return (nullSubst, typeFloat)
+
+inferPrim :: TypeEnv -> [Expr] -> Type -> Infer (Subst, Type)
+inferPrim env l t = do
+  tv <- fresh
+  (s1, tf) <- foldM inferStep (nullSubst, id) l
+  s2 <- unify (apply s1 (tf tv)) t
+  return (s2 `compose` s1, apply s2 tv)
+  where
+  inferStep (s, tf) exp = do
+    (s', t) <- infer (apply s env) exp
+    return (s' `compose` s, tf . (TArr t))
+
+inferExpr :: TypeEnv -> Expr -> Either TypeError Scheme
+inferExpr env = runInfer . infer env
+
+inferTop :: TypeEnv -> [Decl] -> Either TypeError TypeEnv
+inferTop env [] = Right env
+inferTop env ((name, ex):xs) = case inferExpr env ex of
+  Left err -> Left err
+  Right ty -> inferTop (extend env (name, ty)) xs
+
+normalize :: Scheme -> Scheme
+normalize (Forall ts body) = Forall (fmap snd ord) (normtype body)
+  where
+    ord = zip (nub $ fv body) (fmap TV letters)
+
+    fv (TVar a)   = [a]
+    fv (TArr a b) = fv a ++ fv b
+    fv (TCon _)   = []
+
+    normtype (TArr a b) = TArr (normtype a) (normtype b)
+    normtype (TCon a)   = TCon a
+    normtype (TVar a)   =
+      case lookup a ord of
+        Just x -> TVar x
+        Nothing -> error "type variable not in signature"

src/Lexer.hs

@@ -55,6 +55,9 @@ symbol = L.symbol spaceConsumer
 integer :: Parser Integer
 integer = lexeme L.decimal
 
+float :: Parser Float
+float = lexeme L.float
+
 reserved :: String -> Parser ()
 reserved w = (lexeme . try) (string w *> notFollowedBy alphaNumChar)
 

src/Parser.hs

@@ -28,7 +28,12 @@ variable = do
 int :: Parser Expr
 int = do
   n <- L.integer
-  return (Lit (LInt (fromIntegral n)))
+  return (Lit (LInt n))
+
+float :: Parser Expr
+float = do
+  f <- L.float
+  pure . Lit $ LFloat f
 
 bool :: Parser Expr
 bool = (L.reserved "true" >> return (Lit (LBool True)))
@@ -70,14 +75,23 @@ ifthen = do
   fl <- expr
   return (If cond tr fl)
 
+swizzle :: Parser Expr
+swizzle = do
+  name <- L.identifier
+  L.symbol "."
+  op <- L.identifier
+  pure $ Swizzle name op
+
 aexp :: Parser Expr
 aexp =
       L.parens expr
+  <|> try float
   <|> bool
   <|> int
   <|> ifthen
   <|> letin
   <|> lambda
+  <|> try swizzle
   <|> variable
 
 term :: Parser Expr
@@ -87,9 +101,7 @@ term = aexp >>= \x ->
 
 table :: [[Operator Parser Expr]]
 table = 
-    [ [ InfixL (Op Swizzle <$ L.symbol ".")
-      ]
-    , [ InfixL (Op Mul <$ L.symbol "*")
+    [ [ InfixL (Op Mul <$ L.symbol "*")
       , InfixL (Op Div <$ L.symbol "/")
       ]
     , [ InfixL (Op Add <$ L.symbol "+")
@@ -172,5 +184,5 @@ modl = do
 parseExpr :: String -> Either StringError Expr
 parseExpr input = parse expr "<stdin>" input
 
-parseModule ::  FilePath -> String -> Either StringError [(String, Expr)]
+parseModule ::  FilePath -> String -> Either StringError [Decl]
 parseModule fname input = parse modl fname input

src/Syntax.hs

@@ -12,11 +12,14 @@ data Expr
   | If Expr Expr Expr
   | Fix Expr
   | Op Binop Expr Expr
-  | Ty TyLit
+  | Swizzle Var Var
+  | Ty GlslTypes
   deriving (Show, Eq, Ord)
 
-data TyLit
-  = Float
+data GlslTypes
+  = Bool
+  | Int
+  | Float
   | Vec2
   | Vec3
   | Vec4
@@ -28,9 +31,10 @@ data TyLit
 data Lit
   = LInt Integer
   | LBool Bool
+  | LFloat Float
   deriving (Show, Eq, Ord)
 
-data Binop = Add | Sub | Mul | Eql | Div | Swizzle
+data Binop = Add | Sub | Mul | Eql | Div
   deriving (Eq, Ord, Show)
 
 data Program = Program [Decl] Expr deriving (Show, Eq)

src/Type.hs

@@ -0,0 +1,26 @@
+module Type where
+
+import qualified Syntax as S
+
+newtype TVar = TV String
+  deriving (Show, Eq, Ord)
+
+data Type
+  = TVar TVar
+  | TCon S.GlslTypes
+  | TArr Type Type
+  deriving (Show, Eq, Ord)
+
+infixr `TArr`
+
+data Scheme = Forall [TVar] Type
+  deriving (Show, Eq, Ord)
+
+typeInt :: Type
+typeInt  = TCon S.Int
+
+typeFloat :: Type
+typeFloat  = TCon S.Float
+
+typeBool :: Type
+typeBool = TCon S.Bool