sage_codegen/

generator.rs

//! Main code generator.

use crate::emit::Emitter;
use sage_parser::{
    AgentDecl, BinOp, Block, EventKind, Expr, FnDecl, Literal, Program, Stmt, StringPart,
    UnaryOp,
};
use sage_types::TypeExpr;

/// Generated Rust project files.
pub struct GeneratedProject {
    /// The main.rs content.
    pub main_rs: String,
    /// The Cargo.toml content.
    pub cargo_toml: String,
}

/// Generate Rust code from a Sage program.
pub fn generate(program: &Program, project_name: &str) -> GeneratedProject {
    let mut gen = Generator::new();
    let main_rs = gen.generate_program(program);
    let cargo_toml = gen.generate_cargo_toml(project_name);
    GeneratedProject { main_rs, cargo_toml }
}

struct Generator {
    emit: Emitter,
}

impl Generator {
    fn new() -> Self {
        Self {
            emit: Emitter::new(),
        }
    }

    fn generate_program(&mut self, program: &Program) -> String {
        // Prelude
        self.emit.writeln("//! Generated by Sage compiler. Do not edit.");
        self.emit.blank_line();
        self.emit.writeln("use sage_runtime::prelude::*;");
        self.emit.blank_line();

        // Functions
        for func in &program.functions {
            self.generate_function(func);
            self.emit.blank_line();
        }

        // Agents
        for agent in &program.agents {
            self.generate_agent(agent);
            self.emit.blank_line();
        }

        // Entry point
        self.generate_main(&program.run_agent.name);

        std::mem::take(&mut self.emit).finish()
    }

    fn generate_cargo_toml(&self, name: &str) -> String {
        // Use a relative path that works from target/sage/<project>/
        // This assumes the standard project layout
        format!(
            r#"[package]
name = "{name}"
version = "0.1.0"
edition = "2021"

[dependencies]
sage-runtime = {{ path = "../../../crates/sage-runtime" }}
tokio = {{ version = "1", features = ["full"] }}
serde = {{ version = "1", features = ["derive"] }}
serde_json = "1"

# Standalone project, not part of parent workspace
[workspace]
"#
        )
    }

    fn generate_function(&mut self, func: &FnDecl) {
        // Function signature
        self.emit.write("fn ");
        self.emit.write(&func.name.name);
        self.emit.write("(");

        for (i, param) in func.params.iter().enumerate() {
            if i > 0 {
                self.emit.write(", ");
            }
            self.emit.write(&param.name.name);
            self.emit.write(": ");
            self.emit_type(&param.ty);
        }

        self.emit.write(") -> ");
        self.emit_type(&func.return_ty);
        self.emit.write(" ");
        self.generate_block(&func.body);
    }

    fn generate_agent(&mut self, agent: &AgentDecl) {
        let name = &agent.name.name;

        // Struct definition
        self.emit.write("struct ");
        self.emit.write(name);
        if agent.beliefs.is_empty() {
            self.emit.writeln(";");
        } else {
            self.emit.writeln(" {");
            self.emit.indent();
            for belief in &agent.beliefs {
                self.emit.write(&belief.name.name);
                self.emit.write(": ");
                self.emit_type(&belief.ty);
                self.emit.writeln(",");
            }
            self.emit.dedent();
            self.emit.writeln("}");
        }
        self.emit.blank_line();

        // Find the output type from the start handler
        let output_type = self.infer_agent_output_type(agent);

        // Impl block
        self.emit.write("impl ");
        self.emit.write(name);
        self.emit.writeln(" {");
        self.emit.indent();

        // Generate handlers
        for handler in &agent.handlers {
            if let EventKind::Start = &handler.event {
                self.emit.write("async fn on_start(self, ctx: AgentContext<");
                self.emit.write(&output_type);
                self.emit.write(">) -> SageResult<");
                self.emit.write(&output_type);
                self.emit.writeln("> {");
                self.emit.indent();
                self.generate_block_contents(&handler.body);
                self.emit.dedent();
                self.emit.writeln("}");
            }
        }

        self.emit.dedent();
        self.emit.writeln("}");
    }

