ego_tree/sort.rs
1//! Sorting functionality for tree nodes.
2//!
3//! This module provides methods for sorting children of a node in a tree.
4//! The sorting can be done based on the node values or their indices.
5
6use std::cmp::Ordering;
7
8use crate::{NodeMut, NodeRef};
9
10impl<'a, T: 'a> NodeMut<'a, T> {
11 /// Sort children by value in ascending order.
12 ///
13 /// # Examples
14 ///
15 /// ```rust
16 /// use ego_tree::tree;
17 ///
18 /// let mut tree = tree!('a' => { 'd', 'c', 'b' });
19 /// tree.root_mut().sort();
20 /// assert_eq!(
21 /// vec![&'b', &'c', &'d'],
22 /// tree.root()
23 /// .children()
24 /// .map(|n| n.value())
25 /// .collect::<Vec<_>>(),
26 /// );
27 /// ```
28 pub fn sort(&mut self)
29 where
30 T: Ord,
31 {
32 self.sort_by(|a, b| a.value().cmp(b.value()));
33 }
34
35 /// Sort children by `NodeRef` in ascending order using a comparison function.
36 ///
37 /// # Examples
38 ///
39 /// ```rust
40 /// use ego_tree::tree;
41 ///
42 /// let mut tree = tree!('a' => { 'c', 'd', 'b' });
43 /// tree.root_mut().sort_by(|a, b| b.value().cmp(a.value()));
44 /// assert_eq!(
45 /// vec![&'d', &'c', &'b'],
46 /// tree.root()
47 /// .children()
48 /// .map(|n| n.value())
49 /// .collect::<Vec<_>>(),
50 /// );
51 ///
52 /// // Example for sort_by_id.
53 /// tree.root_mut().sort_by(|a, b| a.id().cmp(&b.id()));
54 /// assert_eq!(
55 /// vec![&'c', &'d', &'b'],
56 /// tree.root()
57 /// .children()
58 /// .map(|n| n.value())
59 /// .collect::<Vec<_>>(),
60 /// );
61 /// ```
62 pub fn sort_by<F>(&mut self, mut compare: F)
63 where
64 F: FnMut(NodeRef<T>, NodeRef<T>) -> Ordering,
65 {
66 if !self.has_children() {
67 return;
68 }
69
70 let mut children = {
71 let this = unsafe { self.tree.get_unchecked(self.id) };
72 this.children().map(|child| child.id).collect::<Vec<_>>()
73 };
74
75 children.sort_by(|a, b| {
76 let a = unsafe { self.tree.get_unchecked(*a) };
77 let b = unsafe { self.tree.get_unchecked(*b) };
78 compare(a, b)
79 });
80
81 for id in children {
82 self.append_id(id);
83 }
84 }
85
86 /// Sort children by `NodeRef`'s key in ascending order using a key extraction function.
87 ///
88 /// # Examples
89 ///
90 /// ```rust
91 /// use ego_tree::tree;
92 ///
93 /// let mut tree = tree!("1a" => { "2b", "4c", "3d" });
94 /// tree.root_mut().sort_by_key(|a| a.value().split_at(1).0.parse::<i32>().unwrap());
95 /// assert_eq!(
96 /// vec!["2b", "3d", "4c"],
97 /// tree.root()
98 /// .children()
99 /// .map(|n| *n.value())
100 /// .collect::<Vec<_>>(),
101 /// );
102 ///
103 /// // Example for sort_by_id.
104 /// tree.root_mut().sort_by_key(|n| n.id());
105 /// assert_eq!(
106 /// vec![&"2b", &"4c", &"3d"],
107 /// tree.root()
108 /// .children()
109 /// .map(|n| n.value())
110 /// .collect::<Vec<_>>(),
111 /// );
112 /// ```
113 pub fn sort_by_key<K, F>(&mut self, mut f: F)
114 where
115 F: FnMut(NodeRef<T>) -> K,
116 K: Ord,
117 {
118 self.sort_by(|a, b| f(a).cmp(&f(b)));
119 }
120}