feat: Individual::into_node_id_iter

src/sys/types.rs

@@ -508,6 +508,21 @@ impl<'p> std::cmp::PartialEq for Individual<'p> {
     }
 }
 
+#[repr(transparent)]
+struct IndividualNodeIdIterator<'ind>(&'ind [super::newtypes::NodeId]);
+
+impl<'ind> Iterator for IndividualNodeIdIterator<'ind> {
+    type Item = super::newtypes::NodeId;
+    fn next(&mut self) -> Option<Self::Item> {
+        if let Some((l, r)) = self.0.split_first() {
+            self.0 = r;
+            Some(*l)
+        } else {
+            None
+        }
+    }
+}
+
 impl<'p> Individual<'p> {
     /// Row id
     #[inline(always)]
@@ -557,6 +572,31 @@ impl<'p> Individual<'p> {
     pub fn nodes(&self) -> Option<&[super::newtypes::NodeId]> {
         general_data_body!(self, row, nodes, nodes_length, super::newtypes::NodeId)
     }
+
+    /// Convert into an iterator over node ids.
+    ///
+    /// One use of this function is to flatten an iterator over individuals
+    /// into an iterator over node ids from those individuals.
+    pub fn into_node_id_iter<'iter>(self) -> impl Iterator<Item = super::newtypes::NodeId> + 'iter
+    where
+        'p: 'iter,
+    {
+        let nodes = if self.row.nodes_length > 0 {
+            assert!(!self.row.nodes.is_null());
+            // SAFETY: the pointer is not null.
+            // The cast works b/c NodeId is a transparent newtype for tsk_id_t.
+            unsafe {
+                std::slice::from_raw_parts(
+                    self.row.nodes.cast::<super::newtypes::NodeId>(),
+                    self.row.nodes_length as usize,
+                )
+            }
+        } else {
+            &[]
+        };
+
+        IndividualNodeIdIterator(nodes)
+    }
 }
 
 /// A lifetime-bound Node.

tests/test_trees.rs

@@ -1160,3 +1160,29 @@ mod from_popgen_oxide {
         assert_eq!(site, ts.sites().num_rows().as_usize())
     }
 }
+
+#[test]
+fn flatten_nodes_from_individuals() {
+    let mut tables = tskit::TableCollection::new(100.).unwrap();
+    let ts = tables.deepcopy().unwrap().tree_sequence(tskit::TreeSequenceFlags::default().build_indexes()).unwrap();
+    let nodes: Vec<tskit::NodeId> = ts
+        .individual_iter()
+        .flat_map(|i| i.into_node_id_iter())
+        .collect::<Vec<_>>();
+    assert!(nodes.is_empty());
+    let _ = tables.add_node(0, 0., -1, 0).unwrap();
+    let _ = tables.add_node(0, 0., -1, 0).unwrap();
+    let _ = tables.add_node(0, 0., -1, 1).unwrap();
+    let _ = tables.add_node(0, 0., -1, 1).unwrap();
+    let _ = tables.add_node(0, 0., -1, 1).unwrap();
+    let _ = tables.add_individual(0, None, None).unwrap();
+    let _ = tables.add_individual(0, None, None).unwrap();
+    let ts = tables
+        .tree_sequence(tskit::TreeSequenceFlags::default().build_indexes())
+        .unwrap();
+    let nodes: Vec<tskit::NodeId> = ts
+        .individual_iter()
+        .flat_map(|i| i.into_node_id_iter())
+        .collect::<Vec<_>>();
+    assert_eq!(nodes, &[0, 1, 2, 3, 4]);
+}