plum

// Functions/methods are named camelCase across this project (matching plum's own

// naming convention), not Rust's idiomatic snake_case — silence the resulting lint.

#![allow(non_snake_case)]

pub mod types;

pub mod monomorphize;

use std::collections::BTreeMap;

use types::{PlumType, TypeEnv, TypeScheme, CheckError, CheckResult};

use plum_core::ast;

/// Maps a bare type name (as written in source: a param/return annotation, or

/// a method's receiver name) to its `PlumType`. Builtin primitives get their

/// dedicated variant; anything else is an unmodeled `TNamed`.

pub fn plumTypeFromName(name: &str) -> PlumType {

    match name {

        "Int" => PlumType::TInt,

        "Float" => PlumType::TFloat,

        "Bool" => PlumType::TBool,

        "Str" => PlumType::TStr,

        "Byte" => PlumType::TByte,

        "[]Byte" => PlumType::TByteSlice,

        "Unit" => PlumType::TUnit,

        other => PlumType::TNamed(other.to_string()),

    }

}

pub fn plumTypeFromAst(ty: &ast::Type) -> PlumType {

    plumTypeFromName(&ty.name)

}

/// The `ctx.methods`/`ctx.classes` receiver name for a value of type `ty`, or

/// `None` if `ty` has no methods (e.g. an unresolved `TVar` or a bare

/// function type). Builtin primitive types (`Int`/`Float`/`Bool`/`Str`) are

/// declared as `type Int = fun ...` etc in `libs/std`, exactly like a class —

/// they're just never `TNamed` at the type level, so this maps them back to

/// the same receiver name string `ctx.methods` is keyed by.

pub fn methodReceiverName(ty: &PlumType) -> Option<String> {

    match ty {

        PlumType::TNamed(name) => Some(name.clone()),

        PlumType::TInt => Some("Int".to_string()),

        PlumType::TFloat => Some("Float".to_string()),

        PlumType::TBool => Some("Bool".to_string()),

        PlumType::TStr => Some("Str".to_string()),

        PlumType::TByte => Some("Byte".to_string()),

        PlumType::TByteSlice => Some("ByteSlice".to_string()),

        _ => None,

    }

}

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::TByte, PlumType::TByte) => Ok(()),

        (PlumType::TByteSlice, PlumType::TByteSlice) => 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)),

    }

}

pub fn lookup(env: &TypeEnv, name: &str) -> Result<PlumType, String> {

    env.get(name)

        .map(|s| *s.body.clone())

        .ok_or_else(|| format!("undefined name '{}'", name))

}

/// Field names and types for every `type ClassName = ...` declaration in the source.

pub type ClassEnv = BTreeMap<String, Vec<(String, PlumType)>>;

/// `(receiver type, method name) -> TFun` for every `name<Receiver>(...)` method.

pub type MethodEnv = BTreeMap<(String, String), PlumType>;

/// Info about one `enum` variant: which enum it belongs to, its 0-based runtime tag

/// (numbering is shared across all of that enum's variants), and its payload field

/// types (empty for a payload-free variant like `Red` or `None`).

#[derive(Debug, Clone, PartialEq)]

pub struct EnumVariantInfo {

    pub enum_name: String,

    pub tag: i32,

    pub field_types: Vec<PlumType>,

    pub values: Vec<ast::Expr>,

}

/// Enum variant name -> its info, e.g. `"True" -> { enum_name: "Bool", tag: 1, field_types: [] }`.

pub type EnumVariants = BTreeMap<String, EnumVariantInfo>;

/// Field names and types for every discriminant enum's shared params (`enum Foo(n: Int) = ...`),

/// keyed by the ENUM's name (not a variant name) — e.g. `"Step" -> [("n", TInt)]`. Field

/// access on a value of this type must load/store at offset `(field_idx + 1) * 8`, NOT

/// `field_idx * 8` like a class — slot 0 is always the variant's tag.

pub type EnumParams = BTreeMap<String, Vec<(String, PlumType)>>;

/// Shared, read-only lookup tables built once from the whole source, threaded through

/// every check/infer call alongside the (mutable, scope-local) `TypeEnv`.

pub struct CheckCtx<'a> {

    pub classes: &'a ClassEnv,

    pub methods: &'a MethodEnv,

    pub enum_variants: &'a EnumVariants,

    pub enum_params: &'a EnumParams,

}

/// Builds the global lookup tables (function/const signatures, class fields, method

/// signatures, enum variants) from a whole source. Shared by `checkSource` and by

/// `plum-wasm-codegen`, which needs the same tables to resolve `self`, field access,

