Skip to content

Commit 13f4549

Browse files
molpopgenmolpopgen
authored
feat: support adding nodes with defaults (#488)
1 parent 0ccaa7a commit 13f4549

File tree

4 files changed

+307
-1
lines changed

4 files changed

+307
-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: 281 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,30 @@ where
8383
)
8484
}
8585

86+
pub(crate) fn add_row_with_defaults<T: Into<crate::Time>, D: crate::node_table::DefaultNodeData>(
87+
time: T,
88+
defaults: &D,
89+
table: *mut ll_bindings::tsk_node_table_t,
90+
) -> Result<NodeId, TskitError> {
91+
let md = defaults.metadata()?;
92+
let (ptr, mdlen) = match &md {
93+
Some(value) => (
94+
value.as_ptr().cast::<i8>(),
95+
crate::SizeType::try_from(value.len())?,
96+
),
97+
None => (std::ptr::null(), 0.into()),
98+
};
99+
add_row_details(
100+
defaults.flags().bits(),
101+
time.into().into(),
102+
defaults.population().into(),
103+
defaults.individual().into(),
104+
ptr,
105+
mdlen.into(),
106+
table,
107+
)
108+
}
109+
86110
/// Row of a [`NodeTable`]
87111
#[derive(Debug)]
88112
pub struct NodeTableRow {
@@ -214,6 +238,243 @@ impl<'a> streaming_iterator::StreamingIterator for NodeTableRowView<'a> {
214238
}
215239
}
216240

