header.rs

// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2026 Jonathan D.A. Jewell (hyperpolymath) <j.d.a.jewell@open.ac.uk>
//
// C header generator for Chapeliser.
//
// Generates the C header that declares all functions the user must implement.
// This is the contract between the user's application code and the Chapeliser
// runtime (Chapel wrapper + Zig FFI bridge).

use anyhow::{Context, Result};
use std::fmt::Write as FmtWrite;
use std::fs;
use std::path::Path;

use crate::manifest::Manifest;

/// Generate the C header from a manifest.
/// Writes to `output_dir/include/<name>_chapeliser.h`.
pub fn generate(manifest: &Manifest, output_dir: &Path) -> Result<()> {
    let name = &manifest.workload.name;
    let safe_name = name.replace('-', "_");
    let upper = safe_name.to_uppercase();

    let mut src = String::with_capacity(4096);

    writeln!(src, "/* SPDX-License-Identifier: MPL-2.0 */")?;
    writeln!(
        src,
        "/* Auto-generated by Chapeliser — do not edit manually. */"
    )?;
    writeln!(src, "/* C header for workload: {name} */")?;
    writeln!(src, "/* Regenerate with: chapeliser generate */")?;
    writeln!(src)?;
    writeln!(src, "#ifndef CHAPELISER_{upper}_H")?;
    writeln!(src, "#define CHAPELISER_{upper}_H")?;
    writeln!(src)?;
    writeln!(src, "#include <stddef.h>")?;
    writeln!(src, "#include <stdint.h>")?;
    writeln!(src)?;
    writeln!(src, "#ifdef __cplusplus")?;
    writeln!(src, "extern \"C\" {{")?;
    writeln!(src, "#endif")?;
    writeln!(src)?;

    // Documentation block
    writeln!(src, "/*")?;
    writeln!(src, " * Chapeliser FFI contract for workload: {name}")?;
    writeln!(src, " *")?;
    writeln!(
        src,
        " * Implement these functions with C linkage in your application."
    )?;
    writeln!(
        src,
        " * Chapeliser's generated Chapel + Zig code calls them at runtime."
    )?;
    writeln!(src, " *")?;
    writeln!(src, " * Return values: 0 = success, non-zero = error code.")?;
    writeln!(src, " *")?;
    writeln!(
        src,
        " * For Rust:  #[no_mangle] pub extern \"C\" fn {safe_name}_init() -> i32 {{ ... }}"
    )?;
    writeln!(
        src,
        " * For Zig:   export fn {safe_name}_init() callconv(.C) c_int {{ ... }}"
    )?;
    writeln!(src, " * For C:     int {safe_name}_init(void) {{ ... }}")?;
    writeln!(src, " */")?;
    writeln!(src)?;

    // --- Lifecycle ---
    writeln!(src, "/* ---- Lifecycle ---- */")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Initialise the workload. Called once on locale 0 before processing. */"
    )?;
    writeln!(src, "int {safe_name}_init(void);")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Shut down the workload. Called once on locale 0 after all results stored. */"
    )?;
    writeln!(src, "int {safe_name}_shutdown(void);")?;
    writeln!(src)?;

    // --- Data I/O ---
    writeln!(src, "/* ---- Data I/O ---- */")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Return the total number of input items. Called on locale 0. */"
    )?;
    writeln!(src, "int {safe_name}_get_total_items(void);")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Serialise input item `idx` into `buf`. Set `*len` to bytes written. */"
    )?;
    writeln!(
        src,
        "/* `buf` has capacity {max_bytes} bytes (from manifest max-item-bytes). */",
        max_bytes = manifest.data.max_item_bytes.unwrap_or(1_048_576)
    )?;
    writeln!(
        src,
        "int {safe_name}_load_item(int idx, uint8_t* buf, size_t* len);"
    )?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Receive a processed result. `buf` contains `len` serialised bytes. */"
    )?;
    writeln!(
        src,
        "int {safe_name}_store_result(int idx, const uint8_t* buf, size_t len);"
    )?;
    writeln!(src)?;

    // --- Processing ---
    writeln!(src, "/* ---- Processing ---- */")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Process a single serialised item. Read from in_buf/in_len, */"
    )?;
    writeln!(
        src,
        "/* write result to out_buf, set *out_len. Called on any locale. */"
    )?;
    writeln!(src, "int {safe_name}_process_item(")?;
    writeln!(src, "    const uint8_t* in_buf, size_t in_len,")?;
    writeln!(src, "    uint8_t* out_buf, size_t* out_len")?;
    writeln!(src, ");")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Process a chunk of items (for chunk partition strategy). */"
    )?;
    writeln!(src, "int {safe_name}_process_chunk(")?;
    writeln!(src, "    const uint8_t* items_buf, int item_count,")?;
    writeln!(src, "    const int* item_offsets, const int* item_sizes,")?;
    writeln!(src, "    uint8_t* out_buf, size_t* out_len")?;
    writeln!(src, ");")?;
    writeln!(src)?;

    // --- Reduction ---
    writeln!(
        src,
        "/* ---- Reduction (for reduce/tree-reduce gather) ---- */"
    )?;
    writeln!(src)?;
    writeln!(src, "/* Combine two results into one. */")?;
    writeln!(src, "int {safe_name}_reduce(")?;
    writeln!(src, "    const uint8_t* a_buf, size_t a_len,")?;
    writeln!(src, "    const uint8_t* b_buf, size_t b_len,")?;
    writeln!(src, "    uint8_t* out_buf, size_t* out_len")?;
    writeln!(src, ");")?;
    writeln!(src)?;

    // --- Match predicate ---
    writeln!(src, "/* ---- Match predicate (for 'first' gather) ---- */")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Return 1 if result satisfies the search criterion, 0 otherwise. */"
    )?;
    writeln!(
        src,
        "int {safe_name}_is_match(const uint8_t* buf, size_t len);"
    )?;
    writeln!(src)?;

    // --- Key hash ---
    writeln!(src, "/* ---- Key hash (for keyed partition) ---- */")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Return a hash of the item's key for distribution routing. */"
    )?;
    writeln!(
        src,
        "unsigned int {safe_name}_key_hash(const uint8_t* buf, size_t len);"
    )?;
    writeln!(src)?;

    // --- Checkpoint ---
    writeln!(src, "/* ---- Checkpoint (optional) ---- */")?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Save checkpoint data. Return -1 if not implemented. */"
    )?;
    writeln!(
        src,
        "int {safe_name}_checkpoint_save(const uint8_t* buf, size_t len, const char* tag);"
    )?;
    writeln!(src)?;
    writeln!(
        src,
        "/* Load checkpoint data. Return -1 if not implemented. */"
    )?;
    writeln!(
        src,
        "int {safe_name}_checkpoint_load(uint8_t* buf, size_t* len, const char* tag);"
    )?;
    writeln!(src)?;

    // Close
    writeln!(src, "#ifdef __cplusplus")?;
    writeln!(src, "}}")?;
    writeln!(src, "#endif")?;
    writeln!(src)?;
    writeln!(src, "#endif /* CHAPELISER_{upper}_H */")?;

    let out_path = output_dir.join(format!("{}_chapeliser.h", safe_name));
    fs::write(&out_path, &src)
        .with_context(|| format!("Failed to write C header: {}", out_path.display()))?;
    println!("  C header:       {}", out_path.display());

    // Also emit the c_* ABI header that the generated Chapel `require`s, so chpl
    // can resolve the extern procs (needed for remote 'on'/'coforall' calls).
    // These symbols are provided by the Zig FFI bridge at link time.
    write_abi_header(output_dir, &safe_name, &upper)?;

    Ok(())
}

