1use petgraph::algo::{bellman_ford, min_spanning_tree};
2use petgraph::data::FromElements;
3use petgraph::prelude::*;
4use petgraph::Graph;
5use std::collections::HashMap;
6
7fn convert_shortest_paths(
8 all_paths: &[bellman_ford::Paths<NodeIndex, f64>],
9 nodes: &[NodeIndex],
10) -> Vec<Vec<(usize, Vec<usize>)>> {
11 let mut my_paths: Vec<Vec<(usize, Vec<usize>)>> = Vec::new();
12 for n1 in nodes.iter() {
13 my_paths.push(Vec::new());
14 for n2 in nodes.iter() {
15 let path = get_shortest_path(all_paths, n1, n2);
16 my_paths.last_mut().unwrap().push((path.len() - 1, path));
17 }
18 }
19 my_paths
20}
21
22fn get_shortest_path(
23 all_paths: &[bellman_ford::Paths<NodeIndex, f64>],
24 n1: &NodeIndex,
25 n2: &NodeIndex,
26) -> Vec<usize> {
27 let path = &all_paths[n1.index()];
28 let mut spath: Vec<usize> = Vec::new();
29 spath.push(n2.index());
30 let mut v = path.predecessors[n2.index()];
31 while let Some(ni) = v {
32 spath.push(ni.index());
33 v = path.predecessors[ni.index()];
34 }
35 spath
36}
37
38fn build_closure_mst(
39 all_paths: &[Vec<(usize, Vec<usize>)>],
40 terminals: &[usize],
41) -> Graph<(), f64, Undirected> {
42 let mut g = Graph::new_undirected();
43 let nodes: Vec<_> = (0..terminals.len()).map(|_| g.add_node(())).collect();
44
45 for (i1, n1) in nodes.iter().enumerate() {
46 for (i2, n2) in nodes.iter().enumerate() {
47 if i1 != i2 {
48 g.add_edge(*n1, *n2, all_paths[terminals[i1]][terminals[i2]].0 as f64);
49 }
50 }
51 }
52 let mst: Graph<(), f64, Undirected> = Graph::from_elements(min_spanning_tree(&g));
53 mst
54}
55
56#[derive(Debug, Clone)]
57pub struct SteinerTree {
58 pub graph: UnGraph<(), i32, u32>,
59 pub mapping: HashMap<usize, NodeIndex>,
60 pub inverse_mapping: HashMap<NodeIndex, usize>,
61 pub terminals: Vec<usize>,
62}
63
64impl SteinerTree {
65 pub fn len(&self) -> usize {
66 self.graph.node_count()
67 }
68
69 pub fn is_empty(&self) -> bool {
70 self.graph.node_count() == 0
71 }
72
73 pub fn cnot_cost(&self) -> usize {
74 2 * self.len() - self.terminals.len() - 1
75 }
76
77 pub fn contains(&self, node: usize) -> bool {
78 self.mapping.contains_key(&node)
79 }
80 pub fn is_leaf(&self, node: usize) -> bool {
81 let node_index = self.mapping.get(&node).unwrap();
82 self.graph.neighbors_undirected(*node_index).count() == 1
83 }
84 pub fn neighbors(&self, node: usize) -> Vec<usize> {
85 let mut neighs = Vec::new();
86 for nei in self.graph.neighbors_undirected(self.mapping[&node]) {
87 neighs.push(self.inverse_mapping[&nei]);
88 }
89 neighs
90 }
91 pub fn nodes(&self) -> Vec<usize> {
92 self.mapping.keys().copied().collect()
93 }
94 pub fn add_node(&mut self, node: usize) {
95 let node_index = self.graph.add_node(());
96 self.mapping.insert(node, node_index);
97 self.inverse_mapping.insert(node_index, node);
98 }
99
100 pub fn add_edge(&mut self, n1: usize, n2: usize) {
101 self.graph.add_edge(self.mapping[&n1], self.mapping[&n2], 1);
102 }
103
104 pub fn remove_node(&mut self, node: usize) {
105 if !self.contains(node) {
106 panic!("Node {} is not in the tree", node);
107 }
108 let last_index: NodeIndex<u32> = NodeIndex::new(self.graph.node_count() - 1);
109 let node_index = *self.mapping.get(&node).unwrap();
110 let i = *self.inverse_mapping.get(&last_index).unwrap();
111 if i == node {
112 self.graph.remove_node(node_index);
113 self.mapping.remove(&i);
114 self.inverse_mapping.remove(&node_index);
115 return;
116 }
117 self.inverse_mapping.remove(&last_index);
118 self.mapping.remove(&node);
119 self.mapping.insert(i, node_index);
120 self.inverse_mapping.insert(node_index, i);
121 self.graph.remove_node(node_index);
122 }
123
124 pub fn update_tree(&mut self, new_support: &[usize], qbits: &[usize; 2]) {
125 self.terminals = new_support.to_owned();
126 if self.contains(qbits[0]) && self.contains(qbits[1]) {
127 if new_support.contains(&qbits[0]) && new_support.contains(&qbits[1]) {
128 return;
129 }
130 if !new_support.contains(&qbits[0]) && self.is_leaf(qbits[0]) {
131 self.remove_node(qbits[0]);
132 }
133 if !new_support.contains(&qbits[1]) && self.is_leaf(qbits[1]) {
134 self.remove_node(qbits[1]);
135 }
136 return;
137 }
138
139 if !self.contains(qbits[0]) && new_support.contains(&qbits[0]) {
140 self.add_node(qbits[0]);
141 self.add_edge(qbits[0], qbits[1]);
142 return;
143 }
144
145 if !self.contains(qbits[1]) && new_support.contains(&qbits[1]) {
146 self.add_node(qbits[1]);
147 self.add_edge(qbits[0], qbits[1]);
148 }
149 }
150 pub fn prune_non_terminal_leaves(&mut self) {
151 loop {
152 let mut popped = false;
153 for node in self.graph.node_indices() {
154 if !self.terminals.contains(&self.inverse_mapping[&node])
155 && self.graph.neighbors(node).count() == 1
156 {
157 self.remove_node(self.inverse_mapping[&node]);
158 popped = true;
159 break;
160 }
161 }
162 if !popped {
163 break;
164 }
165 }
166 }
167}
168
169fn build_steiner_from_mst(
170 closure_mst: &Graph<(), f64, Undirected>,
171 terminals: &[usize],
172 all_paths: &[Vec<(usize, Vec<usize>)>],
173) -> SteinerTree {
174 let mut g = Graph::new_undirected();
176 let mut nodes = HashMap::new();
177 let mut inverse_map = HashMap::new();
178 for i in terminals.iter() {
179 let nodei = g.add_node(());
180 nodes.insert(*i, nodei);
181 inverse_map.insert(nodei, *i);
182 }
183 for edge in closure_mst.raw_edges() {
184 let t1 = terminals[edge.source().index()];
185 let t2 = terminals[edge.target().index()];
186 let (_, path) = &all_paths[t1][t2];
187 for i in path.iter() {
188 if !nodes.contains_key(i) {
189 let nodei = g.add_node(());
190 nodes.insert(*i, nodei);
191 inverse_map.insert(nodei, *i);
192 }
193 }
194 let mut left = path[0];
195 for next_v in path.iter().skip(1) {
196 g.add_edge(nodes[&left], nodes[next_v], 1);
197 left = *next_v;
198 }
199 }
200 let mst: Graph<(), i32, Undirected> = Graph::from_elements(min_spanning_tree(&g));
201 let mut as_st = SteinerTree {
202 graph: mst,
203 mapping: nodes,
204 inverse_mapping: inverse_map,
205 terminals: terminals.to_owned(),
206 };
207 as_st.prune_non_terminal_leaves();
208 as_st
209}
210
211fn get_steiner_tree(conv_paths: &[Vec<(usize, Vec<usize>)>], terminals: &[usize]) -> SteinerTree {
212 let closure = build_closure_mst(conv_paths, terminals);
213 build_steiner_from_mst(&closure, terminals, conv_paths)
214}
215#[derive(Debug, Clone)]
216pub struct HardwareGraph {
217 pub graph: UnGraph<(), f64, u32>,
218 shortest_paths: Vec<Vec<(usize, Vec<usize>)>>,
219}
220
221impl HardwareGraph {
222 pub fn from_couplings(couplings: &[(usize, usize)]) -> Self {
223 let mut graph = UnGraph::new_undirected();
224 let n = couplings.iter().map(|c| c.0.max(c.1)).max().unwrap() + 1;
225 let nodes: Vec<_> = (0..n).map(|_| graph.add_node(())).collect();
226 graph.extend_with_edges(couplings.iter().map(|c| (nodes[c.0], nodes[c.1], 1.)));
227 let all_paths: Vec<bellman_ford::Paths<NodeIndex, f64>> = nodes
228 .iter()
229 .map(|node_index| bellman_ford(&graph, *node_index).unwrap())
230 .collect();
231 let shortest_paths = convert_shortest_paths(&all_paths, &nodes);
232 Self {
233 graph,
234 shortest_paths,
235 }
236 }
237 pub fn len(&self) -> usize {
238 self.graph.node_count()
239 }
240
241 pub fn is_empty(&self) -> bool {
242 self.graph.node_count() == 0
243 }
244
245 pub fn get_steiner_tree(&self, terminals: &[usize]) -> SteinerTree {
246 get_steiner_tree(&self.shortest_paths, terminals)
247 }
248}