1use std::collections::{BTreeMap, BTreeSet, VecDeque};
2
3use code_system_graph_model::{Edge, Node, NodeId, TraceReport, TraceSegment};
4use thiserror::Error;
5
6#[derive(Debug, Clone)]
8pub struct FederatedGraph {
9 nodes: BTreeMap<NodeId, Node>,
10 outgoing: BTreeMap<NodeId, Vec<Edge>>,
11}
12
13#[derive(Debug, Error, PartialEq, Eq)]
15pub enum TraceError {
16 #[error("edge `{edge}` references missing node `{node}`")]
18 DanglingEdge {
19 edge: String,
21 node: String,
23 },
24 #[error("trace anchor `{0}` was not observed in the selected snapshot")]
26 UnknownAnchor(String),
27 #[error("trace max_depth must be greater than zero")]
29 InvalidDepth,
30 #[error("trace parent chain is inconsistent at node `{0}`")]
32 InconsistentPath(String),
33}
34
35impl FederatedGraph {
36 pub fn new(nodes: Vec<Node>, edges: Vec<Edge>) -> Result<Self, TraceError> {
42 let nodes = nodes
43 .into_iter()
44 .map(|node| (node.id.clone(), node))
45 .collect::<BTreeMap<_, _>>();
46 let mut outgoing: BTreeMap<NodeId, Vec<Edge>> = BTreeMap::new();
47 for edge in edges {
48 for endpoint in [&edge.source, &edge.target] {
49 if !nodes.contains_key(endpoint) {
50 return Err(TraceError::DanglingEdge {
51 edge: edge.id.as_str().to_owned(),
52 node: endpoint.as_str().to_owned(),
53 });
54 }
55 }
56 outgoing.entry(edge.source.clone()).or_default().push(edge);
57 }
58 for adjacent in outgoing.values_mut() {
59 adjacent.sort_by(|left, right| left.id.cmp(&right.id));
60 }
61 Ok(Self { nodes, outgoing })
62 }
63
64 pub fn trace(
73 &self,
74 from: &NodeId,
75 to: &NodeId,
76 max_depth: usize,
77 ) -> Result<TraceReport, TraceError> {
78 if max_depth == 0 {
79 return Err(TraceError::InvalidDepth);
80 }
81 if !self.nodes.contains_key(from) {
82 return Err(TraceError::UnknownAnchor(from.as_str().to_owned()));
83 }
84 if !self.nodes.contains_key(to) {
85 return Err(TraceError::UnknownAnchor(to.as_str().to_owned()));
86 }
87 if from == to {
88 return Ok(TraceReport {
89 segments: Vec::new(),
90 truncated: false,
91 coverage_gaps: Vec::new(),
92 });
93 }
94
95 let mut queue = VecDeque::from([(from.clone(), 0_usize)]);
96 let mut visited = BTreeSet::from([from.clone()]);
97 let mut parents: BTreeMap<NodeId, (NodeId, Edge)> = BTreeMap::new();
98 let mut truncated = false;
99 let mut found = false;
100
101 while let Some((current, depth)) = queue.pop_front() {
102 if depth >= max_depth {
103 if self
104 .outgoing
105 .get(¤t)
106 .is_some_and(|edges| !edges.is_empty())
107 {
108 truncated = true;
109 }
110 continue;
111 }
112 for edge in self.outgoing.get(¤t).into_iter().flatten() {
113 if !visited.insert(edge.target.clone()) {
114 continue;
115 }
116 parents.insert(edge.target.clone(), (current.clone(), edge.clone()));
117 if &edge.target == to {
118 found = true;
119 break;
120 }
121 queue.push_back((edge.target.clone(), depth + 1));
122 }
123 if found {
124 break;
125 }
126 }
127
128 if !found {
129 return Ok(TraceReport {
130 segments: Vec::new(),
131 truncated,
132 coverage_gaps: vec![
133 "No confirmed path was observed within available coverage and bounds."
134 .to_owned(),
135 ],
136 });
137 }
138
139 self.reconstruct(from, to, &parents, truncated)
140 }
141
142 fn reconstruct(
143 &self,
144 from: &NodeId,
145 to: &NodeId,
146 parents: &BTreeMap<NodeId, (NodeId, Edge)>,
147 truncated: bool,
148 ) -> Result<TraceReport, TraceError> {
149 let mut current = to.clone();
150 let mut segments = Vec::new();
151 while ¤t != from {
152 let (parent, edge) = parents
153 .get(¤t)
154 .ok_or_else(|| TraceError::InconsistentPath(current.as_str().to_owned()))?;
155 let source = self
156 .nodes
157 .get(parent)
158 .ok_or_else(|| TraceError::InconsistentPath(parent.as_str().to_owned()))?;
159 let target = self
160 .nodes
161 .get(¤t)
162 .ok_or_else(|| TraceError::InconsistentPath(current.as_str().to_owned()))?;
163 segments.push(TraceSegment {
164 source: source.clone(),
165 edge: edge.clone(),
166 target: target.clone(),
167 });
168 current = parent.clone();
169 }
170 segments.reverse();
171 Ok(TraceReport {
172 segments,
173 truncated,
174 coverage_gaps: Vec::new(),
175 })
176 }
177}
178
179#[cfg(test)]
180mod tests {
181 use code_system_graph_model::{
182 Edge, EdgeId, EdgeKind, EpistemicStatus, Node, NodeId, NodeKind
183 };
184
185 use super::FederatedGraph;
186
187 fn node(id: &str) -> Node {
188 Node {
189 id: NodeId::new(id),
190 kind: NodeKind::HttpOperation,
191 repo_id: None,
192 stable_key: id.to_owned(),
193 label: id.to_owned(),
194 }
195 }
196
197 #[test]
198 fn trace_should_return_deterministic_shortest_path() {
199 let edge = Edge {
200 id: EdgeId::new("edge:1"),
201 source: NodeId::new("node:web"),
202 target: NodeId::new("node:api"),
203 kind: EdgeKind::CallsRemote,
204 confidence: 1.0,
205 status: EpistemicStatus::Confirmed,
206 evidence: Vec::new(),
207 };
208 let graph = FederatedGraph::new(vec![node("node:web"), node("node:api")], vec![edge]);
209 let result = graph
210 .and_then(|value| value.trace(&NodeId::new("node:web"), &NodeId::new("node:api"), 4));
211 let segment_count = result.map(|report| report.segments.len());
212
213 assert_eq!(segment_count, Ok(1));
214 }
215}