rustytorch/crates/production/rtx-onnx-codegen/src/model_gen.rs

//! Model generation builder API.
//!
//! Provides a fluent API for generating Rust code from ONNX models.

use std::fs;
use std::path::{Path, PathBuf};

use crate::codegen::{CodeGenerator, CodegenConfig, Visibility};
use crate::error::{Error, Result};
use crate::parser::OnnxParser;

/// Builder for generating Rust code from ONNX models.
///
/// # Example
///
/// ```ignore
/// use rtx_onnx_codegen::ModelGen;
///
/// ModelGen::new()
///     .input("models/resnet50.onnx")
///     .out_dir("src/models/")
///     .struct_name("ResNet50")
///     .run()?;
/// ```
pub struct ModelGen {
    /// Input ONNX file path.
    input_path: Option<PathBuf>,
    /// Output directory for generated code.
    out_dir: Option<PathBuf>,
    /// Output file name (defaults to model name).
    out_file: Option<String>,
    /// Configuration for code generation.
    config: CodegenConfig,
    /// Whether to emit cargo rerun-if-changed directives.
    emit_rerun: bool,
}

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

impl ModelGen {
    /// Create a new ModelGen builder.
    pub fn new() -> Self {
        Self {
            input_path: None,
            out_dir: None,
            out_file: None,
            config: CodegenConfig::default(),
            emit_rerun: true,
        }
    }

    /// Set the input ONNX file path.
    pub fn input<P: AsRef<Path>>(mut self, path: P) -> Self {
        self.input_path = Some(path.as_ref().to_path_buf());
        self
    }

    /// Set the output directory.
    pub fn out_dir<P: AsRef<Path>>(mut self, path: P) -> Self {
        self.out_dir = Some(path.as_ref().to_path_buf());
        self
    }

    /// Set the output file name (without .rs extension).
    pub fn out_file(mut self, name: impl Into<String>) -> Self {
        self.out_file = Some(name.into());
        self
    }

    /// Set the name of the generated struct.
    pub fn struct_name(mut self, name: impl Into<String>) -> Self {
        self.config.struct_name = name.into();
        self
    }

    /// Set the documentation string for the generated struct.
    pub fn doc(mut self, doc: impl Into<String>) -> Self {
        self.config.doc = Some(doc.into());
        self
    }

    /// Set whether to derive Debug on the generated struct.
    pub fn derive_debug(mut self, derive: bool) -> Self {
        self.config.derive_debug = derive;
        self
    }

    /// Set whether to derive Clone on the generated struct.
    pub fn derive_clone(mut self, derive: bool) -> Self {
        self.config.derive_clone = derive;
        self
    }

    /// Set whether to use derive(Module) macro.
    pub fn derive_module(mut self, derive: bool) -> Self {
        self.config.use_derive_module = derive;
        self
    }

    /// Set the visibility of the generated struct to public.
    pub fn public(mut self) -> Self {
        self.config.visibility = Visibility::Public;
        self
    }

    /// Set the visibility of the generated struct to crate-level.
    pub fn crate_visible(mut self) -> Self {
        self.config.visibility = Visibility::Crate;
        self
    }

    /// Set the visibility of the generated struct to private.
    pub fn private(mut self) -> Self {
        self.config.visibility = Visibility::Private;
        self
    }

    /// Disable cargo rerun-if-changed directives.
    pub fn no_rerun_if_changed(mut self) -> Self {
        self.emit_rerun = false;
        self
    }

    /// Run the code generation.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No input path was specified
    /// - The input file doesn't exist or can't be read
    /// - The ONNX model is malformed
    /// - Code generation fails
    /// - The output file can't be written
    pub fn run(&self) -> Result<PathBuf> {
        let input_path = self
            .input_path
            .as_ref()
            .ok_or_else(|| Error::CodeGen("No input path specified".into()))?;

        // Determine output path
        let out_dir = self.out_dir.clone().unwrap_or_else(|| {
            std::env::var("OUT_DIR").map_or_else(|_| PathBuf::from("."), PathBuf::from)
        });

        let out_file = self.out_file.clone().unwrap_or_else(|| {
            input_path
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("model")
                .to_string()
        });

        let output_path = out_dir.join(format!("{}.rs", out_file));

        // Parse the ONNX model
        let graph = OnnxParser::parse_file(input_path)?;

        // Update struct name from model name if not explicitly set
        let config = if self.config.struct_name == "Model" {
            let mut config = self.config.clone();
            if !graph.name.is_empty() {
                config.struct_name = to_pascal_case(&graph.name);
            } else {
                config.struct_name = to_pascal_case(&out_file);
            }
            config
        } else {
            self.config.clone()
        };

        // Generate code
        let generator = CodeGenerator::new(config);
        let code = generator.generate(&graph)?;

        // Ensure output directory exists
        if let Some(parent) = output_path.parent() {
            fs::create_dir_all(parent)?;
        }

        // Write output file
        fs::write(&output_path, code)?;

        // Emit cargo rerun directive
        if self.emit_rerun {
            println!("cargo:rerun-if-changed={}", input_path.display());
        }

        Ok(output_path)
    }

    /// Run from a build script context.
    ///
    /// This is a convenience method that:
    /// - Uses OUT_DIR as the default output directory
    /// - Always emits cargo:rerun-if-changed
    /// - Panics on error (appropriate for build scripts)
    pub fn run_from_script(self) {
        match self.run() {
            Ok(path) => {
                eprintln!("Generated: {}", path.display());
            }
            Err(e) => {
                panic!("ONNX code generation failed: {}", e);
            }
        }
    }

    /// Generate code without writing to a file.
    ///
    /// Useful for testing or generating code programmatically.
    pub fn generate_string(&self) -> Result<String> {
        let input_path = self
            .input_path
            .as_ref()
            .ok_or_else(|| Error::CodeGen("No input path specified".into()))?;

        let graph = OnnxParser::parse_file(input_path)?;

        let generator = CodeGenerator::new(self.config.clone());
        generator.generate(&graph)
    }
}

/// Convert a string to PascalCase.
fn to_pascal_case(s: &str) -> String {
    let mut result = String::new();
    let mut capitalize_next = true;

    for c in s.chars() {
        if c == '_' || c == '-' || c == '.' || c == ' ' {
            capitalize_next = true;
        } else if capitalize_next {
            result.extend(c.to_uppercase());
            capitalize_next = false;
        } else {
            result.push(c);
        }
    }

    // Ensure first character is uppercase
    let mut chars: Vec<char> = result.chars().collect();
    if !chars.is_empty() {
        chars[0] = chars[0].to_uppercase().next().unwrap_or(chars[0]);
    }

    chars.into_iter().collect()
}

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

    #[test]
    fn test_to_pascal_case() {
        assert_eq!(to_pascal_case("resnet50"), "Resnet50");
        assert_eq!(to_pascal_case("my_model"), "MyModel");
        assert_eq!(to_pascal_case("bert-base"), "BertBase");
        assert_eq!(to_pascal_case("gpt.2"), "Gpt2");
    }

    #[test]
    fn test_builder() {
        let model_gen = ModelGen::new()
            .input("model.onnx")
            .out_dir("out/")
            .struct_name("MyModel");

        assert_eq!(model_gen.input_path, Some(PathBuf::from("model.onnx")));
        assert_eq!(model_gen.out_dir, Some(PathBuf::from("out/")));
        assert_eq!(model_gen.config.struct_name, "MyModel");
    }
}