use super::edge::EdgeStore;
use super::label_table::{LabelId, LabelTable};
use rustc_hash::{FxHashMap, FxHashSet};
pub trait EdgePredicate: Send + Sync {
fn matches(&self, target: u64, edge_id: u64, label_id: LabelId) -> bool;
}
pub struct LabelFilter {
allowed: FxHashSet<LabelId>,
}
impl LabelFilter {
#[must_use]
pub fn new(allowed: FxHashSet<LabelId>) -> Self {
Self { allowed }
}
}
impl EdgePredicate for LabelFilter {
#[inline]
fn matches(&self, _target: u64, _edge_id: u64, label_id: LabelId) -> bool {
self.allowed.contains(&label_id)
}
}
pub struct NoFilter;
impl EdgePredicate for NoFilter {
#[inline]
fn matches(&self, _target: u64, _edge_id: u64, _label_id: LabelId) -> bool {
true
}
}
pub trait AdjacencySource {
fn neighbors(&self, node_id: u64) -> Vec<u64>;
}
impl AdjacencySource for CsrSnapshot {
#[inline]
fn neighbors(&self, node_id: u64) -> Vec<u64> {
self.neighbors(node_id).to_vec()
}
}
impl AdjacencySource for EdgeStore {
#[inline]
fn neighbors(&self, node_id: u64) -> Vec<u64> {
self.get_outgoing(node_id)
.iter()
.map(|e| e.target())
.collect()
}
}
#[derive(Debug, Clone)]
pub struct CsrSnapshot {
offsets: Vec<usize>,
targets: Vec<u64>,
edge_ids: Vec<u64>,
label_ids: Vec<LabelId>,
node_to_index: FxHashMap<u64, usize>,
index_to_node: Vec<u64>,
label_table: Vec<String>,
label_to_idx: FxHashMap<String, u32>,
}
impl CsrSnapshot {
#[inline]
fn range_of(&self, node_id: u64) -> Option<(usize, usize)> {
let &idx = self.node_to_index.get(&node_id)?;
let start = self.offsets[idx];
let end = self.offsets[idx + 1];
Some((start, end))
}
#[must_use]
#[inline]
pub fn neighbors(&self, node_id: u64) -> &[u64] {
if let Some((start, end)) = self.range_of(node_id) {
&self.targets[start..end]
} else {
&[]
}
}
#[must_use]
#[inline]
pub fn edge_ids(&self, node_id: u64) -> &[u64] {
if let Some((start, end)) = self.range_of(node_id) {
&self.edge_ids[start..end]
} else {
&[]
}
}
#[must_use]
#[inline]
pub fn label_ids(&self, node_id: u64) -> &[LabelId] {
if let Some((start, end)) = self.range_of(node_id) {
&self.label_ids[start..end]
} else {
&[]
}
}
#[must_use]
#[inline]
pub fn label_at(&self, source_id: u64, neighbor_idx: usize) -> Option<&str> {
let (start, end) = self.range_of(source_id)?;
if neighbor_idx >= end - start {
return None;
}
let label_id = self.label_ids[start + neighbor_idx];
self.label_table
.get(label_id.as_u32() as usize)
.map(String::as_str)
}
#[must_use]
#[inline]
pub fn degree(&self, node_id: u64) -> usize {
if let Some((start, end)) = self.range_of(node_id) {
end - start
} else {
0
}
}
#[must_use]
#[inline]
pub fn contains_node(&self, node_id: u64) -> bool {
self.node_to_index.contains_key(&node_id)
}
#[must_use]
#[inline]
pub fn node_count(&self) -> usize {
self.index_to_node.len()
}
#[must_use]
#[inline]
pub fn edge_count(&self) -> usize {
self.targets.len()
}
#[must_use]
#[inline]
pub fn distinct_label_count(&self) -> usize {
self.label_table.len()
}
#[must_use]
#[inline]
pub fn has_label(&self, label: &str) -> bool {
self.label_to_idx.contains_key(label)
}
pub fn neighbors_filtered<'a, P: EdgePredicate>(
&'a self,
node_id: u64,
predicate: &'a P,
) -> impl Iterator<Item = (u64, u64, LabelId)> + 'a {
let (start, end) = self.range_of(node_id).unwrap_or((0, 0));
(start..end).filter_map(move |i| {
let target = self.targets[i];
let eid = self.edge_ids[i];
let lid = self.label_ids[i];
if predicate.matches(target, eid, lid) {
Some((target, eid, lid))
} else {
None
}
})
}
#[cfg(test)]
pub(crate) fn offsets(&self) -> &[usize] {
&self.offsets
}
}
pub(crate) struct SnapshotBuilder;
impl SnapshotBuilder {
pub fn build(edge_store: &EdgeStore, _label_table: &LabelTable) -> CsrSnapshot {
let mut node_ids: Vec<u64> = edge_store.outgoing_keys();
node_ids.sort_unstable();
let node_count = node_ids.len();
let total_edges: usize = edge_store.total_outgoing_edges();
let mut node_to_index =
FxHashMap::with_capacity_and_hasher(node_count, rustc_hash::FxBuildHasher);
for (idx, &nid) in node_ids.iter().enumerate() {
node_to_index.insert(nid, idx);
}
let mut offsets = Vec::with_capacity(node_count + 1);
let mut targets = Vec::with_capacity(total_edges);
let mut edge_ids_buf = Vec::with_capacity(total_edges);
let mut label_ids_buf: Vec<LabelId> = Vec::with_capacity(total_edges);
let mut label_table_vec: Vec<String> = Vec::new();
let mut label_to_idx: FxHashMap<String, u32> =
FxHashMap::with_capacity_and_hasher(16, rustc_hash::FxBuildHasher);
for &nid in &node_ids {
offsets.push(targets.len());
edge_store.for_each_outgoing_edge(nid, |edge| {
targets.push(edge.target());
edge_ids_buf.push(edge.id());
let label_str = edge.label();
let local_idx = *label_to_idx
.entry(label_str.to_string())
.or_insert_with(|| {
let idx = label_table_vec.len();
label_table_vec.push(label_str.to_string());
#[allow(clippy::cast_possible_truncation)]
{
idx as u32
}
});
label_ids_buf.push(LabelId::from_u32(local_idx));
});
}
offsets.push(targets.len());
CsrSnapshot {
offsets,
targets,
edge_ids: edge_ids_buf,
label_ids: label_ids_buf,
node_to_index,
index_to_node: node_ids,
label_table: label_table_vec,
label_to_idx,
}
}
#[must_use]
#[allow(dead_code)] pub fn empty() -> CsrSnapshot {
CsrSnapshot {
offsets: vec![0],
targets: Vec::new(),
edge_ids: Vec::new(),
label_ids: Vec::new(),
node_to_index: FxHashMap::default(),
index_to_node: Vec::new(),
label_table: Vec::new(),
label_to_idx: FxHashMap::default(),
}
}
}