gridoxide-sparse 0.1.0

//! Thin internal wrapper around the vendored SuiteSparse KLU solver
//! (`vendor/suitesparse/`, compiled by `build.rs`), mirroring `sparse.rs`'s
//! shape exactly so it's a drop-in third backend. See
//! `docs/src/solvers/backends.md` for why this exists and
//! `vendor/suitesparse/PROVENANCE.md` for what was vendored and its licensing.
//!
//! All raw-pointer FFI handling is confined to this one file — nothing raw
//! leaks into `solver.rs`/`network.rs`, the same isolation principle
//! `sparse.rs` already applies to `faer` types.
//!
//! KLU represents complex values as interleaved `[re, im, re, im, ...]`
//! `f64` pairs, not a native complex type — `klu_z_factor`/`klu_z_solve`
//! share the exact same `*mut f64` signature as their real counterparts,
//! just interpreting the buffer differently. `pack_complex_values`/the
//! `Vec<f64>` RHS buffers in `klu_solve_complex` do this conversion.

use num_complex::Complex;

#[allow(
    non_camel_case_types,
    non_snake_case,
    non_upper_case_globals,
    dead_code,
    clippy::all
)]
mod bindings {
    include!(concat!(env!("OUT_DIR"), "/klu_bindings.rs"));
}
use bindings::{
    klu_analyze, klu_common, klu_defaults, klu_factor, klu_free_numeric, klu_free_symbolic,
    klu_numeric, klu_refactor, klu_solve, klu_symbolic, klu_z_factor, klu_z_solve,
};

/// Re-exported for `klu_native::ffi_oracle`'s differential tests: `klu_scale`
/// needs the real `klu_common` struct layout (bindgen-generated here from
/// `klu.h`, not hand-declared) — reusing it is simpler and safer than
/// hand-writing a second FFI declaration that must independently match the
/// same ABI, the approach `ffi_oracle.rs` otherwise uses for `btf_order`/
/// `amd_order` (neither of which takes a `klu_common`).
#[cfg(test)]
pub(crate) use bindings::{
    klu_common as klu_common_for_oracle, klu_defaults as klu_defaults_for_oracle,
    klu_scale as klu_scale_for_oracle,
};

const KLU_OK: i32 = 0;

/// Builds a KLU-ready CSC structure (column pointers + sorted row indices)
/// from a set of `(row, col)` index pairs, merging duplicates — the same
/// accumulation semantics used everywhere else in gridoxide's sparse code
/// (e.g. `network::YBus`, parallel branches between the same two buses).
/// Returns `(col_ptr, row_idx, groups)`, where `groups[k]` lists the
/// original `entries` indices contributing to the `k`-th CSC position.
fn build_csc_structure(
    n: usize,
    pairs: &[(usize, usize)],
) -> (Vec<i32>, Vec<i32>, Vec<Vec<usize>>) {
    let mut order: Vec<usize> = (0..pairs.len()).collect();
    order.sort_by_key(|&i| (pairs[i].1, pairs[i].0));

    let mut col_ptr = vec![0i32; n + 1];
    let mut row_idx: Vec<i32> = Vec::new();
    let mut groups: Vec<Vec<usize>> = Vec::new();

    let mut idx = 0;
    while idx < order.len() {
        let (row, col) = pairs[order[idx]];
        let mut group = vec![order[idx]];
        let mut k = idx + 1;
        while k < order.len() && pairs[order[k]] == (row, col) {
            group.push(order[k]);
            k += 1;
        }
        row_idx.push(row as i32);
        groups.push(group);
        col_ptr[col + 1] += 1;
        idx = k;
    }
    for c in 0..n {
        col_ptr[c + 1] += col_ptr[c];
    }
    (col_ptr, row_idx, groups)
}

fn pack_real_values(entries: &[(usize, usize, f64)], groups: &[Vec<usize>]) -> Vec<f64> {
    groups
        .iter()
        .map(|g| g.iter().map(|&i| entries[i].2).sum())
        .collect()
}

