chapel.rs

// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2026 Jonathan D.A. Jewell (hyperpolymath) <j.d.a.jewell@open.ac.uk>
//
// Chapel code generator for Chapeliser.
//
// Produces a complete, compilable Chapel program from a chapeliser.toml manifest.
// The generated .chpl file is a self-contained distributed program that:
//   1. Initialises the user's library via C-ABI (c_init)
//   2. Loads input items as serialised byte buffers (c_load_item)
//   3. Distributes items across Chapel locales using the chosen partition strategy
//   4. Processes items via the user's C-ABI function (c_process_item / c_process_chunk)
//   5. Gathers results using the chosen gather strategy
//   6. Stores output via C-ABI (c_store_result)
//   7. Shuts down (c_shutdown)
//
// The user never writes Chapel code — they implement 6-8 C-ABI functions
// (declared in the generated .h header) and Chapeliser does the rest.

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

use crate::manifest::Manifest;

/// Generate the Chapel wrapper program from a manifest.
/// Writes a single .chpl file to `output_dir/chapel/<name>_distributed.chpl`.
pub fn generate(manifest: &Manifest, output_dir: &Path) -> Result<()> {
    let name = &manifest.workload.name;
    let safe_name = name.replace('-', "_");

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

    write_header(&mut src, manifest)?;
    write_imports(&mut src, manifest)?;
    write_config_constants(&mut src, manifest)?;
    write_ffi_declarations(&mut src, &safe_name)?;
    write_helpers(&mut src, manifest)?;
    write_main_proc(&mut src, manifest, &safe_name)?;
    write_module_close(&mut src)?;

    let out_path = output_dir.join(format!("{}_distributed.chpl", safe_name));
    fs::write(&out_path, &src)
        .with_context(|| format!("Failed to write Chapel wrapper: {}", out_path.display()))?;
    println!("  Chapel wrapper: {}", out_path.display());

    Ok(())
}

/// SPDX header, module declaration, and metadata comment.
fn write_header(src: &mut String, manifest: &Manifest) -> Result<()> {
    let name = &manifest.workload.name;
    let safe_name = name.replace('-', "_");

    writeln!(src, "// SPDX-License-Identifier: MPL-2.0")?;
    writeln!(
        src,
        "// Auto-generated by Chapeliser — do not edit manually."
    )?;
    writeln!(src, "// Workload: {name}")?;
    writeln!(
        src,
        "// Partition: {}, Gather: {}",
        manifest.workload.partition, manifest.workload.gather
    )?;
    writeln!(
        src,
        "// Serialization: {}, Max item bytes: {}",
        manifest.data.serialization,
        manifest.data.max_item_bytes.unwrap_or(1_048_576)
    )?;
    writeln!(src, "// Regenerate with: chapeliser generate")?;
    writeln!(src)?;
    writeln!(src, "module {safe_name}_Distributed {{")?;
    writeln!(src)?;

    Ok(())
}

/// Chapel `use` statements — only include modules actually needed by the
/// chosen partition and gather strategies.
fn write_imports(src: &mut String, manifest: &Manifest) -> Result<()> {
    writeln!(src, "  use CTypes;")?;
    writeln!(src, "  use List;")?;
    writeln!(src, "  use Time;")?;
    writeln!(src, "  use IO;")?;

    // BlockDist needed for spatial partition and general distribution
    if manifest.workload.partition == "spatial" {
        writeln!(src, "  use BlockDist;")?;
    }

    // DynamicIters needed for adaptive partition
    if manifest.workload.partition == "adaptive" {
        writeln!(src, "  use DynamicIters;")?;
    }

    // Atomic variables needed for 'first' gather strategy
    if manifest.workload.gather == "first" {
        writeln!(src, "  use AtomicObjects;")?;
    }

    // C prototypes for the c_* ABI so chpl can resolve the extern procs —
    // needed for calls inside 'on'/'coforall' (potentially remote) contexts.
    let safe_name = manifest.workload.name.replace('-', "_");
    writeln!(src, "  require \"{safe_name}_abi.h\";")?;

    writeln!(src)?;
    Ok(())
}

