From f5860a07a3a7bfe895d897e2b419ba89a715940b Mon Sep 17 00:00:00 2001
From: "Kevin R. Thornton" <krthornt@uci.edu>
Date: Sun, 18 May 2025 11:02:17 -0700
Subject: [PATCH] refactor: TableColumn trait for column access

replaces concrete types added in #726, #730

BREAKING CHANGE: requires MSRV bump
---
 .github/workflows/tests.yml |   2 +-
 Cargo.toml                  |   2 +-
 src/edge_table.rs           |  33 ++++++++---
 src/lib.rs                  |   2 +-
 src/node_table.rs           |  16 +++---
 src/table_column.rs         |  70 +++++------------------
 src/traits.rs               | 108 ++++++++++++++++++++++++++++++++++++
 tests/test_tables.rs        |   4 ++
 8 files changed, 162 insertions(+), 75 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 9fb803212..2c48ea04b 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -83,7 +83,7 @@ jobs:
     strategy:
       matrix:
         rust:
-          - 1.71.0
+          - 1.75.0
     steps:
       - uses: actions/checkout@v4.2.2
       - uses: dtolnay/rust-toolchain@v1
diff --git a/Cargo.toml b/Cargo.toml
index 099b9640a..a43a771d1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,7 +8,7 @@ description = "rust interface to tskit"
 license = "MIT"
 homepage = "https://github.com/tskit-dev/tskit-rust"
 repository = "https://github.com/tskit-dev/tskit-rust"
-rust-version = "1.71.0"
+rust-version = "1.75.0"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [lints.rust]
diff --git a/src/edge_table.rs b/src/edge_table.rs
index 93aefd92f..917fbec61 100644
--- a/src/edge_table.rs
+++ b/src/edge_table.rs
@@ -358,20 +358,37 @@ impl EdgeTable {
         /// Get the child column as a slice of the underlying integer type
         => child, child_slice_raw, ll_bindings::tsk_id_t);
 
-    pub fn parent_column(&self) -> crate::EdgeTableColumn<crate::NodeId> {
-        crate::EdgeTableColumn::new(self.parent_slice())
+    /// Table column with ergonomic indexing.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// use tskit::TableColumn;
+    /// let mut edges = tskit::EdgeTable::new().unwrap();
+    /// // left, right, parent, child
+    /// let edge: tskit::EdgeId = edges.add_row(0., 10., 1, 0).unwrap();
+    /// let p = edges.parent_column();
+    /// assert_eq!(p[edge], 1);
+    /// assert_eq!(p.get_with_id(edge), Some(&tskit::NodeId::from(1)));
+    /// assert!(p.get_with_id(tskit::EdgeId::NULL).is_none())
+    /// ```
+    pub fn parent_column(&self) -> impl crate::TableColumn<EdgeId, NodeId> + '_ {
+        crate::table_column::OpaqueTableColumn(self.parent_slice())
     }
 
-    pub fn child_column(&self) -> crate::EdgeTableColumn<crate::NodeId> {
-        crate::EdgeTableColumn::new(self.child_slice())
+    /// Table column with ergonomic indexing.
+    pub fn child_column(&self) -> impl crate::TableColumn<EdgeId, NodeId> + '_ {
+        crate::table_column::OpaqueTableColumn(self.child_slice())
     }
 
-    pub fn left_column(&self) -> crate::EdgeTableColumn<Position> {
-        crate::EdgeTableColumn::new(self.left_slice())
+    /// Table column with ergonomic indexing.
+    pub fn left_column(&self) -> impl crate::TableColumn<EdgeId, Position> + '_ {
+        crate::table_column::OpaqueTableColumn(self.left_slice())
     }
 
-    pub fn right_column(&self) -> crate::EdgeTableColumn<Position> {
-        crate::EdgeTableColumn::new(self.right_slice())
+    /// Table column with ergonomic indexing.
+    pub fn right_column(&self) -> impl crate::TableColumn<EdgeId, Position> + '_ {
+        crate::table_column::OpaqueTableColumn(self.right_slice())
     }
 
     /// Clear all data from the table