/// `groups` flattened into two arrays: CSC position `k` sums the inputs
/// `src[ptr[k]..ptr[k + 1]]`.
///
/// What a repeated solve packs through, every iteration. A `Vec` per
/// position cost a pointer chase per nonzero, and the packed values a fresh
/// allocation per call; together they were a tenth of a warm Newton solve on
/// case9241pegase, where a Newton Jacobian has no duplicate position at all.
struct FlatGroups {
    ptr: Vec<usize>,
    src: Vec<usize>,
}

impl FlatGroups {
    fn new(groups: &[Vec<usize>]) -> Self {
        let mut ptr = Vec::with_capacity(groups.len() + 1);
        let mut src = Vec::with_capacity(groups.iter().map(Vec::len).sum());
        ptr.push(0);
        for g in groups {
            src.extend_from_slice(g);
            ptr.push(src.len());
        }
        Self { ptr, src }
    }

    /// The CSC position of every input, if each position has exactly one —
    /// `order[k]` the input stored at position `k`.
    fn one_to_one(&self) -> Option<Vec<usize>> {
        (self.src.len() + 1 == self.ptr.len()).then(|| self.src.clone())
    }

    /// Positions that take input `k` as-is.
    fn identity(n: usize) -> Self {
        Self {
            ptr: (0..=n).collect(),
            src: (0..n).collect(),
        }
    }

    /// Packs `value(i)` for every input `i` into `out`, in CSC order.
    fn pack_into(&self, value: impl Fn(usize) -> f64, out: &mut Vec<f64>) {
        out.clear();
        out.extend(self.ptr.windows(2).map(|w| match &self.src[w[0]..w[1]] {
            [only] => value(*only),
            several => several.iter().map(|&i| value(i)).sum(),
        }));
    }
}

fn pack_complex_values(
    entries: &[(usize, usize, Complex<f64>)],
    groups: &[Vec<usize>],
) -> Vec<f64> {
    let mut out = Vec::with_capacity(groups.len() * 2);
    for g in groups {
        let sum: Complex<f64> = g.iter().map(|&i| entries[i].2).sum();
        out.push(sum.re);
        out.push(sum.im);
    }
    out
}

fn new_common() -> klu_common {
    let mut common: klu_common = unsafe { std::mem::zeroed() };
    unsafe { klu_defaults(&mut common) };
    common
}

/// A real sparse system whose sparsity *pattern* is fixed across repeated
/// solves — mirrors `sparse::RealSparseSystem`'s role but backed by KLU
/// instead of `faer`. Caches the symbolic analysis (`klu_analyze`) and a
/// numeric factorization object across calls, using `klu_refactor` for
/// cheap numeric-only re-factorization — the same pattern already validated
/// this session for `RealSparseSystem`/`BlockLu::refactor`.
pub struct KluRealSystem {
    n: usize,
    col_ptr: Vec<i32>,
    row_idx: Vec<i32>,
    groups: FlatGroups,
    /// The packed values of the last call, kept so the next reuses the
    /// allocation.
    packed: Vec<f64>,
    /// Set by [`adopt_storage_order`](Self::adopt_storage_order): values
    /// arrive in CSC order and are handed to KLU as they are.
    in_order: bool,
    common: klu_common,
    symbolic: *mut klu_symbolic,
    numeric: *mut klu_numeric,
}

