assura_codegen/

types_gen.rs

//! Type mapping, Cargo.toml generation, and type definition codegen.

use super::*;

pub(crate) fn generate_cargo_toml_impl(
    crate_name: &str,
    config: &BackendConfig,
    include_proptest: bool,
    include_thiserror: bool,
) -> String {
    let mut toml = format!(
        r#"[package]
name = "{crate_name}"
version = "0.1.0"
edition = "2024"

# Generated by the Assura compiler.
# Do not edit manually.

[dependencies]
"#
    );

    if include_thiserror {
        toml.push_str("thiserror = \"2\"\n");
    }

    if config.runtime_checks {
        // Use crates.io version when published; until then, users must add
        // a [patch] or path override pointing to their local assura-runtime.
        toml.push_str("assura-runtime = \"0.1\"\n");
    }

    if include_proptest {
        toml.push_str("\n[dev-dependencies]\nproptest = \"1\"\n");
    }

    // Add profile configuration based on BackendConfig
    if config.opt_level != 2 || config.debug_info {
        toml.push_str("\n[profile.release]\n");
        toml.push_str("opt-level = ");
        toml.push_str(&config.opt_level.to_string());
        toml.push('\n');
        if config.debug_info {
            toml.push_str("debug = true\n");
        }
    }

    // Add backend-specific configuration
    if matches!(config.backend, CodegenBackend::Cranelift) {
        toml.push_str(
            "\n# Cranelift backend for fast dev builds\n\
             # Install: rustup component add rustc-codegen-cranelift\n\
             # Note: Cranelift requires nightly Rust\n",
        );
        // Force dev-friendly profile settings for Cranelift
        if config.opt_level == 2 && !config.debug_info {
            toml.push_str("\n[profile.dev]\nopt-level = 0\ndebug = true\n");
        }
    }

    // WASM target: add cdylib crate type so cargo produces a .wasm file
    if matches!(config.target, CompileTarget::Wasm) {
        toml.push_str("\n[lib]\ncrate-type = [\"cdylib\"]\n");
        toml.push_str(
            "\n# WASM target: build with `cargo build --target wasm32-wasip1`\n\
             # Install target: `rustup target add wasm32-wasip1`\n",
        );
    }

    toml
}

// ---------------------------------------------------------------------------
// Type mapping: Assura types -> Rust types
// ---------------------------------------------------------------------------

/// Map a single Assura type token to its Rust equivalent.
pub fn map_type_token(tok: &str) -> &str {
    match tok {
        "Int" => "i64",
        "Nat" => "u64",
        "Float" => "f64",
        "Bool" => "bool",
        "String" => "String",
        "Bytes" => "Vec<u8>",
        "Unit" => "()",
        "Never" => "!",
        "U8" => "u8",
        "U16" => "u16",
        "U32" => "u32",
        "U64" => "u64",
        "I8" => "i8",
        "I16" => "i16",
        "I32" => "i32",
        "I64" => "i64",
        "F32" => "f32",
        "F64" => "f64",
        "List" => "Vec",
        "Map" => "std::collections::BTreeMap",
        "Set" => "std::collections::BTreeSet",
        "Sequence" => "Vec",
        // Option, Result map to themselves
        _ => tok,
    }
}

/// Convert an Assura type token sequence (e.g., `["List", "<", "Int", ">"]`)
/// to a valid Rust type string.
///
/// Handles Assura-specific syntax that is not valid Rust:
/// - Refinement types: `{ x : T | P }` -> base type `T`
/// - Taint annotations: `T @ taint : label` -> just `T`
/// - Union error types: `T | E` -> `Result<T, E>`
/// - Smart joining: `Vec < i64 >` -> `Vec<i64>`
pub fn map_type_tokens(tokens: &[String]) -> String {
    if tokens.is_empty() {
        return "()".to_string();
    }

    // Phase 1: Strip annotations, qualifiers, and clause keywords
    // Stops at: "@" (taint), "#" (attribute), "decreases", "where"
    // Filters out type qualifiers with no Rust equivalent
    let clean: Vec<&str> = tokens
        .iter()
        .map(|s| s.as_str())
        .take_while(|t| !matches!(*t, "@" | "#" | "decreases" | "where"))
        .filter(|t| {
            !matches!(
                *t,
                "linear" | "secret" | "tainted" | "taint" | "untrusted" | "validated"
            )
        })
        .collect();
    if clean.is_empty() {
        return "()".to_string();
    }

    // Phase 2: Strip refinement predicates: { x : T | P } -> just the base type
    if clean.first() == Some(&"{") {
        return extract_base_type_from_refined(tokens);
    }

    // Phase 3: Handle union error types: T | E -> Result<T, E>
    // Find a "|" not inside angle brackets
    let mut depth = 0i32;
    let mut pipe_pos = None;
    for (i, tok) in clean.iter().enumerate() {
        match *tok {
            "<" => depth += 1,
            ">" => {
                if depth > 0 {
                    depth -= 1;
                }
            }
            "|" if depth == 0 => {
                pipe_pos = Some(i);
                break;
            }
            _ => {}
        }
    }

    if let Some(pos) = pipe_pos {
        let ok_tokens: Vec<String> = clean[..pos].iter().map(|s| s.to_string()).collect();
        let err_tokens: Vec<String> = clean[pos + 1..].iter().map(|s| s.to_string()).collect();
        let ok_type = map_type_tokens(&ok_tokens);
        let err_type = map_type_tokens(&err_tokens);
        return format!("Result<{ok_type}, {err_type}>");
    }

    // Phase 4: Map each token and join smartly (no extra spaces around < > & etc.)
    // Convert SCREAMING_SNAKE_CASE names inside angle brackets to marker types,
    // since Assura const-generics don't translate directly to Rust generics.
    let mut in_angle = 0i32;
    let mapped: Vec<String> = clean
        .iter()
        .map(|t| {
            match *t {
                "<" => in_angle += 1,
                ">" if in_angle > 0 => {
                    in_angle -= 1;
                }
                _ => {}
            }
            if in_angle > 0 && is_const_name(t) {
                // Const-generic name used as type param: wrap as marker type
                t.to_string()
            } else {
                map_type_token(t).to_string()
            }
        })
        .collect();
    let refs: Vec<&str> = mapped.iter().map(|s| s.as_str()).collect();
    smart_join_type_tokens(&refs)
}

