From b7744efb40f15fe832c70a7986522c567081dfb5 Mon Sep 17 00:00:00 2001
From: Kevin Thornton <krthornt@uci.edu>
Date: Tue, 12 May 2026 14:54:18 -0700
Subject: [PATCH] feat: Individual::into_node_id_iter

---
 src/sys/types.rs    | 40 ++++++++++++++++++++++++++++++++++++++++
 tests/test_trees.rs | 26 ++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/src/sys/types.rs b/src/sys/types.rs
index 2a4a5ebb..0795b5ab 100644
--- a/src/sys/types.rs
+++ b/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.
diff --git a/tests/test_trees.rs b/tests/test_trees.rs
index 9c26017c..ddf0932d 100644
--- a/tests/test_trees.rs
+++ b/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]);
+}