use crate::graph::unified::compaction::CompactionSnapshot;
use crate::graph::unified::edge::EdgeKind;
use crate::graph::unified::node::NodeId;
use crate::graph::unified::storage::CsrGraph;
use anyhow::Result;
use rayon::prelude::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EdgeKindDiscriminant(std::mem::Discriminant<EdgeKind>);
impl EdgeKindDiscriminant {
#[must_use]
pub fn from_edge_kind(kind: &EdgeKind) -> Self {
EdgeKindDiscriminant(std::mem::discriminant(kind))
}
#[must_use]
pub fn matches(&self, kind: &EdgeKind) -> bool {
self.0 == std::mem::discriminant(kind)
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CsrAdjacency {
pub node_count: u32,
pub edge_count: u32,
pub row_offsets: Vec<u32>,
pub col_indices: Vec<u32>,
pub edge_kinds: Vec<EdgeKind>,
}
impl CsrAdjacency {
#[allow(clippy::cast_possible_truncation)]
#[must_use]
pub fn from_csr_graph(csr: &CsrGraph) -> Self {
let node_count = csr.node_count() as u32;
let edge_count = csr.edge_count() as u32;
let row_offsets = csr.row_ptr().to_vec();
let col_indices: Vec<u32> = csr.col_idx().iter().map(|nid| nid.index()).collect();
let edge_kinds = csr.edge_kind().to_vec();
Self {
node_count,
edge_count,
row_offsets,
col_indices,
edge_kinds,
}
}
#[allow(clippy::cast_possible_truncation)] pub fn build_from_snapshot(snapshot: &CompactionSnapshot) -> Result<Self> {
let node_count = snapshot.node_count as u32;
let (merged_edges, _stats) = crate::graph::unified::compaction::merge::merge_with_csr(
&snapshot.csr_edges,
snapshot.delta_edges.clone(),
);
let mut all_edges: Vec<(u32, u32, EdgeKind)> = merged_edges
.into_iter()
.map(|edge| (edge.source.index(), edge.target.index(), edge.kind))
.collect();
all_edges.par_sort_unstable_by_key(|(src, tgt, _)| (*src, *tgt));
let edge_count = all_edges.len() as u32;
let mut out_degree = vec![0u32; node_count as usize];
for &(src, _, _) in &all_edges {
out_degree[src as usize] += 1;
}
let mut row_offsets = Vec::with_capacity(node_count as usize + 1);
row_offsets.push(0);
for &count in &out_degree {
let last = *row_offsets.last().unwrap();
row_offsets.push(last + count);
}
let mut col_indices = Vec::with_capacity(edge_count as usize);
let mut edge_kinds = Vec::with_capacity(edge_count as usize);
for (_, target, kind) in all_edges {
col_indices.push(target);
edge_kinds.push(kind);
}
Ok(Self {
node_count,
edge_count,
row_offsets,
col_indices,
edge_kinds,
})
}
#[must_use]
pub fn neighbors(&self, node: NodeId) -> &[u32] {
let idx = node.index() as usize;
if idx >= self.row_offsets.len() - 1 {
return &[];
}
let start = self.row_offsets[idx] as usize;
let end = self.row_offsets[idx + 1] as usize;
&self.col_indices[start..end]
}
#[must_use]
pub fn neighbors_filtered(&self, node: NodeId, kind: &EdgeKind) -> Vec<u32> {
let idx = node.index() as usize;
if idx >= self.row_offsets.len() - 1 {
return Vec::new();
}
let start = self.row_offsets[idx] as usize;
let end = self.row_offsets[idx + 1] as usize;
let kind_discriminant = std::mem::discriminant(kind);
(start..end)
.filter(|&i| std::mem::discriminant(&self.edge_kinds[i]) == kind_discriminant)
.map(|i| self.col_indices[i])
.collect()
}
#[must_use]
pub fn edge_range(&self, node: NodeId) -> (usize, usize) {
let idx = node.index() as usize;
if idx >= self.row_offsets.len() - 1 {
return (0, 0);
}
let start = self.row_offsets[idx] as usize;
let end = self.row_offsets[idx + 1] as usize;
(start, end)
}
pub fn iter_neighbors_filtered(
&self,
node: NodeId,
kind: &EdgeKind,
) -> impl Iterator<Item = (usize, u32)> + '_ {
let (start, end) = self.edge_range(node);
let kind_discriminant = std::mem::discriminant(kind);
(start..end)
.filter(move |&i| std::mem::discriminant(&self.edge_kinds[i]) == kind_discriminant)
.map(move |i| (i, self.col_indices[i]))
}
}