diff --git a/src/lib.rs b/src/lib.rs
index 1729df519..69f56154f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -117,7 +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::{EdgeTableColumn, NodeTableColumn};
+pub use traits::TableColumn;
 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 2462b20d8..0d4bb3667 100644
--- a/src/node_table.rs
+++ b/src/node_table.rs
@@ -807,20 +807,20 @@ 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 individual_column(&self) -> impl crate::TableColumn<NodeId, IndividualId> + '_ {
+        crate::table_column::OpaqueTableColumn(self.individual_slice())
     }
 
-    pub fn population_column(&self) -> crate::NodeTableColumn<PopulationId> {
-        crate::NodeTableColumn::new(self.population_slice())
+    pub fn population_column(&self) -> impl crate::TableColumn<NodeId, PopulationId> + '_ {
+        crate::table_column::OpaqueTableColumn(self.population_slice())
     }
 
-    pub fn time_column(&self) -> crate::NodeTableColumn<Time> {
-        crate::NodeTableColumn::new(self.time_slice())
+    pub fn time_column(&self) -> impl crate::TableColumn<NodeId, Time> + '_ {
+        crate::table_column::OpaqueTableColumn(self.time_slice())
     }
 
-    pub fn flags_column(&self) -> crate::NodeTableColumn<NodeFlags> {
-        crate::NodeTableColumn::new(self.flags_slice())
+    pub fn flags_column(&self) -> impl crate::TableColumn<NodeId, NodeFlags> + '_ {
+        crate::table_column::OpaqueTableColumn(self.flags_slice())
     }
 
     /// Clear all data from the table
diff --git a/src/table_column.rs b/src/table_column.rs
index 298f367da..dce977f99 100644
--- a/src/table_column.rs
+++ b/src/table_column.rs
@@ -1,61 +1,19 @@
-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]);
+#[repr(transparent)]
+#[derive(Clone)]
+pub(crate) struct OpaqueTableColumn<'table, T>(pub(crate) &'table [T]);
 
-        impl<'table, T> $name<'table, T> {
-            pub(crate) fn new(column: &'table [T]) -> $name<'table, T> {
-                Self(column)
-            }
+impl<T> std::ops::Index<usize> for OpaqueTableColumn<'_, T> {
+    type Output = T;
 
-            /// View the underlying slice
-            pub fn as_slice(&self) -> &[T] {
-                self.0
-            }
-
-            pub fn get_with_id(&self, index: crate::$index) -> Option<&T> {
-                self.get_with_usize(usize::try_from(index).ok()?)
-            }
-
-            pub fn get_with_size_type(&self, index: crate::SizeType) -> Option<&T> {
-                self.get_with_usize(usize::try_from(index).ok()?)
-            }
-
-            pub fn get_with_usize(&self, index: usize) -> Option<&T> {
-                self.0.get(index)
-            }
-        }
-
-        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()]
-            }
-        }
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.0[index]
+    }
+}
 
-        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::ops::Index<crate::SizeType> for OpaqueTableColumn<'_, T> {
+    type Output = T;
 
-        impl<T> std::convert::AsRef<[T]> for $name<'_, T> {
-            fn as_ref(&self) -> &[T] {
-                self.0
-            }
-        }
-    };
+    fn index(&self, index: crate::SizeType) -> &Self::Output {
+        &self.0[usize::try_from(index).unwrap()]
+    }
 }
-
-make_table_column!(NodeTableColumn, NodeId);
-make_table_column!(EdgeTableColumn, EdgeId);
diff --git a/src/traits.rs b/src/traits.rs
index fbc03532d..0466cd737 100644
--- a/src/traits.rs
+++ b/src/traits.rs
@@ -119,3 +119,111 @@ impl_individual_parents!(
 );
 impl_individual_parents!(N, usize, &[crate::IndividualId; N], self, self.as_slice());
 impl_individual_parents!(N, usize, [crate::IndividualId; N], self, self.as_slice());
