spydecy_codegen/

lib.rs

//! Rust Code Generator - Sprint 4+
//!
//! Generates idiomatic Rust code from optimized Unified HIR.
//! This is the final stage of the Spydecy pipeline:
//!
//! ```text
//! Python + C → Unified HIR → Optimizer → Codegen → Rust Code
//! ```

#![warn(missing_docs, clippy::all, clippy::pedantic)]
#![deny(unsafe_code)]
#![allow(clippy::module_name_repetitions)]

use anyhow::{Context, Result};
use spydecy_hir::unified::{UnificationPattern, UnifiedHIR};

/// Rust code generator
///
/// Converts optimized `UnifiedHIR` into idiomatic Rust source code.
pub struct RustCodegen {
    /// Indentation level
    indent_level: usize,
    /// Indentation string (default: 4 spaces)
    indent: String,
}

impl RustCodegen {
    /// Create a new Rust code generator
    #[must_use]
    pub fn new() -> Self {
        Self {
            indent_level: 0,
            indent: "    ".to_owned(), // 4 spaces
        }
    }

    /// Generate Rust code from `UnifiedHIR`
    ///
    /// # Errors
    ///
    /// Returns an error if the HIR cannot be converted to Rust code
    pub fn generate(&mut self, hir: &UnifiedHIR) -> Result<String> {
        match hir {
            UnifiedHIR::Module {
                name, declarations, ..
            } => self.generate_module(name, declarations),
            UnifiedHIR::Function {
                name,
                params,
                return_type,
                body,
                ..
            } => self.generate_function(name, params, return_type, body),
            UnifiedHIR::Call {
                callee,
                args,
                cross_mapping,
                ..
            } => {
                // Check if this is an optimized pattern
                if let Some(mapping) = cross_mapping {
                    if mapping.boundary_eliminated {
                        return Ok(Self::generate_optimized_call(callee, args, mapping.pattern));
                    }
                }
                self.generate_call(callee, args)
            }
            UnifiedHIR::Variable { name, .. } => Ok(name.clone()),
            UnifiedHIR::Return { value, .. } => self.generate_return(value.as_deref()),
            UnifiedHIR::Assign { target, value, .. } => self.generate_assign(target, value),
            _ => Ok("/* Unsupported HIR node */".to_owned()),
        }
    }

    /// Generate a module
    fn generate_module(&mut self, name: &str, declarations: &[UnifiedHIR]) -> Result<String> {
        let mut output = String::new();

        // Module header
        output.push_str("// Module: ");
        output.push_str(name);
        output.push('\n');
        output.push_str("// Generated by Spydecy\n");
        output.push_str("#![allow(dead_code)]\n\n");

        // Generate all declarations
        for decl in declarations {
            let code = self.generate(decl)?;
            output.push_str(&code);
            output.push_str("\n\n");
        }

        Ok(output)
    }

    /// Generate a function
    fn generate_function(
        &mut self,
        name: &str,
        params: &[spydecy_hir::unified::UnifiedParameter],
        return_type: &spydecy_hir::types::Type,
        body: &[UnifiedHIR],
    ) -> Result<String> {
        let mut output = String::new();

        // Function signature
        output.push_str("pub fn ");
        output.push_str(name);
        output.push('(');

        // Parameters
        for (i, param) in params.iter().enumerate() {
            if i > 0 {
                output.push_str(", ");
            }
            output.push_str(&param.name);
            output.push_str(": ");
            output.push_str(&self.generate_type(&param.param_type)?);
        }

        output.push(')');

        // Return type
        if !matches!(
            return_type,
            spydecy_hir::types::Type::Rust(spydecy_hir::types::RustType::Unit)
        ) {
            output.push_str(" -> ");
            output.push_str(&self.generate_type(return_type)?);
        }

        output.push_str(" {\n");

        // Function body
        self.indent_level += 1;
        for stmt in body {
            let code = self.generate(stmt)?;
            output.push_str(&self.indent());
            output.push_str(&code);
            if !code.trim().ends_with('}') && !code.trim().is_empty() {
                output.push(';');
            }
            output.push('\n');
        }
        self.indent_level -= 1;

        output.push('}');

        Ok(output)
    }

    /// Extract the receiver variable name from arguments
    ///
    /// Returns the name of the first Variable argument, or "x" as fallback
    fn extract_receiver_name(args: &[UnifiedHIR]) -> String {
        args.first()
            .and_then(|arg| {
                if let UnifiedHIR::Variable { name, .. } = arg {
                    Some(name.clone())
                } else {
                    None
                }
            })
            .unwrap_or_else(|| "x".to_owned())
    }