impl KluRealSystem {
    /// Builds the symbolic factorization from an initial sparsity pattern.
    /// Subsequent calls to `factor_and_solve` must supply entries with the
    /// exact same `(row, col)` pairs in the exact same order (values may
    /// differ) — the cached CSC position mapping assumes positional
    /// correspondence, not just set equality.
    pub fn new(n: usize, entries: &[(usize, usize, f64)]) -> Option<Self> {
        let pairs: Vec<(usize, usize)> = entries.iter().map(|&(r, c, _)| (r, c)).collect();
        let (mut col_ptr, mut row_idx, groups) = build_csc_structure(n, &pairs);
        let mut values = pack_real_values(entries, &groups);

        let mut common = new_common();
        let symbolic = unsafe {
            klu_analyze(
                n as i32,
                col_ptr.as_mut_ptr(),
                row_idx.as_mut_ptr(),
                &mut common,
            )
        };
        if symbolic.is_null() {
            return None;
        }
        let mut numeric = unsafe {
            klu_factor(
                col_ptr.as_mut_ptr(),
                row_idx.as_mut_ptr(),
                values.as_mut_ptr(),
                symbolic,
                &mut common,
            )
        };
        if numeric.is_null() || common.status != KLU_OK {
            let mut symbolic = symbolic;
            unsafe { klu_free_symbolic(&mut symbolic, &mut common) };
            return None;
        }
        // numeric is non-null and owned from here; suppress "unused mut" if
        // the compiler considers it otherwise (kept `mut` since it's freed
        // via `&mut` in Drop).
        let _ = &mut numeric;

        Some(Self {
            n,
            col_ptr,
            row_idx,
            groups: FlatGroups::new(&groups),
            packed: values,
            in_order: false,
            common,
            symbolic,
            numeric,
        })
    }

    /// Numeric-only refactorization against the cached symbolic pattern,
    /// then solves `A x = b`. Returns `None` if the matrix is singular.
    ///
    /// Success/failure is judged from `common.status` alone (`KLU_OK` = 0,
    /// documented as consistent across every KLU function — `klu.h`:
    /// "`status`: `KLU_OK` if OK, `< 0` if error", with `> 0` a warning
    /// including `KLU_SINGULAR`), not from `klu_refactor`/`klu_solve`'s own
    /// return values — confirmed by direct testing that those follow a
    /// *different*, boolean TRUE(1)/FALSE(0) convention (documented for
    /// `klu_refactor` in `klu.h`: "return TRUE if successful, FALSE
    /// otherwise") rather than `klu_factor`'s/`klu_analyze`'s "0 = OK, < 0 =
    /// error" status-code convention. Mixing the two up here previously
    /// caused every successful refactor to be misread as a failure.
    pub fn factor_and_solve(
        &mut self,
        entries: &[(usize, usize, f64)],
        rhs: &[f64],
    ) -> Option<Vec<f64>> {
        let mut values = std::mem::take(&mut self.packed);
        self.groups.pack_into(|i| entries[i].2, &mut values);
        self.solve_packed(values, rhs)
    }

    /// From now on takes values in KLU's own storage order, and hands them
    /// to KLU without copying: returns `order`, where `order[k]` is the
    /// input index (in the order `new` was given) that position `k` holds.
    /// The caller reorders its values to match, once, and keeps them so.
    ///
    /// `None`, changing nothing, if some position sums several inputs —
    /// duplicate `(row, col)` pairs, which a one-to-one order cannot express.
    pub fn adopt_storage_order(&mut self) -> Option<Vec<usize>> {
        let order = self.groups.one_to_one()?;
        self.groups = FlatGroups::identity(order.len());
        self.in_order = true;
        Some(order)
    }

    /// Shared tail of both `factor_and_solve` paths, taking already-packed
    /// CSC values, whose allocation it keeps for the next call.
    fn solve_packed(&mut self, values: Vec<f64>, rhs: &[f64]) -> Option<Vec<f64>> {
        let x = self.refactor_and_solve(&values, rhs);
        self.packed = values;
        x
    }