/// Runtime-configurable constants. Users override these with Chapel's
/// `--configName=value` command-line syntax.
fn write_config_constants(src: &mut String, manifest: &Manifest) -> Result<()> {
    let max_bytes = manifest.data.max_item_bytes.unwrap_or(1_048_576);
    let grain = manifest.scaling.grain_size;
    let retries = manifest.resilience.retries;
    let checkpoint_enabled = manifest.resilience.checkpoint;
    let checkpoint_interval = manifest.resilience.checkpoint_interval_secs;

    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(
        src,
        "  // Runtime configuration (override with --name=value)"
    )?;
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(src)?;

    // Item count: 0 means "ask the user library via c_get_total_items()"
    writeln!(
        src,
        "  config const totalItems: int = 0;  // 0 = auto-detect via c_get_total_items()"
    )?;
    writeln!(
        src,
        "  config const maxItemBytes: int = {max_bytes};  // max serialised size per item"
    )?;
    writeln!(
        src,
        "  config const grainSize: int = {grain};  // items per task (for chunk strategy)"
    )?;
    writeln!(
        src,
        "  config const maxRetries: int = {retries};  // per-item retry limit"
    )?;
    writeln!(
        src,
        "  config const enableCheckpoint: bool = {checkpoint_enabled};  // periodic checkpoint"
    )?;
    writeln!(
        src,
        "  config const checkpointIntervalSecs: int = {checkpoint_interval};  // seconds between checkpoints"
    )?;
    writeln!(src)?;

    Ok(())
}

/// Declare all C-ABI external functions the Chapel program calls.
/// These are implemented in the Zig FFI layer (or directly by the user).
fn write_ffi_declarations(src: &mut String, _safe_name: &str) -> Result<()> {
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(
        src,
        "  // FFI declarations (implemented in Zig bridge / user code)"
    )?;
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(src)?;

    // Lifecycle
    writeln!(src, "  // Lifecycle: called once on locale 0")?;
    writeln!(src, "  extern proc c_init(): c_int;")?;
    writeln!(src, "  extern proc c_shutdown(): c_int;")?;
    writeln!(src)?;

    // Data I/O
    writeln!(
        src,
        "  // Data I/O: load/store serialised items on locale 0"
    )?;
    writeln!(src, "  extern proc c_get_total_items(): c_int;")?;
    writeln!(
        src,
        "  extern proc c_load_item(idx: c_int, buf: c_ptr(c_uchar), len: c_ptr(c_size_t)): c_int;"
    )?;
    writeln!(
        src,
        "  extern proc c_store_result(idx: c_int, buf: c_ptr(c_uchar), len: c_size_t): c_int;"
    )?;
    writeln!(src)?;

    // Processing
    writeln!(src, "  // Processing: called per item/chunk on any locale")?;
    writeln!(src, "  extern proc c_process_item(")?;
    writeln!(src, "    in_buf: c_ptr(c_uchar), in_len: c_size_t,")?;
    writeln!(src, "    out_buf: c_ptr(c_uchar), out_len: c_ptr(c_size_t)")?;
    writeln!(src, "  ): c_int;")?;
    writeln!(src)?;
    writeln!(src, "  extern proc c_process_chunk(")?;
    writeln!(
        src,
        "    items_buf: c_ptr(c_uchar), item_count: c_int, item_offsets: c_ptr(c_int), item_sizes: c_ptr(c_int),"
    )?;
    writeln!(src, "    out_buf: c_ptr(c_uchar), out_len: c_ptr(c_size_t)")?;
    writeln!(src, "  ): c_int;")?;
    writeln!(src)?;

    // Reduction (for reduce/tree-reduce gather)
    writeln!(src, "  // Reduction: combine two results into one")?;
    writeln!(src, "  extern proc c_reduce(")?;
    writeln!(src, "    a_buf: c_ptr(c_uchar), a_len: c_size_t,")?;
    writeln!(src, "    b_buf: c_ptr(c_uchar), b_len: c_size_t,")?;
    writeln!(src, "    out_buf: c_ptr(c_uchar), out_len: c_ptr(c_size_t)")?;
    writeln!(src, "  ): c_int;")?;
    writeln!(src)?;

    // Match predicate (for 'first' gather)
    writeln!(
        src,
        "  // Match predicate: returns 1 if result satisfies search criterion"
    )?;
    writeln!(
        src,
        "  extern proc c_is_match(buf: c_ptr(c_uchar), len: c_size_t): c_int;"
    )?;
    writeln!(src)?;

    // Key extraction (for keyed partition)
    writeln!(
        src,
        "  // Key hash: returns hash of item's key for keyed distribution"
    )?;
    writeln!(
        src,
        "  extern proc c_key_hash(buf: c_ptr(c_uchar), len: c_size_t): c_uint;"
    )?;
    writeln!(src)?;

    // Checkpoint (optional, only called if enableCheckpoint is true)
    writeln!(src, "  // Checkpoint: save/load progress (optional)")?;
    writeln!(src, "  extern proc c_checkpoint_save(")?;
    writeln!(
        src,
        "    buf: c_ptr(c_uchar), len: c_size_t, tag: c_ptrConst(c_char)"
    )?;
    writeln!(src, "  ): c_int;")?;
    writeln!(src, "  extern proc c_checkpoint_load(")?;
    writeln!(
        src,
        "    buf: c_ptr(c_uchar), len: c_ptr(c_size_t), tag: c_ptrConst(c_char)"
    )?;
    writeln!(src, "  ): c_int;")?;
    writeln!(src)?;

    Ok(())
}

