try to make more efficient, but table integrity checking blocks us

Author: molpopgen

examples/haploid_wright_fisher_simple_edge_buffer.rs

@@ -1,5 +1,3 @@
-#![cfg(feature="bindings")]
-
 // This is a rust implementation of the example
 // found in tskit-c
 
@@ -12,10 +10,19 @@ use rand::distributions::Distribution;
 use rand::prelude::*;
 use rand::SeedableRng;
 
+#[derive(Debug)]
+struct Edge {
+    left: tskit::Position,
+    right: tskit::Position,
+    child: tskit::NodeId,
+    previous: Option<usize>,
+}
+
 #[derive(Default)]
 struct EdgeBuffer {
     parent: Vec<tskit::NodeId>,
-    edges: HashMap<tskit::NodeId, tskit::OwningEdgeTable>,
+    last: HashMap<tskit::NodeId, usize>,
+    edges: Vec<Edge>,
 }
 
 impl EdgeBuffer {
@@ -26,16 +33,34 @@ impl EdgeBuffer {
         parent: tskit::NodeId,
         child: tskit::NodeId,
     ) -> Result<()> {
-        if let Some(ref mut edges) = self.edges.get_mut(&parent) {
-            edges.add_row(left, right, parent, child)?;
+        if let Some(last) = self.last.get_mut(&parent) {
+            let pchild = self.edges[*last].child;
+            assert!(child >= pchild);
+            self.edges.push(Edge {
+                left,
+                right,
+                child,
+                previous: Some(*last),
+            });
+            *last = self.edges.len() - 1;
         } else {
-            let mut edges = tskit::OwningEdgeTable::default();
-            edges.add_row(left, right, parent, child)?;
-            self.edges.insert(parent, edges);
+            self.edges.push(Edge {
+                left,
+                right,
+                child,
+                previous: None,
+            });
+            self.last.insert(parent, self.edges.len() - 1);
             self.parent.push(parent);
         }
         Ok(())
     }
+
+    fn clear(&mut self) {
+        self.parent.clear();
+        self.last.clear();
+        self.edges.clear();
+    }
 }
 
 fn rotate_edges(bookmark: &tskit::types::Bookmark, tables: &mut tskit::TableCollection) {
@@ -97,6 +122,7 @@ fn simulate(
 
     let mut buffer = EdgeBuffer::default();
     for birth_time in (0..num_generations).rev() {
+        println!("bt = {birth_time:?}");
         for c in children.iter_mut() {
             let bt = f64::from(birth_time);
             let child = tables.add_node(0, bt, -1, -1)?;
@@ -115,27 +141,19 @@ fn simulate(
         if birth_time % simplify_interval == 0 {
             //tables.sort(&bookmark, tskit::TableSortOptions::default())?;
             for &parent in buffer.parent.iter().rev() {
-                // remove the buffered edges from the internal hash
-                // NOTE: we unwrap here b/c our impl of EdgeBuffer guarantees that
-                // the object exists in the hash
-                let edges = buffer.edges.remove(&parent).unwrap();
-                println!("{:?}", edges.num_rows());
-                let rv = unsafe {
-                    tskit::bindings::tsk_edge_table_append_columns(
-                        &mut (*tables.as_mut_ptr()).edges,
-                        edges.num_rows().into(),
-                        (*edges.as_ptr()).left,
-                        (*edges.as_ptr()).right,
-                        (*edges.as_ptr()).parent,
-                        (*edges.as_ptr()).child,
-                        (*edges.as_ptr()).metadata,
-                        (*edges.as_ptr()).metadata_offset,
-                    )
-                };
-                assert!(rv >= 0);
+                println!("{parent:?}");
+                let mut last = buffer.last.get(&parent).cloned();
+                while let Some(previous) = last {
+                    let edge = &buffer.edges[previous];
+                    tables.add_edge(edge.left, edge.right, parent, edge.child)?;
+                    println!("adding {edge:?}");
+                    tables.check_integrity(
+                        tskit::TableIntegrityCheckFlags::default().check_edge_ordering(),
+                    )?;
+                    last = edge.previous;
+                }
             }
-            buffer.parent.clear();
-            buffer.edges.clear();
+            buffer.clear();
             rotate_edges(&bookmark, &mut tables);
             if let Some(idmap) =
                 tables.simplify(children, tskit::SimplificationOptions::default(), true)?