    fn generate_main(&mut self, entry_agent: &str) {
        self.emit.writeln("#[tokio::main]");
        self.emit.writeln("async fn main() -> Result<(), Box<dyn std::error::Error>> {");
        self.emit.indent();

        self.emit.write("let handle = sage_runtime::spawn(|ctx| ");
        self.emit.write(entry_agent);
        self.emit.writeln(".on_start(ctx));");
        self.emit.writeln("let result = handle.result().await?;");
        self.emit.writeln("println!(\"{:?}\", result);");
        self.emit.writeln("Ok(())");

        self.emit.dedent();
        self.emit.writeln("}");
    }

    fn generate_block(&mut self, block: &Block) {
        self.emit.open_brace();
        self.generate_block_contents(block);
        self.emit.close_brace();
    }

    fn generate_block_inline(&mut self, block: &Block) {
        self.emit.open_brace();
        self.generate_block_contents(block);
        self.emit.close_brace_inline();
    }

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

    fn generate_stmt(&mut self, stmt: &Stmt) {
        match stmt {
            Stmt::Let { name, ty, value, .. } => {
                self.emit.write("let ");
                if ty.is_some() {
                    self.emit.write(&name.name);
                    self.emit.write(": ");
                    self.emit_type(ty.as_ref().unwrap());
                } else {
                    self.emit.write(&name.name);
                }
                self.emit.write(" = ");
                self.generate_expr(value);
                self.emit.writeln(";");
            }

            Stmt::Assign { name, value, .. } => {
                self.emit.write(&name.name);
                self.emit.write(" = ");
                self.generate_expr(value);
                self.emit.writeln(";");
            }

            Stmt::Return { value, .. } => {
                self.emit.write("return ");
                if let Some(expr) = value {
                    self.generate_expr(expr);
                }
                self.emit.writeln(";");
            }

            Stmt::If {
                condition,
                then_block,
                else_block,
                ..
            } => {
                self.emit.write("if ");
                self.generate_expr(condition);
                self.emit.write(" ");
                if else_block.is_some() {
                    self.generate_block_inline(then_block);
                    self.emit.write(" else ");
                    match else_block.as_ref().unwrap() {
                        sage_parser::ElseBranch::Block(block) => {
                            self.generate_block(block);
                        }
                        sage_parser::ElseBranch::ElseIf(stmt) => {
                            self.generate_stmt(stmt);
                        }
                    }
                } else {
                    self.generate_block(then_block);
                }
            }

            Stmt::For { var, iter, body, .. } => {
                self.emit.write("for ");
                self.emit.write(&var.name);
                self.emit.write(" in ");
                self.generate_expr(iter);
                self.emit.write(" ");
                self.generate_block(body);
            }

            Stmt::While {
                condition, body, ..
            } => {
                self.emit.write("while ");
                self.generate_expr(condition);
                self.emit.write(" ");
                self.generate_block(body);
            }

            Stmt::Expr { expr, .. } => {
                // Handle emit specially
                if let Expr::Emit { value, .. } = expr {
                    self.emit.write("return ctx.emit(");
                    self.generate_expr(value);
                    self.emit.writeln(");");
                } else {
                    self.generate_expr(expr);
                    self.emit.writeln(";");
                }
            }
        }
    }

