kain/codegen/

rust.rs

//! Rust Code Generation - Transpiles KAIN AST to Rust source code
//!
//! This module generates valid Rust code from a typed KAIN program.
//! The generated Rust can be compiled with `rustc` directly or integrated
//! into a Cargo project.

use crate::types::{TypedProgram, TypedItem};
use crate::error::{KainResult, KainError};
use crate::ast::{
    Type, Expr, Stmt, Block, BinaryOp, UnaryOp, Pattern, Function, Struct, Enum,
    Field, Variant, VariantFields, Impl, Param, MatchArm, CallArg, ElseBranch,
    VariantPatternFields, EnumVariantFields,
};
use crate::span::Span;

/// Generate Rust source code from a typed program
pub fn generate(program: &TypedProgram) -> KainResult<String> {
    let mut gen = RustGen::new();
    Ok(gen.gen_program(program))
}

// StringBuilder helper for accumulated output
struct StringBuilder {
    lines: Vec<String>,
}

impl StringBuilder {
    fn new() -> Self {
        Self { lines: Vec::new() }
    }

    fn push(&mut self, text: &str) {
        self.lines.push(text.to_string());
    }

    fn push_line(&mut self, text: &str) {
        self.lines.push(format!("{}\n", text));
    }

    fn build(&self) -> String {
        self.lines.join("")
    }
}

// Main Rust code generator
struct RustGen {
    output: StringBuilder,
    indent: usize,
}

impl RustGen {
    fn new() -> Self {
        Self {
            output: StringBuilder::new(),
            indent: 0,
        }
    }

    fn push_indent(&mut self) {
        self.indent += 1;
    }

    fn pop_indent(&mut self) {
        if self.indent > 0 {
            self.indent -= 1;
        }
    }

    fn indent_str(&self) -> String {
        "    ".repeat(self.indent)
    }

    fn write_line(&mut self, line: &str) {
        let indented = format!("{}{}", self.indent_str(), line);
        self.output.push_line(&indented);
    }

    fn write_blank(&mut self) {
        self.output.push_line("");
    }

    // Generate Rust code for an entire program
    fn gen_program(&mut self, program: &TypedProgram) -> String {
        // Header
        self.write_line("// Generated by KAIN Compiler (Project Ouroboros)");
        self.write_line("// Do not edit - regenerate from .kn source");
        self.write_blank();
        self.write_line("#![allow(unused_variables)]");
        self.write_line("#![allow(unused_mut)]");
        self.write_line("#![allow(dead_code)]");
        self.write_line("#![allow(unused_parens)]");
        self.write_blank();

        // Standard imports
        self.write_line("use std::collections::HashMap;");
        self.write_line("use std::rc::Rc;");
        self.write_line("use std::cell::RefCell;");
        self.write_blank();

        // Generate each item
        for item in &program.items {
            self.gen_item(item);
            self.write_blank();
        }

        self.output.build()
    }

    fn gen_item(&mut self, item: &TypedItem) {
        match item {
            TypedItem::Function(fn_typed) => self.gen_function(&fn_typed.ast),
            TypedItem::Struct(st) => self.gen_struct(&st.ast),
            TypedItem::Enum(en) => self.gen_enum(&en.ast),
            TypedItem::Impl(im) => self.gen_impl(&im.ast),
            _ => {} // Skip shaders, actors, components, traits, etc. for Rust output
        }
    }

    // Generate a function definition
    fn gen_function(&mut self, func: &Function) {
        let vis = match func.visibility {
            crate::ast::Visibility::Public => "pub ",
            _ => "",
        };

        // Parameters
        let params = self.gen_params(&func.params);

        // Return type
        let ret = if let Some(ty) = &func.return_type {
            format!(" -> {}", self.map_type(ty))
        } else {
            String::new()
        };

        // Function signature
        self.write_line(&format!("{}fn {}({}){} {{", vis, func.name, params, ret));
        self.push_indent();

        // Body
        self.gen_block(&func.body);

        self.pop_indent();
        self.write_line("}");
    }

    fn gen_params(&self, params: &[Param]) -> String {
        let parts: Vec<String> = params
            .iter()
            .map(|p| {
                let ty_str = self.map_type(&p.ty);
                if p.mutable {
                    format!("mut {}: {}", p.name, ty_str)
                } else {
                    format!("{}: {}", p.name, ty_str)
                }
            })
            .collect();
        parts.join(", ")
    }

