feat: ergonomic access to immutable NodeTable columns (#726)

Author: molpopgen

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};

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())

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);

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);
+}