vyre-self-substrate 0.4.1

//! Vyre IR Region tree as a string diagram (#53 self-consumer).
//!
//! Closes the recursion thesis for #53 — string-diagram tensor
//! compilation ships to user dialects (quantum circuits, monoidal
//! tensor networks, ZX-calculus) AND IS the substrate semantics for
//! vyre's IR.
//!
//! # The legendary self-use
//!
//! Selinger's (2010) string diagrams are the visual + algebraic
//! language of monoidal categories. Each diagram is built from:
//!
//! - **Boxes**: morphisms (functions f: A → B). In vyre = each Region.
//! - **Wires**: types (objects A in the category). In vyre = buffer
//!   bindings between Regions.
//! - **Composition** ∘: stack boxes vertically (sequential
//!   dependence). In vyre = nested Regions in entry order.
//! - **Tensor product** ⊗: place boxes side-by-side (parallel
//!   independence). In vyre = sibling Regions sharing no buffers.
//!
//! Vyre's Region tree IS a string diagram in
//! `Cat(GPU buffers, Programs)`. Making this explicit means every
//! optimizer rewrite (region_inline, fusion, fission) is a
//! string-diagram rewrite — the equational laws of monoidal
//! categories give us free correctness proofs.
//!
//! # Concrete payoffs
//!
//! 1. **Coherence theorems for free**: associativity of `∘` and `⊗`,
//!    naturality of swap, are baked into the diagram model. Today
//!    these are checked by hand in each pass.
//! 2. **Adjoint pairs as duality**: backward-pass synthesis (gradient
//!    computation) is the dagger-functor in compact closed
//!    categories. Once the IR is a string diagram, `vyre-frontend-c` can
//!    derive backward-pass kernels for free.
//! 3. **Equational rewriting**: the ZX calculus has 7 rewrite rules
//!    that are complete for monoidal-category equivalence. Vyre's
//!    optimizer reduces from ~30 hand-curated passes to 7
//!    algebraic rules + a confluent rewriting strategy.
//!
//! # Algorithm
//!
//! `monoidal_compose(f, g)` is sequential composition `g ∘ f` —
//! exactly the matrix-product semantics over the buffer-passing
//! contract between two Regions. For 0.6 we ship the per-arrow
//! composition step. The full ZX-calculus rewrite engine ships in
//! 1.0.

use vyre_primitives::graph::string_diagram::monoidal_compose_cpu_into;

/// Reusable buffers for string-diagram IR rewrites.
#[derive(Debug, Default)]
pub struct StringDiagramRewriteScratch {
    gf: Vec<f64>,
    h_after_gf: Vec<f64>,
    hg: Vec<f64>,
    hg_after_f: Vec<f64>,
}

