wasm_gen.rs

// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2026 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk>
//
// WASM generation module for affinescriptiser — Produces WASM compilation configuration files
// and entry point declarations from the manifest's WASM settings and the generated AffineScript
// module. The output includes a build configuration (target, optimisation flags, size constraints)
// and an entry point source file that wires up the safe resource wrappers for WASM export.

use crate::abi::{Violation, ViolationType, WasmConfig, WasmTarget};
use crate::codegen::affine_gen::AffineModule;
use std::fmt::Write;

/// Generated WASM build artefacts: a build configuration file and a WASM entry point module.
#[derive(Debug, Clone)]
pub struct WasmOutput {
    /// WASM build configuration (target, flags, constraints) as TOML content.
    pub build_config: String,
    /// WASM entry point source code that imports and re-exports the safe wrappers.
    pub entry_point: String,
    /// Any WASM-specific violations (e.g. size constraint failures).
    pub violations: Vec<Violation>,
}

/// Generate WASM build configuration and entry point from the AffineScript module
/// and WASM settings declared in the manifest.
///
/// The build configuration is emitted as a TOML file suitable for consumption by
/// WASM build tools (wasm-pack, cargo-wasi, or custom AffineScript compiler).
/// The entry point re-exports each resource's safe allocation/deallocation functions
/// as WASM-exported symbols.
pub fn generate_wasm_output(affine_module: &AffineModule, wasm_config: &WasmConfig) -> WasmOutput {
    let build_config = generate_build_config(wasm_config, &affine_module.module_name);
    let entry_point = generate_entry_point(affine_module);
    let violations = Vec::new(); // Size checks happen at build time, not codegen time.

    WasmOutput {
        build_config,
        entry_point,
        violations,
    }
}

/// Generate a TOML build configuration for the WASM target.
fn generate_build_config(config: &WasmConfig, module_name: &str) -> String {
    let mut out = String::new();

    writeln!(
        out,
        "# WASM build configuration generated by affinescriptiser"
    )
    .expect("TODO: handle error");
    writeln!(out, "# SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");
    writeln!(out, "[build]").expect("TODO: handle error");
    writeln!(out, "module = \"{}\"", module_name).expect("TODO: handle error");
    writeln!(out, "target = \"{}\"", config.target).expect("TODO: handle error");
    writeln!(out, "optimize = {}", config.optimize).expect("TODO: handle error");

    if let Some(limit) = config.size_limit_kb {
        writeln!(out, "size-limit-kb = {}", limit).expect("TODO: handle error");
    }

    writeln!(out).expect("TODO: handle error");
    writeln!(out, "[build.flags]").expect("TODO: handle error");

    // Target-specific flags.
    match config.target {
        WasmTarget::Wasm32Unknown => {
            writeln!(out, "# Browser-compatible WASM (no WASI)").expect("TODO: handle error");
            writeln!(out, "no-wasi = true").expect("TODO: handle error");
        }
        WasmTarget::Wasm32Wasi => {
            writeln!(out, "# WASI-enabled WASM (filesystem, env access)")
                .expect("TODO: handle error");
            writeln!(out, "wasi = true").expect("TODO: handle error");
        }
    }

    if config.optimize {
        writeln!(out, "wasm-opt = true").expect("TODO: handle error");
        writeln!(out, "strip-debug = true").expect("TODO: handle error");
        writeln!(out, "lto = true").expect("TODO: handle error");
    }

    out
}

/// Generate the WASM entry point module that imports and re-exports safe resource wrappers.
fn generate_entry_point(affine_module: &AffineModule) -> String {
    let mut out = String::new();

    writeln!(out, "// WASM entry point generated by affinescriptiser").expect("TODO: handle error");
    writeln!(out, "// SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");
    writeln!(
        out,
        "import {{ * }} from \"{}\";",
        affine_module.module_name
    )
    .expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    // Extract and re-export all safe_create_* and safe_destroy_* functions.
    for line in affine_module.source.lines() {
        let trimmed = line.trim();
        if trimmed.starts_with("fn safe_create_") || trimmed.starts_with("fn safe_destroy_") {
            // Extract the function name.
            if let Some(name) = extract_function_name(trimmed) {
                writeln!(out, "@wasm_export").expect("TODO: handle error");
                writeln!(
                    out,
                    "pub fn {}(/* forwarded params */) -> /* forwarded return */ {{",
                    name
                )
                .expect("TODO: handle error");
                writeln!(
                    out,
                    "  {}.{}(/* forwarded args */)",
                    affine_module.module_name, name
                )
                .expect("TODO: handle error");
                writeln!(out, "}}").expect("TODO: handle error");
                writeln!(out).expect("TODO: handle error");
            }
        }
    }

    // Export a health-check function for WASM module validation.
    writeln!(out, "@wasm_export").expect("TODO: handle error");
    writeln!(out, "pub fn affinescriptiser_health() -> i32 {{").expect("TODO: handle error");
    writeln!(out, "  // Returns 0 if all resource invariants hold").expect("TODO: handle error");
    writeln!(out, "  0").expect("TODO: handle error");
    writeln!(out, "}}").expect("TODO: handle error");

    out
}

