safegraph 0.1.0

//! Deeper investigation of the `remove_edges` gap between `sg_vec_scoped`
//! and petgraph. Run: `cargo run --release --example remove_edges_probe`.
//!
//! Decomposes the safegraph cost into: (1) the victim-scan pass, (2) the
//! batched `take_nodes_edges_unchecked`, and (3) a scalar descending loop of
//! `take_edge_unchecked` (the apples-to-apples analog of petgraph's in-place
//! `retain_edges`). A counting global allocator reports allocations per call.

// Dev-only timing probe: uses modern std (`black_box`, `is_multiple_of`) and is
// not bound by the library's 1.56 MSRV.
#![allow(clippy::incompatible_msrv)]

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::time::Instant;

use safegraph::graph::operation::GraphOperation;
use safegraph::graph::Graph;
use safegraph::VecGraph;

// ---- counting allocator -------------------------------------------------

static ALLOCS: AtomicUsize = AtomicUsize::new(0);
static BYTES: AtomicUsize = AtomicUsize::new(0);

struct Counting;
unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        ALLOCS.fetch_add(1, Relaxed);
        BYTES.fetch_add(l.size(), Relaxed);
        System.alloc(l)
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        System.dealloc(p, l)
    }
}
#[global_allocator]
static A: Counting = Counting;

fn reset() {
    ALLOCS.store(0, Relaxed);
    BYTES.store(0, Relaxed);
}
fn snap() -> (usize, usize) {
    (ALLOCS.load(Relaxed), BYTES.load(Relaxed))
}

// ---- workload (mirrors benchmark/common.rs) -----------------------------

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 {
        let mut x = self.0;
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
        self.0 = x;
        x
    }
    fn below(&mut self, n: usize) -> u32 {
        (self.next() % n as u64) as u32
    }
}

fn gen(n: usize, m: usize) -> Vec<(u32, u32)> {
    let mut r = Rng(0x5AFE_6EA9);
    (0..m).map(|_| (r.below(n), r.below(n))).collect()
}

fn victim(e: usize) -> bool {
    e.is_multiple_of(2)
}

// ---- builders -----------------------------------------------------------

type G = VecGraph<usize, usize>;

fn build(n: usize, edges: &[(u32, u32)]) -> G {
    let mut g = G::default();
    let ixs: Vec<_> = (0..n)
        .map(|i| unsafe { g.insert_node_unchecked(i).ok().unwrap() })
        .collect();
    for (j, &(f, t)) in edges.iter().enumerate() {
        unsafe {
            g.insert_edge_unchecked(j, [ixs[f as usize], ixs[t as usize]])
                .ok()
                .unwrap();
        }
    }
    g
}

fn build_pg(n: usize, edges: &[(u32, u32)]) -> petgraph::graph::DiGraph<usize, usize> {
    let mut g = petgraph::graph::DiGraph::<usize, usize>::new();
    let ixs: Vec<_> = (0..n).map(|i| g.add_node(i)).collect();
    for (j, &(f, t)) in edges.iter().enumerate() {
        g.add_edge(ixs[f as usize], ixs[t as usize], j);
    }
    g
}

// ---- the three safegraph removal strategies -----------------------------

/// Collect victims (one O(m) scan + Vec alloc), then batched removal.
fn sg_batched(g: &mut G) -> usize {
    let victims: Vec<_> = unsafe {
        GraphOperation::edge_indices(g)
            .filter(|&e| victim(*Graph::edge_unchecked(g, e)))
            .collect()
    };
    let k = victims.len();
    unsafe { g.remove_nodes_edges_unchecked(None::<u32>, victims) };
    k
}

/// Collect victims, sort descending, then scalar inline-repair removal —
/// the structural analog of petgraph `retain_edges`.
fn sg_scalar(g: &mut G) -> usize {
    let mut victims: Vec<_> = unsafe {
        GraphOperation::edge_indices(g)
            .filter(|&e| victim(*Graph::edge_unchecked(g, e)))
            .collect()
    };
    // Descending so each swap_remove only relocates a survivor we won't touch.
    victims.sort_unstable_by(|a, b| b.cmp(a));
    let k = victims.len();
    for e in victims {
        unsafe { Graph::take_edge_unchecked(g, e) };
    }
    k
}

