From f584c89b50a89937dba73f0c6d1106ab1f962ae5 Mon Sep 17 00:00:00 2001 From: "Kevin R. Thornton" Date: Tue, 31 Jan 2023 11:56:25 -0800 Subject: [PATCH] feat: Allow default arguments when adding nodes. * add serde_json to dev-dependencies --- Cargo.toml | 1 + src/lib.rs | 4 +- src/node_table.rs | 307 +++++++++++++++++++++++++++++++++ src/table_collection.rs | 12 ++ tests/book_table_collection.rs | 142 +++++++++++++++ 5 files changed, 465 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 0ca46ec8f..7e2f28088 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ anyhow = {version = "1.0.66"} clap = {version = "~3.2.8", features = ["derive"]} serde = {version = "1.0.118", features = ["derive"]} serde-pickle = "1.1.0" +serde_json = {version = "1.0.67"} bincode = "1.3.1" rand = "0.8.3" rand_distr = "0.4.0" diff --git a/src/lib.rs b/src/lib.rs index 81e5b1228..db331bb17 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -126,7 +126,9 @@ pub use individual_table::{IndividualTable, IndividualTableRow, OwningIndividual pub use migration_table::{MigrationTable, MigrationTableRow, OwningMigrationTable}; pub use mutation_table::{MutationTable, MutationTableRow, OwningMutationTable}; pub use newtypes::*; -pub use node_table::{NodeTable, NodeTableRow, OwningNodeTable}; +pub use node_table::{ + NodeDefaults, NodeDefaultsWithMetadata, NodeTable, NodeTableRow, OwningNodeTable, +}; pub use population_table::{OwningPopulationTable, PopulationTable, PopulationTableRow}; pub use site_table::{OwningSiteTable, SiteTable, SiteTableRow}; pub use table_collection::TableCollection; diff --git a/src/node_table.rs b/src/node_table.rs index 738a98642..765935063 100644 --- a/src/node_table.rs +++ b/src/node_table.rs @@ -138,6 +138,248 @@ impl<'a> streaming_iterator::StreamingIterator for NodeTableRowView<'a> { } } +/// Defaults for node table rows without metadata +/// +/// # Examples +/// +/// ``` +/// let d = tskit::NodeDefaults::default(); +/// assert_eq!(d.flags, tskit::NodeFlags::default()); +/// assert_eq!(d.population, tskit::PopulationId::NULL); +/// assert_eq!(d.individual, tskit::IndividualId::NULL); +/// ``` +/// +/// [Struct update syntax](https://doc.rust-lang.org/book/ch05-01-defining-structs.html) +/// is your friend here: +/// +/// ``` +/// let d = tskit::NodeDefaults{population: 0.into(), ..Default::default()}; +/// assert_eq!(d.flags, tskit::NodeFlags::default()); +/// assert_eq!(d.population, 0); +/// assert_eq!(d.individual, tskit::IndividualId::NULL); +/// let d2 = tskit::NodeDefaults{flags: tskit::NodeFlags::default().mark_sample(), +/// // update remaining values from d +/// ..d}; +/// assert!(d2.flags.is_sample()); +/// assert_eq!(d2.population, 0); +/// assert_eq!(d2.individual, tskit::IndividualId::NULL); +/// ``` +#[derive(Copy, Clone, Default, Eq, PartialEq, Debug)] +pub struct NodeDefaults { + pub flags: NodeFlags, + pub population: PopulationId, + pub individual: IndividualId, +} + +/// Defaults for node table rows with metadata +/// +/// # Notes +/// +/// This struct derives `Debug` and `Clone`. +/// However, neither is a trait bound on `M`. +/// Therefore, use of `Debug` and/or `Clone` will fail unless `M` +/// also implements the relevant trait. +/// +/// See [the book](https://tskit-dev.github.io/tskit-rust/) +/// for details. +#[derive(Debug, Clone)] +pub struct NodeDefaultsWithMetadata +where + M: crate::metadata::NodeMetadata, +{ + pub flags: NodeFlags, + pub population: PopulationId, + pub individual: IndividualId, + pub metadata: Option, +} + +/// This is a doctest hack as described in the rust book. +/// We do this b/c the specific error messages can change +/// across rust versions, making crates like trybuild +/// less useful. +/// +/// ```compile_fail +/// #[derive(serde::Serialize, serde::Deserialize)] +/// struct NodeMetadata { +/// value: i32, +/// } +/// +/// impl tskit::metadata::MetadataRoundtrip for NodeMetadata { +/// fn encode(&self) -> Result, tskit::metadata::MetadataError> { +/// match serde_json::to_string(self) { +/// Ok(x) => Ok(x.as_bytes().to_vec()), +/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }), +/// } +/// } +/// fn decode(md: &[u8]) -> Result +/// where +/// Self: Sized, +/// { +/// match serde_json::from_slice(md) { +/// Ok(v) => Ok(v), +/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }), +/// } +/// } +/// } +/// +/// impl tskit::metadata::NodeMetadata for NodeMetadata {} +/// +/// type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata; +/// let defaults = DefaultsWithMetadata { +/// metadata: Some(NodeMetadata { value: 42 }), +/// ..Default::default() +/// }; +/// +/// // Fails because metadata type is not Debug +/// println!("{:?}", defaults); +/// ``` +/// +/// ```compile_fail +/// #[derive(serde::Serialize, serde::Deserialize)] +/// struct NodeMetadata { +/// value: i32, +/// } +/// +/// impl tskit::metadata::MetadataRoundtrip for NodeMetadata { +/// fn encode(&self) -> Result, tskit::metadata::MetadataError> { +/// match serde_json::to_string(self) { +/// Ok(x) => Ok(x.as_bytes().to_vec()), +/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }), +/// } +/// } +/// fn decode(md: &[u8]) -> Result +/// where +/// Self: Sized, +/// { +/// match serde_json::from_slice(md) { +/// Ok(v) => Ok(v), +/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }), +/// } +/// } +/// } +/// +/// impl tskit::metadata::NodeMetadata for NodeMetadata {} +/// +/// let mut tables = tskit::TableCollection::new(10.0).unwrap(); +/// type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata; +/// // What if there is default metadata for all rows? +/// let defaults = DefaultsWithMetadata { +/// metadata: Some(NodeMetadata { value: 42 }), +/// ..Default::default() +/// }; +/// // We can scoop all non-metadata fields even though +/// // type is not Copy/Clone +/// let _ = tables +/// .add_node_with_defaults( +/// 0.0, +/// &DefaultsWithMetadata { +/// metadata: Some(NodeMetadata { value: 2 * 42 }), +/// ..defaults +/// }, +/// ) +/// .unwrap(); +/// // But now, we start to cause a problem: +/// // If we don't clone here, our metadata type moves, +/// // so our defaults are moved. +/// let _ = tables +/// .add_node_with_defaults( +/// 0.0, +/// &DefaultsWithMetadata { +/// population: 6.into(), +/// ..defaults +/// }, +/// ) +/// .unwrap(); +/// // Now, we have a use-after-move error +/// // if we hadn't cloned in the last step. +/// let _ = tables +/// .add_node_with_defaults( +/// 0.0, +/// &DefaultsWithMetadata { +/// individual: 7.into(), +/// ..defaults +/// }, +/// ) +/// .unwrap(); +/// ``` +#[cfg(doctest)] +struct NodeDefaultsWithMetadataNotCloneNotDebug; + +// Manual implementation required so that +// we do not force client code to impl Default +// for metadata types. +impl Default for NodeDefaultsWithMetadata +where + M: crate::metadata::NodeMetadata, +{ + fn default() -> Self { + Self { + flags: NodeFlags::default(), + population: PopulationId::default(), + individual: IndividualId::default(), + metadata: None, + } + } +} + +mod private { + pub trait DefaultNodeDataMarker {} + + impl DefaultNodeDataMarker for super::NodeDefaults {} + + impl DefaultNodeDataMarker for super::NodeDefaultsWithMetadata where + M: crate::metadata::NodeMetadata + { + } +} + +/// This trait is sealed. +pub trait DefaultNodeData: private::DefaultNodeDataMarker { + fn flags(&self) -> NodeFlags; + fn population(&self) -> PopulationId; + fn individual(&self) -> IndividualId; + fn metadata(&self) -> Result>, TskitError>; +} + +impl DefaultNodeData for NodeDefaults { + fn flags(&self) -> NodeFlags { + self.flags + } + fn population(&self) -> PopulationId { + self.population + } + fn individual(&self) -> IndividualId { + self.individual + } + fn metadata(&self) -> Result>, TskitError> { + Ok(None) + } +} + +impl DefaultNodeData for NodeDefaultsWithMetadata +where + M: crate::metadata::NodeMetadata, +{ + fn flags(&self) -> NodeFlags { + self.flags + } + fn population(&self) -> PopulationId { + self.population + } + fn individual(&self) -> IndividualId { + self.individual + } + fn metadata(&self) -> Result>, TskitError> { + self.metadata.as_ref().map_or_else( + || Ok(None), + |v| match v.encode() { + Ok(x) => Ok(Some(x)), + Err(e) => Err(e.into()), + }, + ) + } +} + /// A node table. /// /// # Examples @@ -520,6 +762,71 @@ impl NodeTable { ) } + /// Add a row with default values. + /// + /// # Parameters + /// + /// * `time`, the birth time of the node + /// * `defaults`, the default values for remaining fields. + /// + /// ## Notes on parameters + /// + /// The type of `defaults` must be one of: + /// + /// * [`NodeDefaults`] + /// * [`NodeDefaultsWithMetadata`] + /// + /// # Examples + /// + /// ## Without metadata + /// + /// ``` + /// let mut nodes = tskit::NodeTable::default(); + /// let defaults = tskit::NodeDefaults::default(); + /// let id = nodes.add_row_with_defaults(0.0, &defaults).unwrap(); + /// assert_eq!(id, 0); + /// assert_eq!(nodes.individual(id), Some(tskit::IndividualId::NULL)); + /// assert_eq!(nodes.population(id), Some(tskit::PopulationId::NULL)); + /// ``` + /// + /// Use [struct update syntax](https://doc.rust-lang.org/book/ch05-01-defining-structs.html) + /// to customize defaults: + /// + /// ``` + /// # let mut nodes = tskit::NodeTable::default(); + /// let defaults = tskit::NodeDefaults { + /// population: 3.into(), + /// ..Default::default() + /// }; + /// let id = nodes.add_row_with_defaults(0.0, &defaults).unwrap(); + /// # assert_eq!(id, 0); + /// # assert_eq!(nodes.individual(id), Some(tskit::IndividualId::NULL)); + /// assert_eq!(nodes.population(id), Some(tskit::PopulationId::from(3))); + /// ``` + /// + /// ## With metadata + /// + /// See the [book](https://tskit-dev.github.io/tskit-rust) for examples. + pub fn add_row_with_defaults, D: DefaultNodeData>( + &mut self, + time: T, + defaults: &D, + ) -> Result { + let md = defaults.metadata()?; + let (ptr, mdlen) = match &md { + Some(value) => (value.as_ptr(), SizeType::try_from(value.len())?), + None => (std::ptr::null(), 0.into()), + }; + self.add_row_details( + defaults.flags(), + time, + defaults.population(), + defaults.individual(), + ptr.cast::(), + mdlen.into(), + ) + } + build_table_column_slice_getter!( /// Get the time column as a slice => time, time_slice, Time); diff --git a/src/table_collection.rs b/src/table_collection.rs index abf02890f..be4107fe5 100644 --- a/src/table_collection.rs +++ b/src/table_collection.rs @@ -3,6 +3,7 @@ use std::vec; use crate::bindings as ll_bindings; use crate::error::TskitError; +use crate::node_table::DefaultNodeData; use crate::types::Bookmark; use crate::IndividualTableSortOptions; use crate::Position; @@ -545,6 +546,17 @@ impl TableCollection { .add_row_with_metadata(flags, time, population, individual, metadata) } + /// Add a node using defaults. + /// + /// See [`crate::NodeTable::add_row_with_defaults`] for details. + pub fn add_node_with_defaults, D: DefaultNodeData>( + &mut self, + time: T, + defaults: &D, + ) -> Result { + self.nodes_mut().add_row_with_defaults(time, defaults) + } + /// Add a row to the site table pub fn add_site