/// Join type tokens without spurious spaces around angle brackets, ampersands, etc.
pub(crate) fn smart_join_type_tokens(tokens: &[&str]) -> String {
    let mut result = String::new();
    for (i, tok) in tokens.iter().enumerate() {
        if i > 0 {
            let prev = tokens[i - 1];
            let no_space = matches!(*tok, ">" | "<" | "," | ")" | ".")
                || matches!(prev, "<" | "(" | "&" | ".")
                || (*tok == "mut" && prev == "&");
            if !no_space {
                result.push(' ');
            }
        }
        result.push_str(tok);
    }
    result
}

// ---------------------------------------------------------------------------

// Type definitions
// ---------------------------------------------------------------------------

pub(crate) fn generate_type_def(t: &TypeDef, code: &mut String) {
    let items = crate::hir::build_type_def(t);
    for item in &items {
        code.push_str(&crate::hir::render_item_raw(item));
    }
}

// ---------------------------------------------------------------------------
// Type reference collection (for generating stub types)
// ---------------------------------------------------------------------------

/// Check if a name looks like a constant (SCREAMING_SNAKE_CASE).
/// E.g., `TOTAL_TABLE_SIZE`, `MAX_ALPHABET_SIZE`.
pub(crate) fn is_const_name(name: &str) -> bool {
    !name.is_empty()
        && name
            .chars()
            .all(|c| c.is_ascii_uppercase() || c == '_' || c.is_ascii_digit())
        && name.contains('_')
}

/// Info about a generic parameter: either a type param or a const param.
#[derive(Debug, Clone)]
pub(crate) enum GenericParamKind {
    Type,
    Const,
}

/// Scan a type token sequence for `TypeName<Arg1, Arg2, ...>` patterns.
/// Records the generic params (type vs const) for each type name, and
/// collects names that appear as const-generic arguments.
///
/// Superseded in production by `detect_generic_arity_from_type_expr` which
/// operates directly on `TypeExpr`. Retained for unit tests.
#[cfg(test)]
pub(crate) fn detect_generic_arity(
    tokens: &[String],
    param_map: &mut std::collections::HashMap<String, Vec<GenericParamKind>>,
    const_names: &mut std::collections::HashSet<String>,
) {
    let mut i = 0;
    while i < tokens.len() {
        if i + 1 < tokens.len() && tokens[i + 1] == "<" && is_user_type_name(&tokens[i]) {
            let type_name = tokens[i].clone();
            let mut depth = 0;
            let mut params = Vec::new();
            let mut current_is_const = false;
            let mut j = i + 1;
            while j < tokens.len() {
                match tokens[j].as_str() {
                    "<" => {
                        depth += 1;
                        if depth == 1 {
                            current_is_const = false;
                        }
                    }
                    ">" if depth > 0 => {
                        depth -= 1;
                        if depth == 0 {
                            // Record the last parameter
                            params.push(if current_is_const {
                                GenericParamKind::Const
                            } else {
                                GenericParamKind::Type
                            });
                            break;
                        }
                    }
                    "," if depth == 1 => {
                        params.push(if current_is_const {
                            GenericParamKind::Const
                        } else {
                            GenericParamKind::Type
                        });
                        current_is_const = false;
                    }
                    tok if depth == 1 && is_const_name(tok) => {
                        const_names.insert(tok.to_string());
                        current_is_const = true;
                    }
                    tok if depth == 1
                        && !tok.is_empty()
                        && tok.chars().all(|c| c.is_ascii_digit()) =>
                    {
                        // Numeric literal as const generic
                        current_is_const = true;
                    }
                    _ => {}
                }
                j += 1;
            }
            let existing_len = param_map.get(&type_name).map_or(0, |v| v.len());
            if params.len() > existing_len {
                param_map.insert(type_name, params);
            }
            i = j + 1;
        } else {
            i += 1;
        }
    }
}

/// Check if a token looks like a user-defined type name (starts with uppercase,
/// alphanumeric, not a Rust keyword or built-in).
pub(crate) fn is_user_type_name(tok: &str) -> bool {
    !tok.is_empty()
        && tok.chars().next().is_some_and(|c| c.is_uppercase())
        && tok.chars().all(|c| c.is_alphanumeric() || c == '_')
        && !matches!(
            tok,
            "Int"
                | "Nat"
                | "Float"
                | "Bool"
                | "String"
                | "Bytes"
                | "Unit"
                | "Never"
                | "U8"
                | "U16"
                | "U32"
                | "U64"
                | "I8"
                | "I16"
                | "I32"
                | "I64"
                | "F32"
                | "F64"
                | "List"
                | "Vec"
                | "Map"
                | "Set"
                | "Option"
                | "Result"
                | "Sequence"
                | "Self"
                | "Box"
                | "Fn"
                | "FnOnce"
                | "FnMut"
        )
}

