From 71e927a9bb221309afa72cb601038e3232e3e180 Mon Sep 17 00:00:00 2001
From: "Kevin R. Thornton" <krthornt@uci.edu>
Date: Thu, 26 May 2022 10:21:37 -0700
Subject: [PATCH] TreeSequence::new now takes ownership of the TableCollection.
 (#235)

---
 src/table_collection.rs | 15 +++++++++++++++
 src/trees.rs            | 15 +++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/table_collection.rs b/src/table_collection.rs
index 27ff7f08..a0947413 100644
--- a/src/table_collection.rs
+++ b/src/table_collection.rs
@@ -151,6 +151,21 @@ impl TableCollection {
         Ok(tables)
     }
 
+    pub(crate) fn into_raw(self) -> Result<*mut ll_bindings::tsk_table_collection_t, TskitError> {
+        let mut tables = self;
+        // rust won't let use move inner out b/c this type implements Drop.
+        // So we have to replace the existing pointer with a new one.
+        let table_ptr = unsafe {
+            libc::malloc(std::mem::size_of::<ll_bindings::tsk_table_collection_t>())
+                as *mut ll_bindings::tsk_table_collection_t
+        };
+        let rv = unsafe { ll_bindings::tsk_table_collection_init(table_ptr, 0) };
+
+        let mut temp = unsafe { MBox::from_raw(table_ptr) };
+        std::mem::swap(&mut temp, &mut tables.inner);
+        handle_tsk_return_value!(rv, MBox::into_raw(temp))
+    }
+
     /// Load a table collection from a file.
     ///
     /// # Panics
diff --git a/src/trees.rs b/src/trees.rs
index d6bcc045..3e8a1a0e 100644
--- a/src/trees.rs
+++ b/src/trees.rs
@@ -1000,15 +1000,22 @@ impl TreeSequence {
     /// let tree_sequence = tskit::TreeSequence::new(tables,
     /// tskit::TreeSequenceFlags::default()).unwrap();
     /// ```
+    ///
+    /// ## Note
+    ///
+    /// This function makes *no extra copies* of the tables.
+    /// There is, however, a temporary allocation of an empty table collection
+    /// in order to convince rust that we are safely handling all memory.
     pub fn new<F: Into<TreeSequenceFlags>>(
         tables: TableCollection,
         flags: F,
     ) -> Result<Self, TskitError> {
-        let mut t = tables;
         let mut treeseq = Self::wrap();
-        let rv = unsafe {
-            ll_bindings::tsk_treeseq_init(treeseq.as_mut_ptr(), t.as_mut_ptr(), flags.into().bits())
-        };
+        let mut flags: u32 = flags.into().bits();
+        flags |= ll_bindings::TSK_TAKE_OWNERSHIP;
+        let raw_tables_ptr = tables.into_raw()?;
+        let rv =
+            unsafe { ll_bindings::tsk_treeseq_init(treeseq.as_mut_ptr(), raw_tables_ptr, flags) };
         handle_tsk_return_value!(rv, treeseq)
     }