Insertion/removal accessors for tree sequences

Author: molpopgen

src/sys/treeseq.rs

@@ -88,6 +88,12 @@ impl TreeSequence {
         unsafe { (*(self.as_ref()).tables).nodes.num_rows }
     }
 
+    fn num_edges_raw(&self) -> bindings::tsk_size_t {
+        assert!(!self.as_ref().tables.is_null());
+        // SAFETY: none of the pointers are null
+        unsafe { (*(self.as_ref()).tables).edges.num_rows }
+    }
+
     pub fn kc_distance(&self, other: &Self, lambda: f64) -> Result<f64, TskitError> {
         let mut kc: f64 = f64::NAN;
         let kcp: *mut f64 = &mut kc;
@@ -103,4 +109,30 @@ impl TreeSequence {
     pub fn num_samples(&self) -> super::newtypes::SizeType {
         unsafe { bindings::tsk_treeseq_get_num_samples(self.as_ref()) }.into()
     }
+
+    pub fn edge_insertion_order(&self) -> &[super::newtypes::EdgeId] {
+        assert!(!self.as_ref().tables.is_null());
+        // SAFETY: all array lengths are the number of rows in the table
+        // SAFETY: no pointers can be null
+        // SAFETY: tables are indexed in order to create a treeseq
+        unsafe {
+            super::generate_slice(
+                (*self.as_ref().tables).indexes.edge_insertion_order,
+                self.num_edges_raw(),
+            )
+        }
+    }
+
+    pub fn edge_removal_order(&self) -> &[super::newtypes::EdgeId] {
+        assert!(!self.as_ref().tables.is_null());
+        // SAFETY: all array lengths are the number of rows in the table
+        // SAFETY: no pointers can be null
+        // SAFETY: tables are indexed in order to create a treeseq
+        unsafe {
+            super::generate_slice(
+                (*self.as_ref().tables).indexes.edge_removal_order,
+                self.num_edges_raw(),
+            )
+        }
+    }
 }

src/trees/treeseq.rs

@@ -633,6 +633,30 @@ impl TreeSequence {
     ) -> Vec<crate::NodeId> {
         self.tables.create_node_id_vector(f)
     }
+
+    #[allow(unused)]
+    pub(crate) fn edge_insertion_order(&self) -> &[crate::EdgeId] {
+        self.inner.edge_insertion_order()
+    }
+
+    #[allow(unused)]
+    pub(crate) fn edge_insertion_order_column(
+        &self,
+    ) -> impl crate::TableColumn<crate::EdgeId, crate::EdgeId> + '_ {
+        crate::table_column::OpaqueTableColumn(self.edge_insertion_order())
+    }
+
+    #[allow(unused)]
+    pub(crate) fn edge_removal_order(&self) -> &[crate::EdgeId] {
+        self.inner.edge_removal_order()
+    }
+
+    #[allow(unused)]
+    pub(crate) fn edge_removal_order_column(
+        &self,
+    ) -> impl crate::TableColumn<crate::EdgeId, crate::EdgeId> + '_ {
+        crate::table_column::OpaqueTableColumn(self.edge_removal_order())
+    }
 }
 
 impl TryFrom<TableCollection> for TreeSequence {