#![allow(dead_code)]
use super::EdgeStore;
use std::collections::{HashSet, VecDeque};
pub const DEFAULT_MAX_DEPTH: u32 = 3;
pub const SAFETY_MAX_DEPTH: u32 = 100;
#[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,
}
}
}
#[derive(Debug, Clone)]
pub struct TraversalConfig {
pub min_depth: u32,
pub max_depth: u32,
pub limit: usize,
pub rel_types: Vec<String>,
}
impl Default for TraversalConfig {
fn default() -> Self {
Self {
min_depth: 1,
max_depth: DEFAULT_MAX_DEPTH,
limit: 100,
rel_types: Vec::new(),
}
}
}
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
}
}
#[derive(Debug)]
pub(super) struct BfsState {
pub(super) node_id: u64,
pub(super) path: Vec<u64>,
pub(super) depth: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BfsDirection {
Forward,
Reverse,
}
#[must_use]
fn bfs_traverse_directed(
edge_store: &EdgeStore,
source_id: u64,
config: &TraversalConfig,
direction: BfsDirection,
) -> Vec<TraversalResult> {
let mut results = Vec::new();
let mut visited = HashSet::new();
let mut queue = VecDeque::new();
visited.insert(source_id);
queue.push_back(BfsState {
node_id: source_id,
path: Vec::new(),
depth: 0,
});
while let Some(state) = queue.pop_front() {
if results.len() >= config.limit {
break;
}
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,
config,
direction,
&mut results,
&mut visited,
&mut queue,
);
}
results
}
#[inline]
fn process_bfs_neighbors(
edges: &[&super::GraphEdge],
state: &BfsState,
config: &TraversalConfig,
direction: BfsDirection,
results: &mut Vec<TraversalResult>,
visited: &mut HashSet<u64>,
queue: &mut VecDeque<BfsState>,
) {
for edge in edges {
if results.len() >= config.limit {
break;
}
if !config.rel_types.is_empty() && !config.rel_types.contains(&edge.label().to_string()) {
continue;
}
let next_node = match direction {
BfsDirection::Forward => edge.target(),
BfsDirection::Reverse => edge.source(),
};
let new_depth = state.depth + 1;
if new_depth > config.max_depth {
continue;
}
let mut new_path = state.path.clone();
new_path.push(edge.id());
if new_depth >= config.min_depth {
results.push(TraversalResult::new(next_node, new_path.clone(), new_depth));
}
if new_depth < config.max_depth && visited.insert(next_node) {
queue.push_back(BfsState {
node_id: next_node,
path: new_path,
depth: new_depth,
});
}
}
}
#[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)
}
#[must_use]
pub fn bfs_traverse_both(
edge_store: &EdgeStore,
source_id: u64,
config: &TraversalConfig,
) -> Vec<TraversalResult> {
let mut results = Vec::new();
let half_limit = config.limit / 2 + 1;
let config_half = TraversalConfig {
limit: half_limit,
..config.clone()
};
let forward = bfs_traverse(edge_store, source_id, &config_half);
results.extend(forward);
if results.len() < config.limit {
let reverse = bfs_traverse_reverse(edge_store, source_id, &config_half);
for r in reverse {
if results.len() >= config.limit {
break;
}
if !results
.iter()
.any(|existing| existing.target_id == r.target_id && existing.path == r.path)
{
results.push(r);
}
}
}
results.truncate(config.limit);
results
}