/// Extract a function name from a line like `fn safe_create_gpu_buffer(params: ...) -> ... {`.
fn extract_function_name(line: &str) -> Option<String> {
    let after_fn = line.strip_prefix("fn ")?;
    let paren_pos = after_fn.find('(')?;
    Some(after_fn[..paren_pos].to_string())
}

/// Check whether a WASM binary (represented by its size) exceeds the configured limit.
/// Returns a SizeExceeded violation if the limit is exceeded, or None if within bounds.
pub fn check_size_limit(actual_bytes: u64, config: &WasmConfig) -> Option<Violation> {
    if let Some(limit_kb) = config.size_limit_kb {
        let actual_kb = actual_bytes / 1024;
        if actual_kb > limit_kb {
            return Some(Violation::new(ViolationType::SizeExceeded {
                actual_kb,
                limit_kb,
            }));
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::abi::{AffineResource, AffinityKind};
    use crate::codegen::affine_gen::generate_affine_module;

    fn sample_module() -> AffineModule {
        let resources = vec![
            AffineResource {
                name: "GpuBuffer".into(),
                allocator: "create_buffer".into(),
                deallocator: "release_buffer".into(),
                affinity: AffinityKind::Affine,
            },
            AffineResource {
                name: "SessionToken".into(),
                allocator: "create_session".into(),
                deallocator: "revoke_session".into(),
                affinity: AffinityKind::Linear,
            },
        ];
        generate_affine_module("wasm_safe", &resources, &[])
    }

    #[test]
    fn test_generate_build_config() {
        let config = WasmConfig {
            target: WasmTarget::Wasm32Unknown,
            optimize: true,
            size_limit_kb: Some(512),
        };
        let out = generate_build_config(&config, "test");
        assert!(out.contains("target = \"wasm32-unknown-unknown\""));
        assert!(out.contains("optimize = true"));
        assert!(out.contains("size-limit-kb = 512"));
        assert!(out.contains("wasm-opt = true"));
    }

    #[test]
    fn test_generate_build_config_wasi() {
        let config = WasmConfig {
            target: WasmTarget::Wasm32Wasi,
            optimize: false,
            size_limit_kb: None,
        };
        let out = generate_build_config(&config, "test");
        assert!(out.contains("target = \"wasm32-wasi\""));
        assert!(out.contains("wasi = true"));
        assert!(!out.contains("wasm-opt"));
    }

    #[test]
    fn test_generate_entry_point_exports() {
        let module = sample_module();
        let config = WasmConfig::default();
        let output = generate_wasm_output(&module, &config);
        assert!(output.entry_point.contains("@wasm_export"));
        assert!(output.entry_point.contains("affinescriptiser_health"));
    }

    #[test]
    fn test_check_size_limit_within_bounds() {
        let config = WasmConfig {
            target: WasmTarget::default(),
            optimize: false,
            size_limit_kb: Some(512),
        };
        // 256KB = 262144 bytes — within the 512KB limit.
        assert!(check_size_limit(262144, &config).is_none());
    }

    #[test]
    fn test_check_size_limit_exceeded() {
        let config = WasmConfig {
            target: WasmTarget::default(),
            optimize: false,
            size_limit_kb: Some(512),
        };
        // 1MB = 1048576 bytes — exceeds the 512KB limit.
        let violation = check_size_limit(1048576, &config);
        assert!(violation.is_some());
        assert!(
            violation
                .expect("TODO: handle error")
                .violation
                .to_string()
                .contains("SIZE-EXCEEDED")
        );
    }

    #[test]
    fn test_check_size_limit_no_limit() {
        let config = WasmConfig {
            target: WasmTarget::default(),
            optimize: false,
            size_limit_kb: None,
        };
        // No limit configured — any size is fine.
        assert!(check_size_limit(999999999, &config).is_none());
    }

    #[test]
    fn test_extract_function_name() {
        assert_eq!(
            extract_function_name(
                "fn safe_create_gpu_buffer(params: AllocParams) -> GpuBuffer<Affine> {"
            ),
            Some("safe_create_gpu_buffer".into())
        );
        assert_eq!(extract_function_name("let x = 1;"), None);
    }
}