    fn gen_struct(&mut self, struct_def: &Struct) {
        let vis = match struct_def.visibility {
            crate::ast::Visibility::Public => "pub ",
            _ => "",
        };

        self.write_line("#[derive(Debug, Clone)]");
        self.write_line(&format!("{}struct {} {{", vis, struct_def.name));
        self.push_indent();

        for field in &struct_def.fields {
            self.write_line(&format!("pub {}: {},", field.name, self.map_type(&field.ty)));
        }

        self.pop_indent();
        self.write_line("}");
    }

    fn gen_enum(&mut self, enum_def: &Enum) {
        self.write_line("#[derive(Debug, Clone, PartialEq)]");
        self.write_line(&format!("pub enum {} {{", enum_def.name));
        self.push_indent();

        for variant in &enum_def.variants {
            match &variant.fields {
                VariantFields::Unit => {
                    self.write_line(&format!("{},", variant.name));
                }
                VariantFields::Tuple(types) => {
                    let fields: Vec<String> = types.iter().map(|t| self.map_type(t)).collect();
                    self.write_line(&format!("{}({}),", variant.name, fields.join(", ")));
                }
                VariantFields::Struct(fields) => {
                    self.write_line(&format!("{} {{", variant.name));
                    self.push_indent();
                    for f in fields {
                        self.write_line(&format!("{}: {},", f.name, self.map_type(&f.ty)));
                    }
                    self.pop_indent();
                    self.write_line("},");
                }
            }
        }

        self.pop_indent();
        self.write_line("}");
    }

    fn gen_impl(&mut self, impl_def: &Impl) {
        let target = self.map_type(&impl_def.target_type);
        
        if let Some(trait_name) = &impl_def.trait_name {
            self.write_line(&format!("impl {} for {} {{", trait_name, target));
        } else {
            self.write_line(&format!("impl {} {{", target));
        }
        self.push_indent();

        for method in &impl_def.methods {
            self.gen_function(method);
            self.write_blank();
        }

        self.pop_indent();
        self.write_line("}");
    }

    fn gen_block(&mut self, block: &Block) {
        for stmt in &block.stmts {
            self.gen_stmt(stmt);
        }
    }

    fn gen_stmt(&mut self, stmt: &Stmt) {
        match stmt {
            Stmt::Let { pattern, ty, value, .. } => {
                let pat_str = self.gen_pattern(pattern);
                let ty_str = ty.as_ref().map(|t| format!(": {}", self.map_type(t))).unwrap_or_default();
                if let Some(val) = value {
                    self.write_line(&format!("let {}{} = {};", pat_str, ty_str, self.gen_expr(val)));
                } else {
                    self.write_line(&format!("let {}{};", pat_str, ty_str));
                }
            }

            Stmt::Return(maybe_expr, _) => {
                if let Some(expr) = maybe_expr {
                    self.write_line(&format!("return {};", self.gen_expr(expr)));
                } else {
                    self.write_line("return;");
                }
            }

            Stmt::Break(maybe_expr, _) => {
                if let Some(expr) = maybe_expr {
                    self.write_line(&format!("break {};", self.gen_expr(expr)));
                } else {
                    self.write_line("break;");
                }
            }

            Stmt::Continue(_) => {
                self.write_line("continue;");
            }

            Stmt::For { binding, iter, body, .. } => {
                let pat = self.gen_pattern(binding);
                self.write_line(&format!("for {} in {} {{", pat, self.gen_expr(iter)));
                self.push_indent();
                self.gen_block(body);
                self.pop_indent();
                self.write_line("}");
            }

            Stmt::While { condition, body, .. } => {
                self.write_line(&format!("while {} {{", self.gen_expr(condition)));
                self.push_indent();
                self.gen_block(body);
                self.pop_indent();
                self.write_line("}");
            }

            Stmt::Loop { body, .. } => {
                self.write_line("loop {");
                self.push_indent();
                self.gen_block(body);
                self.pop_indent();
                self.write_line("}");
            }

            Stmt::Expr(expr) => {
                // Check for assignment expression
                if let Expr::Assign { target, value, .. } = expr {
                    self.write_line(&format!("{} = {};", self.gen_expr(target), self.gen_expr(value)));
                } else {
                    self.write_line(&format!("{};", self.gen_expr(expr)));
                }
            }

            Stmt::Item(item) => {
                // Nested items - convert to TypedItem and generate
                // For now, skip nested items in Rust output
            }
        }
    }

