use std::collections::{HashSet, VecDeque};
use super::compat::{EntityRef, LinkStore, StorageContext};
use crate::error::Result;
use super::helpers::{get_edge_weight, get_neighbor_entity, get_neighbors, matches_link_type};
use super::types::{PathNode, TraversalOptions, MAX_TRAVERSAL_DEPTH, MAX_TRAVERSAL_RESULTS};
pub async fn bfs_traverse<S: LinkStore>(
store: &S,
ctx: &StorageContext,
start: EntityRef,
options: &TraversalOptions,
) -> Result<Vec<PathNode>> {
let max_depth = options.max_depth.min(MAX_TRAVERSAL_DEPTH);
let limit = options
.limit
.unwrap_or(MAX_TRAVERSAL_RESULTS)
.min(MAX_TRAVERSAL_RESULTS);
let min_weight = options.min_weight.unwrap_or(f64::NEG_INFINITY);
let mut visited: HashSet<EntityRef> = HashSet::new();
let mut results: Vec<PathNode> = Vec::new();
let mut queue: VecDeque<(EntityRef, usize, f64)> = VecDeque::new();
visited.insert(start.clone());
results.push(PathNode::start(start.clone()));
queue.push_back((start, 0, 0.0));
while let Some((current, depth, path_weight)) = queue.pop_front() {
if depth >= max_depth {
continue;
}
if results.len() >= limit {
break;
}
let links = get_neighbors(store, ctx, ¤t, &options.direction).await?;
for link in links {
if !matches_link_type(&link, &options.link_types) {
continue;
}
let edge_weight = get_edge_weight(&link);
if edge_weight < min_weight {
continue;
}
let neighbor = get_neighbor_entity(&link, ¤t, &options.direction);
if visited.contains(&neighbor) {
continue;
}
visited.insert(neighbor.clone());
let new_weight = path_weight + edge_weight;
let node = PathNode {
entity_id: neighbor.clone(),
depth: depth + 1,
via_link: Some(link),
path_weight: new_weight,
};
results.push(node);
if results.len() >= limit {
break;
}
queue.push_back((neighbor, depth + 1, new_weight));
}
}
Ok(results)
}