/// Helper procedures used by the main distribution/gather logic.
fn write_helpers(src: &mut String, manifest: &Manifest) -> Result<()> {
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(src, "  // Helper procedures")?;
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(src)?;

    // Compute the index range for a given locale (even distribution with remainder)
    writeln!(
        src,
        "  // Compute the item range assigned to a locale for even distribution."
    )?;
    writeln!(
        src,
        "  proc localeRange(locId: int, nLocs: int, nItems: int): range {{"
    )?;
    writeln!(src, "    const base = nItems / nLocs;")?;
    writeln!(src, "    const rem  = nItems % nLocs;")?;
    writeln!(src, "    const lo = locId * base + min(locId, rem);")?;
    writeln!(
        src,
        "    const hi = lo + base + (if locId < rem then 1 else 0);"
    )?;
    writeln!(src, "    return lo..#(hi - lo);")?;
    writeln!(src, "  }}")?;
    writeln!(src)?;

    // Process a single item with retry logic
    writeln!(
        src,
        "  // Process one item with retry. Returns 0 on success, last error code on failure."
    )?;
    writeln!(
        src,
        "  proc processWithRetry(ref inBuf: [] uint(8), inLen: c_size_t,"
    )?;
    writeln!(
        src,
        "                        ref outBuf: [] uint(8), ref outLen: c_size_t): c_int {{"
    )?;
    writeln!(src, "    var lastRc: c_int = -1;")?;
    writeln!(src, "    for attempt in 0..#maxRetries {{")?;
    writeln!(
        src,
        "      lastRc = c_process_item(c_ptrTo(inBuf[0]), inLen, c_ptrTo(outBuf[0]), c_ptrTo(outLen));"
    )?;
    writeln!(src, "      if lastRc == 0 then return 0;")?;
    writeln!(
        src,
        "      writeln(\"  WARN: item processing failed (attempt \", attempt + 1, \"/\", maxRetries, \", rc=\", lastRc, \")\");"
    )?;
    writeln!(src, "    }}")?;
    writeln!(src, "    return lastRc;")?;
    writeln!(src, "  }}")?;
    writeln!(src)?;

    // Chunk range helper for chunk strategy
    if manifest.workload.partition == "chunk" {
        writeln!(src, "  // Compute the item range for a given chunk index.")?;
        writeln!(
            src,
            "  proc chunkRange(chunkIdx: int, nItems: int): range {{"
        )?;
        writeln!(src, "    const lo = chunkIdx * grainSize;")?;
        writeln!(src, "    const hi = min(lo + grainSize, nItems);")?;
        writeln!(src, "    return lo..#(hi - lo);")?;
        writeln!(src, "  }}")?;
        writeln!(src)?;
    }

    Ok(())
}

/// The main() proc — orchestrates load, distribute, process, gather, store.
fn write_main_proc(src: &mut String, manifest: &Manifest, safe_name: &str) -> Result<()> {
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(src, "  // Main entry point")?;
    writeln!(
        src,
        "  // ------------------------------------------------------------------"
    )?;
    writeln!(src)?;
    writeln!(src, "  proc main() {{")?;
    writeln!(src, "    const t0 = timeSinceEpoch().totalSeconds();")?;
    writeln!(src)?;

    // --- Init ---
    writeln!(
        src,
        "    writeln(\"Chapeliser [{safe_name}]: initialising on \", numLocales, \" locale(s)\");"
    )?;
    writeln!(src, "    {{")?;
    writeln!(src, "      const rc = c_init();")?;
    writeln!(src, "      if rc != 0 {{")?;
    writeln!(src, "        writeln(\"FATAL: c_init() returned \", rc);")?;
    writeln!(src, "        return;")?;
    writeln!(src, "      }}")?;
    writeln!(src, "    }}")?;
    writeln!(src)?;

    // --- Determine item count ---
    writeln!(
        src,
        "    const nItems: int = if totalItems > 0 then totalItems else c_get_total_items(): int;"
    )?;
    writeln!(src, "    if nItems <= 0 {{")?;
    writeln!(
        src,
        "      writeln(\"Chapeliser: nothing to do (nItems=\", nItems, \")\");"
    )?;
    writeln!(src, "      c_shutdown();")?;
    writeln!(src, "      return;")?;
    writeln!(src, "    }}")?;
    writeln!(
        src,
        "    writeln(\"Chapeliser: distributing \", nItems, \" items (partition={partition}, gather={gather})\");",
        partition = manifest.workload.partition,
        gather = manifest.workload.gather,
    )?;
    writeln!(src)?;

    // --- Load items on locale 0 ---
    write_load_phase(src)?;

    // --- Distribution + processing (strategy-specific) ---
    write_distribution_phase(src, manifest)?;

    // --- Gather (strategy-specific) ---
    write_gather_phase(src, manifest)?;

    // --- Store results ---
    write_store_phase(src, manifest)?;

    // --- Shutdown ---
    writeln!(src, "    c_shutdown();")?;
    writeln!(
        src,
        "    const elapsed = timeSinceEpoch().totalSeconds() - t0;"
    )?;
    writeln!(
        src,
        "    writeln(\"Chapeliser [{safe_name}]: complete in \", elapsed:real(32), \"s\");"
    )?;
    writeln!(src, "  }}")?;
    writeln!(src)?;

    Ok(())
}

