diff --git a/Cargo.toml b/Cargo.toml index c89ff5338..5254781ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ serde_json = {version = "1.0.67", optional = true} bincode = {version = "1.3.1", optional = true} tskit-derive = {version = "0.2.0", path = "tskit-derive", optional = true} mbox = "0.6.0" +delegate = "0.8.0" [dev-dependencies] anyhow = {version = "1.0.66"} diff --git a/src/_macros.rs b/src/_macros.rs index b6a0edfa2..f4b5e5171 100644 --- a/src/_macros.rs +++ b/src/_macros.rs @@ -1191,6 +1191,121 @@ macro_rules! build_table_column_slice_mut_getter { }; } +macro_rules! delegate_table_view_api { + () => { + delegate::delegate! { + to self.views { + /// Get reference to the [``EdgeTable``](crate::EdgeTable). + pub fn edges(&self) -> &crate::EdgeTable; + /// Get reference to the [``IndividualTable``](crate::IndividualTable). + pub fn individuals(&self) -> &crate::IndividualTable; + /// Get reference to the [``MigrationTable``](crate::MigrationTable). + pub fn migrations(&self) -> &crate::MigrationTable; + /// Get reference to the [``MutationTable``](crate::MutationTable). + pub fn mutations(&self) -> &crate::MutationTable; + /// Get reference to the [``NodeTable``](crate::NodeTable). + pub fn nodes(&self) -> &crate::NodeTable; + /// Get reference to the [``PopulationTable``](crate::PopulationTable). + pub fn populations(&self) -> &crate::PopulationTable; + /// Get reference to the [``SiteTable``](crate::SiteTable). + pub fn sites(&self) -> &crate::SiteTable; + + #[cfg(feature = "provenance")] + #[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))] + /// Get reference to the [``ProvenanceTable``](crate::provenance::ProvenanceTable) + pub fn provenances(&self) -> &crate::provenance::ProvenanceTable ; + + /// Return an iterator over the individuals. + pub fn individuals_iter(&self) -> impl Iterator + '_; + /// Return an iterator over the nodes. + pub fn nodes_iter(&self) -> impl Iterator + '_; + /// Return an iterator over the edges. + pub fn edges_iter(&self) -> impl Iterator + '_; + /// Return an iterator over the migrations. + pub fn migrations_iter(&self) -> impl Iterator + '_; + /// Return an iterator over the mutations. + pub fn mutations_iter(&self) -> impl Iterator + '_; + /// Return an iterator over the populations. + pub fn populations_iter(&self) -> impl Iterator + '_; + /// Return an iterator over the sites. + pub fn sites_iter(&self) -> impl Iterator + '_; + + #[cfg(feature = "provenance")] + #[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))] + /// Return an iterator over provenances + pub fn provenances_iter(&self,) -> impl Iterator + '_; + + /// Obtain a vector containing the indexes ("ids") + /// of all nodes for which [`crate::TSK_NODE_IS_SAMPLE`] + /// is `true`. + /// + /// The provided implementation dispatches to + /// [`crate::NodeTable::samples_as_vector`]. + pub fn samples_as_vector(&self) -> Vec; + + /// Obtain a vector containing the indexes ("ids") of all nodes + /// satisfying a certain criterion. + /// + /// The provided implementation dispatches to + /// [`crate::NodeTable::create_node_id_vector`]. + /// + /// # Parameters + /// + /// * `f`: a function. The function is passed the current table + /// collection and each [`crate::node_table::NodeTableRow`]. + /// If `f` returns `true`, the index of that row is included + /// in the return value. + /// + /// # Examples + /// + /// Get all nodes with time > 0.0: + /// + /// ``` + /// use tskit::bindings::tsk_id_t; + /// + /// let mut tables = tskit::TableCollection::new(100.).unwrap(); + /// tables + /// .add_node(tskit::TSK_NODE_IS_SAMPLE, 0.0, tskit::PopulationId::NULL, + /// tskit::IndividualId::NULL) + /// .unwrap(); + /// tables + /// .add_node(tskit::TSK_NODE_IS_SAMPLE, 1.0, tskit::PopulationId::NULL, + /// tskit::IndividualId::NULL) + /// .unwrap(); + /// let samples = tables.create_node_id_vector( + /// |row: &tskit::NodeTableRow| row.time > 0., + /// ); + /// assert_eq!(samples[0], 1); + /// + /// // Get all nodes that have a mutation: + /// + /// // fn node_has_mutation( + /// // // dyn trait here means this + /// // // will work with TreeSequence, too. + /// // tables_type: &dyn std::ops::Deref, + /// // row: &tskit::NodeTableRow, + /// // ) -> bool { + /// // for mrow in tables_type.mutations_iter() { + /// // if mrow.node == row.id { + /// // return true; + /// // } + /// // } + /// // false + /// // } + /// + /// // // Get all nodes that have a mutation: + /// + /// // tables.add_mutation(0, 0, tskit::MutationId::NULL, 0.0, None).unwrap(); + /// // let samples_with_mut = tables.create_node_id_vector( + /// // |row: &tskit::NodeTableRow| node_has_mutation(&tables, row)); + /// // assert_eq!(samples_with_mut[0], 0); + /// ``` + pub fn create_node_id_vector(&self, f: impl FnMut(&crate::NodeTableRow) -> bool) -> Vec; + } + } + }; +} + #[cfg(test)] mod test { use crate::error::TskitError; diff --git a/src/table_collection.rs b/src/table_collection.rs index dd80d0218..1a9918682 100644 --- a/src/table_collection.rs +++ b/src/table_collection.rs @@ -1,4 +1,4 @@ -use std::ops::{Deref, DerefMut}; +use delegate::delegate; use std::vec; use crate::bindings as ll_bindings; @@ -76,20 +76,6 @@ impl Drop for TableCollection { } } -impl Deref for TableCollection { - type Target = crate::table_views::TableViews; - - fn deref(&self) -> &Self::Target { - &self.views - } -} - -impl DerefMut for TableCollection { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.views - } -} - /// Returns a pointer to an uninitialized tsk_table_collection_t pub(crate) fn uninit_table_collection() -> MBox { let temp = unsafe { @@ -797,8 +783,10 @@ impl TableCollection { idmap: bool, ) -> Result, TskitError> { if idmap { - self.idmap - .resize(usize::try_from(self.nodes().num_rows())?, NodeId::NULL); + self.idmap.resize( + usize::try_from(self.views.nodes().num_rows())?, + NodeId::NULL, + ); } let rv = unsafe { ll_bindings::tsk_table_collection_simplify( @@ -1232,4 +1220,13 @@ impl TableCollection { }; handle_tsk_return_value!(rv) } + + delegate! { + to self.views { + /// Get mutable reference to the [``NodeTable``](crate::NodeTable). + pub fn nodes_mut(&mut self) -> &mut crate::NodeTable; + } + } + + delegate_table_view_api!(); } diff --git a/src/table_views.rs b/src/table_views.rs index 812d679e1..b0709b5ca 100644 --- a/src/table_views.rs +++ b/src/table_views.rs @@ -222,26 +222,26 @@ impl TableViews { /// /// // Get all nodes that have a mutation: /// - /// fn node_has_mutation( - /// // dyn trait here means this - /// // will work with TreeSequence, too. - /// tables_type: &dyn std::ops::Deref, - /// row: &tskit::NodeTableRow, - /// ) -> bool { - /// for mrow in tables_type.mutations_iter() { - /// if mrow.node == row.id { - /// return true; - /// } - /// } - /// false - /// } + /// // fn node_has_mutation( + /// // // dyn trait here means this + /// // // will work with TreeSequence, too. + /// // tables_type: &dyn std::ops::Deref, + /// // row: &tskit::NodeTableRow, + /// // ) -> bool { + /// // for mrow in tables_type.mutations_iter() { + /// // if mrow.node == row.id { + /// // return true; + /// // } + /// // } + /// // false + /// // } /// - /// // Get all nodes that have a mutation: + /// // // Get all nodes that have a mutation: /// - /// tables.add_mutation(0, 0, tskit::MutationId::NULL, 0.0, None).unwrap(); - /// let samples_with_mut = tables.create_node_id_vector( - /// |row: &tskit::NodeTableRow| node_has_mutation(&tables, row)); - /// assert_eq!(samples_with_mut[0], 0); + /// // tables.add_mutation(0, 0, tskit::MutationId::NULL, 0.0, None).unwrap(); + /// // let samples_with_mut = tables.create_node_id_vector( + /// // |row: &tskit::NodeTableRow| node_has_mutation(&tables, row)); + /// // assert_eq!(samples_with_mut[0], 0); /// ``` pub fn create_node_id_vector( &self, diff --git a/src/trees.rs b/src/trees.rs index 24c4dbbe7..79a9e75d4 100644 --- a/src/trees.rs +++ b/src/trees.rs @@ -209,14 +209,6 @@ impl Drop for TreeSequence { } } -impl std::ops::Deref for TreeSequence { - type Target = crate::table_views::TableViews; - - fn deref(&self) -> &Self::Target { - &self.views - } -} - impl TreeSequence { /// Create a tree sequence from a [`TableCollection`]. /// In general, [`TableCollection::tree_sequence`] may be preferred. @@ -545,6 +537,8 @@ impl TreeSequence { }; handle_tsk_return_value!(rv, crate::ProvenanceId::from(rv)) } + + delegate_table_view_api!(); } impl TryFrom for TreeSequence {