Skip to main content

weavatrix_graph/
graph.rs

1use crate::{Edge, GraphError, Node, NodeId, Result, SourceSpan};
2use serde::{Deserialize, Deserializer, Serialize, de::Error as _};
3use std::collections::{BTreeMap, BTreeSet};
4
5#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
6pub struct Graph {
7    nodes: Vec<Node>,
8    edges: Vec<Edge>,
9    #[serde(skip)]
10    outgoing_index: BTreeMap<NodeId, Vec<usize>>,
11    #[serde(skip)]
12    incoming_index: BTreeMap<NodeId, Vec<usize>>,
13}
14
15impl Graph {
16    /// Creates a validated graph and canonicalizes its ordering.
17    ///
18    /// # Errors
19    ///
20    /// Returns an error for conflicting nodes, dangling edges, empty extractor
21    /// identities, or invalid source spans.
22    pub fn try_from_parts(
23        nodes: impl IntoIterator<Item = Node>,
24        edges: impl IntoIterator<Item = Edge>,
25    ) -> Result<Self> {
26        let mut builder = GraphBuilder::new();
27        for node in nodes {
28            builder.add_node(node)?;
29        }
30        for edge in edges {
31            builder.add_edge(edge)?;
32        }
33        builder.build()
34    }
35
36    #[must_use]
37    pub fn nodes(&self) -> &[Node] {
38        &self.nodes
39    }
40
41    #[must_use]
42    pub fn edges(&self) -> &[Edge] {
43        &self.edges
44    }
45
46    #[must_use]
47    pub fn node(&self, id: &str) -> Option<&Node> {
48        self.nodes
49            .binary_search_by(|node| node.id.as_str().cmp(id))
50            .ok()
51            .map(|index| &self.nodes[index])
52    }
53
54    pub fn outgoing<'graph>(
55        &'graph self,
56        id: &'graph NodeId,
57    ) -> impl Iterator<Item = &'graph Edge> {
58        self.outgoing_index
59            .get(id)
60            .into_iter()
61            .flatten()
62            .map(|index| &self.edges[*index])
63    }
64
65    pub fn incoming<'graph>(
66        &'graph self,
67        id: &'graph NodeId,
68    ) -> impl Iterator<Item = &'graph Edge> {
69        self.incoming_index
70            .get(id)
71            .into_iter()
72            .flatten()
73            .map(|index| &self.edges[*index])
74    }
75
76    #[must_use]
77    pub const fn node_count(&self) -> usize {
78        self.nodes.len()
79    }
80
81    #[must_use]
82    pub const fn edge_count(&self) -> usize {
83        self.edges.len()
84    }
85
86    #[must_use]
87    pub const fn is_empty(&self) -> bool {
88        self.nodes.is_empty() && self.edges.is_empty()
89    }
90
91    #[must_use]
92    pub fn into_parts(self) -> (Vec<Node>, Vec<Edge>) {
93        (self.nodes, self.edges)
94    }
95}
96
97#[derive(Deserialize)]
98struct GraphWire {
99    nodes: Vec<Node>,
100    edges: Vec<Edge>,
101}
102
103impl<'de> Deserialize<'de> for Graph {
104    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
105    where
106        D: Deserializer<'de>,
107    {
108        let wire = GraphWire::deserialize(deserializer)?;
109        Self::try_from_parts(wire.nodes, wire.edges).map_err(D::Error::custom)
110    }
111}
112
113#[derive(Debug, Default)]
114pub struct GraphBuilder {
115    nodes: BTreeMap<NodeId, Node>,
116    edges: BTreeSet<Edge>,
117}
118
119impl GraphBuilder {
120    #[must_use]
121    pub const fn new() -> Self {
122        Self {
123            nodes: BTreeMap::new(),
124            edges: BTreeSet::new(),
125        }
126    }
127
128    /// Adds a node idempotently.
129    ///
130    /// # Errors
131    ///
132    /// Returns an error when the same identifier already has a different
133    /// definition or the node contains an invalid source span.
134    pub fn add_node(&mut self, node: Node) -> Result<&mut Self> {
135        if let Some(span) = &node.span {
136            validate_span(span)?;
137        }
138        if let Some(language) = &node.language {
139            validate_language(language)?;
140        }
141        if let Some(existing) = self.nodes.get(&node.id) {
142            if existing == &node {
143                return Ok(self);
144            }
145            return Err(GraphError::ConflictingNode {
146                id: node.id.to_string(),
147            });
148        }
149        self.nodes.insert(node.id.clone(), node);
150        Ok(self)
151    }
152
153    /// Adds an edge idempotently. Endpoint existence is validated by `build`,
154    /// so callers may insert edges before nodes.
155    ///
156    /// # Errors
157    ///
158    /// Returns an error when provenance or its source span is invalid.
159    pub fn add_edge(&mut self, edge: Edge) -> Result<&mut Self> {
160        if edge.provenance.extractor.is_empty() {
161            return Err(GraphError::EmptyExtractor);
162        }
163        if let Some(span) = &edge.provenance.span {
164            validate_span(span)?;
165        }
166        self.edges.insert(edge);
167        Ok(self)
168    }
169
170    /// Validates all endpoints and returns an immutable graph.
171    ///
172    /// # Errors
173    ///
174    /// Returns an error when an edge references a missing source or target.
175    pub fn build(self) -> Result<Graph> {
176        for edge in &self.edges {
177            if !self.nodes.contains_key(&edge.source) {
178                return Err(GraphError::MissingEdgeSource {
179                    id: edge.source.to_string(),
180                });
181            }
182            if !self.nodes.contains_key(&edge.target) {
183                return Err(GraphError::MissingEdgeTarget {
184                    id: edge.target.to_string(),
185                });
186            }
187        }
188        let edges = self.edges.into_iter().collect::<Vec<_>>();
189        let (outgoing_index, incoming_index) = build_edge_indexes(&edges);
190        Ok(Graph {
191            nodes: self.nodes.into_values().collect(),
192            edges,
193            outgoing_index,
194            incoming_index,
195        })
196    }
197}
198
199fn build_edge_indexes(
200    edges: &[Edge],
201) -> (BTreeMap<NodeId, Vec<usize>>, BTreeMap<NodeId, Vec<usize>>) {
202    let mut outgoing = BTreeMap::<NodeId, Vec<usize>>::new();
203    let mut incoming = BTreeMap::<NodeId, Vec<usize>>::new();
204    for (index, edge) in edges.iter().enumerate() {
205        outgoing.entry(edge.source.clone()).or_default().push(index);
206        incoming.entry(edge.target.clone()).or_default().push(index);
207    }
208    (outgoing, incoming)
209}
210
211fn validate_language(language: &str) -> Result<()> {
212    if language.is_empty() || language.trim() != language {
213        return Err(GraphError::InvalidKind {
214            category: "language",
215            value: language.to_owned(),
216        });
217    }
218    Ok(())
219}
220
221fn validate_span(span: &SourceSpan) -> Result<()> {
222    if span.file.is_empty() {
223        return Err(GraphError::InvalidSpan {
224            file: span.file.clone(),
225            reason: "file must not be empty",
226        });
227    }
228    if span.start.line == 0 || span.start.column == 0 || span.end.line == 0 || span.end.column == 0
229    {
230        return Err(GraphError::InvalidSpan {
231            file: span.file.clone(),
232            reason: "positions are one-based",
233        });
234    }
235    if span.end < span.start {
236        return Err(GraphError::InvalidSpan {
237            file: span.file.clone(),
238            reason: "end precedes start",
239        });
240    }
241    Ok(())
242}