diff --git a/src/lib.rs b/src/lib.rs index 2832659..6142ca9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,7 +52,7 @@ pub struct Tree { /// Node ID. /// /// Index into a `Tree`-internal `Vec`. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct NodeId(NonZeroUsize); impl NodeId { @@ -935,3 +935,5 @@ impl Display for Tree { Ok(()) } } + +mod sort; diff --git a/src/sort.rs b/src/sort.rs new file mode 100644 index 0000000..ca96e23 --- /dev/null +++ b/src/sort.rs @@ -0,0 +1,120 @@ +//! Sorting functionality for tree nodes. +//! +//! 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. + +use std::cmp::Ordering; + +use crate::{NodeMut, NodeRef}; + +impl<'a, T: 'a> NodeMut<'a, T> { + /// Sort children by value in ascending order. + /// + /// # Examples + /// + /// ```rust + /// use ego_tree::tree; + /// + /// let mut tree = tree!('a' => { 'd', 'c', 'b' }); + /// tree.root_mut().sort(); + /// assert_eq!( + /// vec![&'b', &'c', &'d'], + /// tree.root() + /// .children() + /// .map(|n| n.value()) + /// .collect::>(), + /// ); + /// ``` + pub fn sort(&mut self) + where + T: Ord, + { + self.sort_by(|a, b| a.value().cmp(b.value())); + } + + /// Sort children by `NodeRef` in ascending order using a comparison function. + /// + /// # Examples + /// + /// ```rust + /// use ego_tree::tree; + /// + /// let mut tree = tree!('a' => { 'c', 'd', 'b' }); + /// tree.root_mut().sort_by(|a, b| b.value().cmp(a.value())); + /// assert_eq!( + /// vec![&'d', &'c', &'b'], + /// tree.root() + /// .children() + /// .map(|n| n.value()) + /// .collect::>(), + /// ); + /// + /// // Example for sort_by_id. + /// tree.root_mut().sort_by(|a, b| a.id().cmp(&b.id())); + /// assert_eq!( + /// vec![&'c', &'d', &'b'], + /// tree.root() + /// .children() + /// .map(|n| n.value()) + /// .collect::>(), + /// ); + /// ``` + pub fn sort_by(&mut self, mut compare: F) + where + F: FnMut(NodeRef, NodeRef) -> Ordering, + { + if !self.has_children() { + return; + } + + let mut children = { + let this = unsafe { self.tree.get_unchecked(self.id) }; + this.children().map(|child| child.id).collect::>() + }; + + children.sort_by(|a, b| { + let a = unsafe { self.tree.get_unchecked(*a) }; + let b = unsafe { self.tree.get_unchecked(*b) }; + compare(a, b) + }); + + for id in children { + self.append_id(id); + } + } + + /// Sort children by `NodeRef`'s key in ascending order using a key extraction function. + /// + /// # Examples + /// + /// ```rust + /// use ego_tree::tree; + /// + /// let mut tree = tree!("1a" => { "2b", "4c", "3d" }); + /// tree.root_mut().sort_by_key(|a| a.value().split_at(1).0.parse::().unwrap()); + /// assert_eq!( + /// vec!["2b", "3d", "4c"], + /// tree.root() + /// .children() + /// .map(|n| *n.value()) + /// .collect::>(), + /// ); + /// + /// // Example for sort_by_id. + /// tree.root_mut().sort_by_key(|n| n.id()); + /// assert_eq!( + /// vec![&"2b", &"4c", &"3d"], + /// tree.root() + /// .children() + /// .map(|n| n.value()) + /// .collect::>(), + /// ); + /// ``` + pub fn sort_by_key(&mut self, mut f: F) + where + F: FnMut(NodeRef) -> K, + K: Ord, + { + self.sort_by(|a, b| f(a).cmp(&f(b))); + } +} diff --git a/tests/sort.rs b/tests/sort.rs new file mode 100644 index 0000000..e94e07f --- /dev/null +++ b/tests/sort.rs @@ -0,0 +1,91 @@ +use std::assert_eq; + +use ego_tree::tree; + +#[test] +fn sort() { + let mut tree = tree!('a' => { 'd' => { 'e', 'f' }, 'c', 'b' }); + tree.root_mut().sort(); + assert_eq!( + vec![&'b', &'c', &'d'], + tree.root() + .children() + .map(|n| n.value()) + .collect::>(), + ); + assert_eq!( + tree.to_string(), + tree!('a' => { 'b', 'c', 'd' => { 'e', 'f' } }).to_string() + ); +} + +#[test] +fn sort_by() { + let mut tree = tree!('a' => { 'c', 'd', 'b' }); + tree.root_mut().sort_by(|a, b| b.value().cmp(a.value())); + assert_eq!( + vec![&'d', &'c', &'b'], + tree.root() + .children() + .map(|n| n.value()) + .collect::>(), + ); + + let mut tree = tree!('a' => { 'c','d', 'e', 'b' }); + tree.root_mut().sort_by(|a, b| b.value().cmp(a.value())); + assert_eq!( + vec![&'e', &'d', &'c', &'b'], + tree.root() + .children() + .map(|n| n.value()) + .collect::>(), + ); +} + +#[test] +fn sort_by_key() { + let mut tree = tree!("1a" => { "2b", "4c", "3d" }); + tree.root_mut() + .sort_by_key(|a| a.value().split_at(1).0.parse::().unwrap()); + assert_eq!( + vec!["2b", "3d", "4c"], + tree.root() + .children() + .map(|n| *n.value()) + .collect::>(), + ); +} + +#[test] +fn sort_id() { + let mut tree = tree!('a' => { 'd', 'c', 'b' }); + tree.root_mut().sort(); + assert_ne!( + vec![&'d', &'c', &'b'], + tree.root() + .children() + .map(|n| n.value()) + .collect::>(), + ); + tree.root_mut().sort_by_key(|n| n.id()); + assert_eq!( + vec![&'d', &'c', &'b'], + tree.root() + .children() + .map(|n| n.value()) + .collect::>(), + ); +} + +#[test] +fn sort_by_id() { + let mut tree = tree!('a' => { 'd', 'b', 'c' }); + tree.root_mut().sort_by(|a, b| b.id().cmp(&a.id())); + assert_eq!( + vec![&'c', &'b', &'d'], + tree.root() + .children() + .map(|n| n.value()) + .collect::>(), + ); +}