30f1008 | plum

git clone https://git.pyrossh.dev/plum

A statically typed, imperative programming language inspired by rust, python

30f1008 — Peter John 2026-07-19T19:14:09+05:30

feat(plum-checker): implement type checker with unification

Files changed (2) hide show

plum-checker/src/lib.rs +290 -0
plum-checker/tests/checker_tests.rs +25 -0

plum-checker/src/lib.rs CHANGED Viewed

@@ -1 +1,291 @@
 pub mod types;
+use types::{PlumType, TypeEnv, TypeScheme, CheckError, CheckResult};
+use plum_core::ast;
+pub fn plum_type_from_ast(ty: &ast::Type) -> PlumType {
+    match ty.name.as_str() {
+        "Int" => PlumType::TInt,
+        "Float" => PlumType::TFloat,
+        "Bool" => PlumType::TBool,
+        "Str" => PlumType::TStr,
+        "Unit" => PlumType::TUnit,
+        other => PlumType::TNamed(other.to_string()),
+    }
+}
+pub fn unify(t1: &PlumType, t2: &PlumType) -> Result<(), String> {
+    match (t1, t2) {
+        (PlumType::TVar(_), _) | (_, PlumType::TVar(_)) => Ok(()),
+        (PlumType::TInt, PlumType::TInt) => Ok(()),
+        (PlumType::TFloat, PlumType::TFloat) => Ok(()),
+        (PlumType::TBool, PlumType::TBool) => Ok(()),
+        (PlumType::TStr, PlumType::TStr) => Ok(()),
+        (PlumType::TUnit, PlumType::TUnit) => Ok(()),
+        (PlumType::TNamed(a), PlumType::TNamed(b)) if a == b => Ok(()),
+        (PlumType::TFun(ps1, r1), PlumType::TFun(ps2, r2)) if ps1.len() == ps2.len() => {
+            for (p1, p2) in ps1.iter().zip(ps2.iter()) {
+                unify(p1, p2)?;
+            }
+            unify(r1, r2)
+        }
+        _ => Err(format!("type mismatch: expected {}, found {}", t1, t2)),
+    }
+}
+fn lookup(env: &TypeEnv, name: &str) -> Result<PlumType, String> {
+    env.get(name)
+        .map(|s| *s.body.clone())
+        .ok_or_else(|| format!("undefined name '{}'", name))
+}
+pub fn check_source(source: &ast::Source) -> CheckResult<()> {
+    let mut errors: Vec<CheckError> = Vec::new();
+    let mut global_env: TypeEnv = TypeEnv::new();
+    // First pass: register all top-level function signatures and consts
+    for item in &source.items {
+        match item {
+            ast::Item::Fn(f) => {
+                let param_types: Vec<PlumType> = f.params.iter().map(|p| {
+                    match &p.ty {
+                        ast::ParamType::Type(t) => plum_type_from_ast(t),
+                        ast::ParamType::Variadic(t) => plum_type_from_ast(t),
+                    }
+                }).collect();
+                let ret = f.returns.as_ref()
+                    .map(|r| PlumType::TNamed(r.name.clone()))
+                    .unwrap_or(PlumType::TUnit);
+                let scheme = TypeScheme::mono(PlumType::TFun(param_types, Box::new(ret)));
+                global_env.insert(f.name.clone(), scheme);
+            }
+            ast::Item::Const(c) => {
+                global_env.insert(c.name.clone(), TypeScheme::mono(PlumType::TVar("_".to_string())));
+            }
+            _ => {}
+        }
+    }
+    // Second pass: check each function body
+    for item in &source.items {
+        if let ast::Item::Fn(f) = item {
+            let mut local_errors = check_fn(f, &global_env);
+            errors.append(&mut local_errors);
+        }
+    }
+    if errors.is_empty() { Ok(()) } else { Err(errors) }
+}
+fn check_fn(f: &ast::Fn, global_env: &TypeEnv) -> Vec<CheckError> {
+    let mut errors = Vec::new();
+    let mut env = global_env.clone();
+    // Add params to env
+    for p in &f.params {
+        let ty = match &p.ty {
+            ast::ParamType::Type(t) => plum_type_from_ast(t),
+            ast::ParamType::Variadic(t) => plum_type_from_ast(t),
+        };
+        env.insert(p.name.clone(), TypeScheme::mono(ty));
+    }
+    let declared_ret = f.returns.as_ref()
+        .map(|r| {
+            let ast_ty = ast::Type { name: r.name.clone(), generics: vec![] };
+            plum_type_from_ast(&ast_ty)
+        })
+        .unwrap_or(PlumType::TUnit);
+    match &f.body {
+        ast::FnBody::Expr(e) => {
+            match infer_expr(e, &env) {
+                Ok(t) => {
+                    if let Err(msg) = unify(&declared_ret, &t) {
+                        errors.push(CheckError { message: format!("fn '{}': return type mismatch: {}", f.name, msg) });
+                    }
+                }
+                Err(msg) => errors.push(CheckError { message: format!("fn '{}': {}", f.name, msg) }),
+            }
+        }
+        ast::FnBody::Block(block) => {
+            let mut block_errors = check_block(block, &mut env, &declared_ret, &f.name);
+            errors.append(&mut block_errors);
+        }
+    }
+    errors
+}
+fn check_block(block: &ast::Block, env: &mut TypeEnv, declared_ret: &PlumType, fn_name: &str) -> Vec<CheckError> {
+    let mut errors = Vec::new();
+    for stmt in &block.stmts {
+        let mut stmt_errors = check_stmt(stmt, env, declared_ret, fn_name);
+        errors.append(&mut stmt_errors);
+    }
+    errors
+}
+fn check_stmt(stmt: &ast::Stmt, env: &mut TypeEnv, declared_ret: &PlumType, fn_name: &str) -> Vec<CheckError> {
+    let mut errors = Vec::new();
+    match stmt {
+        ast::Stmt::Assign(a) => {
+            for (target, value) in a.targets.iter().zip(a.values.iter()) {
+                match infer_expr(value, env) {
+                    Ok(t) => { env.insert(target.clone(), TypeScheme::mono(t)); }
+                    Err(msg) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, target, msg) }),
+                }
+            }
+        }
+        ast::Stmt::Return(Some(e)) => {
+            match infer_expr(e, env) {
+                Ok(t) => {
+                    if let Err(msg) = unify(declared_ret, &t) {
+                        errors.push(CheckError { message: format!("fn '{}': return type mismatch: {}", fn_name, msg) });
+                    }
+                }
+                Err(msg) => errors.push(CheckError { message: format!("fn '{}': return: {}", fn_name, msg) }),
+            }
+        }
+        ast::Stmt::Return(None) => {
+            if let Err(msg) = unify(declared_ret, &PlumType::TUnit) {
+                errors.push(CheckError { message: format!("fn '{}': bare return in non-Unit function: {}", fn_name, msg) });
+            }
+        }
+        ast::Stmt::If(if_) => {
+            match infer_expr(&if_.condition, env) {
+                Ok(t) => {
+                    if let Err(msg) = unify(&PlumType::TBool, &t) {
+                        errors.push(CheckError { message: format!("fn '{}': if condition must be Bool: {}", fn_name, msg) });
+                    }
+                }
+                Err(msg) => errors.push(CheckError { message: format!("fn '{}': if condition: {}", fn_name, msg) }),
+            }
+            errors.append(&mut check_block(&if_.body, env, declared_ret, fn_name));
+            for ei in &if_.else_ifs {
+                match infer_expr(&ei.condition, env) {
+                    Ok(t) => {
+                        if let Err(msg) = unify(&PlumType::TBool, &t) {
+                            errors.push(CheckError { message: format!("fn '{}': else if condition must be Bool: {}", fn_name, msg) });
+                        }
+                    }
+                    Err(msg) => errors.push(CheckError { message: format!("fn '{}': else if condition: {}", fn_name, msg) }),
+                }
+                errors.append(&mut check_block(&ei.body, env, declared_ret, fn_name));
+            }
+            if let Some(else_block) = &if_.else_ {
+                errors.append(&mut check_block(else_block, env, declared_ret, fn_name));
+            }
+        }
+        ast::Stmt::While(w) => {
+            match infer_expr(&w.condition, env) {
+                Ok(t) => {
+                    if let Err(msg) = unify(&PlumType::TBool, &t) {
+                        errors.push(CheckError { message: format!("fn '{}': while condition must be Bool: {}", fn_name, msg) });
+                    }
+                }
+                Err(msg) => errors.push(CheckError { message: format!("fn '{}': while condition: {}", fn_name, msg) }),
+            }
+            errors.append(&mut check_block(&w.body, env, declared_ret, fn_name));
+        }
+        ast::Stmt::For(f_stmt) => {
+            match infer_expr(&f_stmt.iter, env) {
+                Ok(_) => {}
+                Err(msg) => errors.push(CheckError { message: format!("fn '{}': for iter: {}", fn_name, msg) }),
+            }
+            let mut inner_env = env.clone();
+            for var in &f_stmt.vars {
+                inner_env.insert(var.clone(), TypeScheme::mono(PlumType::TInt));
+            }
+            errors.append(&mut check_block(&f_stmt.body, &mut inner_env, declared_ret, fn_name));
+        }
+        ast::Stmt::Expr(e) => {
+            if let Err(msg) = infer_expr(e, env) {
+                errors.push(CheckError { message: format!("fn '{}': {}", fn_name, msg) });
+            }
+        }
+        ast::Stmt::Assert(e) => {
+            match infer_expr(e, env) {
+                Ok(t) => {
+                    if let Err(msg) = unify(&PlumType::TBool, &t) {
+                        errors.push(CheckError { message: format!("fn '{}': assert must be Bool: {}", fn_name, msg) });
+                    }
+                }
+                Err(msg) => errors.push(CheckError { message: format!("fn '{}': assert: {}", fn_name, msg) }),
+            }
+        }
+        ast::Stmt::Break | ast::Stmt::Continue | ast::Stmt::Todo => {}
+        ast::Stmt::Match(_) => {}
+    }
+    errors
+}
+fn infer_expr(expr: &ast::Expr, env: &TypeEnv) -> Result<PlumType, String> {
+    match expr {
+        ast::Expr::Int(_) => Ok(PlumType::TInt),
+        ast::Expr::Float(_) => Ok(PlumType::TFloat),
+        ast::Expr::String(_) => Ok(PlumType::TStr),
+        ast::Expr::Var(name) => lookup(env, name),
+        ast::Expr::Self_ => Ok(PlumType::TVar("Self".to_string())),
+        ast::Expr::TypeName(n) => Ok(PlumType::TNamed(n.clone())),
+        ast::Expr::Paren(inner) => infer_expr(inner, env),
+        ast::Expr::Not(inner) => {
+            let t = infer_expr(inner, env)?;
+            unify(&PlumType::TBool, &t)?;
+            Ok(PlumType::TBool)
+        }
+        ast::Expr::Unary(u) => infer_expr(&u.operand, env),
+        ast::Expr::Binary(b) => {
+            let lt = infer_expr(&b.left, env)?;
+            let rt = infer_expr(&b.right, env)?;
+            unify(&lt, &rt).map_err(|e| format!("binary op: {}", e))?;
+            match b.op {
+                ast::BinOp::Range => Ok(PlumType::TNamed("Range".to_string())),
+                _ => Ok(lt),
+            }
+        }
+        ast::Expr::Bool(b) => {
+            let lt = infer_expr(&b.left, env)?;
+            let rt = infer_expr(&b.right, env)?;
+            unify(&PlumType::TBool, &lt).map_err(|e| format!("bool op left: {}", e))?;
+            unify(&PlumType::TBool, &rt).map_err(|e| format!("bool op right: {}", e))?;
+            Ok(PlumType::TBool)
+        }
+        ast::Expr::Compare(c) => {
+            let lt = infer_expr(&c.left, env)?;
+            let rt = infer_expr(&c.right, env)?;
+            unify(&lt, &rt).map_err(|e| format!("compare op: {}", e))?;
+            Ok(PlumType::TBool)
+        }
+        ast::Expr::Ternary(t) => {
+            let ct = infer_expr(&t.condition, env)?;
+            unify(&PlumType::TBool, &ct).map_err(|e| format!("ternary condition: {}", e))?;
+            let tt = infer_expr(&t.then, env)?;
+            let et = infer_expr(&t.else_, env)?;
+            unify(&tt, &et).map_err(|e| format!("ternary branches: {}", e))?;
+            Ok(tt)
+        }
+        ast::Expr::FnCall(call) => {
+            match lookup(env, &call.name) {
+                Ok(PlumType::TFun(param_types, ret)) => {
+                    if call.args.len() != param_types.len() {
+                        return Err(format!("call '{}': expected {} args, got {}", call.name, param_types.len(), call.args.len()));
+                    }
+                    for (i, (arg, expected)) in call.args.iter().zip(param_types.iter()).enumerate() {
+                        let arg_expr = match arg {
+                            ast::Arg::Positional(e) => e,
+                            ast::Arg::Keyword { value, .. } => value,
+                            ast::Arg::Pair { value, .. } => value,
+                        };
+                        let actual = infer_expr(arg_expr, env)?;
+                        unify(expected, &actual).map_err(|e| format!("call '{}' arg {}: {}", call.name, i, e))?;
+                    }
+                    Ok(*ret)
+                }
+                Ok(_) => Err(format!("'{}' is not a function", call.name)),
+                Err(_) => Ok(PlumType::TVar("_".to_string())), // unknown fn: allow, codegen will catch
+            }
+        }
+        ast::Expr::ClassCall(_) => Ok(PlumType::TVar("_".to_string())),
+        ast::Expr::Attribute(_) => Ok(PlumType::TVar("_".to_string())),
+    }
+}

plum-checker/tests/checker_tests.rs CHANGED Viewed

@@ -1,4 +1,29 @@
 use plum_checker::types::*;
+use plum_checker::{plum_type_from_ast, unify};
+use plum_core::ast::Type as AstType;
+#[test]
+fn ast_type_int_maps_to_tint() {
+    let ast_ty = AstType { name: "Int".to_string(), generics: vec![] };
+    assert_eq!(plum_type_from_ast(&ast_ty), PlumType::TInt);
+}
+#[test]
+fn ast_type_unknown_maps_to_named() {
+    let ast_ty = AstType { name: "MyClass".to_string(), generics: vec![] };
+    assert_eq!(plum_type_from_ast(&ast_ty), PlumType::TNamed("MyClass".to_string()));
+}
+#[test]
+fn unify_same_types_ok() {
+    assert!(unify(&PlumType::TInt, &PlumType::TInt).is_ok());
+    assert!(unify(&PlumType::TFloat, &PlumType::TFloat).is_ok());
+}
+#[test]
+fn unify_different_types_err() {
+    assert!(unify(&PlumType::TInt, &PlumType::TFloat).is_err());
+}
 #[test]
 fn fresh_vars_are_unique() {