/// Emit `<name>_abi.h` — C prototypes for the `c_*` bridge functions that the
/// generated Chapel calls via `extern proc`. The Chapel wrapper `require`s this
/// header; the Zig bridge (or, in tests, a stub) provides the definitions.
fn write_abi_header(output_dir: &Path, safe_name: &str, upper: &str) -> Result<()> {
    let mut h = String::with_capacity(2048);
    writeln!(h, "/* SPDX-License-Identifier: MPL-2.0 */")?;
    writeln!(h, "/* Auto-generated by Chapeliser — do not edit manually. */")?;
    writeln!(
        h,
        "/* C-ABI prototypes for the c_* bridge the generated Chapel calls. */"
    )?;
    writeln!(h, "#ifndef CHAPELISER_{upper}_ABI_H")?;
    writeln!(h, "#define CHAPELISER_{upper}_ABI_H")?;
    writeln!(h)?;
    writeln!(h, "#include <stddef.h>")?;
    writeln!(h, "#include <stdint.h>")?;
    writeln!(h)?;
    writeln!(h, "#ifdef __cplusplus")?;
    writeln!(h, "extern \"C\" {{")?;
    writeln!(h, "#endif")?;
    writeln!(h)?;
    writeln!(h, "int c_init(void);")?;
    writeln!(h, "int c_shutdown(void);")?;
    writeln!(h, "int c_get_total_items(void);")?;
    writeln!(h, "int c_load_item(int idx, uint8_t* buf, size_t* len);")?;
    writeln!(h, "int c_store_result(int idx, uint8_t* buf, size_t len);")?;
    writeln!(
        h,
        "int c_process_item(uint8_t* in_buf, size_t in_len, uint8_t* out_buf, size_t* out_len);"
    )?;
    writeln!(
        h,
        "int c_process_chunk(uint8_t* items_buf, int item_count, int* item_offsets, int* item_sizes, uint8_t* out_buf, size_t* out_len);"
    )?;
    writeln!(
        h,
        "int c_reduce(uint8_t* a_buf, size_t a_len, uint8_t* b_buf, size_t b_len, uint8_t* out_buf, size_t* out_len);"
    )?;
    writeln!(h, "int c_is_match(uint8_t* buf, size_t len);")?;
    writeln!(h, "unsigned int c_key_hash(uint8_t* buf, size_t len);")?;
    writeln!(
        h,
        "int c_checkpoint_save(uint8_t* buf, size_t len, const char* tag);"
    )?;
    writeln!(
        h,
        "int c_checkpoint_load(uint8_t* buf, size_t* len, const char* tag);"
    )?;
    writeln!(h)?;
    writeln!(h, "#ifdef __cplusplus")?;
    writeln!(h, "}}")?;
    writeln!(h, "#endif")?;
    writeln!(h)?;
    writeln!(h, "#endif /* CHAPELISER_{upper}_ABI_H */")?;

    let out_path = output_dir.join(format!("{}_abi.h", safe_name));
    fs::write(&out_path, &h)
        .with_context(|| format!("Failed to write ABI header: {}", out_path.display()))?;
    println!("  ABI header:     {}", out_path.display());
    Ok(())
}