+
+mod private {
+    pub trait NewTypeMarker: TryInto<usize, Error = crate::TskitError> {}
+    pub trait TableColumnMarker {}
+}
+
+impl private::NewTypeMarker for crate::EdgeId {}
+impl private::NewTypeMarker for crate::NodeId {}
+impl private::NewTypeMarker for crate::SiteId {}
+impl private::NewTypeMarker for crate::MutationId {}
+impl private::NewTypeMarker for crate::MigrationId {}
+impl private::NewTypeMarker for crate::IndividualId {}
+impl private::NewTypeMarker for crate::PopulationId {}
+#[cfg(feature = "provenance")]
+#[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
+impl private::NewTypeMarker for crate::ProvenanceId {}
+
+/// Interface of a non-ragged table column.
+///
+/// Unlike slice views of table columns, this API
+/// allows indexed via row id types and [`crate::SizeType`].
+///
+/// # Notes
+///
+/// * This trait is sealed.
+///
+/// # For C programmers
+///
+/// The `C` programming language allows implicit casts between
+/// integer types.
+/// This implicit behavior allows one to index a table column
+/// using a row id type ([`crate::bindings::tsk_id_t`]) because
+/// the compiler will cast it to `size_t`.
+///
+/// `rust` does not allow implicit casts, which makes working
+/// with table columns as slices awkward.
+/// One has to manually cast the id type and the resulting code isn't
+/// nice to read.
+///
+/// This trait solves that problem by requiring that [`std::ops::Index`]
+/// by implemented for types that one would like to use as indexes
+/// in the `tskit` world.
+pub trait TableColumn<I, T>:
+    std::ops::Index<I, Output = T>
+    + std::ops::Index<usize, Output = T>
+    + std::ops::Index<crate::SizeType, Output = T>
+    + private::TableColumnMarker
+{
+    /// Get the underlying slice
+    fn as_slice(&self) -> &[T];
+
+    /// Get with a table row identifier such as [`crate::NodeId`]
+    fn get_with_id(&self, at: I) -> Option<&T>;
+
+    /// The "standard" get function
+    fn get(&self, at: usize) -> Option<&T> {
+        self.as_slice().get(at)
+    }
+
+    /// Get with [`crate::SizeType`]
+    fn get_with_size_type(&self, at: crate::SizeType) -> Option<&T> {
+        self.as_slice().get(usize::try_from(at).ok()?)
+    }
+
+    /// Iterator over the data.
+    fn iter<'a, 'b>(&'a self) -> impl Iterator<Item = &'b T>
+    where
+        'a: 'b,
+        T: 'b,
+    {
+        self.as_slice().iter()
+    }
+
+    /// Column length
+    fn len(&self) -> usize {
+        self.as_slice().len()
+    }
+
+    /// Query if column is empty
+    fn is_empty(&self) -> bool {
+        self.as_slice().is_empty()
+    }
+}
+
+impl<T> private::TableColumnMarker for crate::table_column::OpaqueTableColumn<'_, T> {}
+
+impl<I, T> std::ops::Index<I> for crate::table_column::OpaqueTableColumn<'_, T>
+where
+    I: private::NewTypeMarker,
+{
+    type Output = T;
+    fn index(&self, index: I) -> &Self::Output {
+        &self.0[index.try_into().unwrap()]
+    }
+}
+
+impl<I, T> TableColumn<I, T> for crate::table_column::OpaqueTableColumn<'_, T>
+where
+    I: private::NewTypeMarker,
+{
+    fn as_slice(&self) -> &[T] {
+        self.0
+    }
+
+    fn get_with_id(&self, at: I) -> Option<&T> {
+        self.0.get(at.try_into().ok()?)
+    }
+}
diff --git a/tests/test_tables.rs b/tests/test_tables.rs
index b1c68f1fe..da773e2c1 100644
--- a/tests/test_tables.rs
+++ b/tests/test_tables.rs
@@ -1,3 +1,5 @@
+use tskit::TableColumn;
+
 #[test]
 fn test_empty_table_collection() {
     macro_rules! validate_empty_tables {
@@ -610,6 +612,8 @@ fn test_edge_table_column_access() {
         let column = table.parent_column();
         assert_eq!(column[edge], 1);
         assert_eq!(column[edge], tskit::NodeId::from(1));
+        let _: Vec<tskit::NodeId> = table.parent_column().iter().cloned().collect::<Vec<_>>();
+        let _ = table.parent_column().iter().cloned().collect::<Vec<_>>();
     }
 
     {