meta-ast 0.7.0

//! Shard edge serialization, validation, and restoration.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::graph::{CodeGraph, EdgeKind, NodeData};
use crate::output::shard::error::ShardError;
use crate::output::shard::file::SHARD_SCHEMA_VERSION;
use crate::output::shard::name::StableNameIndex;

/// Serialized cross-node edge in a shard file.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ShardEdge {
    pub source_name: String,
    pub target_name: String,
    pub kind: ShardEdgeKind,
    pub confidence: f32,
    pub flow_kind: Option<ShardFlowKind>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ShardEdgeKind {
    Ownership,
    Import,
    Reference,
    Flow,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ShardFlowKind {
    DefUse,
    Argument,
    Return,
    FieldAccess,
}

impl From<EdgeKind> for ShardEdgeKind {
    fn from(kind: EdgeKind) -> Self {
        match kind {
            EdgeKind::Ownership => Self::Ownership,
            EdgeKind::Import => Self::Import,
            EdgeKind::Reference => Self::Reference,
            EdgeKind::Flow => Self::Flow,
        }
    }
}

impl From<ShardEdgeKind> for EdgeKind {
    fn from(kind: ShardEdgeKind) -> Self {
        match kind {
            ShardEdgeKind::Ownership => Self::Ownership,
            ShardEdgeKind::Import => Self::Import,
            ShardEdgeKind::Reference => Self::Reference,
            ShardEdgeKind::Flow => Self::Flow,
        }
    }
}

impl From<crate::model::FlowKind> for ShardFlowKind {
    fn from(kind: crate::model::FlowKind) -> Self {
        match kind {
            crate::model::FlowKind::DefUse => Self::DefUse,
            crate::model::FlowKind::Argument => Self::Argument,
            crate::model::FlowKind::Return => Self::Return,
            crate::model::FlowKind::FieldAccess => Self::FieldAccess,
        }
    }
}

impl From<ShardFlowKind> for crate::model::FlowKind {
    fn from(kind: ShardFlowKind) -> Self {
        match kind {
            ShardFlowKind::DefUse => Self::DefUse,
            ShardFlowKind::Argument => Self::Argument,
            ShardFlowKind::Return => Self::Return,
            ShardFlowKind::FieldAccess => Self::FieldAccess,
        }
    }
}

pub(crate) fn validate_edge(
    edge: &ShardEdge,
    line: usize,
    edge_index: usize,
) -> Result<(), ShardError> {
    let valid_confidence = edge.confidence.is_finite() && (0.0..=1.0).contains(&edge.confidence);
    if !valid_confidence {
        return Err(ShardError::InvalidEdge {
            line,
            edge_index,
            message: "confidence must be finite and in the range 0.0..=1.0".to_string(),
        });
    }
    if edge.kind == ShardEdgeKind::Flow {
        return Err(ShardError::InvalidEdge {
            line,
            edge_index,
            message: format!(
                "schema version {SHARD_SCHEMA_VERSION} does not persist dataflow nodes"
            ),
        });
    }
    if edge.flow_kind.is_some() {
        return Err(ShardError::InvalidEdge {
            line,
            edge_index,
            message: "non-flow edges forbid flow_kind".to_string(),
        });
    }
    Ok(())
}

/// Restore persisted edges after `GraphBuilder::from_extractions` regenerates graph nodes.
pub fn restore_shard_edges(graph: &mut CodeGraph, edges: &[ShardEdge]) -> Result<(), ShardError> {
    let names = StableNameIndex::new(graph)?;
    let endpoint_index: HashMap<String, _> = graph
        .graph()
        .node_indices()
        .filter(|index| !matches!(graph.graph()[*index], NodeData::Data(_)))
        .filter_map(|index| names.name_of(index).map(|name| (name.to_string(), index)))
        .collect();

    for (edge_index, edge) in edges.iter().enumerate() {
        validate_edge(edge, 0, edge_index)?;
        let source = endpoint_index
            .get(&edge.source_name)
            .copied()
            .ok_or_else(|| ShardError::MissingEndpoint {
                name: edge.source_name.clone(),
            })?;
        let target = endpoint_index
            .get(&edge.target_name)
            .copied()
            .ok_or_else(|| ShardError::MissingEndpoint {
                name: edge.target_name.clone(),
            })?;
        graph.add_edge_normalized_with_flow(
            source,
            target,
            edge.kind.into(),
            edge.confidence,
            edge.flow_kind.map(Into::into),
        );
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn flow_edge() -> ShardEdge {
        ShardEdge {
            source_name: "python file a.py . f#function!0 .".to_string(),
            target_name: "python file b.py . g#function!0 .".to_string(),
            kind: ShardEdgeKind::Flow,
            confidence: 1.0,
            flow_kind: None,
        }
    }

    /// The refusal message must name the version that actually applies.
    #[test]
    fn dataflow_refusal_names_the_current_schema_version() {
        let error = validate_edge(&flow_edge(), 1, 0).unwrap_err();
        let message = error.to_string();
        assert!(
            message.contains(&crate::output::shard::file::SHARD_SCHEMA_VERSION.to_string()),
            "message names the schema version in force: {message}"
        );
    }
}