/// Collect user-defined type names from a type token sequence.
///
/// Superseded in production by `collect_type_refs_from_type_expr`. Retained
/// for unit tests.
#[cfg(test)]
pub(crate) fn collect_type_refs_from_tokens(
    tokens: &[String],
    out: &mut std::collections::HashSet<String>,
) {
    for tok in tokens {
        // Skip taint annotations, attributes, and keywords
        if matches!(
            tok.as_str(),
            "@" | "#"
                | "taint"
                | "untrusted"
                | "validated"
                | "secret"
                | "pub"
                | "mut"
                | ":"
                | "|"
                | "&"
                | ">"
                | "<"
                | ","
                | "("
                | ")"
                | "{"
                | "}"
                | "["
                | "]"
                | "decreases"
                | "where"
        ) {
            continue;
        }
        if is_user_type_name(tok) {
            out.insert(tok.clone());
        }
    }
}

/// Collect user-defined type names directly from a `TypeExpr`, avoiding the
/// token roundtrip through `type_expr_to_token_vec`.
pub(crate) fn collect_type_refs_from_type_expr(
    te: Option<&assura_ast::TypeExpr>,
    out: &mut std::collections::HashSet<String>,
) {
    let Some(te) = te else { return };
    match te {
        assura_ast::TypeExpr::Named(name) => {
            // Named may contain spaces from fallback join; check each word.
            for word in name.split_whitespace() {
                if is_user_type_name(word) {
                    out.insert(word.to_string());
                }
            }
        }
        assura_ast::TypeExpr::Generic(name, args) => {
            if is_user_type_name(name) {
                out.insert(name.clone());
            }
            for arg in args {
                collect_type_refs_from_type_expr(Some(arg), out);
            }
        }
        assura_ast::TypeExpr::Tuple(items) => {
            for item in items {
                collect_type_refs_from_type_expr(Some(item), out);
            }
        }
        assura_ast::TypeExpr::Fn { params, ret } => {
            for p in params {
                collect_type_refs_from_type_expr(Some(p), out);
            }
            collect_type_refs_from_type_expr(Some(ret), out);
        }
        assura_ast::TypeExpr::Refined { base, .. } => {
            collect_type_refs_from_type_expr(Some(base), out);
        }
        assura_ast::TypeExpr::Unit => {}
    }
}

/// Detect generic arity directly from `TypeExpr`, replacing the token-based
/// `detect_generic_arity` for codegen purposes.
///
/// Walks the type expression tree looking for `Generic(name, args)` nodes.
/// Each arg is classified as `Const` (SCREAMING_SNAKE_CASE or numeric literal)
/// or `Type` (otherwise). The widest arity seen for each type name wins.
pub(crate) fn detect_generic_arity_from_type_expr(
    te: Option<&assura_ast::TypeExpr>,
    param_map: &mut std::collections::HashMap<String, Vec<GenericParamKind>>,
    const_names: &mut std::collections::HashSet<String>,
) {
    let Some(te) = te else { return };
    match te {
        assura_ast::TypeExpr::Generic(name, args) => {
            if is_user_type_name(name) {
                let params: Vec<GenericParamKind> = args
                    .iter()
                    .map(|arg| {
                        if is_const_type_expr_arg(arg) {
                            // Collect the const name if it's a named constant
                            if let assura_ast::TypeExpr::Named(n) = arg
                                && is_const_name(n)
                            {
                                const_names.insert(n.clone());
                            }
                            GenericParamKind::Const
                        } else {
                            GenericParamKind::Type
                        }
                    })
                    .collect();
                let existing_len = param_map.get(name).map_or(0, |v| v.len());
                if params.len() > existing_len {
                    param_map.insert(name.clone(), params);
                }
            }
            // Recurse into args
            for arg in args {
                detect_generic_arity_from_type_expr(Some(arg), param_map, const_names);
            }
        }
        assura_ast::TypeExpr::Tuple(items) => {
            for item in items {
                detect_generic_arity_from_type_expr(Some(item), param_map, const_names);
            }
        }
        assura_ast::TypeExpr::Fn { params, ret } => {
            for p in params {
                detect_generic_arity_from_type_expr(Some(p), param_map, const_names);
            }
            detect_generic_arity_from_type_expr(Some(ret), param_map, const_names);
        }
        assura_ast::TypeExpr::Refined { base, .. } => {
            detect_generic_arity_from_type_expr(Some(base), param_map, const_names);
        }
        assura_ast::TypeExpr::Named(_) | assura_ast::TypeExpr::Unit => {}
    }
}

/// Check if a `TypeExpr` argument to a generic looks like a const (SCREAMING_SNAKE_CASE
/// or numeric literal embedded in a Named).
fn is_const_type_expr_arg(te: &assura_ast::TypeExpr) -> bool {
    match te {
        assura_ast::TypeExpr::Named(n) => {
            is_const_name(n) || (!n.is_empty() && n.chars().all(|c| c.is_ascii_digit()))
        }
        _ => false,
    }
}

