#[derive(Debug, Clone)]
pub struct CsrIndex {
pub offsets: Vec<u64>,
pub neighbors: Vec<u64>,
pub edge_ids: Vec<u64>,
}
impl CsrIndex {
pub fn build(num_nodes: usize, edges: &mut [(u64, u64, u64)]) -> Self {
edges.sort_by_key(|e| e.0);
let mut offsets = vec![0u64; num_nodes + 1];
let mut neighbors = Vec::with_capacity(edges.len());
let mut edge_ids = Vec::with_capacity(edges.len());
for &(src, dst, eid) in edges.iter() {
if (src as usize) < num_nodes {
offsets[src as usize + 1] += 1;
}
neighbors.push(dst);
edge_ids.push(eid);
}
for i in 1..=num_nodes {
offsets[i] += offsets[i - 1];
}
CsrIndex {
offsets,
neighbors,
edge_ids,
}
}
pub fn neighbors(&self, src: u64) -> &[u64] {
let idx = src as usize;
if idx >= self.offsets.len() - 1 {
return &[] as &[u64];
}
let start = self.offsets[idx] as usize;
let end = self.offsets[idx + 1] as usize;
&self.neighbors[start..end]
}
pub fn edge_ids(&self, src: u64) -> &[u64] {
let idx = src as usize;
if idx >= self.offsets.len() - 1 {
return &[] as &[u64];
}
let start = self.offsets[idx] as usize;
let end = self.offsets[idx + 1] as usize;
&self.edge_ids[start..end]
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_csr_build_and_lookup() {
let mut edges = vec![(0, 1, 0), (0, 2, 1), (1, 3, 2), (2, 3, 3)];
let csr = CsrIndex::build(4, &mut edges);
assert_eq!(csr.neighbors(0), &[1, 2]);
assert_eq!(csr.edge_ids(0), &[0, 1]);
assert_eq!(csr.neighbors(1), &[3]);
assert_eq!(csr.neighbors(2), &[3]);
assert_eq!(csr.neighbors(3), &[] as &[u64]);
}
#[test]
fn test_csr_empty() {
let mut edges: Vec<(u64, u64, u64)> = vec![];
let csr = CsrIndex::build(3, &mut edges);
assert_eq!(csr.neighbors(0), &[] as &[u64]);
assert_eq!(csr.neighbors(1), &[] as &[u64]);
assert_eq!(csr.neighbors(2), &[] as &[u64]);
}
#[test]
fn test_csr_out_of_range() {
let mut edges = vec![(0, 1, 0)];
let csr = CsrIndex::build(2, &mut edges);
assert_eq!(csr.neighbors(99), &[] as &[u64]);
}
}