From 00b82af3d7804f0631d1991fac006e6e17f35606 Mon Sep 17 00:00:00 2001
From: "Kevin R. Thornton" <krthornt@uci.edu>
Date: Tue, 11 Mar 2025 09:10:03 -0700
Subject: [PATCH] feat: ergonomic access to immutable NodeTable columns

---
 src/lib.rs           |  2 ++
 src/node_table.rs    | 16 +++++++++++++++
 src/table_column.rs  | 48 ++++++++++++++++++++++++++++++++++++++++++++
 tests/test_tables.rs | 16 +++++++++++++++
 4 files changed, 82 insertions(+)
 create mode 100644 src/table_column.rs

diff --git a/src/lib.rs b/src/lib.rs
index 3b5fe0e0b..0aad69354 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -97,6 +97,7 @@ pub mod prelude;
 mod site_table;
 mod sys;
 mod table_collection;
+mod table_column;
 mod table_iterator;
 mod traits;
 mod trees;
@@ -116,6 +117,7 @@ pub use site_table::{SiteTable, SiteTableRow};
 pub use sys::flags::*;
 pub use sys::NodeTraversalOrder;
 pub use table_collection::TableCollection;
+pub use table_column::NodeTableColumn;
 pub use traits::IndividualLocation;
 pub use traits::IndividualParents;
 pub use trees::{Tree, TreeSequence};
diff --git a/src/node_table.rs b/src/node_table.rs
index 4ce506690..2462b20d8 100644
--- a/src/node_table.rs
+++ b/src/node_table.rs
@@ -807,6 +807,22 @@ impl NodeTable {
         /// Get the population column as a slice
         => population, population_slice_raw, crate::sys::bindings::tsk_id_t);
 
+    pub fn individual_column(&self) -> crate::table_column::NodeTableColumn<IndividualId> {
+        crate::NodeTableColumn::new(self.individual_slice())
+    }
+
+    pub fn population_column(&self) -> crate::NodeTableColumn<PopulationId> {
+        crate::NodeTableColumn::new(self.population_slice())
+    }
+
+    pub fn time_column(&self) -> crate::NodeTableColumn<Time> {
+        crate::NodeTableColumn::new(self.time_slice())
+    }
+
+    pub fn flags_column(&self) -> crate::NodeTableColumn<NodeFlags> {
+        crate::NodeTableColumn::new(self.flags_slice())
+    }
+
     /// Clear all data from the table
     pub fn clear(&mut self) -> Result<i32, TskitError> {
         handle_tsk_return_value!(self.table_.clear())
diff --git a/src/table_column.rs b/src/table_column.rs
new file mode 100644
index 000000000..1cdc96fbc
--- /dev/null
+++ b/src/table_column.rs
@@ -0,0 +1,48 @@
+macro_rules! make_table_column {
+    ($name: ident, $index: ident) => {
+        /// Immutable view of a column
+        #[derive(Clone, Debug)]
+        #[repr(transparent)]
+        pub struct $name<'table, T>(&'table [T]);
+
+        impl<'table, T> $name<'table, T> {
+            pub(crate) fn new(column: &'table [T]) -> $name<'table, T> {
+                Self(column)
+            }
+
+            /// View the underlying slice
+            pub fn as_slice(&self) -> &[T] {
+                self.0
+            }
+        }
+
+        impl<T> std::ops::Index<usize> for $name<'_, T> {
+            type Output = T;
+            fn index(&self, index: usize) -> &Self::Output {
+                &self.0[index]
+            }
+        }
+
+        impl<T> std::ops::Index<crate::$index> for $name<'_, T> {
+            type Output = T;
+            fn index(&self, index: crate::$index) -> &Self::Output {
+                &self.0[usize::try_from(index).unwrap()]
+            }
+        }
+
+        impl<T> std::ops::Index<crate::SizeType> for $name<'_, T> {
+            type Output = T;
+            fn index(&self, index: crate::SizeType) -> &Self::Output {
+                &self.0[usize::try_from(index).unwrap()]
+            }
+        }
+
+        impl<T> std::convert::AsRef<[T]> for $name<'_, T> {
+            fn as_ref(&self) -> &[T] {
+                self.0
+            }
+        }
+    };
+}
+
+make_table_column!(NodeTableColumn, NodeId);
diff --git a/tests/test_tables.rs b/tests/test_tables.rs
index 964c30130..793b95e9b 100644
--- a/tests/test_tables.rs
+++ b/tests/test_tables.rs
@@ -550,3 +550,19 @@ mod test_metadata_round_trips {
         );
     }
 }
+
+#[test]
+fn test_node_table_column_access() {
+    let mut t = tskit::NodeTable::new().unwrap();
+    let node = t
+        .add_row(tskit::NodeFlags::new_sample(), 0.0, -1, -1)
+        .unwrap();
+    let individual = t.individual_column();
+    assert_eq!(individual[node], tskit::IndividualId::NULL);
+    let population = t.population_column();
+    assert_eq!(population[node], tskit::PopulationId::NULL);
+    let time = t.time_column();
+    assert_eq!(time[node], 0.0);
+    let flags = t.flags_column();
+    assert_eq!(flags[node], tskit::NodeFlags::IS_SAMPLE);
+}