Skip to content

Commit af549cb

Browse files
AuTaAuTa
committed
feature: add sort feature for tree NodeMut
Add a new feature. This module provides methods for sorting children of a node in a tree. The sorting can be done based on the node values or their indices.
1 parent 8127d31 commit af549cb

File tree

3 files changed

+371
-0
lines changed

3 files changed

+371
-0
lines changed

‎src/lib.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -929,3 +929,5 @@ impl<T: Display> Display for Tree<T> {
929929
Ok(())
930930
}
931931
}
932+
933+
mod sort;

‎src/sort.rs‎

Lines changed: 264 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,264 @@
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+
6+
use std::cmp::Ordering;
7+
8+
use crate::{NodeId, NodeMut};
9+
10+
impl<'a, T: 'a> NodeMut<'a, T> {
11+
/// Sort children by value in ascending order.
12+
///
13+
/// This method is a shorthand for calling `sort_by` with the `Ord::cmp` method.
14+
///
15+
/// # Examples
16+
///
17+
/// ```
18+
/// use ego_tree::tree;
19+
///
20+
/// let mut tree = tree!('a' => { 'd', 'c', 'b' });
21+
/// tree.root_mut().sort();
22+
/// assert_eq!(
23+
/// vec![&'b', &'c', &'d'],
24+
/// tree.root()
25+
/// .children()
26+
/// .map(|n| n.value())
27+
/// .collect::<Vec<_>>(),
28+
/// );
29+
/// ```
30+
pub fn sort(&mut self)
31+
where
32+
T: Ord,
33+
{
34+
self.sort_by(|a, b| a.cmp(b));
35+
}
36+
37+
/// Sort children by value in ascending order using a comparison function.
38+
///
39+
/// # Examples
40+
///
41+
/// ```
42+
/// use ego_tree::tree;
43+
///
44+
/// let mut tree = tree!('a' => { 'c', 'd', 'b' });
45+
/// tree.root_mut().sort_by(|a, b| b.cmp(a));
46+
/// assert_eq!(
47+
/// vec![&'d', &'c', &'b'],
48+
/// tree.root()
49+
/// .children()
50+
/// .map(|n| n.value())
51+
/// .collect::<Vec<_>>(),
52+
/// );
53+
/// ```
54+
pub fn sort_by<F>(&mut self, mut compare: F)
55+
where
56+
F: FnMut(&T, &T) -> Ordering,
57+
{
58+
if self.has_children() {
59+
let (unsorted, sorted) = self.sort_handler(|nodes| {
60+
nodes.sort_by(|(_, a), (_, b)| compare(a, b));
61+
});
62+
63+
self.swap(unsorted, sorted);
64+
}
65+
}
66+
67+
/// Sort children by value's key in ascending order using a key extraction function.
68+
///
69+
/// # Examples
70+
///
71+
/// ```
72+
/// use ego_tree::tree;
73+
///
74+
/// let mut tree = tree!("1a" => { "2b", "4c", "3d" });
75+
/// tree.root_mut().sort_by_key(|a| a.split_at(1).0.parse::<i32>().unwrap());
76+
/// assert_eq!(
77+
/// vec!["2b", "3d", "4c"],
78+
/// tree.root()
79+
/// .children()
80+
/// .map(|n| *n.value())
81+
/// .collect::<Vec<_>>(),
82+
/// );
83+
/// ```
84+
pub fn sort_by_key<K, F>(&mut self, mut f: F)
85+
where
86+
F: FnMut(&T) -> K,
87+
K: Ord,
88+
{
89+
if self.has_children() {
90+
let (unsorted, sorted) = self.sort_handler(|nodes| {
91+
nodes.sort_by_key(|(_, value)| f(value));
92+
});
93+
94+
self.swap(unsorted, sorted);
95+
}
96+
}
97+
98+
/// Sort children by their NodeId in ascending order. The purpose is to restore the original order.
99+
///
100+
/// This method is a shorthand for calling `sort_by_id` with the `Ord::cmp` method.
101+
///
102+
/// # Examples
103+
///
104+
/// ```
105+
/// use ego_tree::tree;
106+
///
107+
/// let mut tree = tree!('a' => { 'd', 'c', 'b' });
108+
/// tree.root_mut().sort();
109+
/// assert_ne!(
110+
/// vec![&'d', &'c', &'b'],
111+
/// tree.root()
112+
/// .children()
113+
/// .map(|n| n.value())
114+
/// .collect::<Vec<_>>(),
115+
/// );
116+
/// tree.root_mut().sort_id();
117+
/// assert_eq!(
118+
/// vec![&'d', &'c', &'b'],
119+
/// tree.root()
120+
/// .children()
121+
/// .map(|n| n.value())
122+
/// .collect::<Vec<_>>(),
123+
/// );
124+
/// ```
125+
pub fn sort_id(&mut self) {
126+
self.sort_by_id(|a, b| a.cmp(&b));
127+
}
128+
129+
/// Sort children by their NodeId's index using a comparison function.
130+
///
131+
/// # Examples
132+
///
133+
/// ```
134+
/// use ego_tree::tree;
135+
///
136+
/// let mut tree = tree!('a' => { 'd', 'b', 'c' });
137+
/// tree.root_mut().sort_by_id(|a, b| b.cmp(&a));
138+
/// assert_eq!(
139+
/// vec![&'c', &'b', &'d'],
140+
/// tree.root()
141+
/// .children()
142+
/// .map(|n| n.value())
143+
/// .collect::<Vec<_>>(),
144+
/// );
145+
/// ```
146+
pub fn sort_by_id<F>(&mut self, mut compare: F)
147+
where
148+
F: FnMut(usize, usize) -> Ordering,
149+
{
150+
if self.has_children() {
151+
let (unsorted, sorted) = self.sort_handler(|nodes| {
152+
nodes.sort_by(|(ida, _), (idb, _)| compare(ida.to_index(), idb.to_index()));
153+
});
154+
155+
self.swap(unsorted, sorted);
156+
}
157+
}
158+
159+
/// Sort children by a key function taking a NodeId's index and a `&T` reference
160+
/// returning a key of type `K` that implements `Ord`.
161+
///
162+
/// I don't know how to use this method.
163+
///
164+
/// # Examples
165+
///
166+
/// ```
167+
/// use ego_tree::tree;
168+
/// let mut tree = tree!('a' => { 'd', 'b', 'c' });
169+
/// tree.root_mut()
170+
/// .sort_by_id_key(|id, value| id + *value as usize); // {1+100, 2+98, 3+99}
171+
/// assert_eq!(
172+
/// vec![&'b', &'d', &'c'],
173+
/// tree.root()
174+
/// .children()
175+
/// .map(|n| n.value())
176+
/// .collect::<Vec<_>>(),
177+
/// );
178+
/// ```
179+
pub fn sort_by_id_key<K, F>(&mut self, mut f: F)
180+
where
181+
F: FnMut(usize, &T) -> K,
182+
K: Ord,
183+
{
184+
if self.has_children() {
185+
let (unsorted, sorted) = self.sort_handler(|nodes| {
186+
nodes.sort_by_key(|node| f(node.0.to_index(), node.1));
187+
});
188+
self.swap(unsorted, sorted);
189+
}
190+
}
191+
192+
/// Applies a sorting function to the children of the current node and returns their IDs
193+
/// before and after sorting.
194+
///
195+
/// This function takes a mutable closure `f` that sorts a vector of tuples,
196+
/// where each tuple consists of a `NodeId` and a reference to the node's value `&T`.
197+
///
198+
/// # Returns
199+
///
200+
/// A tuple containing:
201+
/// - `Vec<NodeId>`: The original order of the children's `NodeId`s before sorting.
202+
/// - `Vec<NodeId>`: The order of the children's `NodeId`s after applying the sorting function.
203+
fn sort_handler<F>(&mut self, mut f: F) -> (Vec<NodeId>, Vec<NodeId>)
204+
where
205+
F: FnMut(&mut Vec<(NodeId, &T)>),
206+
{
207+
let children = unsafe { self.tree.get_unchecked(self.id()).children() };
208+
let (unsorted, mut nodes): (Vec<_>, Vec<_>) =
209+
children.map(|n| (n.id(), (n.id(), n.value()))).unzip();
210+
f(&mut nodes);
211+
let sorted = nodes.into_iter().map(|(id, _)| id).collect::<Vec<_>>();
212+
(unsorted, sorted)
213+
}
214+
215+
/// Reorders the children of the current node to match the specified sorted order.
216+
///
217+
/// This method takes two vectors of `NodeId`s: `unsorted`, which represents the original
218+
/// order of the node's children, and `sorted`, which represents the desired order after sorting.
219+
/// It swaps nodes in the tree such that their order in the tree matches the `sorted` vector.
220+
///
221+
/// # Parameters
222+
///
223+
/// - `unsorted`: A vector of `NodeId`s representing the original order of the node's children.
224+
/// - `sorted`: A vector of `NodeId`s representing the desired order of the node's children.
225+
///
226+
/// # Safety
227+
///
228+
/// This function uses unsafe code to access and modify the tree nodes. Ensure that the node
229+
/// indices are valid and that the tree structure remains consistent after the operation.
230+
fn swap(&mut self, unsorted: Vec<NodeId>, sorted: Vec<NodeId>) {
231+
let mut swap = |sorted_id: NodeId, unsorted_id: NodeId| {
232+
let mut node = unsafe { self.tree.get_unchecked_mut(unsorted_id) };
233+
node.insert_id_before(sorted_id);
234+
};
235+
236+
let mut cache = None;
237+
let mut unsorted = unsorted.into_iter();
238+
for (index, &id) in sorted.iter().enumerate() {
239+
match cache {
240+
Some(cache_id) if cache_id != id => {
241+
swap(id, cache_id);
242+
}
243+
Some(_) => cache = None,
244+
None => {
245+
for unsorted_id in unsorted.by_ref() {
246+
// Pass through the swapped elements.
247+
if sorted
248+
.iter()
249+
.position(|&node| node == unsorted_id)
250+
.is_some_and(|uindex| uindex < index)
251+
{
252+
continue;
253+
}
254+
if unsorted_id != id {
255+
swap(id, unsorted_id);
256+
cache = Some(unsorted_id);
257+
break;
258+
}
259+
}
260+
}
261+
}
262+
}
263+
}
264+
}