/// and method dispatch during code generation.

pub fn buildGlobalTables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEnv, EnumVariants, EnumParams) {

    let mut global_env: TypeEnv = TypeEnv::new();

    let mut classes: ClassEnv = BTreeMap::new();

    let mut methods: MethodEnv = BTreeMap::new();

    let mut enum_variants: EnumVariants = BTreeMap::new();

    let mut enum_params: EnumParams = BTreeMap::new();

    // `Bool`'s variants are built in (see `inferExpr`'s TypeName handling) rather

    // than requiring every source file to redeclare `enum Bool = | True | False`.

    enum_variants.insert("True".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 1, field_types: vec![], values: vec![] });

    enum_variants.insert("False".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 0, field_types: vec![], values: vec![] });

    // First pass: register class fields and enum variants so later passes can

    // resolve `self.field`, `ClassName(...)`, and bare enum-tag patterns.

    for item in &source.items {

        match item {

            ast::Item::Class(c) => {

                let fields = c.fields.iter()

                    .map(|f| (f.name.clone(), plumTypeFromAst(&f.ty)))

                    .collect();

                classes.insert(c.name.clone(), fields);

            }

            // `Bool` may be re-"declared" (`enum Bool = | True | False`) purely to

            // give it a nesting site for methods (`and`/`or`/`parse`/...) — the

            // language currently has no other way to attach a method to a builtin

            // type (see `libs/std/int.plum`/`float.plum`'s own `type Int =`/`type

            // Float =` for the same pattern). Re-registering its variants here

            // would silently overwrite the hardcoded tags above with whatever

            // order this declaration happens to list them in, flipping every

            // `True`/`False` tag used throughout the rest of the codegen. Skip.

            ast::Item::Enum(e) if e.name == "Bool" => {}

            ast::Item::Enum(e) => {

                let shared_field_types: Vec<PlumType> = e.params.iter()

                    .map(|p| plumTypeFromAst(&p.ty))

                    .collect();

                if !e.params.is_empty() {

                    let params = e.params.iter()

                        .map(|p| (p.name.clone(), plumTypeFromAst(&p.ty)))

                        .collect();

                    enum_params.insert(e.name.clone(), params);

                }

                for (tag, v) in e.variants.iter().enumerate() {

                    let field_types = if !e.params.is_empty() {

                        shared_field_types.clone()

                    } else {

                        v.fields.iter()

                            .map(|f| plumTypeFromAst(&ast::Type { name: f.clone(), generics: vec![] }))

                            .collect()

                    };

                    enum_variants.insert(v.name.clone(), EnumVariantInfo {

                        enum_name: e.name.clone(),

                        tag: tag as i32,

                        field_types,

                        values: v.values.clone(),

                    });

                }

            }

            _ => {}

        }

    }

    // Second pass: register top-level function/method signatures and consts.

    // Methods (`name<Receiver>(...)`) live in `methods`, keyed by receiver type,

    // so a bare call can't accidentally resolve to some other type's method.

    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) => plumTypeFromAst(t),

                        ast::ParamType::Variadic(t) => PlumType::TVariadic(Box::new(plumTypeFromAst(t))),

                        ast::ParamType::Fn(params, ret) => {

                            let param_types = params.iter().map(plumTypeFromAst).collect();

                            let ret_ty = ret.as_ref().map(|r| plumTypeFromAst(r)).unwrap_or(PlumType::TUnit);

                            PlumType::TFun(param_types, Box::new(ret_ty))

                        }

                    }

                }).collect();

                let ret = f.returns.as_ref()

                    .map(plumTypeFromAst)

                    .unwrap_or(PlumType::TUnit);

                let fn_ty = PlumType::TFun(param_types, Box::new(ret));

                if let Some(recv) = &f.type_param {

                    methods.insert((recv.clone(), f.name.clone()), fn_ty);

                } else {

                    global_env.insert(f.name.clone(), TypeScheme::mono(fn_ty));

                }

            }

            ast::Item::Const(c) => {

                // Top-level consts are always simple literals in practice; infer

                // directly from the literal kind rather than pulling in the full

                // `inferExpr` (which needs a `CheckCtx` this pass hasn't built yet).

                let ty = match &c.value {

                    ast::Expr::Int(_) => PlumType::TInt,

                    ast::Expr::Float(_) => PlumType::TFloat,

                    ast::Expr::String(_) => PlumType::TStr,

                    _ => PlumType::TVar("_".to_string()),

                };

                global_env.insert(c.name.clone(), TypeScheme::mono(ty));

            }

            _ => {}

        }

    }

    (global_env, classes, methods, enum_variants, enum_params)

}

