Skip to content

Commit fbe5b2f

Browse files
molpopgenmolpopgen
authored
feat: Allow default arguments when adding nodes. (#466)
* add serde_json to dev-dependencies
1 parent c556839 commit fbe5b2f

File tree

5 files changed

+465
-1
lines changed

5 files changed

+465
-1
lines changed

‎Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ anyhow = {version = "1.0.66"}
3434
clap = {version = "~3.2.8", features = ["derive"]}
3535
serde = {version = "1.0.118", features = ["derive"]}
3636
serde-pickle = "1.1.0"
37+
serde_json = {version = "1.0.67"}
3738
bincode = "1.3.1"
3839
rand = "0.8.3"
3940
rand_distr = "0.4.0"

‎src/lib.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,9 @@ pub use individual_table::{IndividualTable, IndividualTableRow, OwningIndividual
126126
pub use migration_table::{MigrationTable, MigrationTableRow, OwningMigrationTable};
127127
pub use mutation_table::{MutationTable, MutationTableRow, OwningMutationTable};
128128
pub use newtypes::*;
129-
pub use node_table::{NodeTable, NodeTableRow, OwningNodeTable};
129+
pub use node_table::{
130+
NodeDefaults, NodeDefaultsWithMetadata, NodeTable, NodeTableRow, OwningNodeTable,
131+
};
130132
pub use population_table::{OwningPopulationTable, PopulationTable, PopulationTableRow};
131133
pub use site_table::{OwningSiteTable, SiteTable, SiteTableRow};
132134
pub use table_collection::TableCollection;

‎src/node_table.rs

Lines changed: 307 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,248 @@ impl<'a> streaming_iterator::StreamingIterator for NodeTableRowView<'a> {
138138
}
139139
}
140140

141+
/// Defaults for node table rows without metadata
142+
///
143+
/// # Examples
144+
///
145+
/// ```
146+
/// let d = tskit::NodeDefaults::default();
147+
/// assert_eq!(d.flags, tskit::NodeFlags::default());
148+
/// assert_eq!(d.population, tskit::PopulationId::NULL);
149+
/// assert_eq!(d.individual, tskit::IndividualId::NULL);
150+
/// ```
151+
///
152+
/// [Struct update syntax](https://doc.rust-lang.org/book/ch05-01-defining-structs.html)
153+
/// is your friend here:
154+
///
155+
/// ```
156+
/// let d = tskit::NodeDefaults{population: 0.into(), ..Default::default()};
157+
/// assert_eq!(d.flags, tskit::NodeFlags::default());
158+
/// assert_eq!(d.population, 0);
159+
/// assert_eq!(d.individual, tskit::IndividualId::NULL);
160+
/// let d2 = tskit::NodeDefaults{flags: tskit::NodeFlags::default().mark_sample(),
161+
/// // update remaining values from d
162+
/// ..d};
163+
/// assert!(d2.flags.is_sample());
164+
/// assert_eq!(d2.population, 0);
165+
/// assert_eq!(d2.individual, tskit::IndividualId::NULL);
166+
/// ```
167+
#[derive(Copy, Clone, Default, Eq, PartialEq, Debug)]
168+
pub struct NodeDefaults {
169+
pub flags: NodeFlags,
170+
pub population: PopulationId,
171+
pub individual: IndividualId,
172+
}
173+
174+
/// Defaults for node table rows with metadata
175+
///
176+
/// # Notes
177+
///
178+
/// This struct derives `Debug` and `Clone`.
179+
/// However, neither is a trait bound on `M`.
180+
/// Therefore, use of `Debug` and/or `Clone` will fail unless `M`
181+
/// also implements the relevant trait.
182+
///
183+
/// See [the book](https://tskit-dev.github.io/tskit-rust/)
184+
/// for details.
185+
#[derive(Debug, Clone)]
186+
pub struct NodeDefaultsWithMetadata<M>
187+
where
188+
M: crate::metadata::NodeMetadata,
189+
{
190+
pub flags: NodeFlags,
191+
pub population: PopulationId,
192+
pub individual: IndividualId,
193+
pub metadata: Option<M>,
194+
}
195+
196+
/// This is a doctest hack as described in the rust book.
197+
/// We do this b/c the specific error messages can change
198+
/// across rust versions, making crates like trybuild
199+
/// less useful.
200+
///
201+
/// ```compile_fail
202+
/// #[derive(serde::Serialize, serde::Deserialize)]
203+
/// struct NodeMetadata {
204+
/// value: i32,
205+
/// }
206+
///
207+
/// impl tskit::metadata::MetadataRoundtrip for NodeMetadata {
208+
/// fn encode(&self) -> Result<Vec<u8>, tskit::metadata::MetadataError> {
209+
/// match serde_json::to_string(self) {
210+
/// Ok(x) => Ok(x.as_bytes().to_vec()),
211+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
212+
/// }
213+
/// }
214+
/// fn decode(md: &[u8]) -> Result<Self, tskit::metadata::MetadataError>
215+
/// where
216+
/// Self: Sized,
217+
/// {
218+
/// match serde_json::from_slice(md) {
219+
/// Ok(v) => Ok(v),
220+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
221+
/// }
222+
/// }
223+
/// }
224+
///
225+
/// impl tskit::metadata::NodeMetadata for NodeMetadata {}
226+
///
227+
/// type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata<NodeMetadata>;
228+
/// let defaults = DefaultsWithMetadata {
229+
/// metadata: Some(NodeMetadata { value: 42 }),
230+
/// ..Default::default()
231+
/// };
232+
///
233+
/// // Fails because metadata type is not Debug
234+
/// println!("{:?}", defaults);
235+
/// ```
236+
///
237+
/// ```compile_fail
238+
/// #[derive(serde::Serialize, serde::Deserialize)]
239+
/// struct NodeMetadata {
240+
/// value: i32,
241+
/// }
242+
///
243+
/// impl tskit::metadata::MetadataRoundtrip for NodeMetadata {
244+
/// fn encode(&self) -> Result<Vec<u8>, tskit::metadata::MetadataError> {
245+
/// match serde_json::to_string(self) {
246+
/// Ok(x) => Ok(x.as_bytes().to_vec()),
247+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
248+
/// }
249+
/// }
250+
/// fn decode(md: &[u8]) -> Result<Self, tskit::metadata::MetadataError>
251+
/// where
252+
/// Self: Sized,
253+
/// {
254+
/// match serde_json::from_slice(md) {
255+
/// Ok(v) => Ok(v),
256+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
257+
/// }
258+
/// }
259+
/// }
260+
///
261+
/// impl tskit::metadata::NodeMetadata for NodeMetadata {}
262+
///
263+
/// let mut tables = tskit::TableCollection::new(10.0).unwrap();
264+
/// type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata<NodeMetadata>;
265+
/// // What if there is default metadata for all rows?
266+
/// let defaults = DefaultsWithMetadata {
267+
/// metadata: Some(NodeMetadata { value: 42 }),
268+
/// ..Default::default()
269+
/// };
270+
/// // We can scoop all non-metadata fields even though
271+
/// // type is not Copy/Clone
272+
/// let _ = tables
273+
/// .add_node_with_defaults(
274+
/// 0.0,
275+
/// &DefaultsWithMetadata {
276+
/// metadata: Some(NodeMetadata { value: 2 * 42 }),
277+
/// ..defaults
278+
/// },
279+
/// )
280+
/// .unwrap();
281+
/// // But now, we start to cause a problem:
282+
/// // If we don't clone here, our metadata type moves,
283+
/// // so our defaults are moved.
284+
/// let _ = tables
285+
/// .add_node_with_defaults(
286+
/// 0.0,
287+
/// &DefaultsWithMetadata {
288+
/// population: 6.into(),
289+
/// ..defaults
290+
/// },
291+
/// )
292+
/// .unwrap();
293+
/// // Now, we have a use-after-move error
294+
/// // if we hadn't cloned in the last step.
295+
/// let _ = tables
296+
/// .add_node_with_defaults(
297+
/// 0.0,
298+
/// &DefaultsWithMetadata {
299+
/// individual: 7.into(),
300+
/// ..defaults
301+
/// },
302+
/// )
303+
/// .unwrap();
304+
/// ```
305+
#[cfg(doctest)]
306+
struct NodeDefaultsWithMetadataNotCloneNotDebug;
307+
308+
// Manual implementation required so that
309+
// we do not force client code to impl Default
310+
// for metadata types.
311+
impl<M> Default for NodeDefaultsWithMetadata<M>
312+
where
313+
M: crate::metadata::NodeMetadata,
314+
{
315+
fn default() -> Self {
316+
Self {
317+
flags: NodeFlags::default(),
318+
population: PopulationId::default(),
319+
individual: IndividualId::default(),
320+
metadata: None,
321+
}
322+
}
323+
}
324+
325+
mod private {
326+
pub trait DefaultNodeDataMarker {}
327+
328+
impl DefaultNodeDataMarker for super::NodeDefaults {}
329+
330+
impl<M> DefaultNodeDataMarker for super::NodeDefaultsWithMetadata<M> where
331+
M: crate::metadata::NodeMetadata
332+
{
333+
}
334+
}
335+
336+
/// This trait is sealed.
337+
pub trait DefaultNodeData: private::DefaultNodeDataMarker {
338+
fn flags(&self) -> NodeFlags;
339+
fn population(&self) -> PopulationId;
340+
fn individual(&self) -> IndividualId;
341+
fn metadata(&self) -> Result<Option<Vec<u8>>, TskitError>;
342+
}
343+
344+
impl DefaultNodeData for NodeDefaults {
345+
fn flags(&self) -> NodeFlags {
346+
self.flags
347+
}
348+
fn population(&self) -> PopulationId {
349+
self.population
350+
}
351+
fn individual(&self) -> IndividualId {
352+
self.individual
353+
}
354+
fn metadata(&self) -> Result<Option<Vec<u8>>, TskitError> {
355+
Ok(None)
356+
}
357+
}
358+
359+
impl<M> DefaultNodeData for NodeDefaultsWithMetadata<M>
360+
where
361+
M: crate::metadata::NodeMetadata,
362+
{
363+
fn flags(&self) -> NodeFlags {
364+
self.flags
365+
}
366+
fn population(&self) -> PopulationId {
367+
self.population
368+
}
369+
fn individual(&self) -> IndividualId {
370+
self.individual
371+
}
372+
fn metadata(&self) -> Result<Option<Vec<u8>>, TskitError> {
373+
self.metadata.as_ref().map_or_else(
374+
|| Ok(None),
375+
|v| match v.encode() {
376+
Ok(x) => Ok(Some(x)),
377+
Err(e) => Err(e.into()),
378+
},
379+
)
380+
}
381+
}
382+
141383
/// A node table.
142384
///
143385
/// # Examples
@@ -520,6 +762,71 @@ impl NodeTable {
520762
)
521763
}
522764