‎tests/sort.rs‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
use std::assert_eq;
2+
3+
use ego_tree::tree;
4+
5+
#[test]
6+
fn sort() {
7+
let mut tree = tree!('a' => { 'd' => { 'e', 'f' }, 'c', 'b' });
8+
tree.root_mut().sort();
9+
assert_eq!(
10+
vec![&'b', &'c', &'d'],
11+
tree.root()
12+
.children()
13+
.map(|n| n.value())
14+
.collect::<Vec<_>>(),
15+
);
16+
assert_eq!(
17+
tree.to_string(),
18+
tree!('a' => { 'b', 'c', 'd' => { 'e', 'f' } }).to_string()
19+
);
20+
}
21+
22+
#[test]
23+
fn sort_by() {
24+
let mut tree = tree!('a' => { 'c', 'd', 'b' });
25+
tree.root_mut().sort_by(|a, b| b.cmp(a));
26+
assert_eq!(
27+
vec![&'d', &'c', &'b'],
28+
tree.root()
29+
.children()
30+
.map(|n| n.value())
31+
.collect::<Vec<_>>(),
32+
);
33+
34+
let mut tree = tree!('a' => { 'c','d', 'e', 'b' });
35+
tree.root_mut().sort_by(|a, b| b.cmp(a));
36+
assert_eq!(
37+
vec![&'e', &'d', &'c', &'b'],
38+
tree.root()
39+
.children()
40+
.map(|n| n.value())
41+
.collect::<Vec<_>>(),
42+
);
43+
}
44+
45+
#[test]
46+
fn sort_by_key() {
47+
let mut tree = tree!("1a" => { "2b", "4c", "3d" });
48+
tree.root_mut()
49+
.sort_by_key(|a| a.split_at(1).0.parse::<i32>().unwrap());
50+
assert_eq!(
51+
vec!["2b", "3d", "4c"],
52+
tree.root()
53+
.children()
54+
.map(|n| *n.value())
55+
.collect::<Vec<_>>(),
56+
);
57+
}
58+
59+
#[test]
60+
fn sort_id() {
61+
let mut tree = tree!('a' => { 'd', 'c', 'b' });
62+
tree.root_mut().sort();
63+
assert_ne!(
64+
vec![&'d', &'c', &'b'],
65+
tree.root()
66+
.children()
67+
.map(|n| n.value())
68+
.collect::<Vec<_>>(),
69+
);
70+
tree.root_mut().sort_id();
71+
assert_eq!(
72+
vec![&'d', &'c', &'b'],
73+
tree.root()
74+
.children()
75+
.map(|n| n.value())
76+
.collect::<Vec<_>>(),
77+
);
78+
}
79+
80+
#[test]
81+
fn sort_by_id() {
82+
let mut tree = tree!('a' => { 'd', 'b', 'c' });
83+
tree.root_mut().sort_by_id(|a, b| b.cmp(&a));
84+
assert_eq!(
85+
vec![&'c', &'b', &'d'],
86+
tree.root()
87+
.children()
88+
.map(|n| n.value())
89+
.collect::<Vec<_>>(),
90+
);
91+
}
92+
93+
#[test]
94+
fn sort_by_id_key() {
95+
let mut tree = tree!('a' => { 'd', 'b', 'c' });
96+
tree.root_mut()
97+
.sort_by_id_key(|id, value| id + *value as usize); // {1+100, 2+98, 3+99}
98+
assert_eq!(
99+
vec![&'b', &'d', &'c'],
100+
tree.root()
101+
.children()
102+
.map(|n| n.value())
103+
.collect::<Vec<_>>(),
104+
);
105+
}

0 commit comments

Comments
 (0)