impl StringDiagramRewriteScratch {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

/// Sequential composition of two IR-arrow morphisms. `f` has shape
/// `a × b`, `g` has shape `b × c`. Returns `g ∘ f` with shape
/// `a × c`.
///
/// In vyre IR terms: `f` describes how Region F transforms its
/// `a`-dimensional input buffer into a `b`-dimensional intermediate;
/// `g` describes how Region G transforms the intermediate into the
/// `c`-dimensional output. The composed arrow describes the fused
/// F+G transformation in one step.
///
/// # Panics
///
/// Panics on size mismatches.
#[must_use]
pub fn compose_ir_arrows(f: &[f64], g: &[f64], a: u32, b: u32, c: u32) -> Vec<f64> {
    let mut out = Vec::new();
    compose_ir_arrows_into(f, g, a, b, c, &mut out);
    out
}

/// Sequential composition using caller-owned output storage.
pub fn compose_ir_arrows_into(f: &[f64], g: &[f64], a: u32, b: u32, c: u32, out: &mut Vec<f64>) {
    use crate::observability::{bump, string_diagram_ir_rewrite_calls};
    bump(&string_diagram_ir_rewrite_calls);
    monoidal_compose_cpu_into(f, g, a, b, c, out);
}

/// Identity arrow on dimension `n`. Composes with any arrow as the
/// identity — `id ∘ f = f` and `f ∘ id = f`.
#[must_use]
pub fn identity_arrow(n: u32) -> Vec<f64> {
    let mut out = Vec::new();
    identity_arrow_into(n, &mut out);
    out
}

/// Build an identity arrow using caller-owned output storage.
pub fn identity_arrow_into(n: u32, out: &mut Vec<f64>) {
    let n_us = n as usize;
    out.clear();
    out.resize(n_us * n_us, 0.0);
    for i in 0..n_us {
        out[i * n_us + i] = 1.0;
    }
}

/// Test that composition is associative: `(h ∘ g) ∘ f == h ∘ (g ∘ f)`.
/// Returns true when the two associativities agree to numerical
/// precision. Foundational coherence law for monoidal categories.
#[must_use]
pub fn composition_associates(
    f: &[f64],
    g: &[f64],
    h: &[f64],
    a: u32,
    b: u32,
    c: u32,
    d: u32,
) -> bool {
    let mut scratch = StringDiagramRewriteScratch::new();
    composition_associates_with_scratch(f, g, h, a, b, c, d, &mut scratch)
}

/// Associativity check using caller-owned scratch buffers.
#[must_use]
#[allow(clippy::too_many_arguments)]
pub fn composition_associates_with_scratch(
    f: &[f64],
    g: &[f64],
    h: &[f64],
    a: u32,
    b: u32,
    c: u32,
    d: u32,
    scratch: &mut StringDiagramRewriteScratch,
) -> bool {
    compose_ir_arrows_into(f, g, a, b, c, &mut scratch.gf);
    compose_ir_arrows_into(&scratch.gf, h, a, c, d, &mut scratch.h_after_gf);
    compose_ir_arrows_into(g, h, b, c, d, &mut scratch.hg);
    compose_ir_arrows_into(f, &scratch.hg, a, b, d, &mut scratch.hg_after_f);
    let tol = 1e-9_f64;
    scratch
        .h_after_gf
        .iter()
        .zip(scratch.hg_after_f.iter())
        .all(|(a, b)| (a - b).abs() < tol * (1.0 + a.abs() + b.abs()))
}

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

    fn approx_eq_vec(a: &[f64], b: &[f64]) -> bool {
        if a.len() != b.len() {
            return false;
        }
        a.iter()
            .zip(b.iter())
            .all(|(x, y)| (x - y).abs() < 1e-9 * (1.0 + x.abs() + y.abs()))
    }

    #[test]
    fn identity_left_unit() {
        // id ∘ f = f
        let f = vec![1.0, 2.0, 3.0, 4.0]; // 2x2
        let id = identity_arrow(2);
        let composed = compose_ir_arrows(&f, &id, 2, 2, 2);
        assert!(approx_eq_vec(&composed, &f));
    }

    #[test]
    fn identity_right_unit() {
        // f ∘ id = f
        let f = vec![1.0, 2.0, 3.0, 4.0];
        let id = identity_arrow(2);
        let composed = compose_ir_arrows(&id, &f, 2, 2, 2);
        assert!(approx_eq_vec(&composed, &f));
    }

    #[test]
    fn composition_associativity_holds() {
        // (h ∘ g) ∘ f = h ∘ (g ∘ f) for arbitrary 2x2 matrices.
        let f = vec![1.0, 0.5, -0.25, 0.5];
        let g = vec![0.5, 0.5, 0.5, -0.5];
        let h = vec![1.0, 0.0, 0.0, 1.0];
        assert!(composition_associates(&f, &g, &h, 2, 2, 2, 2));
    }

    #[test]
    fn rectangular_composition_dimensions() {
        // f: 2x3, g: 3x4 → composed: 2x4.
        let f = vec![1.0; 6];
        let g = vec![1.0; 12];
        let composed = compose_ir_arrows(&f, &g, 2, 3, 4);
        assert_eq!(composed.len(), 8);
    }

    #[test]
    fn identity_arrow_size_matches() {
        let id = identity_arrow(3);
        assert_eq!(id.len(), 9);
        // Diagonal = 1.0, off-diagonal = 0.0.
        assert_eq!(id[0], 1.0);
        assert_eq!(id[4], 1.0);
        assert_eq!(id[8], 1.0);
        assert_eq!(id[1], 0.0);
        assert_eq!(id[3], 0.0);
    }

    #[test]
    fn reusable_outputs_preserve_associativity() {
        let f = vec![1.0, 0.5, -0.25, 0.5];
        let g = vec![0.5, 0.5, 0.5, -0.5];
        let h = vec![1.0, 0.0, 0.0, 1.0];
        let mut scratch = StringDiagramRewriteScratch::new();
        assert!(composition_associates_with_scratch(
            &f,
            &g,
            &h,
            2,
            2,
            2,
            2,
            &mut scratch
        ));
    }
}