    /// Generate an optimized call (post-boundary-elimination)
    #[allow(clippy::unnecessary_wraps)]
    fn generate_optimized_call(
        callee: &str,
        args: &[UnifiedHIR],
        pattern: UnificationPattern,
    ) -> String {
        // Helper: Extract variable name from first argument
        let receiver = Self::extract_receiver_name(args);

        // Generate idiomatic Rust based on the pattern
        match pattern {
            UnificationPattern::LenPattern => {
                // Vec::len() becomes <receiver>.len()
                format!("{receiver}.len()")
            }
            UnificationPattern::AppendPattern => {
                // Vec::push() becomes <receiver>.push(item)
                format!("{receiver}.push(item)")
            }
            UnificationPattern::DictGetPattern => {
                // HashMap::get() becomes <receiver>.get(&key)
                format!("{receiver}.get(&key)")
            }
            UnificationPattern::ReversePattern => {
                // Vec::reverse() becomes <receiver>.reverse()
                format!("{receiver}.reverse()")
            }
            UnificationPattern::ClearPattern => {
                // Vec::clear() becomes <receiver>.clear()
                format!("{receiver}.clear()")
            }
            UnificationPattern::PopPattern => {
                // Vec::pop() becomes <receiver>.pop()
                format!("{receiver}.pop()")
            }
            UnificationPattern::InsertPattern => {
                // Vec::insert() becomes <receiver>.insert(index, value)
                format!("{receiver}.insert(index, value)")
            }
            UnificationPattern::ExtendPattern => {
                // Vec::extend() becomes <receiver>.extend(iter)
                format!("{receiver}.extend(iter)")
            }
            UnificationPattern::DictPopPattern => {
                // HashMap::remove() becomes <receiver>.remove(&key)
                format!("{receiver}.remove(&key)")
            }
            UnificationPattern::DictClearPattern => {
                // HashMap::clear() becomes <receiver>.clear()
                format!("{receiver}.clear()")
            }
            UnificationPattern::DictKeysPattern => {
                // HashMap::keys() becomes <receiver>.keys()
                format!("{receiver}.keys()")
            }
            UnificationPattern::Custom => format!("{callee}()"),
        }
    }

    /// Generate a regular call
    fn generate_call(&mut self, callee: &str, args: &[UnifiedHIR]) -> Result<String> {
        let mut output = callee.to_owned();
        output.push('(');

        for (i, arg) in args.iter().enumerate() {
            if i > 0 {
                output.push_str(", ");
            }
            output.push_str(&self.generate(arg)?);
        }

        output.push(')');
        Ok(output)
    }

    /// Generate a return statement
    fn generate_return(&mut self, value: Option<&UnifiedHIR>) -> Result<String> {
        if let Some(val) = value {
            Ok(format!("return {}", self.generate(val)?))
        } else {
            Ok("return".to_owned())
        }
    }

    /// Generate an assignment
    fn generate_assign(&mut self, target: &str, value: &UnifiedHIR) -> Result<String> {
        let val_code = self.generate(value)?;
        Ok(format!("let {target} = {val_code}"))
    }

    /// Generate a type annotation
    #[allow(clippy::only_used_in_recursion)]
    fn generate_type(&self, ty: &spydecy_hir::types::Type) -> Result<String> {
        use spydecy_hir::types::{RustType, Type};

        match ty {
            Type::Rust(rust_ty) => match rust_ty {
                RustType::Int { bits, signed } => {
                    let prefix = if *signed { "i" } else { "u" };
                    let size = match bits {
                        spydecy_hir::types::IntSize::I8 => "8",
                        spydecy_hir::types::IntSize::I16 => "16",
                        spydecy_hir::types::IntSize::I32 => "32",
                        spydecy_hir::types::IntSize::I64 => "64",
                        spydecy_hir::types::IntSize::I128 => "128",
                        spydecy_hir::types::IntSize::ISize => "size",
                    };
                    Ok(format!("{prefix}{size}"))
                }
                RustType::Bool => Ok("bool".to_owned()),
                RustType::String => Ok("String".to_owned()),
                RustType::Str => Ok("&str".to_owned()),
                RustType::Vec(inner) => Ok(format!("Vec<{}>", self.generate_type(inner)?)),
                RustType::Option(inner) => Ok(format!("Option<{}>", self.generate_type(inner)?)),
                RustType::Unit => Ok("()".to_owned()),
                RustType::Reference { mutable, inner } => {
                    let mut_str = if *mutable { "mut " } else { "" };
                    Ok(format!("&{mut_str}{}", self.generate_type(inner)?))
                }
                _ => Ok("/* complex type */".to_owned()),
            },
            Type::Unknown => Ok("/* infer */".to_owned()),
            _ => Ok("/* non-rust type */".to_owned()),
        }
    }

