Skip to content

Commit e663500

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 e663500

File tree

4 files changed

+375
-0
lines changed

4 files changed

+375
-0
lines changed

‎Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ readme = "README.md"
1414

1515
[features]
1616
serde = ["dep:serde"]
17+
sort = []
1718

1819
[dependencies]
1920
serde = { version = "1.0.209", optional = true }

‎src/lib.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ use std::num::NonZeroUsize;
4040

4141
#[cfg(feature = "serde")]
4242
pub mod serde;
43+
#[cfg(feature = "sort")]
44+
pub mod sort;
4345

4446
/// Vec-backed ID-tree.
4547
///
@@ -557,6 +559,7 @@ impl<'a, T: 'a> NodeMut<'a, T> {
557559
);
558560

559561
let last_child_id = self.node().children.map(|(_, id)| id);
562+
560563
{
561564
let mut new_child = self.tree.get_mut(new_child_id).unwrap();
562565
new_child.detach();

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

0 commit comments

Comments
 (0)