/// Phase 1: Load all input items as serialised byte buffers on locale 0.
fn write_load_phase(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(src, "    // Phase 1: Load items on locale 0")?;
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(src)?;

    // Allocate per-item byte buffers and size tracking.
    // Chapel allocates these as local arrays on locale 0 (the launching locale).
    writeln!(
        src,
        "    var itemData: [0..#nItems] [0..#maxItemBytes] uint(8);"
    )?;
    writeln!(src, "    var itemSizes: [0..#nItems] c_size_t;")?;
    writeln!(src)?;

    writeln!(src, "    var loadFailed = false;")?;
    writeln!(src, "    on Locales[0] {{")?;
    writeln!(src, "      for i in 0..#nItems {{")?;
    writeln!(src, "        var sz: c_size_t = maxItemBytes: c_size_t;")?;
    writeln!(
        src,
        "        const rc = c_load_item(i: c_int, c_ptrTo(itemData[i][0]), c_ptrTo(sz));"
    )?;
    writeln!(src, "        if rc != 0 {{")?;
    writeln!(
        src,
        "          writeln(\"FATAL: c_load_item(\", i, \") returned \", rc);"
    )?;
    writeln!(src, "          loadFailed = true;")?;
    writeln!(src, "          break;")?;
    writeln!(src, "        }}")?;
    writeln!(src, "        itemSizes[i] = sz;")?;
    writeln!(src, "      }}")?;
    writeln!(src, "    }}")?;
    writeln!(
        src,
        "    // 'return' is illegal inside an 'on' block; bail out after it."
    )?;
    writeln!(src, "    if loadFailed {{")?;
    writeln!(src, "      c_shutdown();")?;
    writeln!(src, "      return;")?;
    writeln!(src, "    }}")?;
    writeln!(src, "    writeln(\"  Loaded \", nItems, \" items\");")?;
    writeln!(src)?;

    Ok(())
}

/// Phase 2: Distribute items and process them.
/// The output is a set of result buffers — one per processed item (or aggregated,
/// depending on gather strategy).
fn write_distribution_phase(src: &mut String, manifest: &Manifest) -> Result<()> {
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(
        src,
        "    // Phase 2: Distribute & process (strategy: {})",
        manifest.workload.partition
    )?;
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(src)?;

    // Result storage — for merge/stream/first we have one result per item.
    // For reduce/tree-reduce, results are aggregated during gather.
    writeln!(
        src,
        "    var resultData: [0..#nItems] [0..#maxItemBytes] uint(8);"
    )?;
    writeln!(src, "    var resultSizes: [0..#nItems] c_size_t;")?;
    writeln!(
        src,
        "    var resultOk: [0..#nItems] bool;  // track which items succeeded"
    )?;
    writeln!(src)?;

    match manifest.workload.partition.as_str() {
        "per-item" => write_per_item_distribution(src)?,
        "chunk" => write_chunk_distribution(src)?,
        "adaptive" => write_adaptive_distribution(src)?,
        "spatial" => write_spatial_distribution(src)?,
        "keyed" => write_keyed_distribution(src)?,
        other => {
            writeln!(src, "    // ERROR: unknown partition strategy '{other}'")?;
            writeln!(src, "    halt(\"Unknown partition strategy: {other}\");")?;
        }
    }

    writeln!(src)?;
    Ok(())
}

/// Per-item distribution: evenly divide items across locales, each locale
/// processes its slice sequentially.
fn write_per_item_distribution(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Per-item: each locale gets nItems/numLocales items"
    )?;
    writeln!(
        src,
        "    coforall loc in Locales with (ref resultData, ref resultSizes, ref resultOk) do on loc {{"
    )?;
    writeln!(
        src,
        "      const myRange = localeRange(loc.id, numLocales, nItems);"
    )?;
    writeln!(src, "      for i in myRange {{")?;

    write_process_single_item(src, "i")?;

    writeln!(src, "      }}")?;

    write_checkpoint_hook(src)?;

    writeln!(src, "    }}")?;

    Ok(())
}