    /// Get current indentation string
    fn indent(&self) -> String {
        self.indent.repeat(self.indent_level)
    }
}

impl Default for RustCodegen {
    fn default() -> Self {
        Self::new()
    }
}

/// Generate Rust code from `UnifiedHIR` (convenience function)
///
/// # Errors
///
/// Returns an error if code generation fails
pub fn generate_rust(hir: &UnifiedHIR) -> Result<String> {
    let mut codegen = RustCodegen::new();
    codegen
        .generate(hir)
        .context("Failed to generate Rust code")
}

#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
    use super::*;
    use spydecy_hir::{
        metadata::Metadata,
        types::{IntSize, RustType, Type},
        unified::{CrossMapping, UnificationPattern},
        Language, NodeId,
    };

    #[test]
    fn test_generate_len_pattern() {
        // Create UnifiedHIR for optimized len() call
        let hir = UnifiedHIR::Call {
            id: NodeId::new(1),
            target_language: Language::Rust,
            callee: "Vec::len".to_owned(),
            args: vec![],
            inferred_type: Type::Rust(RustType::Int {
                bits: IntSize::ISize,
                signed: false,
            }),
            source_language: Language::Python,
            cross_mapping: Some(CrossMapping {
                python_node: None,
                c_node: None,
                pattern: UnificationPattern::LenPattern,
                boundary_eliminated: true,
            }),
            meta: Metadata::new(),
        };

        let code = generate_rust(&hir).expect("Should generate code");
        assert_eq!(code.trim(), "x.len()");
    }

    #[test]
    fn test_generate_append_pattern() {
        let hir = UnifiedHIR::Call {
            id: NodeId::new(1),
            target_language: Language::Rust,
            callee: "Vec::push".to_owned(),
            args: vec![],
            inferred_type: Type::Rust(RustType::Unit),
            source_language: Language::Python,
            cross_mapping: Some(CrossMapping {
                python_node: None,
                c_node: None,
                pattern: UnificationPattern::AppendPattern,
                boundary_eliminated: true,
            }),
            meta: Metadata::new(),
        };

        let code = generate_rust(&hir).expect("Should generate code");
        assert_eq!(code.trim(), "x.push(item)");
    }

    #[test]
    fn test_generate_dict_get_pattern() {
        let hir = UnifiedHIR::Call {
            id: NodeId::new(1),
            target_language: Language::Rust,
            callee: "HashMap::get".to_owned(),
            args: vec![UnifiedHIR::Variable {
                id: NodeId::new(2),
                name: "map".to_owned(),
                var_type: Type::Unknown,
                source_language: Language::Python,
                meta: Metadata::new(),
            }],
            inferred_type: Type::Unknown,
            source_language: Language::Python,
            cross_mapping: Some(CrossMapping {
                python_node: None,
                c_node: None,
                pattern: UnificationPattern::DictGetPattern,
                boundary_eliminated: true,
            }),
            meta: Metadata::new(),
        };

        let code = generate_rust(&hir).expect("Should generate code");
        assert_eq!(code.trim(), "map.get(&key)");
    }

    #[test]
    fn test_generate_variable() {
        let hir = UnifiedHIR::Variable {
            id: NodeId::new(1),
            name: "my_var".to_owned(),
            var_type: Type::Unknown,
            source_language: Language::Rust,
            meta: Metadata::new(),
        };

        let code = generate_rust(&hir).expect("Should generate code");
        assert_eq!(code, "my_var");
    }

    #[test]
    fn test_generate_type_int() {
        let codegen = RustCodegen::new();
        let ty = Type::Rust(RustType::Int {
            bits: IntSize::I32,
            signed: true,
        });

        let code = codegen.generate_type(&ty).expect("Should generate type");
        assert_eq!(code, "i32");
    }

    #[test]
    fn test_generate_type_vec() {
        let codegen = RustCodegen::new();
        let ty = Type::Rust(RustType::Vec(Box::new(Type::Rust(RustType::Int {
            bits: IntSize::I32,
            signed: true,
        }))));

        let code = codegen.generate_type(&ty).expect("Should generate type");
        assert_eq!(code, "Vec<i32>");
    }
}