cool

Author: molpopgen

src/sys/tree.rs

@@ -265,6 +265,16 @@ impl<'treeseq> LLTree<'treeseq> {
         )
     }
 
+    pub fn next_sample_array(&self) -> Result<&[NodeId], TskitError> {
+        err_if_not_tracking_samples!(
+            self.flags,
+            super::generate_slice(
+                self.as_ll_ref().next_sample,
+                self.treeseq.num_nodes_raw() + 1
+            )
+        )
+    }
+
     pub fn kc_distance(&self, other: &Self, lambda: f64) -> Result<f64, TskitError> {
         let mut kc = f64::NAN;
         let kcp: *mut f64 = &mut kc;

src/trees/tree.rs

@@ -335,8 +335,38 @@ impl<'treeseq> Tree<'treeseq> {
         self.inner.as_ll_ref()
     }
 
-    pub fn next_sample_array(&self) {
-        X
+    /// # Failing examples
+    ///
+    /// An error will be returned if ['crate::TreeFlags::SAMPLE_LISTS`] is not used:
+    ///
+    /// ```should_panic
+    /// use streaming_iterator::StreamingIterator;
+    /// let tables = tskit::TableCollection::new(1.).unwrap();
+    /// let treeseq =
+    /// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
+    /// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap(); // ERROR
+    /// while let Some(tree) = tree_iter.next() {
+    ///     let n = tree.next_sample_array().unwrap();
+    ///     for _ in n {}
+    /// }
+    /// ```
+    ///
+    /// The lifetime of the slice is tied to the parent object:
+    ///
+    /// ```compile_fail
+    /// use streaming_iterator::StreamingIterator;
+    /// let tables = tskit::TableCollection::new(1.).unwrap();
+    /// let treeseq =
+    /// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
+    /// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
+    /// while let Some(tree) = tree_iter.next() {
+    ///     let n = tree.next_sample_array().unwrap();
+    ///     drop(tree_iter);
+    ///     for _ in n {} // ERROR
+    /// }
+    /// ```
+    pub fn next_sample_array(&self) -> Result<&[NodeId], TskitError> {
+        self.inner.next_sample_array()
     }
 
     /// # Failing examples