    /// `klu_refactor` on `values` (CSC order), then `klu_solve`.
    ///
    /// A refactor keeps the last factorization's pivots, and new values can
    /// put an exact zero on one; then `klu_factor` factors afresh, with new
    /// pivots, on the same symbolic analysis — lightsim2grid's fallback —
    /// and only if that fails is the matrix singular.
    fn refactor_and_solve(&mut self, values: &[f64], rhs: &[f64]) -> Option<Vec<f64>> {
        // SAFETY: KLU's `Ax` is declared `double []` but only read —
        // `klu_refactor` and the `klu_scale` it calls never write through it
        // (vendor/suitesparse/KLU/Source/klu_refactor.c, klu_scale.c) — so
        // passing a shared slice's pointer is sound.
        unsafe {
            klu_refactor(
                self.col_ptr.as_mut_ptr(),
                self.row_idx.as_mut_ptr(),
                values.as_ptr().cast_mut(),
                self.symbolic,
                self.numeric,
                &mut self.common,
            )
        };
        if self.common.status != KLU_OK {
            // SAFETY: `numeric` is ours and freed exactly once here, then
            // replaced; `klu_free_numeric` accepts the null a failed
            // `klu_factor` leaves, which `Drop` then frees again harmlessly.
            unsafe {
                klu_free_numeric(&mut self.numeric, &mut self.common);
                self.numeric = klu_factor(
                    self.col_ptr.as_mut_ptr(),
                    self.row_idx.as_mut_ptr(),
                    values.as_ptr().cast_mut(),
                    self.symbolic,
                    &mut self.common,
                );
            }
            if self.numeric.is_null() || self.common.status != KLU_OK {
                return None;
            }
        }

        let mut b = rhs.to_vec();
        unsafe {
            klu_solve(
                self.symbolic,
                self.numeric,
                self.n as i32,
                1,
                b.as_mut_ptr(),
                &mut self.common,
            )
        };
        if self.common.status != KLU_OK || b.iter().any(|v| !v.is_finite()) {
            return None;
        }
        Some(b)
    }
}

impl crate::sparse::LinearSolver for KluRealSystem {
    fn new(n: usize, entries: &[(usize, usize, f64)]) -> Option<Self> {
        KluRealSystem::new(n, entries)
    }

    fn factor_and_solve(
        &mut self,
        entries: &[(usize, usize, f64)],
        rhs: &[f64],
    ) -> Option<Vec<f64>> {
        KluRealSystem::factor_and_solve(self, entries, rhs)
    }

    fn factor_and_solve_values(&mut self, values: &[f64], rhs: &[f64]) -> Option<Vec<f64>> {
        if self.in_order {
            return self.refactor_and_solve(values, rhs);
        }
        let mut packed = std::mem::take(&mut self.packed);
        self.groups.pack_into(|i| values[i], &mut packed);
        self.solve_packed(packed, rhs)
    }
}

impl Drop for KluRealSystem {
    fn drop(&mut self) {
        unsafe {
            klu_free_numeric(&mut self.numeric, &mut self.common);
            klu_free_symbolic(&mut self.symbolic, &mut self.common);
        }
    }
}