/// Chunk distribution: divide items into fixed-size chunks, round-robin
/// chunks across locales. Allows the user's c_process_chunk to batch-process.
fn write_chunk_distribution(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Chunk: fixed-size chunks of grainSize items, round-robin across locales"
    )?;
    writeln!(
        src,
        "    const numChunks = (nItems + grainSize - 1) / grainSize;"
    )?;
    writeln!(src)?;
    writeln!(
        src,
        "    coforall loc in Locales with (ref resultData, ref resultSizes, ref resultOk) do on loc {{"
    )?;
    writeln!(
        src,
        "      for chunkIdx in loc.id..#numChunks by numLocales {{"
    )?;
    writeln!(src, "        const cr = chunkRange(chunkIdx, nItems);")?;
    writeln!(
        src,
        "        // Process each item in the chunk individually through the retry wrapper"
    )?;
    writeln!(src, "        for i in cr {{")?;

    write_process_single_item(src, "i")?;

    writeln!(src, "        }}")?;
    writeln!(src, "      }}")?;

    write_checkpoint_hook(src)?;

    writeln!(src, "    }}")?;

    Ok(())
}

/// Adaptive distribution: work-stealing with Chapel's DynamicIters.
/// Items are drawn from a shared work pool dynamically, so fast locales
/// automatically pick up slack from slow ones.
fn write_adaptive_distribution(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Adaptive: dynamic work-stealing via Chapel's DynamicIters"
    )?;
    writeln!(
        src,
        "    // Items are pulled from a shared iterator — fast locales do more work."
    )?;
    writeln!(
        src,
        "    forall i in dynamic(0..#nItems, chunkSize=grainSize) with (ref resultData, ref resultSizes, ref resultOk) {{"
    )?;

    write_process_single_item(src, "i")?;

    writeln!(src, "    }}")?;

    Ok(())
}

/// Spatial distribution: Block-distributed domain decomposition.
/// Best for 2D/3D data where locality matters (simulations, matrices).
fn write_spatial_distribution(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Spatial: Block distribution — contiguous regions per locale"
    )?;
    writeln!(
        src,
        "    const Space = {{0..#nItems}} dmapped new blockDist({{0..#nItems}});"
    )?;
    writeln!(
        src,
        "    forall i in Space with (ref resultData, ref resultSizes, ref resultOk) {{"
    )?;

    write_process_single_item(src, "i")?;

    writeln!(src, "    }}")?;

    Ok(())
}

/// Keyed distribution: route items to locales by key hash.
/// All items with the same key land on the same locale, enabling
/// locale-local aggregation.
fn write_keyed_distribution(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Keyed: items routed to locales by key hash (same key → same locale)"
    )?;
    writeln!(
        src,
        "    coforall loc in Locales with (ref resultData, ref resultSizes, ref resultOk) do on loc {{"
    )?;
    writeln!(src, "      for i in 0..#nItems {{")?;
    writeln!(
        src,
        "        const h = c_key_hash(c_ptrTo(itemData[i][0]), itemSizes[i]);"
    )?;
    writeln!(
        src,
        "        if (h % numLocales:c_uint) == loc.id:c_uint {{"
    )?;

    write_process_single_item(src, "i")?;

    writeln!(src, "        }}")?;
    writeln!(src, "      }}")?;

    write_checkpoint_hook(src)?;

    writeln!(src, "    }}")?;

    Ok(())
}

/// Emit the code that processes a single item at index `idx_var` with retry.
/// This is inlined into every distribution strategy.
fn write_process_single_item(src: &mut String, idx_var: &str) -> Result<()> {
    writeln!(
        src,
        "        // Copy item to locale-local buffer, process, store result"
    )?;
    writeln!(
        src,
        "        var localIn: [0..#maxItemBytes] uint(8) = itemData[{idx_var}];"
    )?;
    writeln!(src, "        var localOut: [0..#maxItemBytes] uint(8);")?;
    writeln!(src, "        var outLen: c_size_t = 0;")?;
    writeln!(
        src,
        "        const rc = processWithRetry(localIn, itemSizes[{idx_var}], localOut, outLen);"
    )?;
    writeln!(src, "        if rc == 0 {{")?;
    writeln!(src, "          resultData[{idx_var}] = localOut;")?;
    writeln!(src, "          resultSizes[{idx_var}] = outLen;")?;
    writeln!(src, "          resultOk[{idx_var}] = true;")?;
    writeln!(src, "        }} else {{")?;
    writeln!(
        src,
        "          writeln(\"  ERROR: item \", {idx_var}, \" failed after \", maxRetries, \" retries (rc=\", rc, \")\");"
    )?;
    writeln!(src, "          resultOk[{idx_var}] = false;")?;
    writeln!(src, "        }}")?;

    Ok(())
}