fn pg_retain(g: &mut petgraph::graph::DiGraph<usize, usize>) -> usize {
    let before = g.edge_count();
    g.retain_edges(|gg, e| !victim(gg[e]));
    before - g.edge_count()
}

// ---- timing harness -----------------------------------------------------

fn time_sg(label: &str, n: usize, edges: &[(u32, u32)], f: fn(&mut G) -> usize, iters: usize) {
    // warm
    let _ = f(&mut build(n, edges));
    // allocation count for one representative call
    let mut g0 = build(n, edges);
    reset();
    let k = f(&mut g0);
    let (a, b) = snap();
    drop(g0);
    // timing over fresh clones
    let mut total = std::time::Duration::ZERO;
    for _ in 0..iters {
        let mut g = build(n, edges);
        let t = Instant::now();
        std::hint::black_box(f(&mut g));
        total += t.elapsed();
        std::hint::black_box(&g);
    }
    println!(
        "  {label:<22} {:>9.2?}/call  | removed {k:>6} | {a:>5} allocs, {:>7} B",
        total / iters as u32,
        b
    );
}

fn time_pg(n: usize, edges: &[(u32, u32)], iters: usize) {
    let _ = pg_retain(&mut build_pg(n, edges));
    let mut g0 = build_pg(n, edges);
    reset();
    let k = pg_retain(&mut g0);
    let (a, b) = snap();
    drop(g0);
    let mut total = std::time::Duration::ZERO;
    for _ in 0..iters {
        let mut g = build_pg(n, edges);
        let t = Instant::now();
        std::hint::black_box(pg_retain(&mut g));
        total += t.elapsed();
        std::hint::black_box(&g);
    }
    println!(
        "  {:<22} {:>9.2?}/call  | removed {k:>6} | {a:>5} allocs, {:>7} B",
        "pg_retain_edges",
        total / iters as u32,
        b
    );
}

/// Surviving-graph fingerprint: sorted (from_payload, to_payload, edge_payload)
/// triples, walked through the adjacency lists. Detects corrupted adjacency
/// that a bare edge-count would miss.
fn fingerprint(g: &G) -> Vec<(usize, usize, usize)> {
    let mut out = Vec::new();
    unsafe {
        for nix in GraphOperation::node_indices(g) {
            let from = *Graph::node_unchecked(g, nix);
            for (eix, _, to_nix) in GraphOperation::walks_from_unchecked(g, nix).map(|w| w.get()) {
                out.push((from, *Graph::node_unchecked(g, to_nix), *Graph::edge_unchecked(g, eix)));
            }
        }
    }
    out.sort_unstable();
    out
}

fn verify(n: usize, edges: &[(u32, u32)]) {
    let mut a = build(n, edges);
    let mut b = build(n, edges);
    sg_batched(&mut a);
    sg_scalar(&mut b);
    let (fa, fb) = (fingerprint(&a), fingerprint(&b));
    assert_eq!(
        fa, fb,
        "scalar and batched disagree on surviving graph (n={n})"
    );
    // every survivor must be a non-victim, and all of them must survive
    let want = edges.len() - edges.iter().enumerate().filter(|(j, _)| victim(*j)).count();
    assert_eq!(fa.len(), want, "wrong survivor count (n={n})");
    assert!(fa.iter().all(|&(_, _, e)| !victim(e)));
}

fn main() {
    for &(n, m) in &[(100usize, 500usize), (1_000, 5_000)] {
        verify(n, &gen(n, m));
    }
    println!("correctness: scalar-desc fingerprint == batched fingerprint ✓");

    for &(n, m) in &[(100usize, 500usize), (1_000, 5_000), (10_000, 50_000)] {
        let edges = gen(n, m);
        let iters = if n >= 10_000 { 30 } else { 400 };
        println!("\n=== n={n}, m={m} ({iters} iters) ===");
        time_sg("sg_batched", n, &edges, sg_batched, iters);
        time_sg("sg_scalar_desc", n, &edges, sg_scalar, iters);
        time_pg(n, &edges, iters);
    }
}