/// Solves a one-shot general complex sparse linear system `A x = b` via
/// KLU. Mirrors `sparse::solve_complex`'s role for the `Scalar`/`faer`
/// backend. Returns `None` if the matrix is singular.
pub fn klu_solve_complex(
    n: usize,
    entries: &[(usize, usize, Complex<f64>)],
    rhs: &[Complex<f64>],
) -> Option<Vec<Complex<f64>>> {
    let pairs: Vec<(usize, usize)> = entries.iter().map(|&(r, c, _)| (r, c)).collect();
    let (mut col_ptr, mut row_idx, groups) = build_csc_structure(n, &pairs);
    let mut values = pack_complex_values(entries, &groups);

    let mut common = new_common();
    let mut symbolic = unsafe {
        klu_analyze(
            n as i32,
            col_ptr.as_mut_ptr(),
            row_idx.as_mut_ptr(),
            &mut common,
        )
    };
    if symbolic.is_null() {
        return None;
    }
    let mut numeric = unsafe {
        klu_z_factor(
            col_ptr.as_mut_ptr(),
            row_idx.as_mut_ptr(),
            values.as_mut_ptr(),
            symbolic,
            &mut common,
        )
    };

    let result = if numeric.is_null() || common.status != KLU_OK {
        None
    } else {
        let mut b: Vec<f64> = Vec::with_capacity(n * 2);
        for &v in rhs {
            b.push(v.re);
            b.push(v.im);
        }
        unsafe { klu_z_solve(symbolic, numeric, n as i32, 1, b.as_mut_ptr(), &mut common) };
        if common.status != KLU_OK || b.iter().any(|v| !v.is_finite()) {
            None
        } else {
            Some(
                (0..n)
                    .map(|i| Complex::new(b[2 * i], b[2 * i + 1]))
                    .collect(),
            )
        }
    };

    unsafe {
        if !numeric.is_null() {
            klu_free_numeric(&mut numeric, &mut common);
        }
        klu_free_symbolic(&mut symbolic, &mut common);
    }

    result
}

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

    #[test]
    fn solve_complex_simple_system() {
        // [[2, 1], [1, 3]] * [x0, x1] = [5, 10]  =>  x0=1, x1=3
        let entries = vec![
            (0, 0, Complex::new(2.0, 0.0)),
            (0, 1, Complex::new(1.0, 0.0)),
            (1, 0, Complex::new(1.0, 0.0)),
            (1, 1, Complex::new(3.0, 0.0)),
        ];
        let rhs = vec![Complex::new(5.0, 0.0), Complex::new(10.0, 0.0)];
        let x = klu_solve_complex(2, &entries, &rhs).unwrap();
        assert!((x[0] - Complex::new(1.0, 0.0)).norm() < 1e-10);
        assert!((x[1] - Complex::new(3.0, 0.0)).norm() < 1e-10);
    }

    #[test]
    fn solve_complex_duplicate_entries_sum() {
        // (0,0) contributed twice as 1.0+1.0=2.0, matching dense += semantics.
        let entries = vec![
            (0, 0, Complex::new(1.0, 0.0)),
            (0, 0, Complex::new(1.0, 0.0)),
            (0, 1, Complex::new(1.0, 0.0)),
            (1, 0, Complex::new(1.0, 0.0)),
            (1, 1, Complex::new(3.0, 0.0)),
        ];
        let rhs = vec![Complex::new(5.0, 0.0), Complex::new(10.0, 0.0)];
        let x = klu_solve_complex(2, &entries, &rhs).unwrap();
        assert!((x[0] - Complex::new(1.0, 0.0)).norm() < 1e-10);
        assert!((x[1] - Complex::new(3.0, 0.0)).norm() < 1e-10);
    }

    #[test]
    fn real_sparse_system_refactor_reuses_symbolic() {
        let entries_a = vec![(0, 0, 2.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 3.0)];
        let mut sys = KluRealSystem::new(2, &entries_a).unwrap();
        let x = sys.factor_and_solve(&entries_a, &[5.0, 10.0]).unwrap();
        assert!((x[0] - 1.0).abs() < 1e-10);
        assert!((x[1] - 3.0).abs() < 1e-10);

        // Same sparsity pattern, different numeric values (as across NR iterations).
        let entries_b = vec![(0, 0, 4.0), (0, 1, 1.0), (1, 0, 1.0), (1, 1, 2.0)];
        // [[4,1],[1,2]] * x = [6, 5] => x0=1, x1=2
        let x2 = sys.factor_and_solve(&entries_b, &[6.0, 5.0]).unwrap();
        assert!((x2[0] - 1.0).abs() < 1e-10);
        assert!((x2[1] - 2.0).abs() < 1e-10);
    }

    #[test]
    fn singular_matrix_returns_none() {
        let entries = vec![(0, 0, 1.0), (0, 1, 1.0), (1, 0, 2.0), (1, 1, 2.0)];
        let sys = KluRealSystem::new(2, &entries);
        if let Some(mut sys) = sys {
            assert!(sys.factor_and_solve(&entries, &[1.0, 2.0]).is_none());
        }
    }
}