pub fn checkSource(source: &ast::Source) -> CheckResult<()> {

    let source = monomorphize::monomorphizeSource(source).map_err(|e| vec![CheckError { message: e }])?;

    let mut errors: Vec<CheckError> = Vec::new();

    let (global_env, classes, methods, enum_variants, enum_params) = buildGlobalTables(&source);

    let ctx = CheckCtx { classes: &classes, methods: &methods, enum_variants: &enum_variants, enum_params: &enum_params };

    // A name that is both a class and an enum variant is ambiguous: `Name(...)`

    // could mean either construction, and downstream code (both the checker's

    // `inferExpr` and codegen) consults `enum_variants` first, so the class

    // constructor would be silently shadowed with no diagnostic. Reject it.

    for name in classes.keys() {

        if enum_variants.contains_key(name) {

            errors.push(CheckError {

                message: format!("'{}' is declared as both a class and an enum variant", name),

            });

        }

    }

    // A discriminant enum (`enum Foo(n: Int) = ...`) requires every variant to supply

    // exactly one value per declared param, unified against that param's type. An

    // ordinary (param-less) enum must NOT have variants with values — most likely

    // caused by writing `Some(5)` where the generic-payload form `Some[Int]` was meant.

    for item in &source.items {

        if let ast::Item::Enum(e) = item {

            for v in &e.variants {

                if e.params.is_empty() {

                    if !v.values.is_empty() {

                        errors.push(CheckError {

                            message: format!("enum '{}' variant '{}': has discriminant values but '{}' declares no params", e.name, v.name, e.name),

                        });

                    }

                    continue;

                }

                if v.values.len() != e.params.len() {

                    errors.push(CheckError {

                        message: format!(

                            "enum '{}' variant '{}': expected {} discriminant value(s), got {}",

                            e.name, v.name, e.params.len(), v.values.len()

                        ),

                    });

                    continue;

                }

                for (value, param) in v.values.iter().zip(e.params.iter()) {

                    match inferExpr(value, &global_env, &ctx) {

                        Ok(actual) => {

                            let expected = plumTypeFromAst(&param.ty);

                            if let Err(msg) = unify(&expected, &actual) {

                                errors.push(CheckError {

                                    message: format!("enum '{}' variant '{}': param '{}': {}", e.name, v.name, param.name, msg),

                                });

                            }

                        }

                        Err(msg) => errors.push(CheckError {

                            message: format!("enum '{}' variant '{}': param '{}': {}", e.name, v.name, param.name, msg),

                        }),

                    }

                }

            }

        }

    }

    for item in &source.items {

        if let ast::Item::Fn(f) = item {

            let mut local_errors = checkFn(f, &global_env, &ctx);

            errors.append(&mut local_errors);

        }

    }

    if errors.is_empty() { Ok(()) } else { Err(errors) }

}

