From 7c0cba0dd0ee03e3719fdd489b4adda97c047f40 Mon Sep 17 00:00:00 2001
From: Kevin Thornton <krthornt@uci.edu>
Date: Fri, 29 May 2026 03:42:11 -0700
Subject: [PATCH] feat: borrow pointer from python TableCollection

---
 .../maketrees/python/maketrees/__init__.py    |  1 +
 python_example/maketrees/src/lib.rs           | 12 +++++++
 .../maketrees/tests/test_maketrees.py         |  8 +++++
 src/lib.rs                                    | 34 +++++++++++++++++++
 tests/tests.rs                                | 21 ++++++++++++
 5 files changed, 76 insertions(+)

diff --git a/python_example/maketrees/python/maketrees/__init__.py b/python_example/maketrees/python/maketrees/__init__.py
index 192f65e..c2e7de6 100644
--- a/python_example/maketrees/python/maketrees/__init__.py
+++ b/python_example/maketrees/python/maketrees/__init__.py
@@ -2,6 +2,7 @@
 
 from ._maketrees import maketrees
 from ._maketrees import raise_error
+from ._maketrees import clear_shared_tables
 
 
 def make_treeseq_with_metadata() -> tskit.TreeSequence:
diff --git a/python_example/maketrees/src/lib.rs b/python_example/maketrees/src/lib.rs
index b5ef35d..8e717f3 100644
--- a/python_example/maketrees/src/lib.rs
+++ b/python_example/maketrees/src/lib.rs
@@ -70,4 +70,16 @@ mod maketrees {
         let holder = tskit2tskit::SharedTableCollection::new(py, f64::NAN)?;
         Ok(holder.into_python_tables(py)?)
     }
+
+    #[pyfunction]
+    fn clear_shared_tables(py: Python<'_>, pytables: Py<PyAny>) -> PyResult<()> {
+        let mut holder =
+            unsafe { tskit2tskit::SharedTableCollection::new_from_tables(py, pytables) }?;
+        unsafe {
+            holder.with_mut_tables(|tables| {
+                tables.clear(0).unwrap();
+            })
+        };
+        Ok(())
+    }
 }
diff --git a/python_example/maketrees/tests/test_maketrees.py b/python_example/maketrees/tests/test_maketrees.py
index 020438b..725776b 100644
--- a/python_example/maketrees/tests/test_maketrees.py
+++ b/python_example/maketrees/tests/test_maketrees.py
@@ -22,6 +22,14 @@ def test_metadata():
     assert m['data'] == "I am a mutation"
 
 
+def test_mutable_sharing_of_tables():
+    tables = tskit.TableCollection(100)
+    tables.nodes.add_row(0, 0.0, -1, -1)
+    assert tables.nodes.num_rows == 1
+    maketrees.clear_shared_tables(tables)
+    assert tables.nodes.num_rows == 0
+
+
 def test_raise():
     try:
         _ = maketrees.raise_error()
diff --git a/src/lib.rs b/src/lib.rs
index e8b1135..cc5da88 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -111,6 +111,40 @@ impl SharedTableCollection {
         })
     }
 
+    /// Constructor to borrow data from a table collection from tskit-python
+    ///
+    /// # Parameters
+    ///
+    /// * `pytables`: a `tskit.TableCollection` from `tskit-python`
+    ///
+    /// # Safety
+    ///
+    /// * `tskit-python` and `tskit-rust` must share the same memory layout
+    ///   for `tsk_table_collection_t`.
+    pub unsafe fn new_from_tables<'py>(
+        py: Python<'py>,
+        pytables: Py<PyAny>,
+    ) -> Result<Self, Error> {
+        let ll_tables = pytables.getattr(py, "_ll_tables")?;
+        assert!(!ll_tables.as_ptr().is_null());
+        // SAFETY: ll_tables is a valid pointer to the low-level type.
+        // (Else we'd have gotten an error above.)
+        // The pointer is not NULL and it MUST have been initialized
+        // by tskit-python.
+        let ptr = unsafe { read_tsk_ptr(ll_tables.as_ptr()) };
+        assert!(!ptr.is_null());
+        // SAFETY: ptr is not NULL
+        assert!(!unsafe { *ptr }.is_null());
+        // SAFETY: nothing is NULL
+        // The pointee has been initialized w/o error over in Python
+        let tables =
+            unsafe { tskit::TableCollection::new_from_raw(std::ptr::NonNull::new(*ptr).unwrap()) }?;
+        Ok(Self {
+            tables: Some(tables),
+            pytables: Some(pytables),
+        })
+    }
+
     unsafe fn as_rust(&self) -> &tskit::TableCollection {
         self.tables.as_ref().unwrap()
     }
diff --git a/tests/tests.rs b/tests/tests.rs
index 539ce31..fd48000 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -34,3 +34,24 @@ fn test_mutable_holder_into_tree_sequence() {
         let _pyts = holder.into_python_tree_sequence(py).unwrap();
     })
 }
+
+#[test]
+fn test_borrow_tables_from_python() {
+    Python::attach(|py| {
+        py.run(cr"import tskit", None, None).unwrap();
+        py.run(cr"tables = tskit.TableCollection(100) ", None, None)
+            .unwrap();
+        let pytables = py
+            .eval(c"tables", None, None)
+            .map_err(|e| {
+                e.print_and_set_sys_last_vars(py);
+            })
+            .unwrap()
+            .unbind();
+        let mut holder =
+            unsafe { tskit2tskit::SharedTableCollection::new_from_tables(py, pytables) }.unwrap();
+        unsafe { holder.with_tables(|tables| assert_eq!(tables.sequence_length(), 100.)) };
+        unsafe { holder.with_mut_tables(|tables| tables.add_node(0, 0., -1, -1).unwrap()) };
+        py.run(cr"assert tables.nodes.num_rows == 1", None, None).unwrap();
+    });
+}