1use std::collections::HashMap;
16
17use crate::types::{Id, Point, Side, Size};
18
19#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
21pub enum LayoutDirection {
22 #[default]
24 TopToBottom,
25 LeftToRight,
27 BottomToTop,
28 RightToLeft,
29}
30
31impl LayoutDirection {
32 pub fn handle_sides(&self) -> (Side, Side) {
34 match self {
35 LayoutDirection::TopToBottom => (Side::Top, Side::Bottom),
36 LayoutDirection::BottomToTop => (Side::Bottom, Side::Top),
37 LayoutDirection::LeftToRight => (Side::Left, Side::Right),
38 LayoutDirection::RightToLeft => (Side::Right, Side::Left),
39 }
40 }
41}
42
43#[derive(Clone, PartialEq, Debug)]
45pub struct LayoutOptions {
46 pub direction: LayoutDirection,
47 pub node_gap: f64,
49 pub rank_gap: f64,
51 pub update_handle_sides: bool,
54}
55
56impl Default for LayoutOptions {
57 fn default() -> Self {
58 Self {
59 direction: LayoutDirection::default(),
60 node_gap: 50.0,
61 rank_gap: 90.0,
62 update_handle_sides: true,
63 }
64 }
65}
66
67impl LayoutOptions {
68 pub fn direction(mut self, direction: LayoutDirection) -> Self {
69 self.direction = direction;
70 self
71 }
72
73 pub fn gaps(mut self, node_gap: f64, rank_gap: f64) -> Self {
74 self.node_gap = node_gap;
75 self.rank_gap = rank_gap;
76 self
77 }
78}
79
80#[derive(Clone, PartialEq, Debug)]
82pub struct LayoutNode {
83 pub id: Id,
84 pub size: Size,
85}
86
87pub fn compute_layout(
90 nodes: &[LayoutNode],
91 edges: &[(Id, Id)],
92 opts: &LayoutOptions,
93) -> HashMap<Id, Point> {
94 let n = nodes.len();
95 if n == 0 {
96 return HashMap::new();
97 }
98
99 let index: HashMap<&str, usize> = nodes
100 .iter()
101 .enumerate()
102 .map(|(i, node)| (node.id.as_str(), i))
103 .collect();
104
105 let mut arcs: Vec<(usize, usize)> = edges
107 .iter()
108 .filter_map(|(s, t)| {
109 let (s, t) = (*index.get(s.as_str())?, *index.get(t.as_str())?);
110 (s != t).then_some((s, t))
111 })
112 .collect();
113 arcs.sort_unstable();
114 arcs.dedup();
115
116 reverse_back_edges(n, &mut arcs);
117 let ranks = assign_ranks(n, &arcs);
118 let order = order_ranks(n, &arcs, &ranks);
119
120 let horizontal = matches!(
122 opts.direction,
123 LayoutDirection::LeftToRight | LayoutDirection::RightToLeft
124 );
125 let main_size = |i: usize| {
126 if horizontal {
127 nodes[i].size.width
128 } else {
129 nodes[i].size.height
130 }
131 };
132 let cross_size = |i: usize| {
133 if horizontal {
134 nodes[i].size.height
135 } else {
136 nodes[i].size.width
137 }
138 };
139
140 let mut cross = vec![0.0f64; n];
142 for rank in &order {
143 let total: f64 = rank.iter().map(|&i| cross_size(i)).sum::<f64>()
144 + opts.node_gap * (rank.len() - 1) as f64;
145 let mut cursor = -total / 2.0;
146 for &i in rank {
147 cross[i] = cursor + cross_size(i) / 2.0;
148 cursor += cross_size(i) + opts.node_gap;
149 }
150 }
151 let (preds, succs) = neighbor_lists(n, &arcs);
153 for _ in 0..2 {
154 for rank in order.iter().skip(1) {
155 align_rank(rank, &preds, &cross_size, opts.node_gap, &mut cross);
156 }
157 for rank in order.iter().rev().skip(1) {
158 align_rank(rank, &succs, &cross_size, opts.node_gap, &mut cross);
159 }
160 }
161
162 let mut main = vec![0.0f64; n];
164 let mut cursor = 0.0;
165 for rank in &order {
166 let depth = rank.iter().map(|&i| main_size(i)).fold(0.0, f64::max);
167 for &i in rank {
168 main[i] = cursor + depth / 2.0;
169 }
170 cursor += depth + opts.rank_gap;
171 }
172
173 nodes
174 .iter()
175 .enumerate()
176 .map(|(i, node)| {
177 let center = match opts.direction {
178 LayoutDirection::TopToBottom => Point::new(cross[i], main[i]),
179 LayoutDirection::BottomToTop => Point::new(cross[i], -main[i]),
180 LayoutDirection::LeftToRight => Point::new(main[i], cross[i]),
181 LayoutDirection::RightToLeft => Point::new(-main[i], cross[i]),
182 };
183 let top_left = Point::new(
184 center.x - node.size.width / 2.0,
185 center.y - node.size.height / 2.0,
186 );
187 (node.id.clone(), top_left)
188 })
189 .collect()
190}
191
192fn reverse_back_edges(n: usize, arcs: &mut [(usize, usize)]) {
194 let mut out: Vec<Vec<usize>> = vec![Vec::new(); n];
195 for (k, &(s, _)) in arcs.iter().enumerate() {
196 out[s].push(k);
197 }
198 let mut state = vec![0u8; n];
200 let mut back: Vec<usize> = Vec::new();
201 for root in 0..n {
202 if state[root] != 0 {
203 continue;
204 }
205 let mut stack: Vec<(usize, usize)> = vec![(root, 0)];
207 state[root] = 1;
208 while let Some(&(v, cursor)) = stack.last() {
209 if cursor < out[v].len() {
210 stack.last_mut().unwrap().1 += 1;
211 let arc_idx = out[v][cursor];
212 let t = arcs[arc_idx].1;
213 match state[t] {
214 0 => {
215 state[t] = 1;
216 stack.push((t, 0));
217 }
218 1 => back.push(arc_idx),
219 _ => {}
220 }
221 } else {
222 state[v] = 2;
223 stack.pop();
224 }
225 }
226 }
227 for k in back {
228 let (s, t) = arcs[k];
229 arcs[k] = (t, s);
230 }
231}
232
233fn assign_ranks(n: usize, arcs: &[(usize, usize)]) -> Vec<usize> {
235 let mut indeg = vec![0usize; n];
236 let mut out: Vec<Vec<usize>> = vec![Vec::new(); n];
237 for &(s, t) in arcs {
238 indeg[t] += 1;
239 out[s].push(t);
240 }
241 let mut rank = vec![0usize; n];
242 let mut queue: Vec<usize> = (0..n).filter(|&i| indeg[i] == 0).collect();
243 let mut head = 0;
244 while head < queue.len() {
245 let v = queue[head];
246 head += 1;
247 for &t in &out[v] {
248 rank[t] = rank[t].max(rank[v] + 1);
249 indeg[t] -= 1;
250 if indeg[t] == 0 {
251 queue.push(t);
252 }
253 }
254 }
255 rank
256}
257
258fn order_ranks(n: usize, arcs: &[(usize, usize)], ranks: &[usize]) -> Vec<Vec<usize>> {
260 let max_rank = ranks.iter().copied().max().unwrap_or(0);
261 let mut order: Vec<Vec<usize>> = vec![Vec::new(); max_rank + 1];
262 for i in 0..n {
263 order[ranks[i]].push(i);
264 }
265 let (preds, succs) = neighbor_lists(n, arcs);
266
267 let mut pos = vec![0.0f64; n];
268 let write_pos = |order: &[Vec<usize>], pos: &mut [f64]| {
269 for rank in order {
270 for (p, &i) in rank.iter().enumerate() {
271 pos[i] = p as f64;
272 }
273 }
274 };
275 write_pos(&order, &mut pos);
276
277 for _ in 0..4 {
278 for rank in order.iter_mut().skip(1) {
279 barycenter_sort(rank, &preds, &pos);
280 }
281 write_pos(&order, &mut pos);
282 for r in (0..order.len().saturating_sub(1)).rev() {
283 barycenter_sort(&mut order[r], &succs, &pos);
284 write_pos(&order, &mut pos);
285 }
286 }
287 order
288}
289
290fn neighbor_lists(n: usize, arcs: &[(usize, usize)]) -> (Vec<Vec<usize>>, Vec<Vec<usize>>) {
291 let mut preds: Vec<Vec<usize>> = vec![Vec::new(); n];
292 let mut succs: Vec<Vec<usize>> = vec![Vec::new(); n];
293 for &(s, t) in arcs {
294 preds[t].push(s);
295 succs[s].push(t);
296 }
297 (preds, succs)
298}
299
300fn barycenter_sort(rank: &mut [usize], neighbors: &[Vec<usize>], pos: &[f64]) {
301 let mut keyed: Vec<(f64, usize, usize)> = rank
302 .iter()
303 .enumerate()
304 .map(|(current, &i)| {
305 let ns = &neighbors[i];
306 let key = if ns.is_empty() {
307 current as f64
308 } else {
309 ns.iter().map(|&p| pos[p]).sum::<f64>() / ns.len() as f64
310 };
311 (key, current, i)
312 })
313 .collect();
314 keyed.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
315 for (slot, (_, _, i)) in keyed.into_iter().enumerate() {
316 rank[slot] = i;
317 }
318}
319
320fn align_rank(
323 rank: &[usize],
324 neighbors: &[Vec<usize>],
325 cross_size: &impl Fn(usize) -> f64,
326 gap: f64,
327 cross: &mut [f64],
328) {
329 if rank.is_empty() {
330 return;
331 }
332 let desired: Vec<f64> = rank
333 .iter()
334 .map(|&i| {
335 let ns = &neighbors[i];
336 if ns.is_empty() {
337 cross[i]
338 } else {
339 ns.iter().map(|&p| cross[p]).sum::<f64>() / ns.len() as f64
340 }
341 })
342 .collect();
343 let mut placed: Vec<f64> = Vec::with_capacity(rank.len());
346 for (k, &i) in rank.iter().enumerate() {
347 let min_pos = if k == 0 {
348 f64::NEG_INFINITY
349 } else {
350 placed[k - 1] + cross_size(rank[k - 1]) / 2.0 + gap + cross_size(i) / 2.0
351 };
352 placed.push(desired[k].max(min_pos));
353 }
354 let drift: f64 =
356 placed.iter().zip(&desired).map(|(p, d)| p - d).sum::<f64>() / rank.len() as f64;
357 for (k, &i) in rank.iter().enumerate() {
358 cross[i] = placed[k] - drift;
359 }
360}
361
362#[cfg(test)]
363mod tests {
364 use super::*;
365
366 fn nodes(ids: &[&str]) -> Vec<LayoutNode> {
367 ids.iter()
368 .map(|id| LayoutNode {
369 id: id.to_string(),
370 size: Size::new(100.0, 40.0),
371 })
372 .collect()
373 }
374
375 fn e(pairs: &[(&str, &str)]) -> Vec<(Id, Id)> {
376 pairs
377 .iter()
378 .map(|(a, b)| (a.to_string(), b.to_string()))
379 .collect()
380 }
381
382 #[test]
383 fn chain_top_to_bottom() {
384 let ns = nodes(&["a", "b", "c"]);
385 let pos = compute_layout(
386 &ns,
387 &e(&[("a", "b"), ("b", "c")]),
388 &LayoutOptions::default(),
389 );
390 assert_eq!(pos["a"].x, pos["b"].x);
392 assert_eq!(pos["b"].x, pos["c"].x);
393 assert!(pos["b"].y - pos["a"].y >= 40.0 + 90.0);
394 assert!(pos["c"].y - pos["b"].y >= 40.0 + 90.0);
395 }
396
397 #[test]
398 fn chain_left_to_right() {
399 let ns = nodes(&["a", "b"]);
400 let opts = LayoutOptions::default().direction(LayoutDirection::LeftToRight);
401 let pos = compute_layout(&ns, &e(&[("a", "b")]), &opts);
402 assert_eq!(pos["a"].y, pos["b"].y);
403 assert!(pos["b"].x - pos["a"].x >= 100.0 + 90.0);
404 }
405
406 #[test]
407 fn diamond_no_overlap() {
408 let ns = nodes(&["a", "b", "c", "d"]);
409 let pos = compute_layout(
410 &ns,
411 &e(&[("a", "b"), ("a", "c"), ("b", "d"), ("c", "d")]),
412 &LayoutOptions::default(),
413 );
414 assert_eq!(pos["b"].y, pos["c"].y);
416 assert!((pos["b"].x - pos["c"].x).abs() >= 100.0 + 50.0);
417 assert!((pos["a"].x - pos["d"].x).abs() < 1.0);
419 assert!(pos["d"].y > pos["b"].y);
420 }
421
422 #[test]
423 fn cycle_does_not_panic() {
424 let ns = nodes(&["a", "b", "c"]);
425 let pos = compute_layout(
426 &ns,
427 &e(&[("a", "b"), ("b", "c"), ("c", "a")]),
428 &LayoutOptions::default(),
429 );
430 assert_eq!(pos.len(), 3);
431 let mut ys: Vec<i64> = pos.values().map(|p| p.y as i64).collect();
433 ys.sort_unstable();
434 ys.dedup();
435 assert_eq!(ys.len(), 3);
436 }
437
438 #[test]
439 fn disconnected_and_unknown_edges() {
440 let ns = nodes(&["a", "b", "lonely"]);
441 let pos = compute_layout(
442 &ns,
443 &e(&[("a", "b"), ("a", "ghost")]),
444 &LayoutOptions::default(),
445 );
446 assert_eq!(pos.len(), 3);
447 }
448
449 #[test]
450 fn empty() {
451 let pos = compute_layout(&[], &[], &LayoutOptions::default());
452 assert!(pos.is_empty());
453 }
454
455 #[test]
456 fn bottom_to_top_flips() {
457 let ns = nodes(&["a", "b"]);
458 let opts = LayoutOptions::default().direction(LayoutDirection::BottomToTop);
459 let pos = compute_layout(&ns, &e(&[("a", "b")]), &opts);
460 assert!(pos["b"].y < pos["a"].y);
461 }
462}