765+
/// Add a row with default values.
766+
///
767+
/// # Parameters
768+
///
769+
/// * `time`, the birth time of the node
770+
/// * `defaults`, the default values for remaining fields.
771+
///
772+
/// ## Notes on parameters
773+
///
774+
/// The type of `defaults` must be one of:
775+
///
776+
/// * [`NodeDefaults`]
777+
/// * [`NodeDefaultsWithMetadata`]
778+
///
779+
/// # Examples
780+
///
781+
/// ## Without metadata
782+
///
783+
/// ```
784+
/// let mut nodes = tskit::NodeTable::default();
785+
/// let defaults = tskit::NodeDefaults::default();
786+
/// let id = nodes.add_row_with_defaults(0.0, &defaults).unwrap();
787+
/// assert_eq!(id, 0);
788+
/// assert_eq!(nodes.individual(id), Some(tskit::IndividualId::NULL));
789+
/// assert_eq!(nodes.population(id), Some(tskit::PopulationId::NULL));
790+
/// ```
791+
///
792+
/// Use [struct update syntax](https://doc.rust-lang.org/book/ch05-01-defining-structs.html)
793+
/// to customize defaults:
794+
///
795+
/// ```
796+
/// # let mut nodes = tskit::NodeTable::default();
797+
/// let defaults = tskit::NodeDefaults {
798+
/// population: 3.into(),
799+
/// ..Default::default()
800+
/// };
801+
/// let id = nodes.add_row_with_defaults(0.0, &defaults).unwrap();
802+
/// # assert_eq!(id, 0);
803+
/// # assert_eq!(nodes.individual(id), Some(tskit::IndividualId::NULL));
804+
/// assert_eq!(nodes.population(id), Some(tskit::PopulationId::from(3)));
805+
/// ```
806+
///
807+
/// ## With metadata
808+
///
809+
/// See the [book](https://tskit-dev.github.io/tskit-rust) for examples.
810+
pub fn add_row_with_defaults<T: Into<Time>, D: DefaultNodeData>(
811+
&mut self,
812+
time: T,
813+
defaults: &D,
814+
) -> Result<NodeId, TskitError> {
815+
let md = defaults.metadata()?;
816+
let (ptr, mdlen) = match &md {
817+
Some(value) => (value.as_ptr(), SizeType::try_from(value.len())?),
818+
None => (std::ptr::null(), 0.into()),
819+
};
820+
self.add_row_details(
821+
defaults.flags(),
822+
time,
823+
defaults.population(),
824+
defaults.individual(),
825+
ptr.cast::<i8>(),
826+
mdlen.into(),
827+
)
828+
}
829+
523830
build_table_column_slice_getter!(
524831
/// Get the time column as a slice
525832
=> time, time_slice, Time);

‎src/table_collection.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ use std::vec;
33

44
use crate::bindings as ll_bindings;
55
use crate::error::TskitError;
6+
use crate::node_table::DefaultNodeData;
67
use crate::types::Bookmark;
78
use crate::IndividualTableSortOptions;
89
use crate::Position;
@@ -545,6 +546,17 @@ impl TableCollection {
545546
.add_row_with_metadata(flags, time, population, individual, metadata)
546547
}
547548

549+
/// Add a node using defaults.
550+
///
551+
/// See [`crate::NodeTable::add_row_with_defaults`] for details.
552+
pub fn add_node_with_defaults<T: Into<crate::Time>, D: DefaultNodeData>(
553+
&mut self,
554+
time: T,
555+
defaults: &D,
556+
) -> Result<NodeId, TskitError> {
557+
self.nodes_mut().add_row_with_defaults(time, defaults)
558+
}
559+
548560
/// Add a row to the site table
549561
pub fn add_site<P>(
550562
&mut self,

0 commit comments

Comments
 (0)