    fn gen_expr(&self, expr: &Expr) -> String {
        match expr {
            Expr::Int(n, _) => n.to_string(),
            Expr::Float(f, _) => format!("{:.1}", f),
            Expr::String(s, _) => format!("\"{}\".to_string()", self.escape_string(s)),
            Expr::Bool(b, _) => if *b { "true".to_string() } else { "false".to_string() },
            Expr::None(_) => "None".to_string(),
            Expr::Ident(name, _) => name.clone(),

            Expr::Binary { left, op, right, .. } => {
                let l = self.gen_expr(left);
                let r = self.gen_expr(right);
                let rust_op = self.map_binop(op);
                format!("({} {} {})", l, rust_op, r)
            }

            Expr::Unary { op, operand, .. } => {
                let o = self.gen_expr(operand);
                let rust_op = self.map_unaryop(op);
                format!("({}{})", rust_op, o)
            }

            Expr::Call { callee, args, .. } => {
                let fn_name = self.gen_expr(callee);

                // Handle KAIN builtins
                if fn_name == "println" || fn_name == "print" {
                    let arg_strs: Vec<String> = args.iter().map(|a| self.gen_expr(&a.value)).collect();
                    let placeholders: Vec<&str> = arg_strs.iter().map(|_| "{}").collect();
                    let format_str = format!("\"{}\"", placeholders.join(" "));
                    if !arg_strs.is_empty() {
                        return format!("{}!({}, {})", fn_name, format_str, arg_strs.join(", "));
                    }
                    return format!("{}!()", fn_name);
                }

                let arg_strs: Vec<String> = args.iter().map(|a| {
                    if let Some(name) = &a.name {
                        format!("{}: {}", name, self.gen_expr(&a.value))
                    } else {
                        self.gen_expr(&a.value)
                    }
                }).collect();
                format!("{}({})", fn_name, arg_strs.join(", "))
            }

            Expr::MethodCall { receiver, method, args, .. } => {
                let recv = self.gen_expr(receiver);
                let arg_strs: Vec<String> = args.iter().map(|a| self.gen_expr(&a.value)).collect();
                format!("{}.{}({})", recv, method, arg_strs.join(", "))
            }

            Expr::Field { object, field, .. } => {
                format!("{}.{}", self.gen_expr(object), field)
            }

            Expr::Index { object, index, .. } => {
                format!("{}[{}]", self.gen_expr(object), self.gen_expr(index))
            }

            Expr::Array(elements, _) => {
                let elems: Vec<String> = elements.iter().map(|e| self.gen_expr(e)).collect();
                format!("vec![{}]", elems.join(", "))
            }

            Expr::Tuple(elements, _) => {
                let elems: Vec<String> = elements.iter().map(|e| self.gen_expr(e)).collect();
                if elems.len() == 1 {
                    format!("({},)", elems[0])
                } else {
                    format!("({})", elems.join(", "))
                }
            }

            Expr::Struct { name, fields, .. } => {
                let field_strs: Vec<String> = fields
                    .iter()
                    .map(|(fname, fval)| format!("{}: {}", fname, self.gen_expr(fval)))
                    .collect();
                format!("{} {{ {} }}", name, field_strs.join(", "))
            }

            Expr::EnumVariant { enum_name, variant, fields, .. } => {
                match fields {
                    EnumVariantFields::Unit => format!("{}::{}", enum_name, variant),
                    EnumVariantFields::Tuple(args) => {
                        let arg_strs: Vec<String> = args.iter().map(|a| self.gen_expr(a)).collect();
                        format!("{}::{}({})", enum_name, variant, arg_strs.join(", "))
                    }
                    EnumVariantFields::Struct(fields) => {
                        let field_strs: Vec<String> = fields
                            .iter()
                            .map(|(n, v)| format!("{}: {}", n, self.gen_expr(v)))
                            .collect();
                        format!("{}::{} {{ {} }}", enum_name, variant, field_strs.join(", "))
                    }
                }
            }

            Expr::If { condition, then_branch, else_branch, .. } => {
                let cond = self.gen_expr(condition);
                let then_str = self.gen_block_expr(then_branch);
                if let Some(else_b) = else_branch {
                    let else_str = self.gen_else_branch(else_b);
                    format!("if {} {{ {} }} else {}", cond, then_str, else_str)
                } else {
                    format!("if {} {{ {} }}", cond, then_str)
                }
            }

            Expr::Match { scrutinee, arms, .. } => {
                let scrut = self.gen_expr(scrutinee);
                let mut result = format!("match {} {{\n", scrut);
                for arm in arms {
                    let pat = self.gen_pattern(&arm.pattern);
                    let body = self.gen_expr(&arm.body);
                    result.push_str(&format!("    {} => {{ {} }}\n", pat, body));
                }
                result.push_str("}");
                result
            }

            Expr::Lambda { params, body, .. } => {
                let param_strs: Vec<String> = params.iter().map(|p| p.name.clone()).collect();
                format!("|{}| {}", param_strs.join(", "), self.gen_expr(body))
            }

            Expr::Await(future, _) => {
                format!("{}.await", self.gen_expr(future))
            }

            Expr::Ref { mutable, value, .. } => {
                if *mutable {
                    format!("&mut {}", self.gen_expr(value))
                } else {
                    format!("&{}", self.gen_expr(value))
                }
            }

            Expr::Deref(inner, _) => {
                format!("*{}", self.gen_expr(inner))
            }

            Expr::Cast { value, target, .. } => {
                format!("{} as {}", self.gen_expr(value), self.map_type(target))
            }

            Expr::Try(inner, _) => {
                format!("{}?", self.gen_expr(inner))
            }

            Expr::Range { start, end, inclusive, .. } => {
                let s = start.as_ref().map(|e| self.gen_expr(e)).unwrap_or_default();
                let e = end.as_ref().map(|e| self.gen_expr(e)).unwrap_or_default();
                if *inclusive {
                    format!("{}..={}", s, e)
                } else {
                    format!("{}..{}", s, e)
                }
            }

            Expr::Block(block, _) => {
                format!("{{ {} }}", self.gen_block_expr(block))
            }

            Expr::Paren(inner, _) => {
                format!("({})", self.gen_expr(inner))
            }

            Expr::Return(maybe_expr, _) => {
                if let Some(expr) = maybe_expr {
                    format!("return {}", self.gen_expr(expr))
                } else {
                    "return".to_string()
                }
            }

            Expr::Break(maybe_expr, _) => {
                if let Some(expr) = maybe_expr {
                    format!("break {}", self.gen_expr(expr))
                } else {
                    "break".to_string()
                }
            }

            Expr::Continue(_) => "continue".to_string(),

            // FString - format string with interpolation
            Expr::FString(parts, _) => {
                let mut format_str = String::new();
                let mut args = Vec::new();
                for part in parts {
                    if let Expr::String(s, _) = part {
                        format_str.push_str(s);
                    } else {
                        format_str.push_str("{}");
                        args.push(self.gen_expr(part));
                    }
                }
                if args.is_empty() {
                    format!("format!(\"{}\")", format_str)
                } else {
                    format!("format!(\"{}\", {})", format_str, args.join(", "))
                }
            }

            // Fallback for unhandled expressions
            _ => "/* unhandled expr */".to_string(),
        }
    }

