Skip to main content

meta_ast/output/shard/
edge.rs

1//! Shard edge serialization, validation, and restoration.
2
3use 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/// Serialized cross-node edge in a shard file.
13#[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
116/// Restore persisted edges after `GraphBuilder::from_extractions` regenerates graph nodes.
117pub 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    /// The refusal message must name the version that actually applies.
166    #[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}