affine_gen.rs

// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2026 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk>
//
// AffineScript type annotation generator — Takes parsed resource sites and manifest resource
// declarations and produces AffineScript wrapper modules. Each tracked resource gets a wrapper
// type that enforces its affinity discipline (affine = at-most-once, linear = exactly-once)
// at compile time. The generated code also includes violation checks that can be statically
// verified or emitted as runtime assertions for the WASM target.

use crate::abi::{AffineResource, AffinityKind, Violation, ViolationType};
use crate::codegen::parser::{ResourceSite, SiteKind};
use std::collections::HashMap;
use std::fmt::Write;

/// Generated AffineScript module containing type wrappers for a set of tracked resources.
#[derive(Debug, Clone)]
pub struct AffineModule {
    /// Module name (derived from the project name).
    pub module_name: String,
    /// The generated AffineScript source code.
    pub source: String,
    /// Any violations detected during static analysis of the resource sites.
    pub violations: Vec<Violation>,
}

/// Generate an AffineScript wrapper module from the manifest's resource declarations and
/// the parsed resource sites. The module contains:
/// 1. Type declarations for each resource with affine/linear wrappers.
/// 2. Safe allocation functions that return wrapped handles.
/// 3. Safe deallocation functions that consume wrapped handles.
/// 4. Static violation reports for detected misuse patterns.
pub fn generate_affine_module(
    module_name: &str,
    resources: &[AffineResource],
    sites: &[ResourceSite],
) -> AffineModule {
    let mut source = String::new();
    let mut violations = Vec::new();

    // Module header.
    writeln!(
        source,
        "// AffineScript module generated by affinescriptiser"
    )
    .expect("TODO: handle error");
    writeln!(source, "// SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(source, "// Module: {}", module_name).expect("TODO: handle error");
    writeln!(source).expect("TODO: handle error");
    writeln!(source, "module {} {{", module_name).expect("TODO: handle error");
    writeln!(source).expect("TODO: handle error");

    // Generate wrapper types and safe API for each resource.
    for resource in resources {
        generate_resource_wrapper(&mut source, resource);
        writeln!(source).expect("TODO: handle error");
    }

    // Perform static violation analysis on the parsed sites.
    let detected = analyse_sites(resources, sites);
    for v in &detected {
        // Emit violation as a compile-time error annotation.
        writeln!(source, "  // @error: {}", v.violation).expect("TODO: handle error");
    }
    violations.extend(detected);

    writeln!(source, "}}").expect("TODO: handle error");

    AffineModule {
        module_name: module_name.to_string(),
        source,
        violations,
    }
}

/// Generate the AffineScript type wrapper for a single resource. Produces:
/// - An opaque wrapper type with the appropriate affinity marker.
/// - A safe `create_*` function that wraps the raw allocator.
/// - A safe `destroy_*` function that consumes the wrapper.
fn generate_resource_wrapper(out: &mut String, resource: &AffineResource) {
    let type_name = &resource.name;
    let affinity_marker = match resource.affinity {
        AffinityKind::Affine => "Affine",
        AffinityKind::Linear => "Linear",
    };

    // Type declaration with affinity annotation.
    writeln!(
        out,
        "  // {} resource — {} discipline",
        type_name, resource.affinity
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "  // Allocated by: {}  Deallocated by: {}",
        resource.allocator, resource.deallocator
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "  type {}<{}> = opaque handle;",
        type_name, affinity_marker
    )
    .expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    // Safe allocator — returns a wrapped handle.
    writeln!(
        out,
        "  // Safe wrapper for {} — returns an affine-tracked handle",
        resource.allocator
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "  fn safe_create_{name}(params: AllocParams) -> {type_name}<{marker}> {{",
        name = to_snake_case(type_name),
        type_name = type_name,
        marker = affinity_marker
    )
    .expect("TODO: handle error");
    writeln!(out, "    let raw = @ffi::{}(params);", resource.allocator)
        .expect("TODO: handle error");
    writeln!(out, "    {}<{}>::wrap(raw)", type_name, affinity_marker).expect("TODO: handle error");
    writeln!(out, "  }}").expect("TODO: handle error");
    writeln!(out).expect("TODO: handle error");

    // Safe deallocator — consumes the wrapped handle.
    writeln!(
        out,
        "  // Safe wrapper for {} — consumes the handle (cannot be used again)",
        resource.deallocator
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "  fn safe_destroy_{name}(handle: {type_name}<{marker}>) -> () {{",
        name = to_snake_case(type_name),
        type_name = type_name,
        marker = affinity_marker
    )
    .expect("TODO: handle error");
    writeln!(
        out,
        "    let raw = {}<{}>::unwrap(handle);",
        type_name, affinity_marker
    )
    .expect("TODO: handle error");
    writeln!(out, "    @ffi::{}(raw);", resource.deallocator).expect("TODO: handle error");
    writeln!(out, "  }}").expect("TODO: handle error");

    // Linear resources get a must-consume annotation.
    if resource.affinity == AffinityKind::Linear {
        writeln!(out).expect("TODO: handle error");
        writeln!(
            out,
            "  // LINEAR INVARIANT: {} must be consumed exactly once.",
            type_name
        )
        .expect("TODO: handle error");
        writeln!(
            out,
            "  #[must_consume(\"{} is linear — must be explicitly deallocated\")]",
            type_name
        )
        .expect("TODO: handle error");
        writeln!(out, "  impl MustConsume for {}<Linear> {{}}", type_name)
            .expect("TODO: handle error");
    }
}

/// Analyse parsed resource sites to detect affine/linear type violations:
/// - Double-free: two deallocation sites for the same resource within the same file scope.
/// - Leaked linear resources: allocation without a corresponding deallocation.
///
/// This is a conservative heuristic analysis (not a full control-flow analysis).
/// It groups sites by file and resource, then checks for imbalanced alloc/dealloc counts
/// and duplicate deallocation sequences.
fn analyse_sites(resources: &[AffineResource], sites: &[ResourceSite]) -> Vec<Violation> {
    let mut violations = Vec::new();

    // Build a lookup map: resource name -> AffineResource.
    let resource_map: HashMap<&str, &AffineResource> =
        resources.iter().map(|r| (r.name.as_str(), r)).collect();

    // Group sites by (file, resource_name).
    let mut grouped: HashMap<(&str, &str), Vec<&ResourceSite>> = HashMap::new();
    for site in sites {
        grouped
            .entry((site.file.as_str(), site.resource_name.as_str()))
            .or_default()
            .push(site);
    }

    for ((file, res_name), file_sites) in &grouped {
        let resource = match resource_map.get(res_name) {
            Some(r) => r,
            None => continue,
        };

        let allocs: Vec<_> = file_sites
            .iter()
            .filter(|s| s.kind == SiteKind::Allocation)
            .collect();
        let deallocs: Vec<_> = file_sites
            .iter()
            .filter(|s| s.kind == SiteKind::Deallocation)
            .collect();

        // Check for double-free: more than one deallocation site.
        if deallocs.len() > 1 {
            for i in 1..deallocs.len() {
                violations.push(Violation::new(ViolationType::DoubleFree {
                    resource: res_name.to_string(),
                    first_free_at: format!("{}:{}", file, deallocs[0].line),
                    second_free_at: format!("{}:{}", file, deallocs[i].line),
                }));
            }
        }

        // Check for leaked linear resources: allocation with no deallocation.
        if resource.affinity == AffinityKind::Linear && !allocs.is_empty() && deallocs.is_empty() {
            for alloc in &allocs {
                violations.push(Violation::new(ViolationType::Leaked {
                    resource: res_name.to_string(),
                    allocated_at: format!("{}:{}", file, alloc.line),
                }));
            }
        }
    }

    violations
}

/// Convert a PascalCase or camelCase name to snake_case for generated function names.
fn to_snake_case(s: &str) -> String {
    let mut result = String::with_capacity(s.len() + 4);
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() {
            if i > 0 {
                result.push('_');
            }
            result.push(c.to_lowercase().next().expect("TODO: handle error"));
        } else {
            result.push(c);
        }
    }
    result
}

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

    fn gpu_resource() -> AffineResource {
        AffineResource {
            name: "GpuBuffer".into(),
            allocator: "create_buffer".into(),
            deallocator: "release_buffer".into(),
            affinity: AffinityKind::Affine,
        }
    }

    fn session_resource() -> AffineResource {
        AffineResource {
            name: "SessionToken".into(),
            allocator: "create_session".into(),
            deallocator: "revoke_session".into(),
            affinity: AffinityKind::Linear,
        }
    }

    #[test]
    fn test_to_snake_case() {
        assert_eq!(to_snake_case("GpuBuffer"), "gpu_buffer");
        assert_eq!(to_snake_case("SessionToken"), "session_token");
        assert_eq!(to_snake_case("X"), "x");
    }

    #[test]
    fn test_generate_affine_module_produces_output() {
        let resources = vec![gpu_resource()];
        let sites = vec![];
        let module = generate_affine_module("test_module", &resources, &sites);
        assert!(module.source.contains("module test_module"));
        assert!(module.source.contains("type GpuBuffer<Affine>"));
        assert!(module.source.contains("safe_create_gpu_buffer"));
        assert!(module.source.contains("safe_destroy_gpu_buffer"));
    }

    #[test]
    fn test_linear_resource_gets_must_consume() {
        let resources = vec![session_resource()];
        let module = generate_affine_module("test", &resources, &[]);
        assert!(module.source.contains("#[must_consume"));
        assert!(module.source.contains("Linear"));
    }

    #[test]
    fn test_double_free_detection() {
        let resources = vec![gpu_resource()];
        let sites = vec![
            ResourceSite {
                resource_name: "GpuBuffer".into(),
                function: "create_buffer".into(),
                kind: SiteKind::Allocation,
                file: "test.rs".into(),
                line: 1,
                context: "create_buffer(64)".into(),
            },
            ResourceSite {
                resource_name: "GpuBuffer".into(),
                function: "release_buffer".into(),
                kind: SiteKind::Deallocation,
                file: "test.rs".into(),
                line: 3,
                context: "release_buffer(buf)".into(),
            },
            ResourceSite {
                resource_name: "GpuBuffer".into(),
                function: "release_buffer".into(),
                kind: SiteKind::Deallocation,
                file: "test.rs".into(),
                line: 5,
                context: "release_buffer(buf)".into(),
            },
        ];
        let module = generate_affine_module("test", &resources, &sites);
        assert_eq!(module.violations.len(), 1);
        assert!(
            module.violations[0]
                .violation
                .to_string()
                .contains("DOUBLE-FREE")
        );
    }

    #[test]
    fn test_linear_leak_detection() {
        let resources = vec![session_resource()];
        let sites = vec![ResourceSite {
            resource_name: "SessionToken".into(),
            function: "create_session".into(),
            kind: SiteKind::Allocation,
            file: "test.rs".into(),
            line: 1,
            context: "create_session()".into(),
        }];
        // No deallocation site — should detect a leak for the linear resource.
        let module = generate_affine_module("test", &resources, &sites);
        assert_eq!(module.violations.len(), 1);
        assert!(module.violations[0].violation.to_string().contains("LEAK"));
    }

    #[test]
    fn test_affine_resource_no_leak_violation() {
        // Affine resources are allowed to be dropped without deallocation.
        let resources = vec![gpu_resource()];
        let sites = vec![ResourceSite {
            resource_name: "GpuBuffer".into(),
            function: "create_buffer".into(),
            kind: SiteKind::Allocation,
            file: "test.rs".into(),
            line: 1,
            context: "create_buffer(64)".into(),
        }];
        let module = generate_affine_module("test", &resources, &sites);
        assert!(
            module.violations.is_empty(),
            "Affine resources should not report leaks"
        );
    }
}