use super::EdgeStore;
use rustc_hash::{FxHashMap, FxHashSet};
use smallvec::SmallVec;
use std::collections::VecDeque;
use std::time::Instant;
pub(crate) const DEADLINE_CHECK_INTERVAL: u32 = 1024;
#[inline]
pub(super) fn deadline_exceeded(deadline: Option<Instant>) -> bool {
deadline.is_some_and(|d| Instant::now() >= d)
}
#[inline]
pub(crate) fn deadline_reached(deadline: Option<Instant>, nodes_since_check: &mut u32) -> bool {
if deadline.is_none() {
return false;
}
*nodes_since_check += 1;
if *nodes_since_check < DEADLINE_CHECK_INTERVAL {
return false;
}
*nodes_since_check = 0;
deadline_exceeded(deadline)
}
pub const DEFAULT_MAX_DEPTH: u32 = 3;
pub const SAFETY_MAX_DEPTH: u32 = 100;
pub type TraversalPath = SmallVec<[u64; 4]>;
#[derive(Debug, Clone)]
pub struct TraversalResult {
pub target_id: u64,
pub path: Vec<u64>,
pub depth: u32,
}
impl TraversalResult {
#[must_use]
pub fn new(target_id: u64, path: Vec<u64>, depth: u32) -> Self {
Self {
target_id,
path,
depth,
}
}
#[must_use]
#[allow(clippy::needless_pass_by_value)]
#[allow(dead_code)] pub(crate) fn from_smallvec(target_id: u64, path: TraversalPath, depth: u32) -> Self {
Self {
target_id,
path: path.to_vec(),
depth,
}
}
}
#[derive(Debug, Clone)]
pub struct TraversalConfig {
pub min_depth: u32,
pub max_depth: u32,
pub limit: usize,
pub rel_types: Vec<String>,
pub deadline: Option<Instant>,
}
impl Default for TraversalConfig {
fn default() -> Self {
Self {
min_depth: 1,
max_depth: DEFAULT_MAX_DEPTH,
limit: 100,
rel_types: Vec::new(),
deadline: None,
}
}
}
impl TraversalConfig {
#[must_use]
pub fn with_range(min: u32, max: u32) -> Self {
Self {
min_depth: min,
max_depth: max,
..Self::default()
}
}
#[must_use]
pub fn with_unbounded_range(min: u32) -> Self {
Self {
min_depth: min,
max_depth: SAFETY_MAX_DEPTH,
..Self::default()
}
}
#[must_use]
pub fn with_limit(mut self, limit: usize) -> Self {
self.limit = limit;
self
}
#[must_use]
pub fn with_rel_types(mut self, types: Vec<String>) -> Self {
self.rel_types = types;
self
}
#[must_use]
pub fn with_max_depth(mut self, max_depth: u32) -> Self {
self.max_depth = max_depth;
self
}
#[must_use]
pub fn with_deadline(mut self, deadline: Instant) -> Self {
self.deadline = Some(deadline);
self
}
}
#[derive(Debug)]
pub(super) struct BfsState {
pub(super) node_id: u64,
pub(super) depth: u32,
}
#[must_use]
pub(super) fn reconstruct_path(
target: u64,
source: u64,
parent_map: &FxHashMap<u64, (u64, u64)>,
) -> Vec<u64> {
let mut path = Vec::new();
let mut current = target;
while current != source {
if let Some(&(parent, edge_id)) = parent_map.get(¤t) {
path.push(edge_id);
current = parent;
} else {
break;
}
}
path.reverse();
path
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BfsDirection {
Forward,
Reverse,
}
struct BfsWalk<'a> {
config: &'a TraversalConfig,
source_id: u64,
results: Vec<TraversalResult>,
visited: FxHashSet<u64>,
queue: VecDeque<BfsState>,
parent_map: FxHashMap<u64, (u64, u64)>,
}
impl<'a> BfsWalk<'a> {
fn new(config: &'a TraversalConfig, source_id: u64) -> Self {
let mut visited = FxHashSet::default();
visited.insert(source_id);
let mut queue = VecDeque::new();
queue.push_back(BfsState {
node_id: source_id,
depth: 0,
});
Self {
config,
source_id,
results: Vec::new(),
visited,
queue,
parent_map: FxHashMap::default(),
}
}
fn limit_reached(&self) -> bool {
self.results.len() >= self.config.limit
}
#[inline]
fn visit(&mut self, target: u64, edge_id: u64, parent: &BfsState) {
let new_depth = parent.depth + 1;
if new_depth > self.config.max_depth {
return;
}
if self.visited.insert(target) {
self.parent_map.insert(target, (parent.node_id, edge_id));
if new_depth >= self.config.min_depth {
let path = reconstruct_path(target, self.source_id, &self.parent_map);
self.results
.push(TraversalResult::new(target, path, new_depth));
}
if new_depth < self.config.max_depth {
self.queue.push_back(BfsState {
node_id: target,
depth: new_depth,
});
}
}
}
}
#[must_use]
fn bfs_traverse_directed(
edge_store: &EdgeStore,
source_id: u64,
config: &TraversalConfig,
direction: BfsDirection,
) -> Vec<TraversalResult> {
let mut walk = BfsWalk::new(config, source_id);
let rel_filter: FxHashSet<&str> = config.rel_types.iter().map(String::as_str).collect();
let use_csr = direction == BfsDirection::Forward && edge_store.has_csr_snapshot();
let mut nodes_since_check = DEADLINE_CHECK_INTERVAL;
while let Some(state) = walk.queue.pop_front() {
if walk.limit_reached() {
break;
}
if deadline_reached(config.deadline, &mut nodes_since_check) {
break;
}
if use_csr {
process_bfs_csr(edge_store, &state, &rel_filter, &mut walk);
} else {
let edges = match direction {
BfsDirection::Forward => edge_store.get_outgoing(state.node_id),
BfsDirection::Reverse => edge_store.get_incoming(state.node_id),
};
process_bfs_neighbors(&edges, &state, &rel_filter, direction, &mut walk);
}
}
walk.results
}
#[inline]
fn process_bfs_csr(
edge_store: &EdgeStore,
state: &BfsState,
rel_filter: &FxHashSet<&str>,
walk: &mut BfsWalk<'_>,
) {
let Some(snapshot) = edge_store.csr_snapshot() else {
return;
};
let targets = snapshot.neighbors(state.node_id);
let edge_ids = snapshot.edge_ids(state.node_id);
for (i, (&target, &eid)) in targets.iter().zip(edge_ids.iter()).enumerate() {
if walk.limit_reached() {
break;
}
if let Some(label) = snapshot.label_at(state.node_id, i) {
if !rel_filter.is_empty() && !rel_filter.contains(label) {
continue;
}
} else if !rel_filter.is_empty() {
continue;
}
walk.visit(target, eid, state);
}
}
#[inline]
fn process_bfs_neighbors(
edges: &[&super::GraphEdge],
state: &BfsState,
rel_filter: &FxHashSet<&str>,
direction: BfsDirection,
walk: &mut BfsWalk<'_>,
) {
for edge in edges {
if walk.limit_reached() {
break;
}
if !rel_filter.is_empty() && !rel_filter.contains(edge.label()) {
continue;
}
let next_node = match direction {
BfsDirection::Forward => edge.target(),
BfsDirection::Reverse => edge.source(),
};
walk.visit(next_node, edge.id(), state);
}
}
#[must_use]
pub fn bfs_traverse(
edge_store: &EdgeStore,
source_id: u64,
config: &TraversalConfig,
) -> Vec<TraversalResult> {
bfs_traverse_directed(edge_store, source_id, config, BfsDirection::Forward)
}
#[must_use]
pub fn bfs_traverse_reverse(
edge_store: &EdgeStore,
source_id: u64,
config: &TraversalConfig,
) -> Vec<TraversalResult> {
bfs_traverse_directed(edge_store, source_id, config, BfsDirection::Reverse)
}