    fn gen_block_expr(&self, block: &Block) -> String {
        // For expression context, return the last expression value
        if block.stmts.is_empty() {
            return "()".to_string();
        }
        
        let mut parts = Vec::new();
        for stmt in &block.stmts {
            match stmt {
                Stmt::Expr(e) => parts.push(self.gen_expr(e)),
                Stmt::Return(Some(e), _) => parts.push(format!("return {}", self.gen_expr(e))),
                Stmt::Return(None, _) => parts.push("return".to_string()),
                _ => {} // Skip other statements in expression context
            }
        }
        parts.join("; ")
    }

    fn gen_else_branch(&self, branch: &ElseBranch) -> String {
        match branch {
            ElseBranch::Else(block) => {
                format!("{{ {} }}", self.gen_block_expr(block))
            }
            ElseBranch::ElseIf(cond, block, maybe_else) => {
                let cond_str = self.gen_expr(cond);
                let then_str = self.gen_block_expr(block);
                if let Some(else_b) = maybe_else {
                    format!("if {} {{ {} }} else {}", cond_str, then_str, self.gen_else_branch(else_b))
                } else {
                    format!("if {} {{ {} }}", cond_str, then_str)
                }
            }
        }
    }

    fn gen_pattern(&self, pattern: &Pattern) -> String {
        match pattern {
            Pattern::Wildcard(_) => "_".to_string(),
            Pattern::Literal(expr) => self.gen_expr(expr),
            Pattern::Binding { name, mutable, .. } => {
                if *mutable {
                    format!("mut {}", name)
                } else {
                    name.clone()
                }
            }
            Pattern::Struct { name, fields, rest, .. } => {
                let field_pats: Vec<String> = fields
                    .iter()
                    .map(|(n, p)| format!("{}: {}", n, self.gen_pattern(p)))
                    .collect();
                if *rest {
                    format!("{} {{ {}, .. }}", name, field_pats.join(", "))
                } else {
                    format!("{} {{ {} }}", name, field_pats.join(", "))
                }
            }
            Pattern::Tuple(pats, _) => {
                let pat_strs: Vec<String> = pats.iter().map(|p| self.gen_pattern(p)).collect();
                format!("({})", pat_strs.join(", "))
            }
            Pattern::Variant { enum_name, variant, fields, .. } => {
                let full_name = if let Some(en) = enum_name {
                    format!("{}::{}", en, variant)
                } else {
                    variant.clone()
                };

                match fields {
                    VariantPatternFields::Unit => full_name,
                    VariantPatternFields::Tuple(pats) => {
                        let pat_strs: Vec<String> = pats.iter().map(|p| self.gen_pattern(p)).collect();
                        format!("{}({})", full_name, pat_strs.join(", "))
                    }
                    VariantPatternFields::Struct(field_pats) => {
                        let field_strs: Vec<String> = field_pats
                            .iter()
                            .map(|(n, p)| format!("{}: {}", n, self.gen_pattern(p)))
                            .collect();
                        format!("{} {{ {} }}", full_name, field_strs.join(", "))
                    }
                }
            }
            Pattern::Slice { patterns, rest, .. } => {
                let mut pat_strs: Vec<String> = patterns.iter().map(|p| self.gen_pattern(p)).collect();
                if let Some(rest_name) = rest {
                    pat_strs.push(format!("{} @ ..", rest_name));
                }
                format!("[{}]", pat_strs.join(", "))
            }
            Pattern::Or(pats, _) => {
                let pat_strs: Vec<String> = pats.iter().map(|p| self.gen_pattern(p)).collect();
                pat_strs.join(" | ")
            }
            Pattern::Range { start, end, inclusive, .. } => {
                let s = start.as_ref().map(|e| self.gen_expr(e)).unwrap_or_default();
                let e = end.as_ref().map(|e| self.gen_expr(e)).unwrap_or_default();
                if *inclusive {
                    format!("{}..={}", s, e)
                } else {
                    format!("{}..{}", s, e)
                }
            }
        }
    }

