use std::collections::{HashSet, VecDeque};
use crate::types::EdgeType;
use super::code_graph::CodeGraph;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Direction {
Forward,
Backward,
}
#[derive(Debug, Clone)]
pub struct TraversalOptions {
pub max_depth: i32,
pub edge_types: Vec<EdgeType>,
pub direction: Direction,
}
impl Default for TraversalOptions {
fn default() -> Self {
Self {
max_depth: -1,
edge_types: Vec::new(),
direction: Direction::Forward,
}
}
}
pub type TraversalResult = Vec<(u64, u32)>;
pub fn bfs(graph: &CodeGraph, start_id: u64, options: &TraversalOptions) -> TraversalResult {
let mut visited = HashSet::new();
let mut queue = VecDeque::new();
let mut result = Vec::new();
visited.insert(start_id);
queue.push_back((start_id, 0u32));
while let Some((current_id, depth)) = queue.pop_front() {
result.push((current_id, depth));
if options.max_depth >= 0 && depth >= options.max_depth as u32 {
continue;
}
let neighbors = match options.direction {
Direction::Forward => graph.edges_from(current_id),
Direction::Backward => graph.edges_to(current_id),
};
for edge in neighbors {
if !options.edge_types.is_empty() && !options.edge_types.contains(&edge.edge_type) {
continue;
}
let next_id = match options.direction {
Direction::Forward => edge.target_id,
Direction::Backward => edge.source_id,
};
if visited.insert(next_id) {
queue.push_back((next_id, depth + 1));
}
}
}
result
}
pub fn dfs(graph: &CodeGraph, start_id: u64, options: &TraversalOptions) -> TraversalResult {
let mut visited = HashSet::new();
let mut result = Vec::new();
dfs_inner(graph, start_id, 0, options, &mut visited, &mut result);
result
}
fn dfs_inner(
graph: &CodeGraph,
current_id: u64,
depth: u32,
options: &TraversalOptions,
visited: &mut HashSet<u64>,
result: &mut TraversalResult,
) {
if !visited.insert(current_id) {
return;
}
result.push((current_id, depth));
if options.max_depth >= 0 && depth >= options.max_depth as u32 {
return;
}
let neighbors = match options.direction {
Direction::Forward => graph.edges_from(current_id),
Direction::Backward => graph.edges_to(current_id),
};
for edge in neighbors {
if !options.edge_types.is_empty() && !options.edge_types.contains(&edge.edge_type) {
continue;
}
let next_id = match options.direction {
Direction::Forward => edge.target_id,
Direction::Backward => edge.source_id,
};
dfs_inner(graph, next_id, depth + 1, options, visited, result);
}
}
pub fn find_paths(
graph: &CodeGraph,
from: u64,
to: u64,
max_depth: u32,
edge_types: &[EdgeType],
) -> Vec<Vec<u64>> {
let mut paths = Vec::new();
let mut current_path = vec![from];
let mut visited = HashSet::new();
visited.insert(from);
find_paths_inner(
graph,
to,
max_depth,
edge_types,
&mut current_path,
&mut visited,
&mut paths,
);
paths
}
fn find_paths_inner(
graph: &CodeGraph,
target: u64,
max_depth: u32,
edge_types: &[EdgeType],
current_path: &mut Vec<u64>,
visited: &mut HashSet<u64>,
results: &mut Vec<Vec<u64>>,
) {
let current = *current_path.last().unwrap();
if current == target && current_path.len() > 1 {
results.push(current_path.clone());
return;
}
if current_path.len() > max_depth as usize {
return;
}
for edge in graph.edges_from(current) {
if !edge_types.is_empty() && !edge_types.contains(&edge.edge_type) {
continue;
}
if visited.insert(edge.target_id) {
current_path.push(edge.target_id);
find_paths_inner(
graph,
target,
max_depth,
edge_types,
current_path,
visited,
results,
);
current_path.pop();
visited.remove(&edge.target_id);
}
}
}
pub fn shortest_path(
graph: &CodeGraph,
from: u64,
to: u64,
edge_types: &[EdgeType],
) -> Option<Vec<u64>> {
if from == to {
return Some(vec![from]);
}
let mut visited = HashSet::new();
let mut queue = VecDeque::new();
let mut parent: std::collections::HashMap<u64, u64> = std::collections::HashMap::new();
visited.insert(from);
queue.push_back(from);
while let Some(current) = queue.pop_front() {
for edge in graph.edges_from(current) {
if !edge_types.is_empty() && !edge_types.contains(&edge.edge_type) {
continue;
}
if visited.insert(edge.target_id) {
parent.insert(edge.target_id, current);
if edge.target_id == to {
let mut path = vec![to];
let mut node = to;
while let Some(&p) = parent.get(&node) {
path.push(p);
node = p;
}
path.reverse();
return Some(path);
}
queue.push_back(edge.target_id);
}
}
}
None
}