diff --git a/src/sys/table_collection.rs b/src/sys/table_collection.rs index 4973ba79..48810e7e 100644 --- a/src/sys/table_collection.rs +++ b/src/sys/table_collection.rs @@ -105,4 +105,8 @@ impl TableCollection { // SAFETY: self pointer is not null unsafe { &mut (*self.as_mut_ptr()).sites } } + + pub fn into_raw(self) -> *mut tsk_table_collection_t { + unsafe { self.0.into_raw() } + } } diff --git a/src/sys/tree.rs b/src/sys/tree.rs index 6edbcc73..d1d43cc3 100644 --- a/src/sys/tree.rs +++ b/src/sys/tree.rs @@ -25,7 +25,7 @@ impl<'treeseq> LLTree<'treeseq> { let code = unsafe { super::bindings::tsk_tree_set_tracked_samples( inner.as_mut(), - treeseq.num_samples(), + treeseq.num_samples().into(), (inner.as_mut()).samples, ) }; diff --git a/src/sys/treeseq.rs b/src/sys/treeseq.rs index 092ce6aa..89128808 100644 --- a/src/sys/treeseq.rs +++ b/src/sys/treeseq.rs @@ -11,9 +11,10 @@ pub struct TreeSequence(TskBox); impl TreeSequence { pub fn new( - tables: *mut bindings::tsk_table_collection_t, + tables: super::TableCollection, flags: super::flags::TreeSequenceFlags, ) -> Result { + let tables = tables.into_raw(); let inner = TskBox::new(|t: *mut bindings::tsk_treeseq_t| unsafe { tsk_treeseq_init(t, tables, flags.bits() | bindings::TSK_TAKE_OWNERSHIP) })?; @@ -30,9 +31,9 @@ impl TreeSequence { pub fn simplify( &self, - samples: &[bindings::tsk_id_t], + samples: &[super::newtypes::NodeId], options: super::flags::SimplificationOptions, - idmap: *mut bindings::tsk_id_t, + idmap: Option<&mut [super::newtypes::NodeId]>, ) -> Result { // The output is an UNINITIALIZED treeseq, // else we leak memory. @@ -42,11 +43,15 @@ impl TreeSequence { let rv = unsafe { bindings::tsk_treeseq_simplify( self.as_ref(), - samples.as_ptr(), - samples.len() as bindings::tsk_size_t, + // The cast is safe/sound b/c NodeId is repr(transparent) + samples.as_ptr().cast::<_>(), + samples.len().try_into().unwrap(), options.bits(), ts.as_mut_ptr(), - idmap, + match idmap { + Some(s) => s.as_mut_ptr().cast::<_>(), + None => std::ptr::null_mut(), + }, ) }; if rv < 0 { @@ -72,9 +77,9 @@ impl TreeSequence { } } - pub fn num_trees(&self) -> bindings::tsk_size_t { + pub fn num_trees(&self) -> super::newtypes::SizeType { // SAFETY: self pointer is not null - unsafe { bindings::tsk_treeseq_get_num_trees(self.as_ref()) } + unsafe { bindings::tsk_treeseq_get_num_trees(self.as_ref()) }.into() } pub fn num_nodes_raw(&self) -> bindings::tsk_size_t { @@ -95,7 +100,7 @@ impl TreeSequence { } } - pub fn num_samples(&self) -> bindings::tsk_size_t { - unsafe { bindings::tsk_treeseq_get_num_samples(self.as_ref()) } + pub fn num_samples(&self) -> super::newtypes::SizeType { + unsafe { bindings::tsk_treeseq_get_num_samples(self.as_ref()) }.into() } } diff --git a/src/table_collection.rs b/src/table_collection.rs index 0bcfc0aa..971c0cf6 100644 --- a/src/table_collection.rs +++ b/src/table_collection.rs @@ -109,13 +109,8 @@ impl TableCollection { }) } - pub(crate) fn into_raw(self) -> Result<*mut ll_bindings::tsk_table_collection_t, TskitError> { - let mut tables = self; - let mut temp = crate::sys::TableCollection::new(1.)?; - std::mem::swap(&mut temp, &mut tables.inner); - let ptr = temp.as_mut_ptr(); - std::mem::forget(temp); - handle_tsk_return_value!(0, ptr) + pub(crate) fn into_inner(self) -> crate::sys::TableCollection { + self.inner } /// Load a table collection from a file. diff --git a/src/trees/treeseq.rs b/src/trees/treeseq.rs index 75e80774..92516963 100644 --- a/src/trees/treeseq.rs +++ b/src/trees/treeseq.rs @@ -9,7 +9,6 @@ use crate::TableOutputOptions; use crate::TreeFlags; use crate::TreeSequenceFlags; use crate::TskReturnValue; -use ll_bindings::tsk_id_t; use sys::bindings as ll_bindings; use super::Tree; @@ -113,7 +112,7 @@ impl TreeSequence { tables: TableCollection, flags: F, ) -> Result { - let raw_tables_ptr = tables.into_raw()?; + let raw_tables_ptr = tables.into_inner(); let mut inner = sys::TreeSequence::new(raw_tables_ptr, flags.into())?; let views = crate::table_views::TableViews::new_from_tree_sequence(inner.as_mut())?; Ok(Self { inner, views }) @@ -302,7 +301,7 @@ impl TreeSequence { /// Get the number of trees. pub fn num_trees(&self) -> SizeType { - self.inner.num_trees().into() + self.inner.num_trees() } /// Calculate the average Kendall-Colijn (`K-C`) distance between @@ -322,7 +321,7 @@ impl TreeSequence { // FIXME: document pub fn num_samples(&self) -> SizeType { - self.inner.num_samples().into() + self.inner.num_samples() } /// Simplify tables and return a new tree sequence. @@ -348,15 +347,12 @@ impl TreeSequence { if idmap { output_node_map.resize(usize::try_from(self.nodes().num_rows())?, NodeId::NULL); } - let llsamples = unsafe { - std::slice::from_raw_parts(samples.as_ptr().cast::(), samples.len()) - }; let mut inner = self.inner.simplify( - llsamples, + samples, options.into(), match idmap { - true => output_node_map.as_mut_ptr().cast::(), - false => std::ptr::null_mut(), + true => Some(&mut output_node_map), + false => None, }, )?; let views = crate::table_views::TableViews::new_from_tree_sequence(inner.as_mut())?;