use bimap::BiMap; use derive_more::From; use itertools::Itertools; use std::cell::RefCell; use std::collections::HashMap; use std::convert::{TryFrom, TryInto}; use std::fmt::{self, Display}; use std::{mem, result}; use thiserror::Error; use crate::ast::{self, hir, Arg, BinaryOperator, Ident, Literal, Pattern}; use crate::common::env::Env; use crate::common::{Namer, NamerOf}; #[derive(Debug, Error)] pub enum Error { #[error("Undefined variable {0}")] UndefinedVariable(Ident<'static>), #[error("Mismatched types: expected {expected}, but got {actual}")] TypeMismatch { expected: Type, actual: Type }, #[error("Mismatched types, expected numeric type, but got {0}")] NonNumeric(Type), #[error("Ambiguous type {0}")] AmbiguousType(TyVar), } pub type Result<T> = result::Result<T, Error>; #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)] pub struct TyVar(u64); impl Display for TyVar { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "t{}", self.0) } } #[derive(Debug, PartialEq, Eq, Clone, Hash)] pub struct NullaryType(String); impl Display for NullaryType { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(&self.0) } } #[derive(Debug, PartialEq, Eq, Clone, Copy)] pub enum PrimType { Int, Float, Bool, CString, } impl<'a> From<PrimType> for ast::Type<'a> { fn from(pr: PrimType) -> Self { match pr { PrimType::Int => ast::Type::Int, PrimType::Float => ast::Type::Float, PrimType::Bool => ast::Type::Bool, PrimType::CString => ast::Type::CString, } } } impl Display for PrimType { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { PrimType::Int => f.write_str("int"), PrimType::Float => f.write_str("float"), PrimType::Bool => f.write_str("bool"), PrimType::CString => f.write_str("cstring"), } } } #[derive(Debug, PartialEq, Eq, Clone, From)] pub enum Type { #[from(ignore)] Univ(TyVar), #[from(ignore)] Exist(TyVar), Nullary(NullaryType), Prim(PrimType), Tuple(Vec<Type>), Unit, Fun { args: Vec<Type>, ret: Box<Type>, }, } impl<'a> TryFrom<Type> for ast::Type<'a> { type Error = Type; fn try_from(value: Type) -> result::Result<Self, Self::Error> { match value { Type::Unit => Ok(ast::Type::Unit), Type::Univ(_) => todo!(), Type::Exist(_) => Err(value), Type::Nullary(_) => todo!(), Type::Prim(p) => Ok(p.into()), Type::Tuple(members) => Ok(ast::Type::Tuple( members.into_iter().map(|ty| ty.try_into()).try_collect()?, )), Type::Fun { ref args, ref ret } => Ok(ast::Type::Function(ast::FunctionType { args: args .clone() .into_iter() .map(Self::try_from) .try_collect() .map_err(|_| value.clone())?, ret: Box::new((*ret.clone()).try_into().map_err(|_| value.clone())?), })), } } } const INT: Type = Type::Prim(PrimType::Int); const FLOAT: Type = Type::Prim(PrimType::Float); const BOOL: Type = Type::Prim(PrimType::Bool); const CSTRING: Type = Type::Prim(PrimType::CString); impl Display for Type { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Type::Nullary(nt) => nt.fmt(f), Type::Prim(p) => p.fmt(f), Type::Univ(TyVar(n)) => write!(f, "∀{}", n), Type::Exist(TyVar(n)) => write!(f, "∃{}", n), Type::Fun { args, ret } => write!(f, "fn {} -> {}", args.iter().join(", "), ret), Type::Tuple(members) => write!(f, "({})", members.iter().join(", ")), Type::Unit => write!(f, "()"), } } } struct Typechecker<'ast> { ty_var_namer: NamerOf<TyVar>, ctx: HashMap<TyVar, Type>, env: Env<Ident<'ast>, Type>, /// AST type var -> type instantiations: Env<Ident<'ast>, Type>, /// AST type-var -> universal TyVar type_vars: RefCell<(BiMap<Ident<'ast>, TyVar>, NamerOf<Ident<'static>>)>, } impl<'ast> Typechecker<'ast> { fn new() -> Self { Self { ty_var_namer: Namer::new(TyVar).boxed(), type_vars: RefCell::new(( Default::default(), Namer::alphabetic().map(|n| Ident::try_from(n).unwrap()), )), ctx: Default::default(), env: Default::default(), instantiations: Default::default(), } } fn bind_pattern( &mut self, pat: Pattern<'ast>, type_: Type, ) -> Result<hir::Pattern<'ast, Type>> { match pat { Pattern::Id(ident) => { self.env.set(ident.clone(), type_.clone()); Ok(hir::Pattern::Id(ident, type_)) } Pattern::Tuple(members) => { let mut tys = Vec::with_capacity(members.len()); let mut hir_members = Vec::with_capacity(members.len()); for pat in members { let ty = self.fresh_ex(); hir_members.push(self.bind_pattern(pat, ty.clone())?); tys.push(ty); } let tuple_type = Type::Tuple(tys); self.unify(&tuple_type, &type_)?; Ok(hir::Pattern::Tuple(hir_members)) } } } pub(crate) fn tc_expr(&mut self, expr: ast::Expr<'ast>) -> Result<hir::Expr<'ast, Type>> { match expr { ast::Expr::Ident(ident) => { let type_ = self .env .resolve(&ident) .ok_or_else(|| Error::UndefinedVariable(ident.to_owned()))? .clone(); Ok(hir::Expr::Ident(ident, type_)) } ast::Expr::Literal(lit) => { let type_ = match lit { Literal::Int(_) => Type::Prim(PrimType::Int), Literal::Bool(_) => Type::Prim(PrimType::Bool), Literal::String(_) => Type::Prim(PrimType::CString), Literal::Unit => Type::Unit, }; Ok(hir::Expr::Literal(lit.to_owned(), type_)) } ast::Expr::Tuple(members) => { let members = members .into_iter() .map(|expr| self.tc_expr(expr)) .collect::<Result<Vec<_>>>()?; let type_ = Type::Tuple(members.iter().map(|expr| expr.type_().clone()).collect()); Ok(hir::Expr::Tuple(members, type_)) } ast::Expr::UnaryOp { op, rhs } => todo!(), ast::Expr::BinaryOp { lhs, op, rhs } => { let lhs = self.tc_expr(*lhs)?; let rhs = self.tc_expr(*rhs)?; let type_ = match op { BinaryOperator::Equ | BinaryOperator::Neq => { self.unify(lhs.type_(), rhs.type_())?; Type::Prim(PrimType::Bool) } BinaryOperator::Add | BinaryOperator::Sub | BinaryOperator::Mul => { let ty = self.unify(lhs.type_(), rhs.type_())?; // if !matches!(ty, Type::Int | Type::Float) { // return Err(Error::NonNumeric(ty)); // } ty } BinaryOperator::Div => todo!(), BinaryOperator::Pow => todo!(), }; Ok(hir::Expr::BinaryOp { lhs: Box::new(lhs), op, rhs: Box::new(rhs), type_, }) } ast::Expr::Let { bindings, body } => { self.env.push(); let bindings = bindings .into_iter() .map( |ast::Binding { pat, type_, body }| -> Result<hir::Binding<Type>> { let body = self.tc_expr(body)?; if let Some(type_) = type_ { let type_ = self.type_from_ast_type(type_); self.unify(body.type_(), &type_)?; } let pat = self.bind_pattern(pat, body.type_().clone())?; Ok(hir::Binding { pat, body }) }, ) .collect::<Result<Vec<hir::Binding<Type>>>>()?; let body = self.tc_expr(*body)?; self.env.pop(); Ok(hir::Expr::Let { bindings, type_: body.type_().clone(), body: Box::new(body), }) } ast::Expr::If { condition, then, else_, } => { let condition = self.tc_expr(*condition)?; self.unify(&Type::Prim(PrimType::Bool), condition.type_())?; let then = self.tc_expr(*then)?; let else_ = self.tc_expr(*else_)?; let type_ = self.unify(then.type_(), else_.type_())?; Ok(hir::Expr::If { condition: Box::new(condition), then: Box::new(then), else_: Box::new(else_), type_, }) } ast::Expr::Fun(f) => { let ast::Fun { args, body } = *f; self.env.push(); let args: Vec<_> = args .into_iter() .map(|Arg { ident, type_ }| { let ty = match type_ { Some(t) => self.type_from_ast_type(t), None => self.fresh_ex(), }; self.env.set(ident.clone(), ty.clone()); (ident, ty) }) .collect(); let body = self.tc_expr(body)?; self.env.pop(); Ok(hir::Expr::Fun { type_: Type::Fun { args: args.iter().map(|(_, ty)| ty.clone()).collect(), ret: Box::new(body.type_().clone()), }, type_args: vec![], // TODO fill in once we do let generalization args, body: Box::new(body), }) } ast::Expr::Call { fun, args } => { let ret_ty = self.fresh_ex(); let arg_tys = args.iter().map(|_| self.fresh_ex()).collect::<Vec<_>>(); let ft = Type::Fun { args: arg_tys.clone(), ret: Box::new(ret_ty.clone()), }; let fun = self.tc_expr(*fun)?; self.instantiations.push(); self.unify(&ft, fun.type_())?; let args = args .into_iter() .zip(arg_tys) .map(|(arg, ty)| { let arg = self.tc_expr(arg)?; self.unify(&ty, arg.type_())?; Ok(arg) }) .try_collect()?; let type_args = self.commit_instantiations(); Ok(hir::Expr::Call { fun: Box::new(fun), type_args, args, type_: ret_ty, }) } ast::Expr::Ascription { expr, type_ } => { let expr = self.tc_expr(*expr)?; let type_ = self.type_from_ast_type(type_); self.unify(expr.type_(), &type_)?; Ok(expr) } } } pub(crate) fn tc_decl( &mut self, decl: ast::Decl<'ast>, ) -> Result<Option<hir::Decl<'ast, Type>>> { match decl { ast::Decl::Fun { name, body } => { let mut expr = ast::Expr::Fun(Box::new(body)); if let Some(type_) = self.env.resolve(&name) { expr = ast::Expr::Ascription { expr: Box::new(expr), type_: self.finalize_type(type_.clone())?, }; } self.env.push(); let body = self.tc_expr(expr)?; let type_ = body.type_().clone(); self.env.set(name.clone(), type_); self.env.pop(); match body { hir::Expr::Fun { type_args, args, body, type_, } => Ok(Some(hir::Decl::Fun { name, type_args, args, body, type_, })), _ => unreachable!(), } } ast::Decl::Ascription { name, type_ } => { let type_ = self.type_from_ast_type(type_); self.env.set(name.clone(), type_); Ok(None) } ast::Decl::Extern { name, type_ } => { let type_ = self.type_from_ast_type(ast::Type::Function(type_)); self.env.set(name.clone(), type_.clone()); let (arg_types, ret_type) = match type_ { Type::Fun { args, ret } => (args, *ret), _ => unreachable!(), }; Ok(Some(hir::Decl::Extern { name, arg_types, ret_type, })) } } } fn fresh_tv(&mut self) -> TyVar { self.ty_var_namer.make_name() } fn fresh_ex(&mut self) -> Type { Type::Exist(self.fresh_tv()) } fn fresh_univ(&mut self) -> Type { Type::Univ(self.fresh_tv()) } fn unify(&mut self, ty1: &Type, ty2: &Type) -> Result<Type> { match (ty1, ty2) { (Type::Unit, Type::Unit) => Ok(Type::Unit), (Type::Exist(tv), ty) | (ty, Type::Exist(tv)) => match self.resolve_tv(*tv)? { Some(existing_ty) if self.types_match(ty, &existing_ty) => Ok(ty.clone()), Some(var @ ast::Type::Var(_)) => { let var = self.type_from_ast_type(var); self.unify(&var, ty) } Some(existing_ty) => match ty { Type::Exist(_) => { let rhs = self.type_from_ast_type(existing_ty); self.unify(ty, &rhs) } _ => Err(Error::TypeMismatch { expected: ty.clone(), actual: self.type_from_ast_type(existing_ty), }), }, None => match self.ctx.insert(*tv, ty.clone()) { Some(existing) => self.unify(&existing, ty), None => Ok(ty.clone()), }, }, (Type::Univ(u1), Type::Univ(u2)) if u1 == u2 => Ok(ty2.clone()), (Type::Univ(u), ty) | (ty, Type::Univ(u)) => { let ident = self.name_univ(*u); match self.instantiations.resolve(&ident) { Some(existing_ty) if ty == existing_ty => Ok(ty.clone()), Some(existing_ty) => Err(Error::TypeMismatch { expected: ty.clone(), actual: existing_ty.clone(), }), None => { self.instantiations.set(ident, ty.clone()); Ok(ty.clone()) } } } (Type::Prim(p1), Type::Prim(p2)) if p1 == p2 => Ok(ty2.clone()), (Type::Tuple(t1), Type::Tuple(t2)) if t1.len() == t2.len() => { let ts = t1 .iter() .zip(t2.iter()) .map(|(t1, t2)| self.unify(t1, t2)) .try_collect()?; Ok(Type::Tuple(ts)) } ( Type::Fun { args: args1, ret: ret1, }, Type::Fun { args: args2, ret: ret2, }, ) => { let args = args1 .iter() .zip(args2) .map(|(t1, t2)| self.unify(t1, t2)) .try_collect()?; let ret = self.unify(ret1, ret2)?; Ok(Type::Fun { args, ret: Box::new(ret), }) } (Type::Nullary(_), _) | (_, Type::Nullary(_)) => todo!(), _ => Err(Error::TypeMismatch { expected: ty1.clone(), actual: ty2.clone(), }), } } fn finalize_expr( &self, expr: hir::Expr<'ast, Type>, ) -> Result<hir::Expr<'ast, ast::Type<'ast>>> { expr.traverse_type(|ty| self.finalize_type(ty)) } fn finalize_decl( &mut self, decl: hir::Decl<'ast, Type>, ) -> Result<hir::Decl<'ast, ast::Type<'ast>>> { let res = decl.traverse_type(|ty| self.finalize_type(ty))?; if let Some(type_) = res.type_() { let ty = self.type_from_ast_type(type_.clone()); self.env.set(res.name().clone(), ty); } Ok(res) } fn finalize_type(&self, ty: Type) -> Result<ast::Type<'static>> { let ret = match ty { Type::Exist(tv) => self.resolve_tv(tv)?.ok_or(Error::AmbiguousType(tv)), Type::Univ(tv) => Ok(ast::Type::Var(self.name_univ(tv))), Type::Unit => Ok(ast::Type::Unit), Type::Nullary(_) => todo!(), Type::Prim(pr) => Ok(pr.into()), Type::Tuple(members) => Ok(ast::Type::Tuple( members .into_iter() .map(|ty| self.finalize_type(ty)) .try_collect()?, )), Type::Fun { args, ret } => Ok(ast::Type::Function(ast::FunctionType { args: args .into_iter() .map(|ty| self.finalize_type(ty)) .try_collect()?, ret: Box::new(self.finalize_type(*ret)?), })), }; ret } fn resolve_tv(&self, tv: TyVar) -> Result<Option<ast::Type<'static>>> { let mut res = &Type::Exist(tv); Ok(loop { match res { Type::Exist(tv) => { res = match self.ctx.get(tv) { Some(r) => r, None => return Ok(None), }; } Type::Univ(tv) => { let ident = self.name_univ(*tv); if let Some(r) = self.instantiations.resolve(&ident) { res = r; } else { break Some(ast::Type::Var(ident)); } } Type::Nullary(_) => todo!(), Type::Prim(pr) => break Some((*pr).into()), Type::Unit => break Some(ast::Type::Unit), Type::Fun { args, ret } => todo!(), Type::Tuple(_) => break Some(self.finalize_type(res.clone())?), } }) } fn type_from_ast_type(&mut self, ast_type: ast::Type<'ast>) -> Type { match ast_type { ast::Type::Unit => Type::Unit, ast::Type::Int => INT, ast::Type::Float => FLOAT, ast::Type::Bool => BOOL, ast::Type::CString => CSTRING, ast::Type::Tuple(members) => Type::Tuple( members .into_iter() .map(|ty| self.type_from_ast_type(ty)) .collect(), ), ast::Type::Function(ast::FunctionType { args, ret }) => Type::Fun { args: args .into_iter() .map(|t| self.type_from_ast_type(t)) .collect(), ret: Box::new(self.type_from_ast_type(*ret)), }, ast::Type::Var(id) => Type::Univ({ let opt_tv = { self.type_vars.borrow_mut().0.get_by_left(&id).copied() }; opt_tv.unwrap_or_else(|| { let tv = self.fresh_tv(); self.type_vars .borrow_mut() .0 .insert_no_overwrite(id, tv) .unwrap(); tv }) }), } } fn name_univ(&self, tv: TyVar) -> Ident<'static> { let mut vars = self.type_vars.borrow_mut(); vars.0 .get_by_right(&tv) .map(Ident::to_owned) .unwrap_or_else(|| { let name = loop { let name = vars.1.make_name(); if !vars.0.contains_left(&name) { break name; } }; vars.0.insert_no_overwrite(name.clone(), tv).unwrap(); name }) } fn commit_instantiations(&mut self) -> HashMap<Ident<'ast>, Type> { let mut res = HashMap::new(); let mut ctx = mem::take(&mut self.ctx); for (_, v) in ctx.iter_mut() { if let Type::Univ(tv) = v { let tv_name = self.name_univ(*tv); if let Some(concrete) = self.instantiations.resolve(&tv_name) { res.insert(tv_name, concrete.clone()); *v = concrete.clone(); } } } self.ctx = ctx; self.instantiations.pop(); res } fn types_match(&self, type_: &Type, ast_type: &ast::Type<'ast>) -> bool { match (type_, ast_type) { (Type::Univ(u), ast::Type::Var(v)) => { Some(u) == self.type_vars.borrow().0.get_by_left(v) } (Type::Univ(_), _) => false, (Type::Exist(_), _) => false, (Type::Unit, ast::Type::Unit) => true, (Type::Unit, _) => false, (Type::Nullary(_), _) => todo!(), (Type::Prim(pr), ty) => ast::Type::from(*pr) == *ty, (Type::Tuple(members), ast::Type::Tuple(members2)) => members .iter() .zip(members2.iter()) .all(|(t1, t2)| self.types_match(t1, t2)), (Type::Tuple(members), _) => false, (Type::Fun { args, ret }, ast::Type::Function(ft)) => { args.len() == ft.args.len() && args .iter() .zip(&ft.args) .all(|(a1, a2)| self.types_match(a1, &a2)) && self.types_match(&*ret, &*ft.ret) } (Type::Fun { .. }, _) => false, } } } pub fn typecheck_expr(expr: ast::Expr) -> Result<hir::Expr<ast::Type>> { let mut typechecker = Typechecker::new(); let typechecked = typechecker.tc_expr(expr)?; typechecker.finalize_expr(typechecked) } pub fn typecheck_toplevel(decls: Vec<ast::Decl>) -> Result<Vec<hir::Decl<ast::Type>>> { let mut typechecker = Typechecker::new(); let mut res = Vec::with_capacity(decls.len()); for decl in decls { if let Some(hir_decl) = typechecker.tc_decl(decl)? { let hir_decl = typechecker.finalize_decl(hir_decl)?; res.push(hir_decl); } typechecker.ctx.clear(); } Ok(res) } #[cfg(test)] mod tests { use super::*; macro_rules! assert_type { ($expr: expr, $type: expr) => { use crate::parser::{expr, type_}; let parsed_expr = test_parse!(expr, $expr); let parsed_type = test_parse!(type_, $type); let res = typecheck_expr(parsed_expr).unwrap_or_else(|e| panic!("{}", e)); assert!( res.type_().alpha_equiv(&parsed_type), "{} inferred type {}, but expected {}", $expr, res.type_(), $type ); }; (toplevel($program: expr), $($decl: ident => $type: expr),+ $(,)?) => {{ use crate::parser::{toplevel, type_}; let program = test_parse!(toplevel, $program); let res = typecheck_toplevel(program).unwrap_or_else(|e| panic!("{}", e)); $( let parsed_type = test_parse!(type_, $type); let ident = Ident::try_from(::std::stringify!($decl)).unwrap(); let decl = res.iter().find(|decl| { matches!(decl, crate::ast::hir::Decl::Fun { name, .. } if name == &ident) }).unwrap_or_else(|| panic!("Could not find declaration for {}", ident)); assert!( decl.type_().unwrap().alpha_equiv(&parsed_type), "inferred type {} for {}, but expected {}", decl.type_().unwrap(), ident, $type ); )+ }}; } macro_rules! assert_type_error { ($expr: expr) => { use crate::parser::expr; let parsed_expr = test_parse!(expr, $expr); let res = typecheck_expr(parsed_expr); assert!( res.is_err(), "Expected type error, but got type: {}", res.unwrap().type_() ); }; } #[test] fn literal_int() { assert_type!("1", "int"); } #[test] fn conditional() { assert_type!("if 1 == 2 then 3 else 4", "int"); } #[test] #[ignore] fn add_bools() { assert_type_error!("true + false"); } #[test] fn call_generic_function() { assert_type!("(fn x = x) 1", "int"); } #[test] fn call_let_bound_generic() { assert_type!("let id = fn x = x in id 1", "int"); } #[test] fn universal_ascripted_let() { assert_type!("let id: fn a -> a = fn x = x in id 1", "int"); } #[test] fn call_generic_function_toplevel() { assert_type!( toplevel( "ty id : fn a -> a fn id x = x fn main = id 0" ), main => "fn -> int", id => "fn a -> a", ); } #[test] #[ignore] fn let_generalization() { assert_type!("let id = fn x = x in if id true then id 1 else 2", "int"); } #[test] fn concrete_function() { assert_type!("fn x = x + 1", "fn int -> int"); } #[test] fn arg_ascriptions() { assert_type!("fn (x: int) = x", "fn int -> int"); } #[test] fn call_concrete_function() { assert_type!("(fn x = x + 1) 2", "int"); } #[test] fn conditional_non_bool() { assert_type_error!("if 3 then true else false"); } #[test] fn let_int() { assert_type!("let x = 1 in x", "int"); } }