use crate::symbolic::codegen::CodegenIR::{CodegenLanguage, CodegenModule};
use crate::symbolic::codegen::codegen_manifest::PreparedProblemManifest;
use crate::symbolic::codegen::codegen_provider_api::{
PreparedBandedProblem, PreparedDenseProblem, PreparedProblem, PreparedSparseProblem,
};
use log::info;
use std::fs;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone)]
pub struct GeneratedCAotLibrary {
pub library_name: String,
pub c_source: String,
pub c_header: String,
pub manifest: PreparedProblemManifest,
pub compiler: String,
pub optimization: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WrittenCAotLibrary {
pub library_dir: PathBuf,
pub makefile: PathBuf,
pub generated_c: PathBuf,
pub generated_h: PathBuf,
pub manifest_h: PathBuf,
pub aot_interface_c: PathBuf,
pub aot_interface_h: PathBuf,
}
impl GeneratedCAotLibrary {
pub fn new(
library_name: impl Into<String>,
c_source: impl Into<String>,
c_header: impl Into<String>,
manifest: PreparedProblemManifest,
) -> Self {
let library_name = library_name.into();
validate_library_name(&library_name);
Self {
library_name,
c_source: c_source.into(),
c_header: c_header.into(),
manifest,
compiler: "gcc".to_string(),
optimization: "-O3".to_string(),
}
}
pub fn from_codegen_module(
library_name: impl Into<String>,
module: &CodegenModule,
manifest: PreparedProblemManifest,
) -> Self {
let mut module = module.clone();
module.set_language(CodegenLanguage::C);
Self::new(
library_name,
module.emit_source(),
module.emit_c_header(),
manifest,
)
}
pub fn from_prepared_dense_problem(
library_name: impl Into<String>,
problem: &PreparedDenseProblem<'_>,
module: &CodegenModule,
) -> Self {
Self::from_codegen_module(library_name, module, PreparedProblemManifest::from(problem))
}
pub fn from_prepared_sparse_problem(
library_name: impl Into<String>,
problem: &PreparedSparseProblem<'_>,
module: &CodegenModule,
) -> Self {
Self::from_codegen_module(library_name, module, PreparedProblemManifest::from(problem))
}
pub fn from_prepared_banded_problem(
library_name: impl Into<String>,
problem: &PreparedBandedProblem<'_>,
module: &CodegenModule,
) -> Self {
Self::from_codegen_module(library_name, module, PreparedProblemManifest::from(problem))
}
pub fn from_prepared_problem(
library_name: impl Into<String>,
problem: &PreparedProblem<'_>,
module: &CodegenModule,
) -> Self {
Self::from_codegen_module(library_name, module, PreparedProblemManifest::from(problem))
}
pub fn with_compiler(mut self, compiler: impl Into<String>) -> Self {
self.compiler = compiler.into();
self
}
pub fn with_optimization(mut self, optimization: impl Into<String>) -> Self {
self.optimization = optimization.into();
self
}
pub fn write_to_dir<P: AsRef<Path>>(&self, parent_dir: P) -> io::Result<WrittenCAotLibrary> {
let library_dir = parent_dir.as_ref().join(&self.library_name);
info!(
"Writing generated C AOT library '{}' into '{}'",
self.library_name,
library_dir.display()
);
fs::create_dir_all(&library_dir)?;
let makefile = library_dir.join("Makefile");
let generated_c = library_dir.join("generated.c");
let generated_h = library_dir.join("generated.h");
let manifest_h = library_dir.join("manifest.h");
let aot_interface_c = library_dir.join("aot_interface.c");
let aot_interface_h = library_dir.join("aot_interface.h");
fs::write(&makefile, self.emit_makefile())?;
self.write_generated_c(&generated_c)?;
self.write_generated_h(&generated_h)?;
fs::write(&manifest_h, self.emit_manifest_h())?;
fs::write(&aot_interface_c, self.emit_aot_interface_c())?;
fs::write(&aot_interface_h, self.emit_aot_interface_h())?;
let written = WrittenCAotLibrary {
library_dir,
makefile,
generated_c,
generated_h,
manifest_h,
aot_interface_c,
aot_interface_h,
};
info!(
"Generated C AOT library '{}' written successfully",
self.library_name
);
Ok(written)
}
fn emit_makefile(&self) -> String {
let target_name = if cfg!(target_os = "windows") {
format!("lib{}.dll", self.library_name)
} else if cfg!(target_os = "macos") {
format!("lib{}.dylib", self.library_name)
} else {
format!("lib{}.so", self.library_name)
};
format!(
"# Auto-generated Makefile for C AOT library\n\
CC = {}\n\
CFLAGS = {} -fPIC -Wall -Wextra\n\
LDFLAGS = -shared -lm\n\
\n\
TARGET = {}\n\
OBJS = generated.o aot_interface.o\n\
\n\
all: $(TARGET)\n\
\n\
$(TARGET): $(OBJS)\n\
\t$(CC) $(LDFLAGS) -o $@ $^\n\
\n\
generated.o: generated.c generated.h manifest.h\n\
\t$(CC) $(CFLAGS) -c generated.c\n\
\n\
aot_interface.o: aot_interface.c aot_interface.h generated.h manifest.h\n\
\t$(CC) $(CFLAGS) -c aot_interface.c\n\
\n\
clean:\n\
\trm -f $(OBJS) $(TARGET)\n\
\n\
.PHONY: all clean\n",
self.compiler, self.optimization, target_name
)
}
fn write_generated_c(&self, path: &Path) -> io::Result<()> {
let mut file = fs::File::create(path)?;
file.write_all(b"/* AUTO-GENERATED C AOT SOURCE */\n\n")?;
file.write_all(b"#include \"generated.h\"\n")?;
file.write_all(b"#include \"manifest.h\"\n\n")?;
file.write_all(self.c_source.as_bytes())?;
Ok(())
}
fn write_generated_h(&self, path: &Path) -> io::Result<()> {
let mut file = fs::File::create(path)?;
file.write_all(b"/* AUTO-GENERATED C AOT HEADER */\n\n")?;
file.write_all(b"#ifndef GENERATED_H\n")?;
file.write_all(b"#define GENERATED_H\n\n")?;
file.write_all(self.c_header.as_bytes())?;
file.write_all(b"\n#endif /* GENERATED_H */\n")?;
Ok(())
}
fn emit_manifest_h(&self) -> String {
let jacobian_nnz = match self.manifest.io.jacobian_nnz {
Some(nnz) => nnz.to_string(),
None => "0".to_string(),
};
format!(
"/* AUTO-GENERATED C AOT MANIFEST */\n\n\
#ifndef MANIFEST_H\n\
#define MANIFEST_H\n\n\
#define BACKEND_KIND \"{}\"\n\
#define MATRIX_BACKEND \"{}\"\n\
#define RESIDUAL_LEN {}\n\
#define JACOBIAN_ROWS {}\n\
#define JACOBIAN_COLS {}\n\
#define JACOBIAN_NNZ {}\n\
#define RESIDUAL_FN_NAME \"{}\"\n\
#define JACOBIAN_FN_NAME \"{}\"\n\n\
#endif /* MANIFEST_H */\n",
self.manifest.backend_kind.as_str(),
self.manifest.matrix_backend.as_str(),
self.manifest.io.residual_len,
self.manifest.io.jacobian_rows,
self.manifest.io.jacobian_cols,
jacobian_nnz,
escape_c_string(&self.manifest.functions.residual_fn_name),
escape_c_string(&self.manifest.functions.jacobian_fn_name)
)
}
fn emit_aot_interface_c(&self) -> String {
let jacobian_zero_init = if self.manifest.matrix_backend.as_str() == "dense" {
" for (size_t i = 0; i < out_len; ++i) out[i] = 0.0;\n"
} else {
""
};
let residual_chunk_exports = self
.manifest
.functions
.residual_chunks
.iter()
.map(|chunk| {
format!(
"\nRUSTEDSCITHE_AOT_EXPORT int rustedscithe_aot_chunk_{fn_name}(\n\
const double* args_ptr,\n\
size_t args_len,\n\
double* out_ptr,\n\
size_t out_len\n\
) {{\n\
if (args_ptr == NULL || out_ptr == NULL || out_len != {out_len}) {{\n\
return 0;\n\
}}\n\
const double* args = args_ptr;\n\
double* out = out_ptr;\n\
(void)args_len;\n\
{fn_name}(args, out);\n\
return 1;\n\
}}\n",
fn_name = chunk.fn_name,
out_len = chunk.len,
)
})
.collect::<Vec<_>>()
.join("");
let jacobian_chunk_exports = self
.manifest
.functions
.jacobian_chunks
.iter()
.map(|chunk| {
format!(
"\nRUSTEDSCITHE_AOT_EXPORT int rustedscithe_aot_chunk_{fn_name}(\n\
const double* args_ptr,\n\
size_t args_len,\n\
double* out_ptr,\n\
size_t out_len\n\
) {{\n\
if (args_ptr == NULL || out_ptr == NULL || out_len != {out_len}) {{\n\
return 0;\n\
}}\n\
const double* args = args_ptr;\n\
double* out = out_ptr;\n\
(void)args_len;\n\
{jacobian_zero_init}\
{fn_name}(args, out);\n\
return 1;\n\
}}\n",
fn_name = chunk.fn_name,
out_len = chunk.len,
jacobian_zero_init = jacobian_zero_init,
)
})
.collect::<Vec<_>>()
.join("");
let residual_dispatch = if self.manifest.functions.residual_chunks.is_empty() {
format!(
" {}(args, out);",
self.manifest.functions.residual_fn_name
)
} else {
self.manifest
.functions
.residual_chunks
.iter()
.map(|chunk| format!(" {}(args, &out[{}]);", chunk.fn_name, chunk.offset))
.collect::<Vec<_>>()
.join("\n")
};
let jacobian_dispatch = if self.manifest.functions.jacobian_fn_name.is_empty() {
String::new()
} else if self.manifest.functions.jacobian_chunks.is_empty() {
format!(
" {}(args, out);",
self.manifest.functions.jacobian_fn_name
)
} else {
self.manifest
.functions
.jacobian_chunks
.iter()
.map(|chunk| format!(" {}(args, &out[{}]);", chunk.fn_name, chunk.offset))
.collect::<Vec<_>>()
.join("\n")
};
format!(
"/* AUTO-GENERATED C AOT FFI INTERFACE */\n\n\
#include \"aot_interface.h\"\n\
#include \"generated.h\"\n\
#include \"manifest.h\"\n\
#include <stddef.h>\n\
#include <stdint.h>\n\n\
RUSTEDSCITHE_AOT_EXPORT int rustedscithe_aot_eval_residual(\n\
const double* args_ptr,\n\
size_t args_len,\n\
double* out_ptr,\n\
size_t out_len\n\
) {{\n\
if (args_ptr == NULL || out_ptr == NULL || out_len != RESIDUAL_LEN) {{\n\
return 0;\n\
}}\n\
const double* args = args_ptr;\n\
double* out = out_ptr;\n\
(void)args_len;\n\
{}\n\
return 1;\n\
}}\n\n\
RUSTEDSCITHE_AOT_EXPORT int rustedscithe_aot_eval_jacobian_values(\n\
const double* args_ptr,\n\
size_t args_len,\n\
double* out_ptr,\n\
size_t out_len\n\
) {{\n\
if (args_ptr == NULL || out_ptr == NULL) {{\n\
return 0;\n\
}}\n\
size_t expected_out_len = JACOBIAN_NNZ > 0 ? JACOBIAN_NNZ : (JACOBIAN_ROWS * JACOBIAN_COLS);\n\
if (out_len != expected_out_len) {{\n\
return 0;\n\
}}\n\
const double* args = args_ptr;\n\
double* out = out_ptr;\n\
(void)args_len;\n\
{}{}\n\
return 1;\n\
}}\n",
residual_dispatch, jacobian_zero_init, jacobian_dispatch
) + residual_chunk_exports.as_str()
+ jacobian_chunk_exports.as_str()
}
fn emit_aot_interface_h(&self) -> String {
"/* AUTO-GENERATED C AOT FFI INTERFACE HEADER */\n\n\
#ifndef AOT_INTERFACE_H\n\
#define AOT_INTERFACE_H\n\n\
#include <stddef.h>\n\
#include <stdint.h>\n\n\
#if defined(_WIN32) || defined(__CYGWIN__)\n\
#define RUSTEDSCITHE_AOT_EXPORT __declspec(dllexport)\n\
#elif defined(__GNUC__)\n\
#define RUSTEDSCITHE_AOT_EXPORT __attribute__((visibility(\"default\")))\n\
#else\n\
#define RUSTEDSCITHE_AOT_EXPORT\n\
#endif\n\n\
#ifdef __cplusplus\n\
extern \"C\" {\n\
#endif\n\n\
RUSTEDSCITHE_AOT_EXPORT int rustedscithe_aot_eval_residual(\n\
const double* args_ptr,\n\
size_t args_len,\n\
double* out_ptr,\n\
size_t out_len\n\
);\n\n\
RUSTEDSCITHE_AOT_EXPORT int rustedscithe_aot_eval_jacobian_values(\n\
const double* args_ptr,\n\
size_t args_len,\n\
double* out_ptr,\n\
size_t out_len\n\
);\n\n\
#ifdef __cplusplus\n\
}\n\
#endif\n\n\
#endif /* AOT_INTERFACE_H */\n"
.to_string()
}
}
fn validate_library_name(library_name: &str) {
assert!(
!library_name.is_empty(),
"generated C AOT library name must not be empty"
);
assert!(
library_name
.chars()
.all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '_'),
"generated C AOT library name must be snake_case ASCII"
);
assert!(
library_name
.chars()
.next()
.is_some_and(|ch| ch.is_ascii_lowercase()),
"generated C AOT library name must start with a lowercase ASCII letter"
);
}
fn escape_c_string(value: &str) -> String {
value.replace('\\', "\\\\").replace('\"', "\\\"")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::symbolic::codegen::CodegenIR::CodegenModule;
use crate::symbolic::codegen::codegen_provider_api::{
BackendKind, MatrixBackend, PreparedBandedProblem, PreparedDenseProblem, PreparedProblem,
};
use crate::symbolic::codegen::codegen_runtime_api::{
DenseJacobianChunkingStrategy, ResidualChunkingStrategy,
};
use crate::symbolic::codegen::codegen_tasks::{
BandedChunkingStrategy, BandedExprEntry, BandedJacobianTask, JacobianTask, ResidualTask,
};
use crate::symbolic::symbolic_engine::Expr;
use tempfile::tempdir;
fn sample_prepared_problem() -> PreparedProblem<'static> {
let residuals = Box::leak(Box::new(vec![Expr::parse_expression("x + 1")]));
let jacobian = Box::leak(Box::new(vec![vec![Expr::parse_expression("1")]]));
let vars = Box::leak(Box::new(vec!["x"]));
PreparedProblem::dense(PreparedDenseProblem::new(
BackendKind::Aot,
MatrixBackend::Dense,
ResidualTask {
fn_name: "eval_residual",
residuals,
variables: vars,
params: None,
}
.runtime_plan(ResidualChunkingStrategy::Whole),
JacobianTask {
fn_name: "eval_jacobian",
jacobian,
variables: vars,
params: None,
}
.runtime_plan(DenseJacobianChunkingStrategy::Whole),
))
}
fn sample_chunked_prepared_problem() -> PreparedProblem<'static> {
let residuals = Box::leak(Box::new(vec![
Expr::parse_expression("x + 1"),
Expr::parse_expression("y + 2"),
]));
let jacobian = Box::leak(Box::new(vec![
vec![Expr::parse_expression("1"), Expr::parse_expression("0")],
vec![Expr::parse_expression("0"), Expr::parse_expression("1")],
]));
let vars = Box::leak(Box::new(vec!["x", "y"]));
PreparedProblem::dense(PreparedDenseProblem::new(
BackendKind::Aot,
MatrixBackend::Dense,
ResidualTask {
fn_name: "eval_residual",
residuals,
variables: vars,
params: None,
}
.runtime_plan(ResidualChunkingStrategy::ByOutputCount {
max_outputs_per_chunk: 1,
}),
JacobianTask {
fn_name: "eval_jacobian",
jacobian,
variables: vars,
params: None,
}
.runtime_plan(DenseJacobianChunkingStrategy::ByRowCount { rows_per_chunk: 1 }),
))
}
#[test]
fn generated_c_library_writes_expected_file_layout() {
let prepared = sample_prepared_problem();
let module = CodegenModule::new("generated_fixture")
.with_language(CodegenLanguage::C)
.add_block("eval_residual", &[Expr::parse_expression("x + 1")], &["x"]);
let library_spec =
GeneratedCAotLibrary::from_prepared_problem("generated_c_fixture", &prepared, &module);
let dir = tempdir().expect("tempdir should exist");
let written = library_spec
.write_to_dir(dir.path())
.expect("generated C library should write");
assert!(written.makefile.exists());
assert!(written.generated_c.exists());
assert!(written.generated_h.exists());
assert!(written.manifest_h.exists());
assert!(written.aot_interface_c.exists());
assert!(written.aot_interface_h.exists());
let makefile = fs::read_to_string(&written.makefile).expect("Makefile should be readable");
let generated_c =
fs::read_to_string(&written.generated_c).expect("generated.c should be readable");
let manifest_h =
fs::read_to_string(&written.manifest_h).expect("manifest.h should be readable");
let aot_interface_c = fs::read_to_string(&written.aot_interface_c)
.expect("aot_interface.c should be readable");
let expected_target = if cfg!(target_os = "windows") {
"libgenerated_c_fixture.dll"
} else if cfg!(target_os = "macos") {
"libgenerated_c_fixture.dylib"
} else {
"libgenerated_c_fixture.so"
};
assert!(makefile.contains(expected_target));
assert!(generated_c.contains("#include <math.h>"));
assert!(manifest_h.contains("#define BACKEND_KIND \"aot\""));
assert!(manifest_h.contains("#define RESIDUAL_FN_NAME \"eval_residual\""));
assert!(aot_interface_c.contains("rustedscithe_aot_eval_residual"));
assert!(aot_interface_c.contains("rustedscithe_aot_eval_jacobian_values"));
assert!(aot_interface_c.contains("const double* args = args_ptr;"));
assert!(aot_interface_c.contains("double* out = out_ptr;"));
}
#[test]
fn generated_c_library_defers_compiler_choice_to_build_configuration() {
let prepared = sample_prepared_problem();
let module = CodegenModule::new("generated_fixture").with_language(CodegenLanguage::C);
let library_spec =
GeneratedCAotLibrary::from_prepared_problem("generated_c_fixture", &prepared, &module);
assert_eq!(library_spec.compiler, "gcc");
assert_eq!(library_spec.with_compiler("tcc").compiler, "tcc");
}
#[test]
fn generated_c_library_exports_chunk_ffi_symbols() {
let prepared = sample_chunked_prepared_problem();
let module = CodegenModule::new("generated_fixture").with_language(CodegenLanguage::C);
let library_spec =
GeneratedCAotLibrary::from_prepared_problem("generated_c_fixture", &prepared, &module);
let aot_interface_c = library_spec.emit_aot_interface_c();
assert!(aot_interface_c.contains("rustedscithe_aot_chunk_eval_residual_chunk_0"));
assert!(aot_interface_c.contains("rustedscithe_aot_chunk_eval_residual_chunk_1"));
assert!(aot_interface_c.contains("rustedscithe_aot_chunk_eval_jacobian_chunk_0"));
assert!(aot_interface_c.contains("rustedscithe_aot_chunk_eval_jacobian_chunk_1"));
assert!(aot_interface_c.contains("for (size_t i = 0; i < out_len; ++i) out[i] = 0.0;"));
}
#[test]
#[should_panic(expected = "snake_case ASCII")]
fn generated_c_library_rejects_invalid_names() {
let prepared = sample_prepared_problem();
let module = CodegenModule::new("test").with_language(CodegenLanguage::C);
let _ = GeneratedCAotLibrary::from_prepared_problem("Bad-Name", &prepared, &module);
}
#[test]
fn generated_c_library_forces_c_source_even_if_module_defaults_to_rust() {
let prepared = sample_prepared_problem();
let module = CodegenModule::new("generated_fixture").add_block(
"eval_residual",
&[Expr::parse_expression("x + 1")],
&["x"],
);
let library_spec =
GeneratedCAotLibrary::from_prepared_problem("generated_c_fixture", &prepared, &module);
assert!(library_spec.c_source.contains("#include <math.h>"));
assert!(!library_spec.c_source.contains("pub mod generated_fixture"));
assert!(!library_spec.c_source.contains("#![allow(clippy::all)]"));
}
#[test]
fn generated_c_library_supports_banded_prepared_problem() {
let residuals = vec![Expr::parse_expression("x + p")];
let vars = vec!["x"];
let params = vec!["p"];
let jac = Expr::parse_expression("2");
let entries = vec![BandedExprEntry {
row: 0,
col: 0,
diag_offset: 0,
diag_position: 0,
expr: &jac,
}];
let prepared = PreparedProblem::banded(PreparedBandedProblem::new(
BackendKind::Aot,
MatrixBackend::Banded,
ResidualTask {
fn_name: "eval_residual",
residuals: &residuals,
variables: &vars,
params: Some(¶ms),
}
.runtime_plan(ResidualChunkingStrategy::Whole),
BandedJacobianTask {
fn_name: "eval_banded_values",
shape: (1, 1),
kl: 0,
ku: 0,
compact: false,
entries: &entries,
variables: &vars,
params: Some(¶ms),
}
.runtime_plan(BandedChunkingStrategy::Whole),
));
let module = CodegenModule::new("generated_banded_fixture")
.with_language(CodegenLanguage::C)
.add_residual_block_plan(match &prepared {
PreparedProblem::Banded(problem) => &problem.residual_plan.chunks[0].plan,
PreparedProblem::Dense(_) | PreparedProblem::Sparse(_) => {
unreachable!("test uses banded prepared problem")
}
})
.add_sparse_values_plan(match &prepared {
PreparedProblem::Banded(problem) => &problem.jacobian_plan.chunks[0].plan,
PreparedProblem::Dense(_) | PreparedProblem::Sparse(_) => {
unreachable!("test uses banded prepared problem")
}
});
let library_spec = GeneratedCAotLibrary::from_prepared_problem(
"generated_banded_c_fixture",
&prepared,
&module,
);
assert_eq!(library_spec.manifest.matrix_backend, MatrixBackend::Banded);
assert!(library_spec.c_source.contains("eval_banded_values"));
}
}