fn checkFn(f: &ast::Fn, global_env: &TypeEnv, ctx: &CheckCtx) -> Vec<CheckError> {

    let mut errors = Vec::new();

    // `is_extern` and `body == FnBody::Extern` must agree — either both (a genuine

    // host-backed declaration) or neither (a normal Plum function). A mismatch here

    // is always a parser artifact of a malformed declaration, not user intent, so

    // report it plainly rather than trying to guess which side was "right".

    match (f.is_extern, &f.body) {

        (true, ast::FnBody::Extern) => {

            if f.type_param.is_some() {

                errors.push(CheckError { message: format!("fn '{}': extern functions can't be methods (no receiver)", f.name) });

            }

            return errors;

        }

        (true, _) => {

            errors.push(CheckError { message: format!("fn '{}': extern fn must not have a body", f.name) });

            return errors;

        }

        (false, ast::FnBody::Extern) => {

            errors.push(CheckError { message: format!("fn '{}': missing a body (or mark it `extern`)", f.name) });

            return errors;

        }

        (false, _) => {}

    }

    let mut env = global_env.clone();

    let variadic_positions: Vec<usize> = f.params.iter().enumerate()

        .filter(|(_, p)| matches!(p.ty, ast::ParamType::Variadic(_)))

        .map(|(i, _)| i)

        .collect();

    if variadic_positions.len() > 1 {

        errors.push(CheckError { message: format!("fn '{}': at most one variadic parameter is allowed", f.name) });

    } else if let Some(&pos) = variadic_positions.first() {

        if pos != f.params.len() - 1 {

            errors.push(CheckError { message: format!("fn '{}': a variadic parameter must be last", f.name) });

        }

    }

    // Methods (`name<Receiver>(...)`) get an implicit `self: Receiver` binding.

    if let Some(recv) = &f.type_param {

        let recv_ty = ast::Type { name: recv.clone(), generics: vec![] };

        env.insert("self".to_string(), TypeScheme::mono(plumTypeFromAst(&recv_ty)));

    }

    // Add params to env

    for p in &f.params {

        let ty = match &p.ty {

            ast::ParamType::Type(t) => plumTypeFromAst(t),

            ast::ParamType::Variadic(t) => PlumType::TVariadic(Box::new(plumTypeFromAst(t))),

            ast::ParamType::Fn(params, ret) => {

                let param_types = params.iter().map(plumTypeFromAst).collect();

                let ret_ty = ret.as_ref().map(|r| plumTypeFromAst(r)).unwrap_or(PlumType::TUnit);

                PlumType::TFun(param_types, Box::new(ret_ty))

            }

        };

        env.insert(p.name.clone(), TypeScheme::mono(ty));

    }

    let declared_ret = f.returns.as_ref()

        .map(plumTypeFromAst)

        .unwrap_or(PlumType::TUnit);

    match &f.body {

        ast::FnBody::Expr(e) => {

            match inferExpr(e, &env, ctx) {

                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 = checkBlock(block, &mut env, &declared_ret, &f.name, ctx);

            errors.append(&mut block_errors);

            // Check the type of the last expression statement against the declared return type

            if let Some(ast::Stmt::Expr(last_expr)) = block.stmts.last() {

                match inferExpr(last_expr, &env, ctx) {

                    Ok(t) => {

                        if let Err(msg) = unify(&declared_ret, &t) {

                            errors.push(CheckError { message: format!("fn '{}': return type mismatch: {}", f.name, msg) });

                        }

                    }

                    Err(_) => {} // already reported by checkBlock

                }

            }

        }

        // Ruled out (and returned early) by the `is_extern`/body coherence check above.

        ast::FnBody::Extern => unreachable!("extern fn bodies are handled earlier in checkFn"),

    }

    errors

}

fn checkBlock(block: &ast::Block, env: &mut TypeEnv, declared_ret: &PlumType, fn_name: &str, ctx: &CheckCtx) -> Vec<CheckError> {

    let mut errors = Vec::new();

    for stmt in &block.stmts {

        let mut stmt_errors = checkStmt(stmt, env, declared_ret, fn_name, ctx);

        errors.append(&mut stmt_errors);

    }

    errors

}

fn describeTargetObject(expr: &ast::Expr) -> String {

    match expr {

        ast::Expr::Self_ => "self".to_string(),

        ast::Expr::Var(n) => n.clone(),

        ast::Expr::Attribute(a) => {

            if let ast::AttrKind::Field(f) = &a.attr {

                format!("{}.{}", describeTargetObject(&a.object), f)

            } else {

                "<expr>".to_string()

            }

        }

        _ => "<expr>".to_string(),

    }

}

fn checkStmt(stmt: &ast::Stmt, env: &mut TypeEnv, declared_ret: &PlumType, fn_name: &str, ctx: &CheckCtx) -> 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 target {

                    ast::AssignTarget::Var(name) => {

                        let already_declared = lookup(env, name).is_ok();

                        if a.declare && already_declared {

                            errors.push(CheckError {

                                message: format!("fn '{}': assign '{}': already declared — use '=' to reassign, not ':='", fn_name, name),

                            });

                        } else if !a.declare && already_declared {

                            // Reassignment: the new value's type must match what

                            // `name` was already bound to — this is the actual

                            // point of `:=` existing at all (catches `x := 5`

                            // then later `x = 3.14` at compile time instead of

                            // silently miscompiling or silently changing type).

                            match inferExpr(value, env, ctx) {

                                Ok(t) => {

                                    let existing = lookup(env, name).expect("already_declared just confirmed this succeeds");

                                    if let Err(msg) = unify(&existing, &t) {

                                        errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, name, msg) });

                                    } else if matches!(existing, PlumType::TVar(_)) {

                                        // `unify` never narrows a `TVar` itself —

                                        // once a concrete type is available, record

                                        // it so later uses of `name` see it too.

                                        env.insert(name.clone(), TypeScheme::mono(t));

                                    }

                                }

                                Err(msg) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, name, msg) }),

                            }

                        } else {

                            // First declaration — via `:=`, or (for backward

                            // compatibility with every existing bare `x = ...`

                            // first-use) via a plain `=` on a name not yet bound.

                            match inferExpr(value, env, ctx) {

                                Ok(t) => { env.insert(name.clone(), TypeScheme::mono(t)); }

                                Err(msg) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, name, msg) }),

                            }

                        }

                    }

                    ast::AssignTarget::Field(object, field_name) => {

                        let label = format!("{}.{}", describeTargetObject(object), field_name);

                        match (inferExpr(object, env, ctx), inferExpr(value, env, ctx)) {

                            (Ok(PlumType::TNamed(class_name)), Ok(value_ty)) => {

                                match ctx.classes.get(&class_name).and_then(|fields| {

                                    fields.iter().find(|(n, _)| n == field_name).map(|(_, ty)| ty.clone())

                                }) {

                                    Some(field_ty) => {

                                        if let Err(msg) = unify(&field_ty, &value_ty) {

                                            errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, label, msg) });

                                        }

                                    }

                                    None => errors.push(CheckError { message: format!("fn '{}': assign '{}': no field '{}' on class '{}'", fn_name, label, field_name, class_name) }),

                                }

                            }

                            (Ok(other), Ok(_)) => errors.push(CheckError { message: format!("fn '{}': assign '{}': cannot access field on non-class type {}", fn_name, label, other) }),

                            (Err(msg), _) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, label, msg) }),

                            (_, Err(msg)) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, label, msg) }),

                        }

                    }

                }

            }

        }

        ast::Stmt::Return(Some(e)) => {

            match inferExpr(e, env, ctx) {

                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 inferExpr(&if_.condition, env, ctx) {

                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 checkBlock(&if_.body, env, declared_ret, fn_name, ctx));

            for ei in &if_.else_ifs {

                match inferExpr(&ei.condition, env, ctx) {

                    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 checkBlock(&ei.body, env, declared_ret, fn_name, ctx));

            }

            if let Some(else_block) = &if_.else_ {

                errors.append(&mut checkBlock(else_block, env, declared_ret, fn_name, ctx));

            }

        }

        ast::Stmt::While(w) => {

            match inferExpr(&w.condition, env, ctx) {

                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 checkBlock(&w.body, env, declared_ret, fn_name, ctx));

        }

        ast::Stmt::For(f_stmt) => {

            let iter_ty = inferExpr(&f_stmt.iter, env, ctx);

            let mut inner_env = env.clone();

            match &iter_ty {

                Ok(PlumType::TVariadic(elem)) => {

                    if f_stmt.vars.len() != 1 {

                        errors.push(CheckError { message: format!("fn '{}': for-loop over a variadic param must bind exactly one variable", fn_name) });

                    }

                    for var in &f_stmt.vars {

                        inner_env.insert(var.clone(), TypeScheme::mono((**elem).clone()));

                    }

                }

                Ok(_) => {

                    for var in &f_stmt.vars {

                        inner_env.insert(var.clone(), TypeScheme::mono(PlumType::TInt));

                    }

                }

                Err(msg) => {

                    errors.push(CheckError { message: format!("fn '{}': for iter: {}", fn_name, msg) });

                    for var in &f_stmt.vars {

                        inner_env.insert(var.clone(), TypeScheme::mono(PlumType::TInt));

                    }

                }

            }

            errors.append(&mut checkBlock(&f_stmt.body, &mut inner_env, declared_ret, fn_name, ctx));

        }

        ast::Stmt::Expr(e) => {

            if let Err(msg) = inferExpr(e, env, ctx) {

                errors.push(CheckError { message: format!("fn '{}': {}", fn_name, msg) });

            }

        }

        ast::Stmt::Assert(e) => {

            match inferExpr(e, env, ctx) {

                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::Match(m) => {

            errors.append(&mut checkMatch(m, env, declared_ret, fn_name, ctx));

        }

        ast::Stmt::Break | ast::Stmt::Continue | ast::Stmt::Todo => {}

    }

    errors

}

fn checkMatch(m: &ast::Match, env: &TypeEnv, declared_ret: &PlumType, fn_name: &str, ctx: &CheckCtx) -> Vec<CheckError> {

    let mut errors = Vec::new();

    let mut subject_types: Vec<PlumType> = Vec::new();

    for s in &m.subjects {

        match inferExpr(s, env, ctx) {

            Ok(t) => subject_types.push(t),

            Err(msg) => errors.push(CheckError { message: format!("fn '{}': match subject: {}", fn_name, msg) }),

        }

    }

    if subject_types.len() != m.subjects.len() {

        return errors; // a subject failed to type — cases can't be checked meaningfully

    }

    for case in &m.cases {

        let mut case_env = env.clone();

        if case.patterns.len() == subject_types.len() {

            for (pat, sty) in case.patterns.iter().zip(subject_types.iter()) {

                if let Err(msg) = checkPattern(pat, sty, &mut case_env, ctx) {

                    errors.push(CheckError { message: format!("fn '{}': match case: {}", fn_name, msg) });

                }

            }

        } else {

            errors.push(CheckError {

                message: format!(

                    "fn '{}': match case has {} pattern(s), expected {}",

                    fn_name, case.patterns.len(), subject_types.len()

                ),

            });

        }

        errors.append(&mut checkBlock(&case.body, &mut case_env, declared_ret, fn_name, ctx));

    }

    errors

}

/// Checks a single case pattern against the type of the subject it matches, binding any

/// new names it introduces into `env`. Constructor-payload sub-patterns (`Some(x)`) bind

/// against that variant's declared field types (see `EnumVariantInfo::field_types`).

fn checkPattern(pat: &ast::CasePattern, subject_ty: &PlumType, env: &mut TypeEnv, ctx: &CheckCtx) -> Result<(), String> {

    match pat {

        ast::CasePattern::Wildcard => Ok(()),

        ast::CasePattern::Int(_) => unify(subject_ty, &PlumType::TInt),

        ast::CasePattern::Float(_) => unify(subject_ty, &PlumType::TFloat),

        ast::CasePattern::String(_) => unify(subject_ty, &PlumType::TStr),

        ast::CasePattern::Name(n) => {

            let is_known_variant = n.chars().next().map(|c| c.is_uppercase()).unwrap_or(false)

                && ctx.enum_variants.contains_key(n);

            if is_known_variant {

                Ok(()) // equality check against a known enum tag, e.g. `True`

            } else {

                env.insert(n.clone(), TypeScheme::mono(subject_ty.clone()));

                Ok(())

            }

        }

        ast::CasePattern::Class { name, fields } => match ctx.enum_variants.get(name) {

            Some(info) => {

                if fields.len() != info.field_types.len() {

                    return Err(format!(

                        "constructor pattern '{}' expects {} field(s), got {}",

                        name, info.field_types.len(), fields.len()

                    ));

                }

                for (f, fty) in fields.iter().zip(info.field_types.iter()) {

                    checkPattern(f, fty, env, ctx)?;

                }

                Ok(())

            }

            // Unmodeled/builtin variant: allow, codegen will catch.

            None => {

                for f in fields {

                    checkPattern(f, &PlumType::TVar("_".to_string()), env, ctx)?;

                }

                Ok(())

            }

        },

    }

}

pub fn inferExpr(expr: &ast::Expr, env: &TypeEnv, ctx: &CheckCtx) -> 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_ => lookup(env, "self"),

        // A bare `NAME` lexes as a type_identifier whenever it starts with an

        // uppercase letter — which includes every SCREAMING_CASE top-level

        // const (`PI`, `MAX_VALUE`, ...), since there's no separate "constant"

        // token. Prefer a matching const's real type over the enum-variant/

        // unmodeled-type-name fallbacks below.

        ast::Expr::TypeName(n) => match n.as_str() {

            "True" | "False" => Ok(PlumType::TBool),

            _ => match lookup(env, n) {

                Ok(ty) if !matches!(ty, PlumType::TVar(_)) => Ok(ty),

                _ => match ctx.enum_variants.get(n) {

                    Some(info) => Ok(PlumType::TNamed(info.enum_name.clone())),

                    // Unmodeled/builtin type name: allow, codegen will catch.

                    None => Ok(PlumType::TNamed(n.to_string())),

                },

            },

        },

        ast::Expr::Paren(inner) => inferExpr(inner, env, ctx),

        ast::Expr::Closure(cl) => {

            let mut closure_env = env.clone();

            let param_types: Vec<PlumType> = cl.params.iter().map(|p| {

                let t = PlumType::TVar(format!("_closure_{}", p));

                closure_env.insert(p.clone(), TypeScheme::mono(t.clone()));

                t

            }).collect();

            let body_ty = match &cl.body.stmts.last() {

                Some(ast::Stmt::Expr(e)) => inferExpr(e, &closure_env, ctx)?,

                Some(ast::Stmt::Return(Some(e))) => inferExpr(e, &closure_env, ctx)?,

                _ => PlumType::TUnit,

            };

            Ok(PlumType::TFun(param_types, Box::new(body_ty)))

        }

        ast::Expr::Not(inner) => {

            let t = inferExpr(inner, env, ctx)?;

            unify(&PlumType::TBool, &t)?;

            Ok(PlumType::TBool)

        }

        ast::Expr::Unary(u) => inferExpr(&u.operand, env, ctx),

        ast::Expr::Binary(b) => {

            let lt = inferExpr(&b.left, env, ctx)?;

            let rt = inferExpr(&b.right, env, ctx)?;

            unify(&lt, &rt).map_err(|e| format!("binary op: {}", e))?;

            Ok(lt)

        }

        ast::Expr::Bool(b) => {

            let lt = inferExpr(&b.left, env, ctx)?;

            let rt = inferExpr(&b.right, env, ctx)?;

            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 = inferExpr(&c.left, env, ctx)?;

            let rt = inferExpr(&c.right, env, ctx)?;

            unify(&lt, &rt).map_err(|e| format!("compare op: {}", e))?;

            Ok(PlumType::TBool)

        }

        ast::Expr::Ternary(t) => {

            let ct = inferExpr(&t.condition, env, ctx)?;

            unify(&PlumType::TBool, &ct).map_err(|e| format!("ternary condition: {}", e))?;

            let tt = inferExpr(&t.then, env, ctx)?;

            let et = inferExpr(&t.else_, env, ctx)?;

            unify(&tt, &et).map_err(|e| format!("ternary branches: {}", e))?;

            Ok(tt)

        }

        ast::Expr::FnCall(call) => {

            // `Int(x)`/`Float(x)`/`Byte(x)` are builtin numeric conversions, not

            // ordinary calls — handled here so `y = Float(x)` unifies against

            // `TFloat` rather than falling through to the permissive

            // "unknown fn" case.

            if (call.name == "Int" || call.name == "Float" || call.name == "Byte") && call.args.len() == 1 {

                let arg_expr = match &call.args[0] {

                    ast::Arg::Positional(e) => e,

                    ast::Arg::Keyword { value, .. } => value,

                    ast::Arg::Pair { value, .. } => value,

                };

                let actual = inferExpr(arg_expr, env, ctx)?;

                return match (call.name.as_str(), &actual) {

                    ("Float", PlumType::TInt) => Ok(PlumType::TFloat),

                    ("Int", PlumType::TFloat) => Ok(PlumType::TInt),

                    ("Float", PlumType::TFloat) | ("Int", PlumType::TInt) => Ok(actual),

                    ("Byte", PlumType::TInt) | ("Byte", PlumType::TByte) => Ok(PlumType::TByte),

                    ("Int", PlumType::TByte) => Ok(PlumType::TInt),

                    (name, other) => Err(format!("call '{}': cannot convert {} to {}", name, other, name)),

                };

            }

            if let Some(info) = ctx.enum_variants.get(&call.name) {

                if call.args.len() != info.field_types.len() {

                    return Err(format!(

                        "variant '{}': expected {} arg(s), got {}",

                        call.name, info.field_types.len(), call.args.len()

                    ));

                }

                for (i, (arg, expected)) in call.args.iter().zip(info.field_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 = inferExpr(arg_expr, env, ctx)?;

                    unify(expected, &actual).map_err(|e| format!("variant '{}' arg {}: {}", call.name, i, e))?;

                }

                return Ok(PlumType::TNamed(info.enum_name.clone()));

            }

            match lookup(env, &call.name) {

                Ok(PlumType::TFun(param_types, ret)) => {

                    match param_types.last() {

                        Some(PlumType::TVariadic(elem)) => {

                            let fixed = &param_types[..param_types.len() - 1];

                            if call.args.len() < fixed.len() {

                                return Err(format!("call '{}': expected at least {} arg(s), got {}", call.name, fixed.len(), call.args.len()));

                            }

                            for (i, (arg, expected)) in call.args.iter().zip(fixed.iter()).enumerate() {

                                let arg_expr = match arg {

                                    ast::Arg::Positional(e) => e,

                                    ast::Arg::Keyword { value, .. } => value,

                                    ast::Arg::Pair { value, .. } => value,

                                };

                                let actual = inferExpr(arg_expr, env, ctx)?;

                                unify(expected, &actual).map_err(|e| format!("call '{}' arg {}: {}", call.name, i, e))?;

                            }

                            for (i, arg) in call.args.iter().enumerate().skip(fixed.len()) {

                                let arg_expr = match arg {

                                    ast::Arg::Positional(e) => e,

                                    ast::Arg::Keyword { value, .. } => value,

                                    ast::Arg::Pair { value, .. } => value,

                                };

                                let actual = inferExpr(arg_expr, env, ctx)?;

                                unify(elem, &actual).map_err(|e| format!("call '{}' variadic arg {}: {}", call.name, i, e))?;

                            }

                            Ok(*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 = inferExpr(arg_expr, env, ctx)?;

                                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(call) => {

            match ctx.classes.get(&call.type_name) {

                Some(fields) => {

                    for fa in &call.fields {

                        match fields.iter().find(|(n, _)| n == &fa.name) {

                            Some((_, expected)) => {

                                let actual = inferExpr(&fa.value, env, ctx)?;

                                unify(expected, &actual)

                                    .map_err(|e| format!("class '{}' field '{}': {}", call.type_name, fa.name, e))?;

                            }

                            None => return Err(format!("unknown field '{}' on class '{}'", fa.name, call.type_name)),

                        }

                    }

                    Ok(PlumType::TNamed(call.type_name.clone()))

                }

                // Unmodeled (e.g. builtin/std) type: allow, codegen will catch.

                None => Ok(PlumType::TNamed(call.type_name.clone())),

            }

        }

        ast::Expr::Attribute(attr) => {

            let obj_ty = inferExpr(&attr.object, env, ctx)?;

            match &attr.attr {

                ast::AttrKind::Field(field_name) => match &obj_ty {

                    PlumType::TNamed(class_name) => match ctx.classes.get(class_name) {

                        Some(fields) => fields.iter()

                            .find(|(n, _)| n == field_name)

                            .map(|(_, t)| t.clone())

                            .ok_or_else(|| format!("no field '{}' on type '{}'", field_name, class_name)),

                        None => match ctx.enum_params.get(class_name) {

                            Some(params) => params.iter()

                                .find(|(n, _)| n == field_name)

                                .map(|(_, t)| t.clone())

                                .ok_or_else(|| format!("no field '{}' on type '{}'", field_name, class_name)),

                            // Unmodeled type: allow, codegen will catch.

                            None => Ok(PlumType::TVar("_".to_string())),

                        },

                    },

                    _ => Err(format!("cannot access field '{}' on non-class type {}", field_name, obj_ty)),

                },

                ast::AttrKind::Method(call) => match methodReceiverName(&obj_ty) {

                    Some(class_name) => match ctx.methods.get(&(class_name.clone(), call.name.clone())) {

                        Some(PlumType::TFun(param_types, ret)) => {

                            match param_types.last() {

                                Some(PlumType::TVariadic(elem)) => {

                                    let fixed = &param_types[..param_types.len() - 1];

                                    if call.args.len() < fixed.len() {

                                        return Err(format!(

                                            "method '{}.{}': expected at least {} arg(s), got {}",

                                            class_name, call.name, fixed.len(), call.args.len()

                                        ));

                                    }

                                    for (i, (arg, expected)) in call.args.iter().zip(fixed.iter()).enumerate() {

                                        let arg_expr = match arg {

                                            ast::Arg::Positional(e) => e,

                                            ast::Arg::Keyword { value, .. } => value,

                                            ast::Arg::Pair { value, .. } => value,

                                        };

                                        let actual = inferExpr(arg_expr, env, ctx)?;

                                        unify(expected, &actual)

                                            .map_err(|e| format!("method '{}.{}' arg {}: {}", class_name, call.name, i, e))?;

                                    }

                                    for (i, arg) in call.args.iter().enumerate().skip(fixed.len()) {

                                        let arg_expr = match arg {

                                            ast::Arg::Positional(e) => e,

                                            ast::Arg::Keyword { value, .. } => value,

                                            ast::Arg::Pair { value, .. } => value,

                                        };

                                        let actual = inferExpr(arg_expr, env, ctx)?;

                                        unify(elem, &actual)

                                            .map_err(|e| format!("method '{}.{}' variadic arg {}: {}", class_name, call.name, i, e))?;

                                    }

                                    Ok(*ret.clone())

                                }

                                _ => {

                                    if call.args.len() != param_types.len() {

                                        return Err(format!(

                                            "method '{}.{}': expected {} args, got {}",

                                            class_name, 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 = inferExpr(arg_expr, env, ctx)?;

                                        unify(expected, &actual)

                                            .map_err(|e| format!("method '{}.{}' arg {}: {}", class_name, call.name, i, e))?;

                                    }

                                    Ok(*ret.clone())

                                }

                            }

                        }

                        // Unmodeled method (e.g. builtin/std): allow, codegen will catch.

                        _ => Ok(PlumType::TVar("_".to_string())),

                    },

                    None => Ok(PlumType::TVar("_".to_string())),

                },

            }

        }

    }

}