1use std::collections::HashMap;
4
5use serde::{Deserialize, Serialize};
6
7use crate::graph::{CodeGraph, EdgeKind, NodeData};
8use crate::output::shard::error::ShardError;
9use crate::output::shard::file::SHARD_SCHEMA_VERSION;
10use crate::output::shard::name::StableNameIndex;
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14pub struct ShardEdge {
15 pub source_name: String,
16 pub target_name: String,
17 pub kind: ShardEdgeKind,
18 pub confidence: f32,
19 pub flow_kind: Option<ShardFlowKind>,
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
23#[serde(rename_all = "snake_case")]
24pub enum ShardEdgeKind {
25 Ownership,
26 Import,
27 Reference,
28 Flow,
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
32#[serde(rename_all = "snake_case")]
33pub enum ShardFlowKind {
34 DefUse,
35 Argument,
36 Return,
37 FieldAccess,
38}
39
40impl From<EdgeKind> for ShardEdgeKind {
41 fn from(kind: EdgeKind) -> Self {
42 match kind {
43 EdgeKind::Ownership => Self::Ownership,
44 EdgeKind::Import => Self::Import,
45 EdgeKind::Reference => Self::Reference,
46 EdgeKind::Flow => Self::Flow,
47 }
48 }
49}
50
51impl From<ShardEdgeKind> for EdgeKind {
52 fn from(kind: ShardEdgeKind) -> Self {
53 match kind {
54 ShardEdgeKind::Ownership => Self::Ownership,
55 ShardEdgeKind::Import => Self::Import,
56 ShardEdgeKind::Reference => Self::Reference,
57 ShardEdgeKind::Flow => Self::Flow,
58 }
59 }
60}
61
62impl From<crate::model::FlowKind> for ShardFlowKind {
63 fn from(kind: crate::model::FlowKind) -> Self {
64 match kind {
65 crate::model::FlowKind::DefUse => Self::DefUse,
66 crate::model::FlowKind::Argument => Self::Argument,
67 crate::model::FlowKind::Return => Self::Return,
68 crate::model::FlowKind::FieldAccess => Self::FieldAccess,
69 }
70 }
71}
72
73impl From<ShardFlowKind> for crate::model::FlowKind {
74 fn from(kind: ShardFlowKind) -> Self {
75 match kind {
76 ShardFlowKind::DefUse => Self::DefUse,
77 ShardFlowKind::Argument => Self::Argument,
78 ShardFlowKind::Return => Self::Return,
79 ShardFlowKind::FieldAccess => Self::FieldAccess,
80 }
81 }
82}
83
84pub(crate) fn validate_edge(
85 edge: &ShardEdge,
86 line: usize,
87 edge_index: usize,
88) -> Result<(), ShardError> {
89 let valid_confidence = edge.confidence.is_finite() && (0.0..=1.0).contains(&edge.confidence);
90 if !valid_confidence {
91 return Err(ShardError::InvalidEdge {
92 line,
93 edge_index,
94 message: "confidence must be finite and in the range 0.0..=1.0".to_string(),
95 });
96 }
97 if edge.kind == ShardEdgeKind::Flow {
98 return Err(ShardError::InvalidEdge {
99 line,
100 edge_index,
101 message: format!(
102 "schema version {SHARD_SCHEMA_VERSION} does not persist dataflow nodes"
103 ),
104 });
105 }
106 if edge.flow_kind.is_some() {
107 return Err(ShardError::InvalidEdge {
108 line,
109 edge_index,
110 message: "non-flow edges forbid flow_kind".to_string(),
111 });
112 }
113 Ok(())
114}
115
116pub fn restore_shard_edges(graph: &mut CodeGraph, edges: &[ShardEdge]) -> Result<(), ShardError> {
118 let names = StableNameIndex::new(graph)?;
119 let endpoint_index: HashMap<String, _> = graph
120 .graph()
121 .node_indices()
122 .filter(|index| !matches!(graph.graph()[*index], NodeData::Data(_)))
123 .filter_map(|index| names.name_of(index).map(|name| (name.to_string(), index)))
124 .collect();
125
126 for (edge_index, edge) in edges.iter().enumerate() {
127 validate_edge(edge, 0, edge_index)?;
128 let source = endpoint_index
129 .get(&edge.source_name)
130 .copied()
131 .ok_or_else(|| ShardError::MissingEndpoint {
132 name: edge.source_name.clone(),
133 })?;
134 let target = endpoint_index
135 .get(&edge.target_name)
136 .copied()
137 .ok_or_else(|| ShardError::MissingEndpoint {
138 name: edge.target_name.clone(),
139 })?;
140 graph.add_edge_normalized_with_flow(
141 source,
142 target,
143 edge.kind.into(),
144 edge.confidence,
145 edge.flow_kind.map(Into::into),
146 );
147 }
148 Ok(())
149}
150
151#[cfg(test)]
152mod tests {
153 use super::*;
154
155 fn flow_edge() -> ShardEdge {
156 ShardEdge {
157 source_name: "python file a.py . f#function!0 .".to_string(),
158 target_name: "python file b.py . g#function!0 .".to_string(),
159 kind: ShardEdgeKind::Flow,
160 confidence: 1.0,
161 flow_kind: None,
162 }
163 }
164
165 #[test]
167 fn dataflow_refusal_names_the_current_schema_version() {
168 let error = validate_edge(&flow_edge(), 1, 0).unwrap_err();
169 let message = error.to_string();
170 assert!(
171 message.contains(&crate::output::shard::file::SHARD_SCHEMA_VERSION.to_string()),
172 "message names the schema version in force: {message}"
173 );
174 }
175}