/// Detect feature_max names used as type arguments directly from `TypeExpr`.
pub(crate) fn detect_feature_max_as_type_from_type_expr(
    te: Option<&assura_ast::TypeExpr>,
    feature_max_set: &std::collections::HashSet<&str>,
    out: &mut std::collections::HashSet<String>,
) {
    let Some(te) = te else { return };
    match te {
        assura_ast::TypeExpr::Generic(_, args) => {
            for arg in args {
                // Each arg is in a type-argument position (inside `<>`)
                if let assura_ast::TypeExpr::Named(n) = arg
                    && feature_max_set.contains(n.as_str())
                {
                    out.insert(n.clone());
                }
                // Recurse for nested generics
                detect_feature_max_as_type_from_type_expr(Some(arg), feature_max_set, out);
            }
        }
        assura_ast::TypeExpr::Tuple(items) => {
            for item in items {
                detect_feature_max_as_type_from_type_expr(Some(item), feature_max_set, out);
            }
        }
        assura_ast::TypeExpr::Fn { params, ret } => {
            for p in params {
                detect_feature_max_as_type_from_type_expr(Some(p), feature_max_set, out);
            }
            detect_feature_max_as_type_from_type_expr(Some(ret), feature_max_set, out);
        }
        assura_ast::TypeExpr::Refined { base, .. } => {
            detect_feature_max_as_type_from_type_expr(Some(base), feature_max_set, out);
        }
        assura_ast::TypeExpr::Named(_) | assura_ast::TypeExpr::Unit => {}
    }
}

/// Collect type names referenced in expressions (e.g., constructor calls).
pub(crate) fn collect_type_refs_from_expr(
    expr: &SpExpr,
    out: &mut std::collections::HashSet<String>,
) {
    match &expr.node {
        Expr::Ident(name) => {
            if is_user_type_name(name) {
                out.insert(name.clone());
            }
        }
        Expr::Call { func, args } => {
            collect_type_refs_from_expr(func, out);
            for a in args {
                collect_type_refs_from_expr(a, out);
            }
        }
        Expr::MethodCall {
            receiver,
            method: _,
            args,
        } => {
            collect_type_refs_from_expr(receiver, out);
            for a in args {
                collect_type_refs_from_expr(a, out);
            }
        }
        Expr::Field(recv, _) => collect_type_refs_from_expr(recv, out),
        Expr::Index { expr: e, index } => {
            collect_type_refs_from_expr(e, out);
            collect_type_refs_from_expr(index, out);
        }
        Expr::BinOp { lhs, rhs, .. } => {
            collect_type_refs_from_expr(lhs, out);
            collect_type_refs_from_expr(rhs, out);
        }
        Expr::UnaryOp { expr: e, .. }
        | Expr::Old(e)
        | Expr::Cast { expr: e, .. }
        | Expr::Ghost(e) => {
            collect_type_refs_from_expr(e, out);
        }
        Expr::If {
            cond,
            then_branch,
            else_branch,
        } => {
            collect_type_refs_from_expr(cond, out);
            collect_type_refs_from_expr(then_branch, out);
            if let Some(eb) = else_branch {
                collect_type_refs_from_expr(eb, out);
            }
        }
        Expr::Forall { domain, body, .. } | Expr::Exists { domain, body, .. } => {
            collect_type_refs_from_expr(domain, out);
            collect_type_refs_from_expr(body, out);
        }
        Expr::Let { value, body, .. } => {
            collect_type_refs_from_expr(value, out);
            collect_type_refs_from_expr(body, out);
        }
        Expr::Match { scrutinee, arms } => {
            collect_type_refs_from_expr(scrutinee, out);
            for arm in arms {
                collect_type_refs_from_expr(&arm.body, out);
            }
        }
        Expr::Apply { args, .. } => {
            for a in args {
                collect_type_refs_from_expr(a, out);
            }
        }
        Expr::List(items) | Expr::Tuple(items) | Expr::Block(items) => {
            for item in items {
                collect_type_refs_from_expr(item, out);
            }
        }
        Expr::Raw(tokens) => {
            for tok in tokens {
                if is_user_type_name(tok) {
                    out.insert(tok.clone());
                }
            }
        }
        // Literals never contain type references
        Expr::Literal(_) => {}
    }
}

/// Find the value for a feature_max constant from the AST.
pub(crate) fn find_feature_max_value(source: &assura_ast::SourceFile, name: &str) -> String {
    for decl in &source.decls {
        if let Decl::Block {
            kind,
            name: n,
            value,
            body,
        } = &decl.node
            && *kind == BlockKind::FeatureMax
            && n == name
        {
            // First, try the inline value field.
            // The value tokens may include a type annotation: `["Nat", "=", "280"]`
            // from `feature_max X: Nat = 280`. Extract after the `=`.
            if let Some(val_tokens) = value {
                if let Some(eq_pos) = val_tokens.iter().position(|t| t == "=") {
                    let after_eq = &val_tokens[eq_pos + 1..];
                    if !after_eq.is_empty() {
                        return after_eq.join(" ");
                    }
                } else if val_tokens.len() == 1 {
                    // Single-token value without `=` (e.g., just a number)
                    return val_tokens[0].clone();
                }
            }
            // Fallback: try to extract from body clauses
            for clause in body {
                let val = expr_to_rust_static(&clause.body);
                if !val.is_empty() && val != "()" {
                    return val;
                }
            }
        }
    }
    // No value found: emit a compile_error! so generated code fails to
    // build rather than silently using 0, which hides missing definitions.
    format!("compile_error!(\"feature_max `{name}` has no value\")")
}

// expr_to_rust_static is now defined in expr.rs (unified RustCodegenFolder)

