Skip to main content

tensor4all_treetn/
named_graph.rs

1//! Named graph wrapper inspired by NamedGraphs.jl
2//! (<https://github.com/mtfishman/NamedGraphs.jl>)
3//!
4//! Provides a mapping between arbitrary node name types (NodeName) and internal NodeIndex.
5//! This allows using meaningful identifiers (coordinates, strings, etc.) instead of raw indices.
6
7use petgraph::stable_graph::{EdgeIndex, NodeIndex, StableGraph};
8use petgraph::EdgeType;
9use petgraph::Undirected;
10use std::collections::{HashMap, HashSet};
11use std::fmt::Debug;
12use std::hash::Hash;
13
14/// Generic named graph wrapper (inspired by NamedGraphs.jl)
15///
16/// Provides a mapping between arbitrary node name types (NodeName) and internal NodeIndex.
17/// This allows using meaningful identifiers (coordinates, strings, etc.) instead of raw indices.
18///
19/// # Type Parameters
20/// - `NodeName`: Node name type (must be Clone, Hash, Eq, Send, Sync, Debug)
21/// - `NodeData`: Node weight/data type
22/// - `EdgeData`: Edge weight/data type
23/// - `Ty`: Edge type (Directed or Undirected, default: Undirected)
24pub struct NamedGraph<NodeName, NodeData, EdgeData, Ty = Undirected>
25where
26    NodeName: Clone + Hash + Eq + Send + Sync + Debug,
27    Ty: EdgeType,
28{
29    /// Internal graph structure (uses NodeIndex)
30    graph: StableGraph<NodeData, EdgeData, Ty>,
31
32    /// Mapping: node name (NodeName) -> NodeIndex
33    node_name_to_index: HashMap<NodeName, NodeIndex>,
34
35    /// Reverse mapping: NodeIndex -> node name (NodeName)
36    index_to_node_name: HashMap<NodeIndex, NodeName>,
37}
38
39impl<NodeName, NodeData, EdgeData, Ty> NamedGraph<NodeName, NodeData, EdgeData, Ty>
40where
41    NodeName: Clone + Hash + Eq + Send + Sync + Debug,
42    Ty: EdgeType,
43{
44    /// Create a new empty NamedGraph.
45    pub fn new() -> Self {
46        Self {
47            graph: StableGraph::with_capacity(0, 0),
48            node_name_to_index: HashMap::new(),
49            index_to_node_name: HashMap::new(),
50        }
51    }
52
53    /// Create a new NamedGraph with initial capacity.
54    pub fn with_capacity(nodes: usize, edges: usize) -> Self {
55        Self {
56            graph: StableGraph::with_capacity(nodes, edges),
57            node_name_to_index: HashMap::with_capacity(nodes),
58            index_to_node_name: HashMap::with_capacity(nodes),
59        }
60    }
61
62    /// Add a node with the given name and data.
63    ///
64    /// Returns an error if the node already exists.
65    pub fn add_node(&mut self, node_name: NodeName, data: NodeData) -> Result<NodeIndex, String> {
66        if self.node_name_to_index.contains_key(&node_name) {
67            return Err(format!("Node already exists: {:?}", node_name));
68        }
69        let node = self.graph.add_node(data);
70        self.node_name_to_index.insert(node_name.clone(), node);
71        self.index_to_node_name.insert(node, node_name);
72        Ok(node)
73    }
74
75    /// Check if a node exists.
76    pub fn has_node(&self, node_name: &NodeName) -> bool {
77        self.node_name_to_index.contains_key(node_name)
78    }
79
80    /// Get the NodeIndex for a node name.
81    pub fn node_index(&self, node_name: &NodeName) -> Option<NodeIndex> {
82        self.node_name_to_index.get(node_name).copied()
83    }
84
85    /// Get the node name for a NodeIndex.
86    pub fn node_name(&self, node: NodeIndex) -> Option<&NodeName> {
87        self.index_to_node_name.get(&node)
88    }
89
90    /// Rename an existing node without changing its data or incident edges.
91    ///
92    /// Returns an error if the old node doesn't exist or the new name is already in use.
93    pub fn rename_node(&mut self, old_name: &NodeName, new_name: NodeName) -> Result<(), String> {
94        if old_name == &new_name {
95            return Ok(());
96        }
97        if self.node_name_to_index.contains_key(&new_name) {
98            return Err(format!("Node already exists: {:?}", new_name));
99        }
100
101        let node_idx = self
102            .node_name_to_index
103            .remove(old_name)
104            .ok_or_else(|| format!("Node not found: {:?}", old_name))?;
105
106        self.node_name_to_index.insert(new_name.clone(), node_idx);
107        self.index_to_node_name.insert(node_idx, new_name);
108        Ok(())
109    }
110
111    /// Get a reference to the data of a node (by node name).
112    pub fn node_data(&self, node_name: &NodeName) -> Option<&NodeData> {
113        self.node_name_to_index
114            .get(node_name)
115            .and_then(|node| self.graph.node_weight(*node))
116    }
117
118    /// Get a mutable reference to the data of a node (by node name).
119    pub fn node_data_mut(&mut self, node_name: &NodeName) -> Option<&mut NodeData> {
120        self.node_name_to_index
121            .get(node_name)
122            .and_then(|node| self.graph.node_weight_mut(*node))
123    }
124
125    /// Get a reference to the data of a node (by NodeIndex).
126    pub fn node_weight(&self, node: NodeIndex) -> Option<&NodeData> {
127        self.graph.node_weight(node)
128    }
129
130    /// Get a mutable reference to the data of a node (by NodeIndex).
131    pub fn node_weight_mut(&mut self, node: NodeIndex) -> Option<&mut NodeData> {
132        self.graph.node_weight_mut(node)
133    }
134
135    /// Add an edge between two nodes.
136    ///
137    /// Returns an error if either node doesn't exist.
138    pub fn add_edge(
139        &mut self,
140        n1: &NodeName,
141        n2: &NodeName,
142        weight: EdgeData,
143    ) -> Result<EdgeIndex, String>
144    where
145        EdgeData: Clone,
146    {
147        let node1 = self
148            .node_name_to_index
149            .get(n1)
150            .ok_or_else(|| format!("Node not found: {:?}", n1))?;
151        let node2 = self
152            .node_name_to_index
153            .get(n2)
154            .ok_or_else(|| format!("Node not found: {:?}", n2))?;
155        Ok(self.graph.add_edge(*node1, *node2, weight))
156    }
157
158    /// Get the weight of an edge between two nodes.
159    pub fn edge_weight(&self, n1: &NodeName, n2: &NodeName) -> Option<&EdgeData> {
160        let node1 = self.node_name_to_index.get(n1)?;
161        let node2 = self.node_name_to_index.get(n2)?;
162        self.graph
163            .find_edge(*node1, *node2)
164            .and_then(|edge| self.graph.edge_weight(edge))
165    }
166
167    /// Get a mutable reference to the weight of an edge between two nodes.
168    pub fn edge_weight_mut(&mut self, n1: &NodeName, n2: &NodeName) -> Option<&mut EdgeData> {
169        let node1 = self.node_name_to_index.get(n1)?;
170        let node2 = self.node_name_to_index.get(n2)?;
171        self.graph
172            .find_edge(*node1, *node2)
173            .and_then(|edge| self.graph.edge_weight_mut(edge))
174    }
175
176    /// Get all neighbors of a node.
177    pub fn neighbors(&self, node_name: &NodeName) -> Vec<&NodeName> {
178        self.node_name_to_index
179            .get(node_name)
180            .map(|node| {
181                self.graph
182                    .neighbors(*node)
183                    .filter_map(|n| self.index_to_node_name.get(&n))
184                    .collect()
185            })
186            .unwrap_or_default()
187    }
188
189    /// Get all node names.
190    pub fn node_names(&self) -> Vec<&NodeName> {
191        self.node_name_to_index.keys().collect()
192    }
193
194    /// Get the number of nodes.
195    pub fn node_count(&self) -> usize {
196        self.node_name_to_index.len()
197    }
198
199    /// Get the number of edges.
200    pub fn edge_count(&self) -> usize {
201        self.graph.edge_count()
202    }
203
204    /// Remove a node and all its edges.
205    ///
206    /// Returns the node data if the node existed.
207    pub fn remove_node(&mut self, node_name: &NodeName) -> Option<NodeData> {
208        let node = self.node_name_to_index.remove(node_name)?;
209        self.index_to_node_name.remove(&node);
210        self.graph.remove_node(node)
211    }
212
213    /// Remove an edge between two nodes.
214    ///
215    /// Returns the edge weight if the edge existed.
216    pub fn remove_edge(&mut self, n1: &NodeName, n2: &NodeName) -> Option<EdgeData> {
217        let node1 = self.node_name_to_index.get(n1)?;
218        let node2 = self.node_name_to_index.get(n2)?;
219        self.graph
220            .find_edge(*node1, *node2)
221            .and_then(|edge| self.graph.remove_edge(edge))
222    }
223
224    /// Check if a node exists in the internal graph.
225    pub fn contains_node(&self, node: NodeIndex) -> bool {
226        self.graph.contains_node(node)
227    }
228
229    /// Get a reference to the internal graph.
230    ///
231    /// This allows direct access to petgraph algorithms that work with NodeIndex.
232    pub fn graph(&self) -> &StableGraph<NodeData, EdgeData, Ty> {
233        &self.graph
234    }
235
236    /// Get a mutable reference to the internal graph.
237    ///
238    /// **Warning**: Directly modifying the internal graph can break the node-name-to-index mapping.
239    /// Use the provided methods instead.
240    pub fn graph_mut(&mut self) -> &mut StableGraph<NodeData, EdgeData, Ty> {
241        &mut self.graph
242    }
243
244    /// Perform an Euler tour traversal starting from the given root node.
245    ///
246    /// The Euler tour visits each edge exactly twice (once in each direction),
247    /// forming a closed walk that covers all edges. This is useful for sweep
248    /// algorithms where we need to visit all nodes while maintaining the
249    /// canonical center at a single point.
250    ///
251    /// # Algorithm
252    /// Uses DFS with backtracking. When visiting a node:
253    /// 1. For each unvisited neighbor, record the edge and recurse
254    /// 2. When backtracking, record the edge back to the parent
255    ///
256    /// # Returns
257    /// A vector of directed edges `(from, to)` in traversal order.
258    /// Each undirected edge appears twice: once as (u, v) and once as (v, u).
259    ///
260    /// # Example
261    /// For a Y-shaped tree with root at center C:
262    /// ```text
263    ///     A
264    ///     |
265    /// B - C - D
266    /// ```
267    /// Starting from C, the tour might be:
268    /// `[(C,A), (A,C), (C,B), (B,C), (C,D), (D,C)]`
269    pub fn euler_tour_edges(&self, root: &NodeName) -> Option<Vec<(NodeIndex, NodeIndex)>> {
270        let root_idx = self.node_index(root)?;
271        Some(self.euler_tour_edges_by_index(root_idx))
272    }
273
274    /// Perform an Euler tour traversal starting from the given root NodeIndex.
275    ///
276    /// See [`Self::euler_tour_edges`] for details.
277    pub fn euler_tour_edges_by_index(&self, root: NodeIndex) -> Vec<(NodeIndex, NodeIndex)> {
278        let mut visited_edges: HashSet<(NodeIndex, NodeIndex)> = HashSet::new();
279        let mut tour = Vec::new();
280        let mut stack = vec![root];
281
282        while let Some(u) = stack.last().copied() {
283            let mut pushed = false;
284
285            // Try to find an unvisited neighbor
286            for v in self.graph.neighbors(u) {
287                if !visited_edges.contains(&(u, v)) {
288                    // Mark both directions as visited (undirected edge)
289                    visited_edges.insert((u, v));
290                    visited_edges.insert((v, u));
291                    // Record the forward edge
292                    tour.push((u, v));
293                    // Push neighbor onto stack
294                    stack.push(v);
295                    pushed = true;
296                    break; // Handle one neighbor at a time
297                }
298            }
299
300            if !pushed {
301                // No unvisited neighbors, backtrack
302                stack.pop();
303                if let Some(&parent) = stack.last() {
304                    // Record the backtracking edge
305                    tour.push((u, parent));
306                }
307            }
308        }
309
310        tour
311    }
312
313    /// Perform an Euler tour traversal and return the vertex sequence.
314    ///
315    /// This returns the sequence of vertices visited, including repeated visits.
316    /// Each vertex appears multiple times based on its degree in the tree.
317    ///
318    /// # Returns
319    /// A vector of NodeIndex in visitation order, or None if root doesn't exist.
320    ///
321    /// # Example
322    /// For a chain A - B - C starting from B:
323    /// - Edges: `[(B,A), (A,B), (B,C), (C,B)]`
324    /// - Vertices: `[B, A, B, C, B]`
325    pub fn euler_tour_vertices(&self, root: &NodeName) -> Option<Vec<NodeIndex>> {
326        let root_idx = self.node_index(root)?;
327        Some(self.euler_tour_vertices_by_index(root_idx))
328    }
329
330    /// Perform an Euler tour traversal and return the vertex sequence by NodeIndex.
331    ///
332    /// See [`Self::euler_tour_vertices`] for details.
333    pub fn euler_tour_vertices_by_index(&self, root: NodeIndex) -> Vec<NodeIndex> {
334        let edges = self.euler_tour_edges_by_index(root);
335        if edges.is_empty() {
336            // Single node case
337            return vec![root];
338        }
339        // Start with the first vertex, then append destinations
340        let mut vertices = vec![edges[0].0];
341        for (_, to) in &edges {
342            vertices.push(*to);
343        }
344        vertices
345    }
346}
347
348impl<NodeName, NodeData, EdgeData, Ty> Default for NamedGraph<NodeName, NodeData, EdgeData, Ty>
349where
350    NodeName: Clone + Hash + Eq + Send + Sync + Debug,
351    Ty: EdgeType,
352{
353    fn default() -> Self {
354        Self::new()
355    }
356}
357
358impl<NodeName, NodeData, EdgeData, Ty> Clone for NamedGraph<NodeName, NodeData, EdgeData, Ty>
359where
360    NodeName: Clone + Hash + Eq + Send + Sync + Debug,
361    NodeData: Clone,
362    EdgeData: Clone,
363    Ty: EdgeType,
364{
365    fn clone(&self) -> Self {
366        Self {
367            graph: self.graph.clone(),
368            node_name_to_index: self.node_name_to_index.clone(),
369            index_to_node_name: self.index_to_node_name.clone(),
370        }
371    }
372}
373
374impl<NodeName, NodeData, EdgeData, Ty> Debug for NamedGraph<NodeName, NodeData, EdgeData, Ty>
375where
376    NodeName: Clone + Hash + Eq + Send + Sync + Debug,
377    NodeData: Debug,
378    EdgeData: Debug,
379    Ty: EdgeType,
380{
381    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
382        f.debug_struct("NamedGraph")
383            .field("graph", &self.graph)
384            .field("node_name_to_index", &self.node_name_to_index)
385            .field("index_to_node_name", &self.index_to_node_name)
386            .finish()
387    }
388}
389
390#[cfg(test)]
391mod tests;