    fn generate_expr(&mut self, expr: &Expr) {
        match expr {
            Expr::Literal { value, .. } => {
                self.emit_literal(value);
            }

            Expr::Var { name, .. } => {
                self.emit.write(&name.name);
            }

            Expr::Binary {
                op, left, right, ..
            } => {
                // Handle string concatenation specially
                if matches!(op, BinOp::Concat) {
                    self.emit.write("format!(\"{}{}\", ");
                    self.generate_expr(left);
                    self.emit.write(", ");
                    self.generate_expr(right);
                    self.emit.write(")");
                } else {
                    self.emit.write("(");
                    self.generate_expr(left);
                    self.emit.write(" ");
                    self.emit_binop(op);
                    self.emit.write(" ");
                    self.generate_expr(right);
                    self.emit.write(")");
                }
            }

            Expr::Unary { op, operand, .. } => {
                self.emit_unaryop(op);
                self.generate_expr(operand);
            }

            Expr::Call { name, args, .. } => {
                let fn_name = &name.name;

                // Handle builtins
                match fn_name.as_str() {
                    "print" => {
                        self.emit.write("println!(\"{}\", ");
                        self.generate_expr(&args[0]);
                        self.emit.write(")");
                    }
                    "str" => {
                        self.generate_expr(&args[0]);
                        self.emit.write(".to_string()");
                    }
                    "len" => {
                        self.generate_expr(&args[0]);
                        self.emit.write(".len() as i64");
                    }
                    _ => {
                        self.emit.write(fn_name);
                        self.emit.write("(");
                        for (i, arg) in args.iter().enumerate() {
                            if i > 0 {
                                self.emit.write(", ");
                            }
                            self.generate_expr(arg);
                        }
                        self.emit.write(")");
                    }
                }
            }

            Expr::SelfField { field, .. } => {
                self.emit.write("self.");
                self.emit.write(&field.name);
            }

            Expr::SelfMethodCall { method, args, .. } => {
                self.emit.write("self.");
                self.emit.write(&method.name);
                self.emit.write("(");
                for (i, arg) in args.iter().enumerate() {
                    if i > 0 {
                        self.emit.write(", ");
                    }
                    self.generate_expr(arg);
                }
                self.emit.write(")");
            }

            Expr::List { elements, .. } => {
                self.emit.write("vec![");
                for (i, elem) in elements.iter().enumerate() {
                    if i > 0 {
                        self.emit.write(", ");
                    }
                    self.generate_expr(elem);
                }
                self.emit.write("]");
            }

            Expr::Paren { inner, .. } => {
                self.emit.write("(");
                self.generate_expr(inner);
                self.emit.write(")");
            }

            Expr::Infer { template, .. } => {
                self.emit.write("ctx.infer_string(&");
                self.emit_string_template(template);
                self.emit.write(").await?");
            }

            Expr::Spawn { agent, fields, .. } => {
                self.emit.write("sage_runtime::spawn(|ctx| ");
                self.emit.write(&agent.name);
                if fields.is_empty() {
                    self.emit.write(".on_start(ctx))");
                } else {
                    self.emit.write(" { ");
                    for (i, field) in fields.iter().enumerate() {
                        if i > 0 {
                            self.emit.write(", ");
                        }
                        self.emit.write(&field.name.name);
                        self.emit.write(": ");
                        self.generate_expr(&field.value);
                    }
                    self.emit.write(" }.on_start(ctx))");
                }
            }

            Expr::Await { handle, .. } => {
                self.generate_expr(handle);
                self.emit.write(".result().await?");
            }

            Expr::Send { handle, message, .. } => {
                self.generate_expr(handle);
                self.emit.write(".send(sage_runtime::Message::new(");
                self.generate_expr(message);
                self.emit.write(")?).await?");
            }

            Expr::Emit { value, .. } => {
                self.emit.write("ctx.emit(");
                self.generate_expr(value);
                self.emit.write(")");
            }

            Expr::StringInterp { template, .. } => {
                self.emit_string_template(template);
            }
        }
    }

    fn emit_literal(&mut self, lit: &Literal) {
        match lit {
            Literal::Int(n) => {
                self.emit.write(&format!("{n}_i64"));
            }
            Literal::Float(f) => {
                self.emit.write(&format!("{f}_f64"));
            }
            Literal::Bool(b) => {
                self.emit.write(if *b { "true" } else { "false" });
            }
            Literal::String(s) => {
                // Escape the string for Rust
                self.emit.write("\"");
                for c in s.chars() {
                    match c {
                        '"' => self.emit.write_raw("\\\""),
                        '\\' => self.emit.write_raw("\\\\"),
                        '\n' => self.emit.write_raw("\\n"),
                        '\r' => self.emit.write_raw("\\r"),
                        '\t' => self.emit.write_raw("\\t"),
                        _ => self.emit.write_raw(&c.to_string()),
                    }
                }
                self.emit.write("\".to_string()");
            }
        }
    }

