more

Author: molpopgen

src/sys/tree.rs

@@ -5,6 +5,7 @@ use super::bindings::tsk_size_t;
 use super::bindings::tsk_tree_t;
 use super::flags::TreeFlags;
 use super::newtypes::NodeId;
+use super::newtypes::SizeType;
 use super::tskbox::TskBox;
 use super::TreeSequence;
 use super::TskitError;
@@ -69,7 +70,7 @@ impl<'treeseq> LLTree<'treeseq> {
         self.inner.as_mut()
     }
 
-    pub fn as_ptr(&mut self) -> *const tsk_tree_t {
+    pub fn as_ptr(&self) -> *const tsk_tree_t {
         self.inner.as_ptr()
     }
 
@@ -112,6 +113,17 @@ impl<'treeseq> LLTree<'treeseq> {
             .num_rows,
         )
     }
+
+    pub fn num_tracked_samples(&self, u: NodeId) -> Result<SizeType, TskitError> {
+        let mut n = tsk_size_t::MAX;
+        let np: *mut tsk_size_t = &mut n;
+        assert!(!self.as_ptr().is_null());
+        // SAFETY: internal pointer not null and is initialized.
+        let code = unsafe {
+            super::bindings::tsk_tree_get_num_tracked_samples(self.as_ptr(), u.into(), np)
+        };
+        handle_tsk_return_value!(code, n.into())
+    }
 }
 
 // Trait defining iteration over nodes.

src/trees/tree.rs

@@ -2,6 +2,7 @@ use crate::sys::bindings as ll_bindings;
 use crate::sys::{LLTree, TreeSequence};
 use crate::NodeId;
 use crate::Position;
+use crate::SizeType;
 use crate::TreeFlags;
 use crate::TskitError;
 
@@ -115,6 +116,15 @@ impl<'treeseq> Tree<'treeseq> {
     pub fn right_child<N: Into<NodeId> + Copy>(&self, u: N) -> Option<NodeId> {
         self.inner.left_sib(u.into())
     }
+
+    /// Get the number of samples below node `u`.
+    ///
+    /// # Errors
+    ///
+    /// * [`TskitError`] if [`TreeFlags::NO_SAMPLE_COUNTS`].
+    pub fn num_tracked_samples(&self, u: NodeId) -> Result<SizeType, TskitError> {
+        self.inner.num_tracked_samples(u)
+    }
 }
 
 impl<'ts> streaming_iterator::StreamingIterator for Tree<'ts> {