refactor: TableColumn trait for column access

replaces concrete types added in #726, #730

Author: molpopgen

.github/workflows/tests.yml

@@ -83,7 +83,7 @@ jobs:
     strategy:
       matrix:
         rust:
-          - 1.71.0
+          - 1.75.0
     steps:
       - uses: actions/[email protected]
       - uses: dtolnay/rust-toolchain@v1

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]

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

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

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

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

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()?)
+    }
+}

tests/test_tables.rs

@@ -1,3 +1,5 @@
+use tskit::TableColumn;
+
 #[test]
 fn test_empty_table_collection() {
     macro_rules! validate_empty_tables {