use std::collections::{HashSet, VecDeque};
use std::rc::Rc;
use crate::model::{EdgeType, Graph, NodeId};
use super::{TraceEdge, TraceNode, TracePath};
pub(crate) struct WorkItem {
node_id: NodeId,
node: TraceNode,
edge: Option<TraceEdge>,
depth: usize,
parent: Option<Rc<WorkItem>>,
path_set: Rc<HashSet<NodeId>>,
}
impl WorkItem {
fn new_root(node_id: NodeId, node: TraceNode) -> Self {
Self {
path_set: Rc::new(HashSet::from([node_id.clone()])),
node_id,
node,
edge: None,
depth: 0,
parent: None,
}
}
fn child(parent: &Rc<WorkItem>, node_id: NodeId, node: TraceNode, edge: TraceEdge) -> Self {
let mut path_set = (*parent.path_set).clone();
path_set.insert(node_id.clone());
Self {
path_set: Rc::new(path_set),
node_id,
node,
edge: Some(edge),
depth: parent.depth + 1,
parent: Some(Rc::clone(parent)),
}
}
fn has_parent_edge(&self) -> bool {
self.edge.is_some()
}
fn path_contains(&self, id: &NodeId) -> bool {
self.path_set.contains(id)
}
fn build_path(&self) -> TracePath {
let mut nodes = Vec::new();
let mut edges = Vec::new();
let mut current: &WorkItem = self;
loop {
nodes.push(current.node.clone());
if let Some(edge) = ¤t.edge {
edges.push(edge.clone());
}
match ¤t.parent {
Some(parent) => current = parent,
None => break,
}
}
nodes.reverse();
edges.reverse();
TracePath {
nodes,
edges,
depth: self.depth,
}
}
}
pub(crate) fn bfs_trace(
graph: &Graph,
start: &NodeId,
max_depth: usize,
edge_filter: impl Fn(&EdgeType) -> bool,
sink: Option<&NodeId>,
) -> Vec<TracePath> {
let Some(start_node) = graph.get_node(start) else {
return Vec::new();
};
if let Some(sink_id) = sink {
if graph.get_node(sink_id).is_none() {
return Vec::new();
}
}
let mut queue: VecDeque<WorkItem> = VecDeque::new();
queue.push_back(WorkItem::new_root(
start.clone(),
TraceNode::from(start_node),
));
let mut results = Vec::new();
while let Some(work) = queue.pop_front() {
let has_parent_edge = work.has_parent_edge();
let current_id = work.node_id.clone();
if let Some(sink_id) = sink {
if ¤t_id == sink_id && has_parent_edge {
results.push(work.build_path());
continue;
}
}
if work.depth >= max_depth {
if sink.is_none() && has_parent_edge {
results.push(work.build_path());
}
continue;
}
let work_rc = Rc::new(work);
for edge in graph.edges_from(¤t_id) {
if !edge_filter(&edge.edge_type) {
continue;
}
let Some(target_node) = graph.get_node(&edge.target) else {
continue;
};
if work_rc.path_contains(&edge.target) {
continue;
}
let child = WorkItem::child(
&work_rc,
edge.target.clone(),
TraceNode::from(target_node),
TraceEdge {
edge_type: edge.edge_type.as_db_type().to_string(),
reason: edge.reason.clone(),
confidence: edge.confidence,
},
);
queue.push_back(child);
}
if sink.is_none() && has_parent_edge {
results.push(work_rc.build_path());
}
}
results
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{Edge, Node, NodeLabel};
fn abc_graph() -> Graph {
let mut g = Graph::new();
g.add_node(
Node::builder(NodeLabel::Function, "a", "proj.a")
.id("a")
.project("proj")
.build(),
);
g.add_node(
Node::builder(NodeLabel::Function, "b", "proj.b")
.id("b")
.project("proj")
.build(),
);
g.add_node(
Node::builder(NodeLabel::Function, "c", "proj.c")
.id("c")
.project("proj")
.build(),
);
g.add_edge(Edge::new("a", "b", EdgeType::Calls, "proj"));
g.add_edge(Edge::new("b", "c", EdgeType::Calls, "proj"));
g
}
fn trace_edge() -> TraceEdge {
TraceEdge {
edge_type: "Calls".to_string(),
reason: None,
confidence: 1.0,
}
}
#[test]
fn has_parent_edge_false_for_root_true_for_child() {
let g = abc_graph();
let a = g.get_node(&"a".to_string()).unwrap();
let b = g.get_node(&"b".to_string()).unwrap();
let root = WorkItem::new_root("a".to_string(), TraceNode::from(a));
assert!(!root.has_parent_edge());
let rc_root = Rc::new(root);
let child = WorkItem::child(&rc_root, "b".to_string(), TraceNode::from(b), trace_edge());
assert!(child.has_parent_edge());
}
#[test]
fn build_path_root_returns_single_node_no_edges() {
let g = abc_graph();
let a = g.get_node(&"a".to_string()).unwrap();
let root = WorkItem::new_root("a".to_string(), TraceNode::from(a));
let path = root.build_path();
assert_eq!(path.depth, 0);
assert_eq!(path.nodes.len(), 1);
assert!(path.edges.is_empty());
}
#[test]
fn build_path_reconstructs_root_to_leaf_order() {
let g = abc_graph();
let a = g.get_node(&"a".to_string()).unwrap();
let b = g.get_node(&"b".to_string()).unwrap();
let c = g.get_node(&"c".to_string()).unwrap();
let root = WorkItem::new_root("a".to_string(), TraceNode::from(a));
let rc_root = Rc::new(root);
let mid = WorkItem::child(&rc_root, "b".to_string(), TraceNode::from(b), trace_edge());
let rc_mid = Rc::new(mid);
let leaf = WorkItem::child(&rc_mid, "c".to_string(), TraceNode::from(c), trace_edge());
let path = leaf.build_path();
assert_eq!(path.depth, 2);
assert_eq!(path.nodes.len(), 3);
assert_eq!(path.nodes[0].name, "a");
assert_eq!(path.nodes[1].name, "b");
assert_eq!(path.nodes[2].name, "c");
assert_eq!(path.edges.len(), 2);
assert_eq!(path.edges[0].edge_type, "Calls");
assert_eq!(path.edges[1].edge_type, "Calls");
}
#[test]
fn path_contains_returns_true_for_self() {
let g = abc_graph();
let a = g.get_node(&"a".to_string()).unwrap();
let root = WorkItem::new_root("a".to_string(), TraceNode::from(a));
assert!(root.path_contains(&"a".to_string()));
}
#[test]
fn path_contains_returns_true_for_ancestor() {
let g = abc_graph();
let a = g.get_node(&"a".to_string()).unwrap();
let b = g.get_node(&"b".to_string()).unwrap();
let root = WorkItem::new_root("a".to_string(), TraceNode::from(a));
let rc_root = Rc::new(root);
let mid = WorkItem::child(&rc_root, "b".to_string(), TraceNode::from(b), trace_edge());
assert!(mid.path_contains(&"a".to_string()));
assert!(mid.path_contains(&"b".to_string()));
}
#[test]
fn path_contains_returns_false_for_outside_node() {
let g = abc_graph();
let a = g.get_node(&"a".to_string()).unwrap();
let root = WorkItem::new_root("a".to_string(), TraceNode::from(a));
assert!(!root.path_contains(&"zzz".to_string()));
}
#[test]
fn bfs_trace_start_not_in_graph_returns_empty() {
let g = abc_graph();
let paths = bfs_trace(&g, &"zzz".to_string(), 5, |_| true, None);
assert!(paths.is_empty());
}
#[test]
fn bfs_trace_sink_not_in_graph_returns_empty() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 5, |_| true, Some(&"zzz".to_string()));
assert!(paths.is_empty());
}
#[test]
fn bfs_trace_sink_mode_skips_trivial_start_eq_sink() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 5, |_| true, Some(&"a".to_string()));
assert!(paths.is_empty());
}
#[test]
fn bfs_trace_sink_mode_records_complete_path() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 5, |_| true, Some(&"c".to_string()));
assert_eq!(paths.len(), 1);
let path = &paths[0];
assert_eq!(path.depth, 2);
assert_eq!(path.nodes.len(), 3);
assert_eq!(path.nodes[0].name, "a");
assert_eq!(path.nodes[2].name, "c");
}
#[test]
fn bfs_trace_sink_mode_records_path_at_max_depth_boundary() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 2, |_| true, Some(&"c".to_string()));
assert_eq!(paths.len(), 1);
assert_eq!(paths[0].nodes.last().unwrap().name, "c");
}
#[test]
fn bfs_trace_sink_mode_max_depth_below_sink_returns_empty() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 1, |_| true, Some(&"c".to_string()));
assert!(paths.is_empty());
}
#[test]
fn bfs_trace_non_sink_returns_all_prefixes() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 5, |_| true, None);
assert_eq!(paths.len(), 2);
let depths: Vec<usize> = paths.iter().map(|p| p.depth).collect();
assert!(depths.contains(&1));
assert!(depths.contains(&2));
}
#[test]
fn bfs_trace_edge_filter_skips_unmatched() {
let g = abc_graph();
let paths = bfs_trace(&g, &"a".to_string(), 5, |et| *et != EdgeType::Calls, None);
assert!(paths.is_empty());
}
#[test]
fn bfs_trace_cycle_does_not_revisit_node() {
let mut g = abc_graph();
g.add_edge(Edge::new("c", "a", EdgeType::Calls, "proj"));
let paths = bfs_trace(&g, &"a".to_string(), 5, |_| true, Some(&"c".to_string()));
assert_eq!(paths.len(), 1);
assert_eq!(paths[0].nodes.len(), 3);
}
}