    fn emit_string_template(&mut self, template: &sage_parser::StringTemplate) {
        if !template.has_interpolations() {
            // Simple string literal
            if let Some(StringPart::Literal(s)) = template.parts.first() {
                self.emit.write("\"");
                self.emit.write_raw(s);
                self.emit.write("\".to_string()");
            }
            return;
        }

        // Build format string and args
        self.emit.write("format!(\"");
        for part in &template.parts {
            match part {
                StringPart::Literal(s) => {
                    // Escape braces for format string
                    let escaped = s.replace('{', "{{").replace('}', "}}");
                    self.emit.write_raw(&escaped);
                }
                StringPart::Interpolation(_) => {
                    self.emit.write_raw("{}");
                }
            }
        }
        self.emit.write("\"");

        // Add the interpolation args
        for part in &template.parts {
            if let StringPart::Interpolation(ident) = part {
                self.emit.write(", ");
                self.emit.write(&ident.name);
            }
        }
        self.emit.write(")");
    }

    fn emit_type(&mut self, ty: &TypeExpr) {
        match ty {
            TypeExpr::Int => self.emit.write("i64"),
            TypeExpr::Float => self.emit.write("f64"),
            TypeExpr::Bool => self.emit.write("bool"),
            TypeExpr::String => self.emit.write("String"),
            TypeExpr::Unit => self.emit.write("()"),
            TypeExpr::List(inner) => {
                self.emit.write("Vec<");
                self.emit_type(inner);
                self.emit.write(">");
            }
            TypeExpr::Option(inner) => {
                self.emit.write("Option<");
                self.emit_type(inner);
                self.emit.write(">");
            }
            TypeExpr::Inferred(inner) => {
                // Inferred<T> just becomes T at runtime
                self.emit_type(inner);
            }
            TypeExpr::Agent(agent_name) => {
                // Agent handles use the agent's output type, but we don't know it here
                // For now, just use a generic output type
                self.emit.write("AgentHandle<");
                self.emit.write(&agent_name.name);
                self.emit.write("Output>");
            }
            TypeExpr::Named(name) => {
                self.emit.write(&name.name);
            }
        }
    }

    fn emit_binop(&mut self, op: &BinOp) {
        let s = match op {
            BinOp::Add => "+",
            BinOp::Sub => "-",
            BinOp::Mul => "*",
            BinOp::Div => "/",
            BinOp::Eq => "==",
            BinOp::Ne => "!=",
            BinOp::Lt => "<",
            BinOp::Gt => ">",
            BinOp::Le => "<=",
            BinOp::Ge => ">=",
            BinOp::And => "&&",
            BinOp::Or => "||",
            BinOp::Concat => "++", // Handled specially above
        };
        self.emit.write(s);
    }

    fn emit_unaryop(&mut self, op: &UnaryOp) {
        let s = match op {
            UnaryOp::Neg => "-",
            UnaryOp::Not => "!",
        };
        self.emit.write(s);
    }

    fn infer_agent_output_type(&self, agent: &AgentDecl) -> String {
        // Look for emit expression in start handler to infer return type
        // For now, default to i64
        for handler in &agent.handlers {
            if let EventKind::Start = &handler.event {
                if let Some(ty) = self.find_emit_type(&handler.body) {
                    return ty;
                }
            }
        }
        "i64".to_string()
    }

    fn find_emit_type(&self, block: &Block) -> Option<String> {
        for stmt in &block.stmts {
            if let Stmt::Expr { expr, .. } = stmt {
                if let Expr::Emit { value, .. } = expr {
                    return Some(self.infer_expr_type(value));
                }
            }
            // Check nested blocks
            if let Stmt::If {
                then_block,
                else_block,
                ..
            } = stmt
            {
                if let Some(ty) = self.find_emit_type(then_block) {
                    return Some(ty);
                }
                if let Some(else_branch) = else_block {
                    if let sage_parser::ElseBranch::Block(block) = else_branch {
                        if let Some(ty) = self.find_emit_type(block) {
                            return Some(ty);
                        }
                    }
                }
            }
        }
        None
    }

