refactor: LLTreeSeq impl now takes flag types directly (#514)

tskit-dev · May 22, 2023 · a365a7e · a365a7e
1 parent eaecd5a
commit a365a7e
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/src/sys/treeseq.rs b/src/sys/treeseq.rs
@@ -9,11 +9,10 @@ pub struct LLTreeSeq(bindings::tsk_treeseq_t);
 impl LLTreeSeq {
     pub fn new(
         tables: *mut bindings::tsk_table_collection_t,
-        flags: bindings::tsk_flags_t,
+        flags: super::flags::TreeSequenceFlags,
     ) -> Result<Self, Error> {
         let mut inner = std::mem::MaybeUninit::<bindings::tsk_treeseq_t>::uninit();
-        let mut flags = flags;
-        flags |= bindings::TSK_TAKE_OWNERSHIP;
+        let flags = flags.bits() | bindings::TSK_TAKE_OWNERSHIP;
         match unsafe { bindings::tsk_treeseq_init(inner.as_mut_ptr(), tables, flags) } {
             code if code < 0 => Err(Error::Code(code)),
             _ => Ok(Self(unsafe { inner.assume_init() })),
@@ -35,7 +34,7 @@ impl LLTreeSeq {
     pub fn simplify(
         &self,
         samples: &[bindings::tsk_id_t],
-        options: bindings::tsk_flags_t,
+        options: super::flags::SimplificationOptions,
         idmap: *mut bindings::tsk_id_t,
     ) -> Result<Self, Error> {
         // The output is an UNINITIALIZED treeseq,
@@ -48,7 +47,7 @@ impl LLTreeSeq {
                 self.as_ptr(),
                 samples.as_ptr(),
                 samples.len() as bindings::tsk_size_t,
-                options,
+                options.bits(),
                 ts.as_mut_ptr(),
                 idmap,
             )

diff --git a/src/trees.rs b/src/trees.rs
@@ -248,7 +248,7 @@ impl TreeSequence {
         flags: F,
     ) -> Result<Self, TskitError> {
         let raw_tables_ptr = tables.into_raw()?;
-        let mut inner = sys::LLTreeSeq::new(raw_tables_ptr, flags.into().bits())?;
+        let mut inner = sys::LLTreeSeq::new(raw_tables_ptr, flags.into())?;
         let views = crate::table_views::TableViews::new_from_tree_sequence(inner.as_mut_ptr())?;
         Ok(Self { inner, views })
     }
@@ -460,7 +460,7 @@ impl TreeSequence {
         };
         let mut inner = self.inner.simplify(
             llsamples,
-            options.into().bits(),
+            options.into(),
             match idmap {
                 true => output_node_map.as_mut_ptr().cast::<tsk_id_t>(),
                 false => std::ptr::null_mut(),