/// Optional checkpoint hook — inserted after each locale finishes its batch.
fn write_checkpoint_hook(src: &mut String) -> Result<()> {
    writeln!(src)?;
    writeln!(
        src,
        "      // Checkpoint: save locale-local progress if enabled"
    )?;
    writeln!(src, "      if enableCheckpoint {{")?;
    writeln!(
        src,
        "        // Serialise locale's result indices into a compact buffer for checkpoint."
    )?;
    writeln!(
        src,
        "        // The tag encodes the locale id so checkpoints don't collide."
    )?;
    writeln!(
        src,
        "        var tag: string = \"locale-\" + loc.id:string;"
    )?;
    writeln!(
        src,
        "        var ckBuf: [0..#8] uint(8);  // minimal: just store completion marker"
    )?;
    writeln!(src, "        ckBuf[0] = 1;  // 1 = this locale finished")?;
    writeln!(
        src,
        "        c_checkpoint_save(c_ptrTo(ckBuf[0]), 1: c_size_t, tag.c_str());"
    )?;
    writeln!(src, "      }}")?;

    Ok(())
}

/// Phase 3: Gather results according to the chosen gather strategy.
fn write_gather_phase(src: &mut String, manifest: &Manifest) -> Result<()> {
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(
        src,
        "    // Phase 3: Gather results (strategy: {})",
        manifest.workload.gather
    )?;
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(src)?;

    match manifest.workload.gather.as_str() {
        "merge" => write_gather_merge(src)?,
        "reduce" => write_gather_reduce(src)?,
        "tree-reduce" => write_gather_tree_reduce(src)?,
        "stream" => write_gather_stream(src)?,
        "first" => write_gather_first(src)?,
        other => {
            writeln!(src, "    // ERROR: unknown gather strategy '{other}'")?;
            writeln!(src, "    halt(\"Unknown gather strategy: {other}\");")?;
        }
    }

    writeln!(src)?;
    Ok(())
}

/// Merge gather: all results are already in resultData[] from the distribution
/// phase. Count successes and report.
fn write_gather_merge(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Merge: results already stored in resultData[i] by distribution phase."
    )?;
    writeln!(src, "    // Count successes for reporting.")?;
    writeln!(src, "    var nSuccess = + reduce resultOk:int;")?;
    writeln!(
        src,
        "    writeln(\"  Gathered \", nSuccess, \"/\", nItems, \" results (merge)\");"
    )?;

    Ok(())
}

/// Reduce gather: fold all results into a single value using c_reduce.
/// Processes results sequentially on locale 0.
fn write_gather_reduce(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Reduce: fold all results into a single accumulator using c_reduce."
    )?;
    writeln!(
        src,
        "    var accumBuf: [0..#maxItemBytes] uint(8) = resultData[0];"
    )?;
    writeln!(src, "    var accumLen: c_size_t = resultSizes[0];")?;
    writeln!(src)?;
    writeln!(src, "    on Locales[0] {{")?;
    writeln!(src, "      for i in 1..#(nItems - 1) {{")?;
    writeln!(src, "        if !resultOk[i] then continue;")?;
    writeln!(src, "        var tmpBuf: [0..#maxItemBytes] uint(8);")?;
    writeln!(src, "        var tmpLen: c_size_t = 0;")?;
    writeln!(src, "        const rc = c_reduce(")?;
    writeln!(src, "          c_ptrTo(accumBuf[0]), accumLen,")?;
    writeln!(src, "          c_ptrTo(resultData[i][0]), resultSizes[i],")?;
    writeln!(src, "          c_ptrTo(tmpBuf[0]), c_ptrTo(tmpLen)")?;
    writeln!(src, "        );")?;
    writeln!(src, "        if rc == 0 {{")?;
    writeln!(src, "          accumBuf = tmpBuf;")?;
    writeln!(src, "          accumLen = tmpLen;")?;
    writeln!(src, "        }}")?;
    writeln!(src, "      }}")?;
    writeln!(src, "    }}")?;
    writeln!(src)?;
    writeln!(src, "    // Store reduced result as item 0")?;
    writeln!(src, "    resultData[0] = accumBuf;")?;
    writeln!(src, "    resultSizes[0] = accumLen;")?;
    writeln!(src, "    resultOk[0] = true;")?;
    writeln!(
        src,
        "    writeln(\"  Reduced \", nItems, \" results into 1\");"
    )?;

    Ok(())
}

