Skip to content

Commit 204889b

Browse files
AuTaAuTa
committed
feature: add sort for tree NodeMut
Add sorted_ids for NodeRef. 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 b8adcd1 commit 204889b

File tree

3 files changed

+268
-1
lines changed

3 files changed

+268
-1
lines changed

‎src/lib.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ pub struct Tree<T> {
5252
/// Node ID.
5353
///
5454
/// Index into a `Tree`-internal `Vec`.
55-
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
55+
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
5656
pub struct NodeId(NonZeroUsize);
5757

5858
impl NodeId {
@@ -935,3 +935,5 @@ impl<T: Display> Display for Tree<T> {
935935
Ok(())
936936
}
937937
}
938+
939+
mod sort;

‎src/sort.rs

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
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, NodeRef};
9+
10+
impl<'a, T: 'a> NodeMut<'a, T> {
11+
/// Sort children by value in ascending order.
12+
///
13+
/// # Examples
14+
///
15+
/// ```
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+
/// ```
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+
let ids = self.sorted_ids(|nodes| {
68+
nodes.sort_by(|a, b| compare(a, b));
69+
});
70+
71+
for id in ids {
72+
self.append_id(id);
73+
}
74+
}
75+
}
76+
77+
/// Sort children by `NodeRef`'s key in ascending order using a key extraction function.
78+
///
79+
/// # Examples
80+
///
81+
/// ```
82+
/// use ego_tree::tree;
83+
///
84+
/// let mut tree = tree!("1a" => { "2b", "4c", "3d" });
85+
/// tree.root_mut().sort_by_key(|a| a.value().split_at(1).0.parse::<i32>().unwrap());
86+
/// assert_eq!(
87+
/// vec!["2b", "3d", "4c"],
88+
/// tree.root()
89+
/// .children()
90+
/// .map(|n| *n.value())
91+
/// .collect::<Vec<_>>(),
92+
/// );
93+
///
94+
/// // Example for sort_by_id.
95+
/// tree.root_mut().sort_by_key(|n| n.id());
96+
/// assert_eq!(
97+
/// vec![&"2b", &"4c", &"3d"],
98+
/// tree.root()
99+
/// .children()
100+
/// .map(|n| n.value())
101+
/// .collect::<Vec<_>>(),
102+
/// );
103+
/// ```
104+
pub fn sort_by_key<K, F>(&mut self, mut f: F)
105+
where
106+
F: FnMut(&NodeRef<T>) -> K,
107+
K: Ord,
108+
{
109+
if self.has_children() {
110+
let ids = self.sorted_ids(|nodes| {
111+
nodes.sort_by_key(|node| f(node));
112+
});
113+
114+
for id in ids {
115+
self.append_id(id);
116+
}
117+
}
118+
}
119+
120+
/// Returns a vector of `NodeId`s of the node's children, sorted according to
121+
/// the provided closure.
122+
///
123+
/// The closure `f` is applied to a mutable vector of `NodeRef`s representing
124+
/// the children of the node. This allows you to define the sorting logic
125+
/// within the closure, which will then determine the order of the returned
126+
/// `NodeId`s.
127+
///
128+
/// # Parameters
129+
///
130+
/// - `f`: A closure that takes a mutable reference to a vector of `NodeRef`s.
131+
/// The closure is responsible for sorting the vector in-place.
132+
///
133+
/// # Returns
134+
///
135+
/// A vector of `NodeId`s representing the children of the node, sorted
136+
/// according to the logic defined in the closure `f`.
137+
pub fn sorted_ids<F>(&self, mut f: F) -> Vec<NodeId>
138+
where
139+
F: FnMut(&mut Vec<NodeRef<T>>),
140+
{
141+
let node_ref = unsafe { self.tree.get_unchecked(self.id) };
142+
node_ref.sorted_ids(&mut f)
143+
}
144+
}
145+
146+
impl<'a, T: 'a> NodeRef<'a, T> {
147+
/// Returns a vector of `NodeId`s of the node's children, sorted according to
148+
/// the provided closure.
149+
///
150+
/// The closure `f` is applied to a mutable vector of `NodeRef`s representing
151+
/// the children of the node. This allows you to define the sorting logic
152+
/// within the closure, which will then determine the order of the returned
153+
/// `NodeId`s.
154+
///
155+
/// # Parameters
156+
///
157+
/// - `f`: A closure that takes a mutable reference to a vector of `NodeRef`s.
158+
/// The closure is responsible for sorting the vector in-place.
159+
///
160+
/// # Returns
161+
///
162+
/// A vector of `NodeId`s representing the children of the node, sorted
163+
/// according to the logic defined in the closure `f`.
164+
pub fn sorted_ids<F>(&self, mut f: F) -> Vec<NodeId>
165+
where
166+
F: FnMut(&mut Vec<NodeRef<T>>),
167+
{
168+
let mut nodes = self.children().collect::<Vec<_>>();
169+
170+
f(&mut nodes);
171+
172+
nodes.into_iter().map(|n| n.id()).collect()
173+
}
174+
}

‎tests/sort.rs

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
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.value().cmp(a.value()));
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.value().cmp(a.value()));
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.value().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_by_key(|n| n.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(|a, b| b.id().cmp(&a.id()));
84+
assert_eq!(
85+
vec![&'c', &'b', &'d'],
86+
tree.root()
87+
.children()
88+
.map(|n| n.value())
89+
.collect::<Vec<_>>(),
90+
);
91+
}

0 commit comments

Comments
 (0)