/// Extract the base type from a refined type token sequence like
/// `["n", ":", "Int", "|", "n", ">", "0"]` -> "i64".
pub(crate) fn extract_base_type_from_refined(tokens: &[String]) -> String {
    // Look for a token after ':' that starts with uppercase
    let mut after_colon = false;
    for tok in tokens {
        if tok == ":" {
            after_colon = true;
            continue;
        }
        if after_colon {
            if tok.chars().next().is_some_and(|c| c.is_ascii_uppercase()) {
                return map_type_token(tok).to_string();
            }
            // Not a type, stop looking
            after_colon = false;
        }
    }
    // Fallback: just use the first type-looking token
    for tok in tokens {
        if tok.chars().next().is_some_and(|c| c.is_ascii_uppercase()) {
            return map_type_token(tok).to_string();
        }
    }
    "i64".to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use assura_ast::Spanned;

    // ---- map_type_token ----

    #[test]
    fn map_basic_types() {
        assert_eq!(map_type_token("Int"), "i64");
        assert_eq!(map_type_token("Nat"), "u64");
        assert_eq!(map_type_token("Float"), "f64");
        assert_eq!(map_type_token("Bool"), "bool");
        assert_eq!(map_type_token("String"), "String");
        assert_eq!(map_type_token("Bytes"), "Vec<u8>");
        assert_eq!(map_type_token("Unit"), "()");
        assert_eq!(map_type_token("Never"), "!");
    }

    #[test]
    fn map_fixed_width_types() {
        assert_eq!(map_type_token("U8"), "u8");
        assert_eq!(map_type_token("U16"), "u16");
        assert_eq!(map_type_token("U32"), "u32");
        assert_eq!(map_type_token("U64"), "u64");
        assert_eq!(map_type_token("I8"), "i8");
        assert_eq!(map_type_token("I16"), "i16");
        assert_eq!(map_type_token("I32"), "i32");
        assert_eq!(map_type_token("I64"), "i64");
        assert_eq!(map_type_token("F32"), "f32");
        assert_eq!(map_type_token("F64"), "f64");
    }

    #[test]
    fn map_collection_types() {
        assert_eq!(map_type_token("List"), "Vec");
        assert_eq!(map_type_token("Map"), "std::collections::BTreeMap");
        assert_eq!(map_type_token("Set"), "std::collections::BTreeSet");
        assert_eq!(map_type_token("Sequence"), "Vec");
    }

    #[test]
    fn map_passthrough() {
        assert_eq!(map_type_token("Option"), "Option");
        assert_eq!(map_type_token("Result"), "Result");
        assert_eq!(map_type_token("MyCustomType"), "MyCustomType");
    }

    // ---- map_type_tokens ----

    #[test]
    fn map_tokens_empty() {
        assert_eq!(map_type_tokens(&[]), "()");
    }

    #[test]
    fn map_tokens_single() {
        let tokens = vec!["Int".to_string()];
        assert_eq!(map_type_tokens(&tokens), "i64");
    }

    #[test]
    fn map_tokens_generic() {
        let tokens: Vec<String> = vec!["List", "<", "Int", ">"]
            .into_iter()
            .map(String::from)
            .collect();
        assert_eq!(map_type_tokens(&tokens), "Vec<i64>");
    }

    #[test]
    fn map_tokens_strips_taint() {
        let tokens: Vec<String> = vec!["Int", "@", "taint", ":", "untrusted"]
            .into_iter()
            .map(String::from)
            .collect();
        assert_eq!(map_type_tokens(&tokens), "i64");
    }

    #[test]
    fn map_tokens_union_error() {
        let tokens: Vec<String> = vec!["Int", "|", "MyError"]
            .into_iter()
            .map(String::from)
            .collect();
        assert_eq!(map_type_tokens(&tokens), "Result<i64, MyError>");
    }

    #[test]
    fn map_tokens_refinement() {
        let tokens: Vec<String> = vec!["{", "x", ":", "Int", "|", "x", ">", "0", "}"]
            .into_iter()
            .map(String::from)
            .collect();
        assert_eq!(map_type_tokens(&tokens), "i64");
    }

    // ---- smart_join_type_tokens ----

    #[test]
    fn smart_join_no_space_after_open_angle() {
        let tokens = vec!["Vec", "<", "i64", ">"];
        // No space before `<`, after `<`, or before `>`
        assert_eq!(smart_join_type_tokens(&tokens), "Vec<i64>");
    }

    #[test]
    fn smart_join_no_space_after_ampersand() {
        let tokens = vec!["&", "str"];
        assert_eq!(smart_join_type_tokens(&tokens), "&str");
    }

    #[test]
    fn smart_join_mut_ref() {
        let tokens = vec!["&", "mut", "i64"];
        assert_eq!(smart_join_type_tokens(&tokens), "&mut i64");
    }

    // ---- is_const_name ----

    #[test]
    fn const_name_true() {
        assert!(is_const_name("MAX_SIZE"));
        assert!(is_const_name("TOTAL_TABLE_SIZE"));
    }

    #[test]
    fn const_name_false() {
        assert!(!is_const_name("MyType"));
        assert!(!is_const_name("x"));
        assert!(!is_const_name(""));
        assert!(!is_const_name("ALLCAPS")); // no underscore
    }

    // ---- is_user_type_name ----

    #[test]
    fn user_type_name_true() {
        assert!(is_user_type_name("MyStruct"));
        assert!(is_user_type_name("Point"));
        assert!(is_user_type_name("HuffmanTree"));
    }

    #[test]
    fn user_type_name_false_builtins() {
        assert!(!is_user_type_name("Int"));
        assert!(!is_user_type_name("Bool"));
        assert!(!is_user_type_name("String"));
        assert!(!is_user_type_name("List"));
        assert!(!is_user_type_name("Option"));
        assert!(!is_user_type_name("Result"));
    }

    #[test]
    fn user_type_name_false_lowercase() {
        assert!(!is_user_type_name("myvar"));
        assert!(!is_user_type_name(""));
    }

    // ---- collect_type_refs_from_tokens ----

    #[test]
    fn collect_type_refs_basic() {
        let tokens: Vec<String> = vec!["List", "<", "Point", ">"]
            .into_iter()
            .map(String::from)
            .collect();
        let mut out = std::collections::HashSet::new();
        collect_type_refs_from_tokens(&tokens, &mut out);
        assert!(out.contains("Point"));
        assert!(!out.contains("List")); // built-in, excluded
    }

    #[test]
    fn collect_type_refs_skips_annotations() {
        let tokens: Vec<String> = vec!["MyType", "@", "taint"]
            .into_iter()
            .map(String::from)
            .collect();
        let mut out = std::collections::HashSet::new();
        collect_type_refs_from_tokens(&tokens, &mut out);
        assert!(out.contains("MyType"));
        assert!(!out.contains("@"));
        assert!(!out.contains("taint"));
    }

    // ---- collect_type_refs_from_expr ----

    #[test]
    fn collect_refs_from_ident() {
        let mut out = std::collections::HashSet::new();
        collect_type_refs_from_expr(&Spanned::no_span(Expr::Ident("MyType".into())), &mut out);
        assert!(out.contains("MyType"));
    }

    #[test]
    fn collect_refs_from_ident_lowercase() {
        let mut out = std::collections::HashSet::new();
        collect_type_refs_from_expr(&Spanned::no_span(Expr::Ident("x".into())), &mut out);
        assert!(out.is_empty());
    }

    #[test]
    fn collect_refs_nested() {
        let mut out = std::collections::HashSet::new();
        let e = Spanned::no_span(Expr::Call {
            func: Box::new(Spanned::no_span(Expr::Ident("Create".into()))),
            args: vec![Spanned::no_span(Expr::Ident("Config".into()))],
        });
        collect_type_refs_from_expr(&e, &mut out);
        assert!(out.contains("Create"));
        assert!(out.contains("Config"));
    }

    // ---- detect_generic_arity ----

    #[test]
    fn detect_generic_one_type_param() {
        let tokens: Vec<String> = vec!["Wrapper", "<", "Int", ">"]
            .into_iter()
            .map(String::from)
            .collect();
        let mut map = std::collections::HashMap::new();
        let mut consts = std::collections::HashSet::new();
        detect_generic_arity(&tokens, &mut map, &mut consts);
        assert_eq!(map.get("Wrapper").map(|v| v.len()), Some(1));
    }

    #[test]
    fn detect_generic_const_param() {
        let tokens: Vec<String> = vec!["Buffer", "<", "MAX_SIZE", ">"]
            .into_iter()
            .map(String::from)
            .collect();
        let mut map = std::collections::HashMap::new();
        let mut consts = std::collections::HashSet::new();
        detect_generic_arity(&tokens, &mut map, &mut consts);
        assert!(consts.contains("MAX_SIZE"));
        assert!(matches!(
            map.get("Buffer").and_then(|v| v.first()),
            Some(GenericParamKind::Const)
        ));
    }

    // ---- generate_cargo_toml_impl ----

    #[test]
    fn cargo_toml_default_config() {
        let cfg = BackendConfig::default();
        let toml = generate_cargo_toml_impl("my_crate", &cfg, false, false);
        assert!(toml.contains("name = \"my_crate\""));
        assert!(toml.contains("edition = \"2024\""));
        assert!(!toml.contains("[profile.release]"));
    }

    #[test]
    fn cargo_toml_with_proptest() {
        let cfg = BackendConfig::default();
        let toml = generate_cargo_toml_impl("test_crate", &cfg, true, false);
        assert!(toml.contains("[dev-dependencies]"));
        assert!(toml.contains("proptest"));
    }

    #[test]
    fn cargo_toml_with_thiserror() {
        let cfg = BackendConfig::default();
        let toml = generate_cargo_toml_impl("err_crate", &cfg, false, true);
        assert!(toml.contains("thiserror"));
    }

    #[test]
    fn cargo_toml_cranelift() {
        let cfg = BackendConfig {
            backend: CodegenBackend::Cranelift,
            ..Default::default()
        };
        let toml = generate_cargo_toml_impl("fast", &cfg, false, false);
        assert!(toml.contains("Cranelift"));
    }

    #[test]
    fn cargo_toml_wasm() {
        let cfg = BackendConfig {
            target: CompileTarget::Wasm,
            ..Default::default()
        };
        let toml = generate_cargo_toml_impl("wasm_mod", &cfg, false, false);
        assert!(toml.contains("cdylib"));
        assert!(toml.contains("wasm32-wasip1"));
    }

    #[test]
    fn cargo_toml_runtime_checks() {
        let cfg = BackendConfig {
            runtime_checks: true,
            ..Default::default()
        };
        let toml = generate_cargo_toml_impl("rt_crate", &cfg, false, false);
        assert!(
            toml.contains("assura-runtime"),
            "runtime checks should add assura-runtime dep: {toml}"
        );
    }

    #[test]
    fn cargo_toml_no_runtime_checks() {
        let cfg = BackendConfig::default();
        let toml = generate_cargo_toml_impl("rt_crate", &cfg, false, false);
        assert!(
            !toml.contains("assura-runtime"),
            "default should not include assura-runtime: {toml}"
        );
    }

    // ---- generate_type_def ----

    #[test]
    fn type_def_struct() {
        let t = TypeDef {
            name: "Point".into(),
            type_params: vec![],
            body: TypeBody::Struct(vec![
                assura_ast::FieldDef {
                    name: "x".into(),
                    ty: Some(assura_ast::TypeExpr::named("Int")),
                    is_pub: true,
                },
                assura_ast::FieldDef {
                    name: "y".into(),
                    ty: Some(assura_ast::TypeExpr::named("Int")),
                    is_pub: true,
                },
            ]),
        };
        let mut code = String::new();
        generate_type_def(&t, &mut code);
        assert!(code.contains("pub struct Point"));
        assert!(code.contains("pub x: i64"));
        assert!(code.contains("pub y: i64"));
    }

    #[test]
    fn type_def_alias() {
        let t = TypeDef {
            name: "Identifier".into(),
            type_params: vec![],
            body: TypeBody::Alias(vec!["String".into()]),
        };
        let mut code = String::new();
        generate_type_def(&t, &mut code);
        assert!(code.contains("pub type Identifier = String"));
    }

    #[test]
    fn type_def_generic() {
        let t = TypeDef {
            name: "Wrapper".into(),
            type_params: vec!["T".into()],
            body: TypeBody::Empty,
        };
        let mut code = String::new();
        generate_type_def(&t, &mut code);
        assert!(code.contains("pub struct Wrapper<T>"));
    }

    #[test]
    fn type_def_refined() {
        let t = TypeDef {
            name: "PositiveInt".into(),
            type_params: vec![],
            body: TypeBody::Refined(vec![
                "n".into(),
                ":".into(),
                "Int".into(),
                "|".into(),
                "n".into(),
                ">".into(),
                "0".into(),
            ]),
        };
        let mut code = String::new();
        generate_type_def(&t, &mut code);
        assert!(code.contains("pub struct PositiveInt(pub i64)"));
    }

    // ---- extract_base_type_from_refined ----

    #[test]
    fn extract_base_type_after_colon() {
        let tokens: Vec<String> = vec!["n", ":", "Int", "|", "n", ">", "0"]
            .into_iter()
            .map(String::from)
            .collect();
        assert_eq!(extract_base_type_from_refined(&tokens), "i64");
    }

    #[test]
    fn extract_base_type_fallback() {
        let tokens: Vec<String> = vec!["Nat"].into_iter().map(String::from).collect();
        assert_eq!(extract_base_type_from_refined(&tokens), "u64");
    }

    // ---- expr_to_rust_static ----

    #[test]
    fn static_int_literal() {
        assert_eq!(
            expr_to_rust_static(&Spanned::no_span(Expr::Literal(Literal::Int("42".into())))),
            "42"
        );
    }

    #[test]
    fn static_binop() {
        let e = Spanned::no_span(Expr::BinOp {
            lhs: Box::new(Spanned::no_span(Expr::Ident("a".into()))),
            op: BinOp::Add,
            rhs: Box::new(Spanned::no_span(Expr::Ident("b".into()))),
        });
        assert_eq!(expr_to_rust_static(&e), "(a + b)");
    }

    #[test]
    fn static_ghost_erased() {
        let e = Spanned::no_span(Expr::Ghost(Box::new(Spanned::no_span(Expr::Ident(
            "x".into(),
        )))));
        let result = expr_to_rust_static(&e);
        assert!(result.contains("ghost:"));
        assert!(result.ends_with("()"));
    }

    #[test]
    fn static_forall_comment() {
        let e = Spanned::no_span(Expr::Forall {
            var: "i".into(),
            domain: Box::new(Spanned::no_span(Expr::Ident("items".into()))),
            body: Box::new(Spanned::no_span(Expr::Literal(Literal::Bool(true)))),
        });
        let result = expr_to_rust_static(&e);
        assert!(result.contains("forall"));
        assert!(result.ends_with("true"));
    }

    #[test]
    fn static_raw_single_token() {
        let e = Spanned::no_span(Expr::Raw(vec!["42".into()]));
        assert_eq!(expr_to_rust_static(&e), "42");
    }

    #[test]
    fn static_raw_eq_value() {
        let e = Spanned::no_span(Expr::Raw(vec!["=".into(), "100".into()]));
        assert_eq!(expr_to_rust_static(&e), "100");
    }

    // ---- collect_type_refs_from_type_expr ----

    #[test]
    fn collect_type_refs_type_expr_named() {
        let mut out = std::collections::HashSet::new();
        let te = assura_ast::TypeExpr::Named("MyStruct".into());
        collect_type_refs_from_type_expr(Some(&te), &mut out);
        assert!(out.contains("MyStruct"));
    }

    #[test]
    fn collect_type_refs_type_expr_named_builtin() {
        let mut out = std::collections::HashSet::new();
        let te = assura_ast::TypeExpr::Named("Int".into());
        collect_type_refs_from_type_expr(Some(&te), &mut out);
        assert!(out.is_empty());
    }

    #[test]
    fn collect_type_refs_type_expr_generic() {
        let mut out = std::collections::HashSet::new();
        let te = assura_ast::TypeExpr::Generic(
            "List".into(),
            vec![assura_ast::TypeExpr::Named("Point".into())],
        );
        collect_type_refs_from_type_expr(Some(&te), &mut out);
        assert!(out.contains("Point"));
        assert!(!out.contains("List")); // built-in
    }

    #[test]
    fn collect_type_refs_type_expr_none() {
        let mut out = std::collections::HashSet::new();
        collect_type_refs_from_type_expr(None, &mut out);
        assert!(out.is_empty());
    }

    // ---- detect_generic_arity_from_type_expr ----

    #[test]
    fn detect_generic_arity_type_expr_one_type_param() {
        let te = assura_ast::TypeExpr::Generic(
            "Buffer".into(),
            vec![assura_ast::TypeExpr::Named("Byte".into())],
        );
        let mut map = std::collections::HashMap::new();
        let mut consts = std::collections::HashSet::new();
        detect_generic_arity_from_type_expr(Some(&te), &mut map, &mut consts);
        assert_eq!(map.get("Buffer").map(|v| v.len()), Some(1));
        assert!(matches!(
            map.get("Buffer").and_then(|v| v.first()),
            Some(GenericParamKind::Type)
        ));
    }

    #[test]
    fn detect_generic_arity_type_expr_const_param() {
        let te = assura_ast::TypeExpr::Generic(
            "Buffer".into(),
            vec![assura_ast::TypeExpr::Named("MAX_SIZE".into())],
        );
        let mut map = std::collections::HashMap::new();
        let mut consts = std::collections::HashSet::new();
        detect_generic_arity_from_type_expr(Some(&te), &mut map, &mut consts);
        assert!(consts.contains("MAX_SIZE"));
        assert!(matches!(
            map.get("Buffer").and_then(|v| v.first()),
            Some(GenericParamKind::Const)
        ));
    }

    #[test]
    fn detect_generic_arity_type_expr_mixed() {
        let te = assura_ast::TypeExpr::Generic(
            "Table".into(),
            vec![
                assura_ast::TypeExpr::Named("Entry".into()),
                assura_ast::TypeExpr::Named("MAX_ENTRIES".into()),
            ],
        );
        let mut map = std::collections::HashMap::new();
        let mut consts = std::collections::HashSet::new();
        detect_generic_arity_from_type_expr(Some(&te), &mut map, &mut consts);
        assert_eq!(map.get("Table").map(|v| v.len()), Some(2));
        assert!(consts.contains("MAX_ENTRIES"));
    }

    #[test]
    fn detect_generic_arity_type_expr_widest_wins() {
        // First usage with 1 param
        let te1 = assura_ast::TypeExpr::Generic(
            "Container".into(),
            vec![assura_ast::TypeExpr::Named("Foo".into())],
        );
        // Second usage with 2 params
        let te2 = assura_ast::TypeExpr::Generic(
            "Container".into(),
            vec![
                assura_ast::TypeExpr::Named("Foo".into()),
                assura_ast::TypeExpr::Named("Bar".into()),
            ],
        );
        let mut map = std::collections::HashMap::new();
        let mut consts = std::collections::HashSet::new();
        detect_generic_arity_from_type_expr(Some(&te1), &mut map, &mut consts);
        detect_generic_arity_from_type_expr(Some(&te2), &mut map, &mut consts);
        assert_eq!(map.get("Container").map(|v| v.len()), Some(2));
    }

    // ---- detect_feature_max_as_type_from_type_expr ----

    #[test]
    fn feature_max_as_type_expr_found() {
        let te = assura_ast::TypeExpr::Generic(
            "Buffer".into(),
            vec![assura_ast::TypeExpr::Named("MAX_BUF_SIZE".into())],
        );
        let fm_set: std::collections::HashSet<&str> = ["MAX_BUF_SIZE"].iter().copied().collect();
        let mut out = std::collections::HashSet::new();
        detect_feature_max_as_type_from_type_expr(Some(&te), &fm_set, &mut out);
        assert!(out.contains("MAX_BUF_SIZE"));
    }

    #[test]
    fn feature_max_as_type_expr_not_in_generic() {
        // Named at top level (not inside Generic args) should NOT be found
        let te = assura_ast::TypeExpr::Named("MAX_BUF_SIZE".into());
        let fm_set: std::collections::HashSet<&str> = ["MAX_BUF_SIZE"].iter().copied().collect();
        let mut out = std::collections::HashSet::new();
        detect_feature_max_as_type_from_type_expr(Some(&te), &fm_set, &mut out);
        assert!(out.is_empty());
    }

    #[test]
    fn feature_max_as_type_expr_nested() {
        // List<Buffer<MAX_SIZE>> should find MAX_SIZE
        let te = assura_ast::TypeExpr::Generic(
            "List".into(),
            vec![assura_ast::TypeExpr::Generic(
                "Buffer".into(),
                vec![assura_ast::TypeExpr::Named("MAX_SIZE".into())],
            )],
        );
        let fm_set: std::collections::HashSet<&str> = ["MAX_SIZE"].iter().copied().collect();
        let mut out = std::collections::HashSet::new();
        detect_feature_max_as_type_from_type_expr(Some(&te), &fm_set, &mut out);
        assert!(out.contains("MAX_SIZE"));
    }

    // ---- qualifier stripping in map_type_tokens ----

    #[test]
    fn map_type_tokens_strips_secret_qualifier() {
        let tokens: Vec<String> = vec!["secret".into(), "Int".into()];
        assert_eq!(map_type_tokens(&tokens), "i64");
    }

    #[test]
    fn map_type_tokens_strips_tainted_qualifier() {
        let tokens: Vec<String> = vec!["tainted".into(), "String".into()];
        assert_eq!(map_type_tokens(&tokens), "String");
    }

    #[test]
    fn map_type_tokens_strips_linear_qualifier() {
        let tokens: Vec<String> = vec!["linear".into(), "Int".into()];
        assert_eq!(map_type_tokens(&tokens), "i64");
    }

    #[test]
    fn map_type_tokens_strips_untrusted_qualifier() {
        let tokens: Vec<String> = vec!["untrusted".into(), "Bytes".into()];
        assert_eq!(map_type_tokens(&tokens), "Vec<u8>");
    }

    #[test]
    fn map_type_tokens_strips_validated_qualifier() {
        let tokens: Vec<String> = vec!["validated".into(), "String".into()];
        assert_eq!(map_type_tokens(&tokens), "String");
    }
}