    // Type mapping from KAIN to Rust
    fn map_type(&self, ty: &Type) -> String {
        match ty {
            Type::Named { name, generics, .. } => {
                let rust_name = match name.as_str() {
                    "Int" => "i64",
                    "Float" => "f64",
                    "Bool" => "bool",
                    "String" => "String",
                    "Unit" => "()",
                    "Array" => "Vec",
                    _ => name,
                };

                if generics.is_empty() {
                    rust_name.to_string()
                } else {
                    let gen_strs: Vec<String> = generics.iter().map(|g| self.map_type(g)).collect();
                    format!("{}<{}>", rust_name, gen_strs.join(", "))
                }
            }
            Type::Tuple(types, _) => {
                let type_strs: Vec<String> = types.iter().map(|t| self.map_type(t)).collect();
                format!("({})", type_strs.join(", "))
            }
            Type::Array(inner, size, _) => {
                format!("[{}; {}]", self.map_type(inner), size)
            }
            Type::Slice(inner, _) => {
                format!("[{}]", self.map_type(inner))
            }
            Type::Ref { mutable, inner, .. } => {
                if *mutable {
                    format!("&mut {}", self.map_type(inner))
                } else {
                    format!("&{}", self.map_type(inner))
                }
            }
            Type::Function { params, return_type, .. } => {
                let param_strs: Vec<String> = params.iter().map(|p| self.map_type(p)).collect();
                format!("fn({}) -> {}", param_strs.join(", "), self.map_type(return_type))
            }
            Type::Option(inner, _) => {
                format!("Option<{}>", self.map_type(inner))
            }
            Type::Result(ok, err, _) => {
                format!("Result<{}, {}>", self.map_type(ok), self.map_type(err))
            }
            Type::Infer(_) => "_".to_string(),
            Type::Never(_) => "!".to_string(),
            Type::Unit(_) => "()".to_string(),
            Type::Impl { trait_name, generics, .. } => {
                if generics.is_empty() {
                    format!("impl {}", trait_name)
                } else {
                    let gen_strs: Vec<String> = generics.iter().map(|g| self.map_type(g)).collect();
                    format!("impl {}<{}>", trait_name, gen_strs.join(", "))
                }
            }
        }
    }

