Skip to content

Commit fbce6ca

Browse files
authored
API fixes. (#104)
* Add Tree::parents. Closes #103 * Rename parent node iterator * Change return type of Tree::traverse_nodes so that it is easier to add other traversal orders. Closes #102. * Add Tree/TreeSequence ::sample_nodes to give non-owning access to sample nodes via a slice. Deprecate ::samples_to_vec.
1 parent c431880 commit fbce6ca

File tree

3 files changed

+51
-22
lines changed

3 files changed

+51
-22
lines changed

examples/tree_traversals.rs

+4-8
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,8 @@ use streaming_iterator::StreamingIterator; // Required for tree iteration
33

44
// "Manual" traversal from samples to root
55
fn traverse_upwards(tree: &tskit::Tree) {
6-
let samples = tree.samples_to_vec();
7-
8-
for s in samples.iter() {
9-
let mut u = *s;
6+
for &s in tree.sample_nodes() {
7+
let mut u = s;
108
while u != tskit::TSK_NULL {
119
u = tree.parent(u).unwrap();
1210
}
@@ -15,12 +13,10 @@ fn traverse_upwards(tree: &tskit::Tree) {
1513

1614
// Iterate from each node up to its root.
1715
fn traverse_upwards_with_iterator(tree: &tskit::Tree) {
18-
let samples = tree.samples_to_vec();
19-
20-
for s in samples.iter() {
16+
for &s in tree.sample_nodes() {
2117
// _steps_to_root counts the number of steps,
2218
// including the starting node s.
23-
for (_steps_to_root, _) in tree.path_to_root(*s).unwrap().enumerate() {}
19+
for (_steps_to_root, _) in tree.parents(s).unwrap().enumerate() {}
2420
}
2521
}
2622

src/test_simplification.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -30,14 +30,14 @@ mod tests {
3030
#[test]
3131
fn test_simplify_treeseq() {
3232
let ts = treeseq_from_small_table_collection_two_trees();
33-
let samples = ts.samples_to_vec();
33+
let samples = ts.sample_nodes();
3434
let (_, idmap_option) = ts
3535
.simplify(&samples, SimplificationOptions::default(), true)
3636
.unwrap();
3737
assert!(idmap_option.is_some());
3838
let idmap = idmap_option.unwrap();
39-
for i in samples.iter() {
40-
assert_ne!(idmap[*i as usize], TSK_NULL);
39+
for &i in samples {
40+
assert_ne!(idmap[i as usize], TSK_NULL);
4141
}
4242
}
4343
}

src/trees.rs

+44-11
Original file line numberDiff line numberDiff line change
@@ -391,6 +391,7 @@ impl Tree {
391391

392392
/// Obtain the list of samples for the current tree/tree sequence
393393
/// as a vector.
394+
#[deprecated(since = "0.2.3", note = "Please use Tree::sample_nodes instead")]
394395
pub fn samples_to_vec(&self) -> Vec<tsk_id_t> {
395396
let num_samples =
396397
unsafe { ll_bindings::tsk_treeseq_get_num_samples((*self.as_ptr()).tree_sequence) };
@@ -403,16 +404,34 @@ impl Tree {
403404
rv
404405
}
405406

407+
/// Get the list of sample nodes as a slice.
408+
pub fn sample_nodes(&self) -> &[tsk_id_t] {
409+
let num_samples =
410+
unsafe { ll_bindings::tsk_treeseq_get_num_samples((*self.as_ptr()).tree_sequence) };
411+
tree_array_slice!(self, samples, num_samples)
412+
}
413+
406414
/// Return an [`Iterator`] from the node `u` to the root of the tree.
407415
///
408416
/// # Errors
409417
///
410418
/// [`TskitError::IndexError`] if `u` is out of range.
419+
#[deprecated(since = "0.2.3", note = "Please use Tree::parents instead")]
411420
pub fn path_to_root(
412421
&self,
413422
u: tsk_id_t,
414423
) -> Result<impl Iterator<Item = tsk_id_t> + '_, TskitError> {
415-
PathToRootIterator::new(self, u)
424+
self.parents(u)
425+
}
426+
427+
/// Return an [`Iterator`] from the node `u` to the root of the tree,
428+
/// travering all parent nodes.
429+
///
430+
/// # Errors
431+
///
432+
/// [`TskitError::IndexError`] if `u` is out of range.
433+
pub fn parents(&self, u: tsk_id_t) -> Result<impl Iterator<Item = tsk_id_t> + '_, TskitError> {
434+
ParentsIterator::new(self, u)
416435
}
417436

418437
/// Return an [`Iterator`] over the children of node `u`.
@@ -466,9 +485,13 @@ impl Tree {
466485
///
467486
/// * `order`: A value from [`NodeTraversalOrder`] specifying the
468487
/// iteration order.
469-
pub fn traverse_nodes(&self, order: NodeTraversalOrder) -> impl Iterator<Item = tsk_id_t> + '_ {
488+
// Return value is dyn for later addition of other traversal orders
489+
pub fn traverse_nodes(
490+
&self,
491+
order: NodeTraversalOrder,
492+
) -> Box<dyn Iterator<Item = tsk_id_t> + '_> {
470493
match order {
471-
NodeTraversalOrder::Preorder => PreorderNodeIterator::new(&self),
494+
NodeTraversalOrder::Preorder => Box::new(PreorderNodeIterator::new(&self)),
472495
}
473496
}
474497

@@ -719,17 +742,17 @@ impl NodeIterator for ChildIterator<'_> {
719742

720743
iterator_for_nodeiterator!(ChildIterator<'_>);
721744

722-
struct PathToRootIterator<'a> {
745+
struct ParentsIterator<'a> {
723746
current_node: Option<tsk_id_t>,
724747
next_node: tsk_id_t,
725748
tree: &'a Tree,
726749
}
727750

728-
impl<'a> PathToRootIterator<'a> {
751+
impl<'a> ParentsIterator<'a> {
729752
fn new(tree: &'a Tree, u: tsk_id_t) -> Result<Self, TskitError> {
730753
match u >= tree.num_nodes as tsk_id_t {
731754
true => Err(TskitError::IndexError),
732-
false => Ok(PathToRootIterator {
755+
false => Ok(ParentsIterator {
733756
current_node: None,
734757
next_node: u,
735758
tree,
@@ -738,7 +761,7 @@ impl<'a> PathToRootIterator<'a> {
738761
}
739762
}
740763

741-
impl NodeIterator for PathToRootIterator<'_> {
764+
impl NodeIterator for ParentsIterator<'_> {
742765
fn next_node(&mut self) {
743766
self.current_node = match self.next_node {
744767
TSK_NULL => None,
@@ -756,7 +779,7 @@ impl NodeIterator for PathToRootIterator<'_> {
756779
}
757780
}
758781

759-
iterator_for_nodeiterator!(PathToRootIterator<'_>);
782+
iterator_for_nodeiterator!(ParentsIterator<'_>);
760783

761784
struct SamplesIterator<'a> {
762785
current_node: Option<tsk_id_t>,
@@ -978,6 +1001,10 @@ impl TreeSequence {
9781001
}
9791002

9801003
/// Get the list of samples as a vector.
1004+
#[deprecated(
1005+
since = "0.2.3",
1006+
note = "Please use TreeSequence::sample_nodes instead"
1007+
)]
9811008
pub fn samples_to_vec(&self) -> Vec<tsk_id_t> {
9821009
let num_samples = unsafe { ll_bindings::tsk_treeseq_get_num_samples(self.as_ptr()) };
9831010
let mut rv = vec![];
@@ -989,6 +1016,12 @@ impl TreeSequence {
9891016
rv
9901017
}
9911018

1019+
/// Get the list of sample nodes as a slice.
1020+
pub fn sample_nodes(&self) -> &[tsk_id_t] {
1021+
let num_samples = unsafe { ll_bindings::tsk_treeseq_get_num_samples(self.as_ptr()) };
1022+
tree_array_slice!(self, samples, num_samples)
1023+
}
1024+
9921025
/// Get the number of trees.
9931026
pub fn num_trees(&self) -> tsk_size_t {
9941027
unsafe { ll_bindings::tsk_treeseq_get_num_trees(self.as_ptr()) }
@@ -1146,7 +1179,7 @@ pub(crate) mod test_trees {
11461179
fn test_create_treeseq_new_from_tables() {
11471180
let tables = make_small_table_collection();
11481181
let treeseq = TreeSequence::new(tables, TreeSequenceFlags::default()).unwrap();
1149-
let samples = treeseq.samples_to_vec();
1182+
let samples = treeseq.sample_nodes();
11501183
assert_eq!(samples.len(), 2);
11511184
for i in 1..3 {
11521185
assert_eq!(samples[i - 1], i as tsk_id_t);
@@ -1168,13 +1201,13 @@ pub(crate) mod test_trees {
11681201
while let Some(tree) = tree_iter.next() {
11691202
ntrees += 1;
11701203
assert_eq!(tree.current_tree, ntrees);
1171-
let samples = tree.samples_to_vec();
1204+
let samples = tree.sample_nodes();
11721205
assert_eq!(samples.len(), 2);
11731206
for i in 1..3 {
11741207
assert_eq!(samples[i - 1], i as tsk_id_t);
11751208

11761209
let mut nsteps = 0;
1177-
for _ in tree.path_to_root(samples[i - 1]).unwrap() {
1210+
for _ in tree.parents(samples[i - 1]).unwrap() {
11781211
nsteps += 1;
11791212
}
11801213
assert_eq!(nsteps, 2);

0 commit comments

Comments
 (0)