1use std::collections::BTreeMap;
2
3use code_system_graph_model::{
4 Edge, EdgeId, EdgeKind, EpistemicStatus, Evidence, EvidenceId, Node, NodeId, NodeKind, Provenance, RepoId, stable_id
5};
6
7use crate::{ProtoGeneratedMarker, ProtoGeneratedRole, ProtoRpcMethod, ProtobufDocument};
8
9#[derive(Debug, Clone, Default)]
11pub struct ProtobufGraphFacts {
12 pub nodes: Vec<Node>,
14 pub edges: Vec<Edge>,
16 pub evidence: Vec<Evidence>,
18}
19
20struct RpcProvider {
21 service: String,
22 method: String,
23 node: Node,
24 evidence: Evidence,
25}
26
27#[must_use]
32pub fn protobuf_documents_to_graph(
33 inputs: &[(&RepoId, &str, &str, &ProtobufDocument)],
34) -> ProtobufGraphFacts {
35 let mut result = ProtobufGraphFacts::default();
36 let mut providers = Vec::new();
37 let mut markers = Vec::new();
38 for (repo_id, source_path, content_hash, document) in inputs {
39 let repository = repository_node(repo_id);
40 let artifact = artifact_node(repo_id, source_path);
41 let artifact_evidence =
42 protobuf_evidence(repo_id, source_path, content_hash, 1, "protobuf artifact");
43 result.edges.push(edge(
44 &repository.id,
45 &artifact.id,
46 EdgeKind::Contains,
47 vec![artifact_evidence.id.clone()],
48 ));
49 result.nodes.extend([repository, artifact.clone()]);
50 result.evidence.push(artifact_evidence);
51 match document {
52 ProtobufDocument::File(file) => {
53 for message in &file.messages {
54 let stable_key = format!(
55 "protobuf-message:{}:{}",
56 repo_id.as_str(),
57 message.full_name
58 );
59 let node = Node {
60 id: NodeId::new(stable_id("node", &stable_key)),
61 kind: NodeKind::Artifact,
62 repo_id: Some((*repo_id).clone()),
63 stable_key,
64 label: message.full_name.clone(),
65 };
66 let evidence = protobuf_evidence(
67 repo_id,
68 source_path,
69 content_hash,
70 message.line,
71 "protobuf message",
72 );
73 result.edges.push(edge(
74 &artifact.id,
75 &node.id,
76 EdgeKind::Contains,
77 vec![evidence.id.clone()],
78 ));
79 result.nodes.push(node);
80 result.evidence.push(evidence);
81 }
82 for service in &file.services {
83 for method in &service.methods {
84 providers.push(rpc_provider(
85 repo_id,
86 source_path,
87 content_hash,
88 &service.full_name,
89 method,
90 ));
91 }
92 }
93 }
94 ProtobufDocument::Generated(generated) => {
95 markers.extend(generated.iter().map(|marker| {
96 (
97 (*repo_id).clone(),
98 (*source_path).to_owned(),
99 (*content_hash).to_owned(),
100 marker.clone(),
101 )
102 }));
103 }
104 }
105 }
106 link_markers(&providers, &markers, &mut result);
107 for provider in providers {
108 result.nodes.push(provider.node);
109 result.evidence.push(provider.evidence);
110 }
111 finish(&mut result);
112 result
113}
114
115fn rpc_provider(
116 repo_id: &RepoId,
117 source_path: &str,
118 content_hash: &str,
119 service: &str,
120 method: &ProtoRpcMethod,
121) -> RpcProvider {
122 let stable_key = format!(
123 "rpc:{}:provider:{service}/{}",
124 repo_id.as_str(),
125 method.name
126 );
127 RpcProvider {
128 service: service.to_owned(),
129 method: method.name.clone(),
130 node: Node {
131 id: NodeId::new(stable_id("node", &stable_key)),
132 kind: NodeKind::RpcMethod,
133 repo_id: Some(repo_id.clone()),
134 stable_key,
135 label: format!("{service}/{}", method.name),
136 },
137 evidence: protobuf_evidence(
138 repo_id,
139 source_path,
140 content_hash,
141 method.line,
142 "protobuf RPC method",
143 ),
144 }
145}
146
147fn link_markers(
148 providers: &[RpcProvider],
149 markers: &[(RepoId, String, String, ProtoGeneratedMarker)],
150 result: &mut ProtobufGraphFacts,
151) {
152 let mut candidates: BTreeMap<(&str, &str), Vec<&RpcProvider>> = BTreeMap::new();
153 for provider in providers {
154 candidates
155 .entry((provider.service.as_str(), provider.method.as_str()))
156 .or_default()
157 .push(provider);
158 }
159 for (repo_id, source_path, content_hash, marker) in markers {
160 let Some(matches) = candidates.get(&(marker.service.as_str(), marker.method.as_str()))
161 else {
162 continue;
163 };
164 if matches.len() != 1 {
165 continue;
166 }
167 let provider = matches[0];
168 let source = generated_node(repo_id, source_path, marker);
169 let source_evidence = protobuf_evidence(
170 repo_id,
171 source_path,
172 content_hash,
173 marker.line,
174 "generated gRPC method marker",
175 );
176 match marker.role {
177 ProtoGeneratedRole::Client => result.edges.push(edge(
178 &source.id,
179 &provider.node.id,
180 EdgeKind::CallsRemote,
181 vec![source_evidence.id.clone(), provider.evidence.id.clone()],
182 )),
183 ProtoGeneratedRole::Server => result.edges.push(edge(
184 &provider.node.id,
185 &source.id,
186 EdgeKind::ImplementedBy,
187 vec![provider.evidence.id.clone(), source_evidence.id.clone()],
188 )),
189 ProtoGeneratedRole::Unknown => {}
190 }
191 result.nodes.push(source);
192 result.evidence.push(source_evidence);
193 }
194}
195
196fn generated_node(repo_id: &RepoId, source_path: &str, marker: &ProtoGeneratedMarker) -> Node {
197 let stable_key = format!(
198 "rpc:{}:generated:{source_path}:{}/{}:{:?}",
199 repo_id.as_str(),
200 marker.service,
201 marker.method,
202 marker.role
203 );
204 Node {
205 id: NodeId::new(stable_id("node", &stable_key)),
206 kind: NodeKind::SymbolRef,
207 repo_id: Some(repo_id.clone()),
208 stable_key,
209 label: format!("{}/{}", marker.service, marker.method),
210 }
211}
212
213fn artifact_node(repo_id: &RepoId, source_path: &str) -> Node {
214 let stable_key = format!("protobuf-artifact:{}:{source_path}", repo_id.as_str());
215 Node {
216 id: NodeId::new(stable_id("node", &stable_key)),
217 kind: NodeKind::Artifact,
218 repo_id: Some(repo_id.clone()),
219 stable_key,
220 label: source_path.to_owned(),
221 }
222}
223
224fn repository_node(repo_id: &RepoId) -> Node {
225 let stable_key = format!("repository:{}", repo_id.as_str());
226 Node {
227 id: NodeId::new(stable_id("node", &stable_key)),
228 kind: NodeKind::Repository,
229 repo_id: Some(repo_id.clone()),
230 stable_key,
231 label: repo_id.as_str().to_owned(),
232 }
233}
234
235fn protobuf_evidence(
236 repo_id: &RepoId,
237 source_path: &str,
238 content_hash: &str,
239 line: u32,
240 note: &str,
241) -> Evidence {
242 let key = format!(
243 "{}:{source_path}:{line}:{note}:{content_hash}",
244 repo_id.as_str()
245 );
246 Evidence {
247 id: EvidenceId::new(stable_id("evidence", &key)),
248 repo_id: Some(repo_id.clone()),
249 file_path: Some(source_path.to_owned()),
250 start_line: Some(line),
251 end_line: Some(line),
252 extractor: "code-system-graph.protobuf".to_owned(),
253 extractor_version: "1.0.0".to_owned(),
254 provenance: Provenance::Extracted,
255 confidence: 1.0,
256 observed_at_commit: None,
257 content_hash: Some(content_hash.to_owned()),
258 note: Some(note.to_owned()),
259 }
260}
261
262fn edge(source: &NodeId, target: &NodeId, kind: EdgeKind, evidence: Vec<EvidenceId>) -> Edge {
263 let key = format!("{}:{kind:?}:{}", source.as_str(), target.as_str());
264 Edge {
265 id: EdgeId::new(stable_id("edge", &key)),
266 source: source.clone(),
267 target: target.clone(),
268 kind,
269 confidence: 1.0,
270 status: EpistemicStatus::Confirmed,
271 evidence,
272 }
273}
274
275fn finish(result: &mut ProtobufGraphFacts) {
276 result.nodes.sort_by(|left, right| left.id.cmp(&right.id));
277 result.nodes.dedup_by(|left, right| left.id == right.id);
278 result.edges.sort_by(|left, right| left.id.cmp(&right.id));
279 result.edges.dedup_by(|left, right| left.id == right.id);
280 result
281 .evidence
282 .sort_by(|left, right| left.id.cmp(&right.id));
283 result.evidence.dedup_by(|left, right| left.id == right.id);
284}
285
286#[cfg(test)]
287mod tests {
288 use code_system_graph_model::{EdgeKind, RepoId};
289
290 use super::protobuf_documents_to_graph;
291 use crate::{
292 ProtoGeneratedRole, ProtobufDocument, SourceLanguage, extract_protobuf, parse_protobuf_generated_source
293 };
294
295 #[test]
296 fn protobuf_graph_should_link_generated_client_to_exact_method() {
297 let contract = extract_protobuf(
298 "orders.proto",
299 r#"
300syntax = "proto3";
301package commerce.orders.v1;
302message Request {}
303message Reply {}
304service Orders { rpc GetOrder(Request) returns (Reply); }
305"#,
306 );
307 let markers = parse_protobuf_generated_source(
308 SourceLanguage::Python,
309 "orders_pb2_grpc.py",
310 "# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!\nchannel.unary_unary(\"/commerce.orders.v1.Orders/GetOrder\")",
311 );
312 assert_eq!(
313 markers.first().map(|marker| marker.role),
314 Some(ProtoGeneratedRole::Client)
315 );
316 let api = RepoId::new("repo:api");
317 let worker = RepoId::new("repo:worker");
318 let facts = contract.as_ref().ok().map(|contract| {
319 let contract = ProtobufDocument::File(Box::new(contract.clone()));
320 let generated = ProtobufDocument::Generated(markers);
321 protobuf_documents_to_graph(&[
322 (&api, "orders.proto", "contract", &contract),
323 (&worker, "orders_pb2_grpc.py", "generated", &generated),
324 ])
325 });
326
327 assert!(matches!(
328 facts,
329 Some(facts)
330 if facts
331 .edges
332 .iter()
333 .any(|edge| edge.kind == EdgeKind::CallsRemote)
334 ));
335 }
336}