/// Tree-reduce: logarithmic reduction across locales. Each round halves the
/// number of active participants — much faster than linear reduce for large
/// locale counts.
fn write_gather_tree_reduce(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Tree-reduce: logarithmic pairwise reduction across locales."
    )?;
    writeln!(
        src,
        "    // First, each locale reduces its local slice into one value."
    )?;
    writeln!(src)?;

    // Per-locale local reduction into localeResult[loc.id]
    writeln!(
        src,
        "    var localeResultBuf: [0..#numLocales] [0..#maxItemBytes] uint(8);"
    )?;
    writeln!(src, "    var localeResultLen: [0..#numLocales] c_size_t;")?;
    writeln!(src, "    var localeHasResult: [0..#numLocales] bool;")?;
    writeln!(src)?;

    writeln!(
        src,
        "    coforall loc in Locales with (ref localeResultBuf, ref localeResultLen, ref localeHasResult) do on loc {{"
    )?;
    writeln!(
        src,
        "      const myRange = localeRange(loc.id, numLocales, nItems);"
    )?;
    writeln!(src, "      var first = true;")?;
    writeln!(src, "      var localAccum: [0..#maxItemBytes] uint(8);")?;
    writeln!(src, "      var localAccumLen: c_size_t = 0;")?;
    writeln!(src)?;
    writeln!(src, "      for i in myRange {{")?;
    writeln!(src, "        if !resultOk[i] then continue;")?;
    writeln!(src, "        if first {{")?;
    writeln!(src, "          localAccum = resultData[i];")?;
    writeln!(src, "          localAccumLen = resultSizes[i];")?;
    writeln!(src, "          first = false;")?;
    writeln!(src, "        }} else {{")?;
    writeln!(src, "          var tmpBuf: [0..#maxItemBytes] uint(8);")?;
    writeln!(src, "          var tmpLen: c_size_t = 0;")?;
    writeln!(
        src,
        "          c_reduce(c_ptrTo(localAccum[0]), localAccumLen,"
    )?;
    writeln!(
        src,
        "                   c_ptrTo(resultData[i][0]), resultSizes[i],"
    )?;
    writeln!(
        src,
        "                   c_ptrTo(tmpBuf[0]), c_ptrTo(tmpLen));"
    )?;
    writeln!(src, "          localAccum = tmpBuf;")?;
    writeln!(src, "          localAccumLen = tmpLen;")?;
    writeln!(src, "        }}")?;
    writeln!(src, "      }}")?;
    writeln!(src)?;
    writeln!(src, "      localeResultBuf[loc.id] = localAccum;")?;
    writeln!(src, "      localeResultLen[loc.id] = localAccumLen;")?;
    writeln!(src, "      localeHasResult[loc.id] = !first;")?;
    writeln!(src, "    }}")?;
    writeln!(src)?;

    // Tree reduction across locale-level accumulators
    writeln!(src, "    // Now tree-reduce across locale accumulators")?;
    writeln!(src, "    var active = numLocales;")?;
    writeln!(src, "    while active > 1 {{")?;
    writeln!(src, "      const half = active / 2;")?;
    writeln!(
        src,
        "      coforall i in 0..#half with (ref localeResultBuf, ref localeResultLen, ref localeHasResult) do on Locales[i] {{"
    )?;
    writeln!(src, "        const partner = i + half;")?;
    writeln!(src, "        if localeHasResult[partner] {{")?;
    writeln!(src, "          if !localeHasResult[i] {{")?;
    writeln!(
        src,
        "            localeResultBuf[i] = localeResultBuf[partner];"
    )?;
    writeln!(
        src,
        "            localeResultLen[i] = localeResultLen[partner];"
    )?;
    writeln!(src, "            localeHasResult[i] = true;")?;
    writeln!(src, "          }} else {{")?;
    writeln!(src, "            var tmpBuf: [0..#maxItemBytes] uint(8);")?;
    writeln!(src, "            var tmpLen: c_size_t = 0;")?;
    writeln!(
        src,
        "            c_reduce(c_ptrTo(localeResultBuf[i][0]), localeResultLen[i],"
    )?;
    writeln!(
        src,
        "                     c_ptrTo(localeResultBuf[partner][0]), localeResultLen[partner],"
    )?;
    writeln!(
        src,
        "                     c_ptrTo(tmpBuf[0]), c_ptrTo(tmpLen));"
    )?;
    writeln!(src, "            localeResultBuf[i] = tmpBuf;")?;
    writeln!(src, "            localeResultLen[i] = tmpLen;")?;
    writeln!(src, "          }}")?;
    writeln!(src, "        }}")?;
    writeln!(src, "      }}")?;
    writeln!(src, "      active = (active + 1) / 2;")?;
    writeln!(src, "    }}")?;
    writeln!(src)?;

    // Final result is in localeResultBuf[0]
    writeln!(src, "    // Store tree-reduced result as item 0")?;
    writeln!(src, "    resultData[0] = localeResultBuf[0];")?;
    writeln!(src, "    resultSizes[0] = localeResultLen[0];")?;
    writeln!(src, "    resultOk[0] = localeHasResult[0];")?;
    writeln!(
        src,
        "    writeln(\"  Tree-reduced \", nItems, \" results into 1 (\", numLocales, \" locales)\");"
    )?;

    Ok(())
}

