Skip to main content

code_system_graph_core/
protobuf_graph.rs

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/// Graph facts assembled from protobuf files and exact generated gRPC markers.
10#[derive(Debug, Clone, Default)]
11pub struct ProtobufGraphFacts {
12    /// RPC methods, generated symbols, messages, artifacts, and repositories.
13    pub nodes: Vec<Node>,
14    /// Exact containment, remote call, and implementation relationships.
15    pub edges: Vec<Edge>,
16    /// Direct protobuf and generated-source evidence.
17    pub evidence: Vec<Evidence>,
18}
19
20struct RpcProvider {
21    service: String,
22    method: String,
23    node: Node,
24    evidence: Evidence,
25}
26
27/// Converts protobuf contracts and generated markers into ambiguity-safe graph facts.
28///
29/// Generated clients and servers link only when one provider defines the exact
30/// package-qualified service and method.
31#[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}