    fn infer_expr_type(&self, expr: &Expr) -> String {
        match expr {
            Expr::Literal { value, .. } => match value {
                Literal::Int(_) => "i64".to_string(),
                Literal::Float(_) => "f64".to_string(),
                Literal::Bool(_) => "bool".to_string(),
                Literal::String(_) => "String".to_string(),
            },
            Expr::Var { .. } => "i64".to_string(), // Conservative default
            Expr::Binary { op, .. } => {
                if matches!(
                    op,
                    BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Gt | BinOp::Le | BinOp::Ge
                ) {
                    "bool".to_string()
                } else if matches!(op, BinOp::Concat) {
                    "String".to_string()
                } else {
                    "i64".to_string()
                }
            }
            Expr::Infer { .. } | Expr::StringInterp { .. } => "String".to_string(),
            Expr::Call { name, .. } if name.name == "str" => "String".to_string(),
            Expr::Call { name, .. } if name.name == "len" => "i64".to_string(),
            _ => "i64".to_string(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use sage_lexer::lex;
    use sage_parser::parse;
    use std::sync::Arc;

    fn generate_source(source: &str) -> String {
        let lex_result = lex(source).expect("lexing failed");
        let source_arc: Arc<str> = Arc::from(source);
        let (program, errors) = parse(lex_result.tokens(), source_arc);
        assert!(errors.is_empty(), "parse errors: {errors:?}");
        let program = program.expect("should parse");
        generate(&program, "test").main_rs
    }

    #[test]
    fn generate_minimal_program() {
        let source = r#"
            agent Main {
                on start {
                    emit(42);
                }
            }
            run Main;
        "#;

        let output = generate_source(source);
        assert!(output.contains("struct Main;"));
        assert!(output.contains("async fn on_start"));
        assert!(output.contains("ctx.emit(42_i64)"));
        assert!(output.contains("#[tokio::main]"));
    }

    #[test]
    fn generate_function() {
        let source = r#"
            fn add(a: Int, b: Int) -> Int {
                return a + b;
            }
            agent Main {
                on start {
                    emit(add(1, 2));
                }
            }
            run Main;
        "#;

        let output = generate_source(source);
        assert!(output.contains("fn add(a: i64, b: i64) -> i64"));
        assert!(output.contains("return (a + b);"));
    }

    #[test]
    fn generate_agent_with_beliefs() {
        let source = r#"
            agent Worker {
                belief value: Int

                on start {
                    emit(self.value * 2);
                }
            }
            agent Main {
                on start {
                    emit(0);
                }
            }
            run Main;
        "#;

        let output = generate_source(source);
        assert!(output.contains("struct Worker {"));
        assert!(output.contains("value: i64,"));
        assert!(output.contains("self.value"));
    }

    #[test]
    fn generate_string_interpolation() {
        let source = r#"
            agent Main {
                on start {
                    let name = "World";
                    let msg = "Hello, {name}!";
                    print(msg);
                    emit(0);
                }
            }
            run Main;
        "#;

        let output = generate_source(source);
        assert!(output.contains("format!(\"Hello, {}!\", name)"));
    }

    #[test]
    fn generate_control_flow() {
        let source = r#"
            agent Main {
                on start {
                    let x = 10;
                    if x > 5 {
                        emit(1);
                    } else {
                        emit(0);
                    }
                }
            }
            run Main;
        "#;

        let output = generate_source(source);
        assert!(output.contains("if (x > 5_i64)"), "output:\n{output}");
        // else is on the same line after close brace
        assert!(output.contains("else"), "output:\n{output}");
    }

    #[test]
    fn generate_loops() {
        let source = r#"
            agent Main {
                on start {
                    for x in [1, 2, 3] {
                        print(str(x));
                    }
                    let n = 0;
                    while n < 5 {
                        n = n + 1;
                    }
                    emit(n);
                }
            }
            run Main;
        "#;

        let output = generate_source(source);
        assert!(output.contains("for x in vec![1_i64, 2_i64, 3_i64]"));
        assert!(output.contains("while (n < 5_i64)"));
    }
}