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}