    fn map_binop(&self, op: &BinaryOp) -> &'static str {
        match op {
            BinaryOp::Add => "+",
            BinaryOp::Sub => "-",
            BinaryOp::Mul => "*",
            BinaryOp::Div => "/",
            BinaryOp::Mod => "%",
            BinaryOp::Pow => ".pow", // Rust uses method call, but we approximate
            BinaryOp::And => "&&",
            BinaryOp::Or => "||",
            BinaryOp::Eq => "==",
            BinaryOp::Ne => "!=",
            BinaryOp::Lt => "<",
            BinaryOp::Le => "<=",
            BinaryOp::Gt => ">",
            BinaryOp::Ge => ">=",
            BinaryOp::BitAnd => "&",
            BinaryOp::BitOr => "|",
            BinaryOp::BitXor => "^",
            BinaryOp::Shl => "<<",
            BinaryOp::Shr => ">>",
            BinaryOp::Assign => "=",
            BinaryOp::AddAssign => "+=",
            BinaryOp::SubAssign => "-=",
            BinaryOp::MulAssign => "*=",
            BinaryOp::DivAssign => "/=",
            BinaryOp::Range => "..",
            BinaryOp::RangeInclusive => "..=",
        }
    }

    fn map_unaryop(&self, op: &UnaryOp) -> &'static str {
        match op {
            UnaryOp::Not => "!",
            UnaryOp::Neg => "-",
            UnaryOp::Ref => "&",
            UnaryOp::RefMut => "&mut ",
            UnaryOp::BitNot => "!", // Rust uses ! for bitwise not too
            UnaryOp::Deref => "*",
        }
    }

    fn escape_string(&self, s: &str) -> String {
        let mut result = String::new();
        for c in s.chars() {
            match c {
                '\n' => result.push_str("\\n"),
                '\t' => result.push_str("\\t"),
                '\\' => result.push_str("\\\\"),
                '"' => result.push_str("\\\""),
                _ => result.push(c),
            }
        }
        result
    }
}

/// Generate a Cargo.toml for the transpiled Rust project
pub fn gen_cargo_toml(name: &str, deps: &[&str]) -> String {
    let mut sb = StringBuilder::new();

    sb.push_line("[package]");
    sb.push_line(&format!("name = \"{}\"", name));
    sb.push_line("version = \"0.1.0\"");
    sb.push_line("edition = \"2021\"");
    sb.push_line("");
    sb.push_line("# Generated by KAIN Compiler");
    sb.push_line("");
    sb.push_line("[dependencies]");

    for dep in deps {
        sb.push_line(&format!("{} = \"*\"", dep));
    }

    sb.build()
}

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

    #[test]
    fn test_type_mapping() {
        let gen = RustGen::new();
        let int_ty = Type::Named {
            name: "Int".to_string(),
            generics: vec![],
            span: Span::default(),
        };
        assert_eq!(gen.map_type(&int_ty), "i64");
    }
}