/// Stream gather: results are stored as they complete (same as merge in this
/// buffer-based model). The distinction is semantic — stream mode implies
/// the user may read partial results before all items complete.
fn write_gather_stream(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // Stream: results stored incrementally as each item completes."
    )?;
    writeln!(
        src,
        "    // In buffer mode, this is equivalent to merge — results are in resultData[]."
    )?;
    writeln!(
        src,
        "    // The user's c_store_result will be called as items complete in the next phase."
    )?;
    writeln!(src, "    var nSuccess = + reduce resultOk:int;")?;
    writeln!(
        src,
        "    writeln(\"  Streamed \", nSuccess, \"/\", nItems, \" results\");"
    )?;

    Ok(())
}

/// First gather: return as soon as any item matches the predicate.
/// Uses an atomic flag so all locales stop work once a match is found.
fn write_gather_first(src: &mut String) -> Result<()> {
    writeln!(
        src,
        "    // First: return the first result that matches c_is_match."
    )?;
    writeln!(
        src,
        "    // The distribution phase ran all items; now scan for the first match."
    )?;
    writeln!(src, "    var foundIdx: int = -1;")?;
    writeln!(src, "    for i in 0..#nItems {{")?;
    writeln!(
        src,
        "      if resultOk[i] && c_is_match(c_ptrTo(resultData[i][0]), resultSizes[i]) == 1 {{"
    )?;
    writeln!(src, "        foundIdx = i;")?;
    writeln!(src, "        break;")?;
    writeln!(src, "      }}")?;
    writeln!(src, "    }}")?;
    writeln!(src)?;
    writeln!(src, "    if foundIdx >= 0 {{")?;
    writeln!(src, "      writeln(\"  Found match at index \", foundIdx);")?;
    writeln!(src, "    }} else {{")?;
    writeln!(
        src,
        "      writeln(\"  No match found in \", nItems, \" items\");"
    )?;
    writeln!(src, "    }}")?;

    Ok(())
}

/// Phase 4: Store results back via C-ABI.
fn write_store_phase(src: &mut String, manifest: &Manifest) -> Result<()> {
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(
        src,
        "    // Phase 4: Store results via c_store_result on locale 0"
    )?;
    writeln!(
        src,
        "    // ================================================================"
    )?;
    writeln!(src)?;

    match manifest.workload.gather.as_str() {
        // Reduce strategies produce a single result in slot 0
        "reduce" | "tree-reduce" => {
            writeln!(src, "    on Locales[0] {{")?;
            writeln!(src, "      if resultOk[0] {{")?;
            writeln!(
                src,
                "        const rc = c_store_result(0: c_int, c_ptrTo(resultData[0][0]), resultSizes[0]);"
            )?;
            writeln!(
                src,
                "        if rc != 0 then writeln(\"  WARN: c_store_result(0) returned \", rc);"
            )?;
            writeln!(src, "      }}")?;
            writeln!(src, "    }}")?;
        }

        // First produces a single result at foundIdx
        "first" => {
            writeln!(src, "    on Locales[0] {{")?;
            writeln!(src, "      if foundIdx >= 0 && resultOk[foundIdx] {{")?;
            writeln!(
                src,
                "        const rc = c_store_result(0: c_int, c_ptrTo(resultData[foundIdx][0]), resultSizes[foundIdx]);"
            )?;
            writeln!(
                src,
                "        if rc != 0 then writeln(\"  WARN: c_store_result(0) returned \", rc);"
            )?;
            writeln!(src, "      }}")?;
            writeln!(src, "    }}")?;
        }

        // Merge and stream store all results
        _ => {
            writeln!(src, "    on Locales[0] {{")?;
            writeln!(src, "      for i in 0..#nItems {{")?;
            writeln!(src, "        if !resultOk[i] then continue;")?;
            writeln!(
                src,
                "        const rc = c_store_result(i: c_int, c_ptrTo(resultData[i][0]), resultSizes[i]);"
            )?;
            writeln!(
                src,
                "        if rc != 0 then writeln(\"  WARN: c_store_result(\", i, \") returned \", rc);"
            )?;
            writeln!(src, "      }}")?;
            writeln!(src, "    }}")?;
        }
    }

    writeln!(src)?;
    Ok(())
}

/// Close the module block.
fn write_module_close(src: &mut String) -> Result<()> {
    writeln!(src, "}}")?;
    Ok(())
}