241+
/// Defaults for node table rows without metadata
242+
///
243+
/// # Examples
244+
///
245+
/// ```
246+
/// let d = tskit::NodeDefaults::default();
247+
/// assert_eq!(d.flags, tskit::NodeFlags::default());
248+
/// assert_eq!(d.population, tskit::PopulationId::NULL);
249+
/// assert_eq!(d.individual, tskit::IndividualId::NULL);
250+
/// ```
251+
///
252+
/// [Struct update syntax](https://doc.rust-lang.org/book/ch05-01-defining-structs.html)
253+
/// is your friend here:
254+
///
255+
/// ```
256+
/// let d = tskit::NodeDefaults{population: 0.into(), ..Default::default()};
257+
/// assert_eq!(d.flags, tskit::NodeFlags::default());
258+
/// assert_eq!(d.population, 0);
259+
/// assert_eq!(d.individual, tskit::IndividualId::NULL);
260+
/// let d2 = tskit::NodeDefaults{flags: tskit::NodeFlags::default().mark_sample(),
261+
/// // update remaining values from d
262+
/// ..d};
263+
/// assert!(d2.flags.is_sample());
264+
/// assert_eq!(d2.population, 0);
265+
/// assert_eq!(d2.individual, tskit::IndividualId::NULL);
266+
/// ```
267+
#[derive(Copy, Clone, Default, Eq, PartialEq, Debug)]
268+
pub struct NodeDefaults {
269+
pub flags: NodeFlags,
270+
pub population: PopulationId,
271+
pub individual: IndividualId,
272+
}
273+
274+
// Defaults for node table rows with metadata
275+
///
276+
/// # Notes
277+
///
278+
/// This struct derives `Debug` and `Clone`.
279+
/// However, neither is a trait bound on `M`.
280+
/// Therefore, use of `Debug` and/or `Clone` will fail unless `M`
281+
/// also implements the relevant trait.
282+
///
283+
/// See [the book](https://tskit-dev.github.io/tskit-rust/)
284+
/// for details.
285+
#[derive(Debug, Clone, Default)]
286+
pub struct NodeDefaultsWithMetadata<M>
287+
where
288+
M: crate::metadata::NodeMetadata,
289+
{
290+
pub flags: NodeFlags,
291+
pub population: PopulationId,
292+
pub individual: IndividualId,
293+
pub metadata: Option<M>,
294+
}
295+
296+
mod private {
297+
pub trait DefaultNodeDataMarker {}
298+
299+
impl DefaultNodeDataMarker for super::NodeDefaults {}
300+
301+
impl<M> DefaultNodeDataMarker for super::NodeDefaultsWithMetadata<M> where
302+
M: crate::metadata::NodeMetadata
303+
{
304+
}
305+
}
306+
307+
/// This trait is sealed.
308+
pub trait DefaultNodeData: private::DefaultNodeDataMarker {
309+
fn flags(&self) -> NodeFlags;
310+
fn population(&self) -> PopulationId;
311+
fn individual(&self) -> IndividualId;
312+
fn metadata(&self) -> Result<Option<Vec<u8>>, TskitError>;
313+
}
314+
315+
impl DefaultNodeData for NodeDefaults {
316+
fn flags(&self) -> NodeFlags {
317+
self.flags
318+
}
319+
fn population(&self) -> PopulationId {
320+
self.population
321+
}
322+
fn individual(&self) -> IndividualId {
323+
self.individual
324+
}
325+
fn metadata(&self) -> Result<Option<Vec<u8>>, TskitError> {
326+
Ok(None)
327+
}
328+
}
329+
330+
impl<M> DefaultNodeData for NodeDefaultsWithMetadata<M>
331+
where
332+
M: crate::metadata::NodeMetadata,
333+
{
334+
fn flags(&self) -> NodeFlags {
335+
self.flags
336+
}
337+
fn population(&self) -> PopulationId {
338+
self.population
339+
}
340+
fn individual(&self) -> IndividualId {
341+
self.individual
342+
}
343+
fn metadata(&self) -> Result<Option<Vec<u8>>, TskitError> {
344+
self.metadata.as_ref().map_or_else(
345+
|| Ok(None),
346+
|v| match v.encode() {
347+
Ok(x) => Ok(Some(x)),
348+
Err(e) => Err(e.into()),
349+
},
350+
)
351+
}
352+
}
353+
354+
/// This is a doctest hack as described in the rust book.
355+
/// We do this b/c the specific error messages can change
356+
/// across rust versions, making crates like trybuild
357+
/// less useful.
358+
///
359+
/// ```compile_fail
360+
/// #[derive(serde::Serialize, serde::Deserialize)]
361+
/// struct NodeMetadata {
362+
/// value: i32,
363+
/// }
364+
///
365+
/// impl Default for NodeMetadata {
366+
/// fn default() -> Self {
367+
/// Self{value: 0}
368+
/// }
369+
/// }
370+
///
371+
/// impl tskit::metadata::MetadataRoundtrip for NodeMetadata {
372+
/// fn encode(&self) -> Result<Vec<u8>, tskit::metadata::MetadataError> {
373+
/// match serde_json::to_string(self) {
374+
/// Ok(x) => Ok(x.as_bytes().to_vec()),
375+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
376+
/// }
377+
/// }
378+
/// fn decode(md: &[u8]) -> Result<Self, tskit::metadata::MetadataError>
379+
/// where
380+
/// Self: Sized,
381+
/// {
382+
/// match serde_json::from_slice(md) {
383+
/// Ok(v) => Ok(v),
384+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
385+
/// }
386+
/// }
387+
/// }
388+
///
389+
/// impl tskit::metadata::NodeMetadata for NodeMetadata {}
390+
///
391+
/// type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata<NodeMetadata>;
392+
/// let defaults = DefaultsWithMetadata {
393+
/// metadata: Some(NodeMetadata { value: 42 }),
394+
/// ..Default::default()
395+
/// };
396+
///
397+
/// // Fails because metadata type is not Debug
398+
/// println!("{:?}", defaults);
399+
/// ```
400+
///
401+
/// ```compile_fail
402+
/// #[derive(serde::Serialize, serde::Deserialize)]
403+
/// struct NodeMetadata {
404+
/// value: i32,
405+
/// }
406+
///
407+
/// impl Default for NodeMetadata {
408+
/// fn default() -> Self {
409+
/// Self{value: 0}
410+
/// }
411+
/// }
412+
///
413+
/// impl tskit::metadata::MetadataRoundtrip for NodeMetadata {
414+
/// fn encode(&self) -> Result<Vec<u8>, tskit::metadata::MetadataError> {
415+
/// match serde_json::to_string(self) {
416+
/// Ok(x) => Ok(x.as_bytes().to_vec()),
417+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
418+
/// }
419+
/// }
420+
/// fn decode(md: &[u8]) -> Result<Self, tskit::metadata::MetadataError>
421+
/// where
422+
/// Self: Sized,
423+
/// {
424+
/// match serde_json::from_slice(md) {
425+
/// Ok(v) => Ok(v),
426+
/// Err(e) => Err(::tskit::metadata::MetadataError::RoundtripError { value: Box::new(e) }),
427+
/// }
428+
/// }
429+
/// }
430+
///
431+
/// impl tskit::metadata::NodeMetadata for NodeMetadata {}
432+
///
433+
/// let mut tables = tskit::TableCollection::new(10.0).unwrap();
434+
/// type DefaultsWithMetadata = tskit::NodeDefaultsWithMetadata<NodeMetadata>;
435+
/// // What if there is default metadata for all rows?
436+
/// let defaults = DefaultsWithMetadata {
437+
/// metadata: Some(NodeMetadata { value: 42 }),
438+
/// ..Default::default()
439+
/// };
440+
/// // We can scoop all non-metadata fields even though
441+
/// // type is not Copy/Clone
442+
/// let _ = tables
443+
/// .add_node_with_defaults(
444+
/// 0.0,
445+
/// &DefaultsWithMetadata {
446+
/// metadata: Some(NodeMetadata { value: 2 * 42 }),
447+
/// ..defaults
448+
/// },
449+
/// )
450+
/// .unwrap();
451+
/// // But now, we start to cause a problem:
452+
/// // If we don't clone here, our metadata type moves,
453+
/// // so our defaults are moved.
454+
/// let _ = tables
455+
/// .add_node_with_defaults(
456+
/// 0.0,
457+
/// &DefaultsWithMetadata {
458+
/// population: 6.into(),
459+
/// ..defaults
460+
/// },
461+
/// )
462+
/// .unwrap();
463+
/// // Now, we have a use-after-move error
464+
/// // if we hadn't cloned in the last step.
465+
/// let _ = tables
466+
/// .add_node_with_defaults(
467+
/// 0.0,
468+
/// &DefaultsWithMetadata {
469+
/// individual: 7.into(),
470+
/// ..defaults
471+
/// },
472+
/// )
473+
/// .unwrap();
474+
/// ```
475+
#[cfg(doctest)]
476+
struct NodeDefaultsWithMetadataNotCloneNotDebug;
477+
217478
/// An immtable view of a node table.
218479
#[derive(Debug)]
219480
#[repr(transparent)]
@@ -700,6 +961,26 @@ impl OwningNodeTable {
700961
self.as_mut_ptr(),
701962
)
702963
}
964+
965+
/// Add row with defaults
966+
///
967+
/// # Examples
968+
///
969+
/// ```
970+
/// # let mut nodes = tskit::OwningNodeTable::default();
971+
/// let node_defaults = tskit::NodeDefaults::default();
972+
/// let rv = nodes.add_row_with_defaults(1.0, &node_defaults).unwrap();
973+
/// assert_eq!(rv, 0);
974+
/// let rv = nodes.add_row_with_defaults(1.0, &node_defaults).unwrap();
975+
/// assert_eq!(rv, 1);
976+
/// ```
977+
pub fn add_row_with_defaults<T: Into<crate::Time> + Copy, D: DefaultNodeData>(
978+
&mut self,
979+
time: T,
980+
defaults: &D,
981+
) -> Result<NodeId, TskitError> {
982+
add_row_with_defaults(time, defaults, self.as_mut_ptr())
983+
}
703984
}
704985

705986
#[cfg(test)]

‎src/table_collection.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,28 @@ impl TableCollection {
418418
})
419419
}
420420

421+
/// Add a node using default values
422+
///
423+
/// # Examples
424+
///
425+
/// ```
426+
/// # let mut tables = tskit::TableCollection::new(1.).unwrap();
427+
/// let node_defaults = tskit::NodeDefaults::default();
428+
/// let rv = tables.add_node_with_defaults(1.0, &node_defaults).unwrap();
429+
/// assert_eq!(rv, 0);
430+
/// let rv = tables.add_node_with_defaults(2.0, &node_defaults).unwrap();
431+
/// assert_eq!(rv, 1);
432+
/// ```
433+
pub fn add_node_with_defaults<T: Into<crate::Time>, D: crate::node_table::DefaultNodeData>(
434+
&mut self,
435+
time: T,
436+
defaults: &D,
437+
) -> Result<NodeId, TskitError> {
438+
crate::node_table::add_row_with_defaults(time, defaults, unsafe {
439+
&mut (*self.as_mut_ptr()).nodes
440+
})
441+
}
442+
421443
/// Add a row with optional metadata to the node table
422444
///
423445
/// # Examples

0 commit comments

Comments
 (0)