( &mut self, diff --git a/tests/book_table_collection.rs b/tests/book_table_collection.rs index 98ed1fbcf..78c4b3e4e 100644 --- a/tests/book_table_collection.rs +++ b/tests/book_table_collection.rs @@ -165,3 +165,145 @@ fn get_data_from_edge_table() { .check_integrity(tskit::TableIntegrityCheckFlags::default()) .is_ok()); } + +#[test] +fn test_adding_node_table_row_with_defaults() { + let mut tables = tskit::TableCollection::new(10.).unwrap(); + let defaults = tskit::NodeDefaults::default(); + let node = tables.add_node_with_defaults(0.0, &defaults).unwrap(); + assert_eq!(node, 0); + let node = tables + .add_node_with_defaults( + 0.0, + // Create a new, temporary defaults instance + &tskit::NodeDefaults { + // Mark the new node as a sample + flags: tskit::NodeFlags::new_sample(), + // Use remaining values from our current defaults + ..defaults + }, + ) + .unwrap(); + assert!(tables.nodes().flags(node).unwrap().is_sample()); +} + +macro_rules! impl_node_metadata_traits { + () => { + impl tskit::metadata::MetadataRoundtrip for NodeMetadata { + fn encode(&self) -> Result, tskit::metadata::MetadataError> { + match serde_json::to_string(self) { + Ok(x) => Ok(x.as_bytes().to_vec()), + Err(e) => { + Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }) + } + } + } + fn decode(md: &[u8]) -> Result + where + Self: Sized, + { + match serde_json::from_slice(md) { + Ok(v) => Ok(v), + Err(e) => { + Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }) + } + } + } + } + impl tskit::metadata::NodeMetadata for NodeMetadata {} + }; +} + +mod node_metadata { + #[derive(serde::Serialize, serde::Deserialize)] + pub struct NodeMetadata { + pub value: i32, + } + impl_node_metadata_traits!(); +} + +mod node_metadata_clone { + #[derive(Clone, serde::Serialize, serde::Deserialize)] + pub struct NodeMetadata { + pub value: i32, + } + impl_node_metadata_traits!(); +} + +#[test] +fn test_adding_node_table_row_with_defaults_and_metadata() { + use node_metadata::NodeMetadata; + let mut tables = tskit::TableCollection::new(10.0).unwrap(); + type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata; + let defaults = DefaultsWithMetadata::default(); + let _ = tables.add_node_with_defaults(0.0, &defaults).unwrap(); + let _ = tables + .add_node_with_defaults( + 0.0, + &DefaultsWithMetadata { + population: 3.into(), + metadata: Some(NodeMetadata { value: 42 }), + ..defaults + }, + ) + .unwrap(); + let _ = tables + .add_node_with_defaults( + 0.0, + &DefaultsWithMetadata { + population: 3.into(), + metadata: Some(NodeMetadata { value: 42 }), + ..defaults + }, + ) + .unwrap(); +} + +#[test] +fn test_adding_node_table_row_with_defaults_and_metadata_requiring_clone() { + use node_metadata_clone::NodeMetadata; + let mut tables = tskit::TableCollection::new(10.0).unwrap(); + type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata; + // What if there is default metadata for all rows? + let defaults = DefaultsWithMetadata { + metadata: Some(NodeMetadata { value: 42 }), + ..Default::default() + }; + + // We can scoop all non-metadata fields even though + // type is not Copy/Clone + let _ = tables + .add_node_with_defaults( + 0.0, + &DefaultsWithMetadata { + metadata: Some(NodeMetadata { value: 2 * 42 }), + ..defaults + }, + ) + .unwrap(); + + // But now, we start to cause a problem: + // If we don't clone here, our metadata type moves, + // so our defaults are moved. + let _ = tables + .add_node_with_defaults( + 0.0, + &DefaultsWithMetadata { + population: 6.into(), + ..defaults.clone() + }, + ) + .unwrap(); + + // Now, we have a use-after-move error + // if we hadn't cloned in the last step. + let _ = tables + .add_node_with_defaults( + 0.0, + &DefaultsWithMetadata { + individual: 7.into(), + ..defaults + }, + ) + .unwrap(); +}