stlc

main.rs

use std::convert::TryFrom;
use std::rc::Rc;

use symbol::Symbol;

#[derive(Debug)]
enum Error {
    UnboundName(Symbol),
    Mismatch(Rc<Type>, Rc<Type>),
    NotArrow,
}

#[derive(Debug)]
enum Expr {
    Abs(Symbol, Rc<Type>, Rc<Expr>),
    App(Rc<Expr>, Rc<Expr>),
    Atom(Symbol),
    Var(Symbol),
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum Type {
    Atom,
    Arrow(Rc<Type>, Rc<Type>),
}

#[derive(Debug)]
enum TypedExpr {
    Abs(Rc<TypedExpr>, Rc<Type>),
    App(Rc<TypedExpr>, Rc<TypedExpr>, Rc<Type>),
    Atom(Symbol, Rc<Type>),
    Var(usize, Rc<Type>),
}

impl TypedExpr {
    pub fn get_type(&self) -> Rc<Type> {
        match self {
            TypedExpr::Abs(_, ty)
            | TypedExpr::App(_, _, ty)
            | TypedExpr::Atom(_, ty)
            | TypedExpr::Var(_, ty) => ty.clone(),
        }
    }
}

impl TryFrom<Expr> for TypedExpr {
    type Error = Error;

    fn try_from(expr: Expr) -> Result<Self, Self::Error> {
        fn typeck(expr: Rc<Expr>, ctx: &mut Vec<(Symbol, Rc<Type>)>) -> Result<TypedExpr, Error> {
            match expr.as_ref() {
                Expr::Abs(arg, arg_ty, body) => {
                    ctx.push((*arg, arg_ty.clone()));
                    let typed_body = typeck(body.clone(), ctx)?;
                    ctx.pop();
                    let ty = Type::Arrow(arg_ty.clone(), typed_body.get_type());
                    Ok(TypedExpr::Abs(Rc::new(typed_body), Rc::new(ty)))
                }
                Expr::App(func, arg) => {
                    let typed_func = typeck(func.clone(), ctx)?;
                    let (func_arg_ty, func_body_ty) = match typed_func.get_type().as_ref() {
                        Type::Atom => return Err(Error::NotArrow),
                        Type::Arrow(left, right) => (left.clone(), right.clone()),
                    };
                    let typed_arg = typeck(arg.clone(), ctx)?;
                    let arg_ty = typed_arg.get_type();
                    if func_arg_ty.as_ref() != arg_ty.as_ref() {
                        return Err(Error::Mismatch(func_arg_ty.clone(), arg_ty.clone()));
                    }
                    Ok(TypedExpr::App(
                        Rc::new(typed_func),
                        Rc::new(typed_arg),
                        func_body_ty.clone(),
                    ))
                }
                Expr::Atom(atom_name) => Ok(TypedExpr::Atom(*atom_name, Rc::new(Type::Atom))),
                Expr::Var(var_name) => ctx
                    .iter()
                    .rev()
                    .enumerate()
                    .find(|(_, (name, _))| name == var_name)
                    .map(|(position, (_, ty))| TypedExpr::Var(position, ty.clone()))
                    .ok_or_else(|| Error::UnboundName(*var_name)),
            }
        }

        let mut ctx = Vec::new();
        typeck(Rc::new(expr), &mut ctx)
    }
}

fn main() {
    let expr = Expr::App(
        Rc::new(Expr::Abs(
            Symbol::from("x"),
            Rc::new(Type::Atom),
            Rc::new(Expr::Var(Symbol::from("x"))),
        )),
        Rc::new(Expr::Atom(Symbol::from("foo"))),
    );

    println!("{:?}", TypedExpr::try_from(expr));
}