Add batched generation

5 files changed, 55 insertions, 32 deletions

diff --git a/build.rs b/build.rs
index e6f102c..f9bb718 100644
--- a/build.rs
+++ b/build.rs
@@ -10,7 +10,7 @@ fn main() {
     cc::Build::new()
         .file("nauty/geng.c")
         .file("nauty/geng-iter.c")
-        // .flag("-O3")
+        .flag("-O3")
         .flag("-Wno-unused-parameter")
         .flag("-Wno-sign-compare")
         .flag("-Wno-unused-variable")
diff --git a/nauty/geng-iter.c b/nauty/geng-iter.c
index cec9834..b5efe21 100644
--- a/nauty/geng-iter.c
+++ b/nauty/geng-iter.c
@@ -30,21 +30,12 @@ void
 OUTPROC(__attribute__((unused)) FILE *outfile,
         graph *g, int n, struct geng_iterator *iter)
 {
-    iter->cur = g;
+    memcpy(&iter->batch[iter->batch_size * n], g, sizeof(set) * n);
     iter->graph_size = n;
+    iter->batch_size++;
 
-    // TODO: add support for generating graphs in batches (for speed)
-#if 0
-    static int i = 1;
-    i++;
-    if (i % 100 == 0)
-    {
-        swapcontext(&geng_worker, &geng_user);
-        i = 0;
-    }
-#else
-    swapcontext(&iter->geng_worker, &iter->geng_user);
-#endif
+    if (iter->batch_size == iter->batch_capacity)
+        swapcontext(&iter->geng_worker, &iter->geng_user);
 }
 
 // TODO: probably don't need to name this function with macro
@@ -58,10 +49,13 @@ geng_iterator_create(struct geng_iterator **iterator_ptr,
                      size_t graph_size,
                      size_t batch_capacity)
 {
-    // TODO: malloc geng_iterator batch
     *iterator_ptr = malloc(sizeof(struct geng_iterator));
     
     struct geng_iterator *iterator = *iterator_ptr;
+    iterator->batch = malloc(sizeof(set) * graph_size * batch_capacity);
+    iterator->batch_capacity = batch_capacity;
+    iterator->batch_size = 0;
+
     iterator->iteration_done = false;
 
     // TODO: add support for more arguments
@@ -75,6 +69,7 @@ geng_iterator_create(struct geng_iterator **iterator_ptr,
     geng_argv[2] = n_str;
     geng_argv[3] = NULL;
 
+    // TODO: make macro
     uint32_t p_argv[2];
     p_argv[0] = (uint32_t) (((size_t) &geng_argv) & ((1llu << 32) - 1llu));
     p_argv[1] = ((size_t) &geng_argv) >> 32;
@@ -98,18 +93,38 @@ geng_iterator_create(struct geng_iterator **iterator_ptr,
 }
 
 bool
-geng_iterator_next(struct geng_iterator *iter, graph *g)
+geng_iterator_next(struct geng_iterator *iter, set *g)
 {
+    static int i = 0;
+
+    if (i == iter->batch_size && iter->generation_done)
+        iter->iteration_done = true;
+
     if (iter->iteration_done) return false;
     else
     {
-        memcpy(g, iter->cur, sizeof(set) * iter->graph_size);
-        swapcontext(&iter->geng_user, &iter->geng_worker);
-        if (!iter->iteration_done && iter->generation_done)
-            iter->iteration_done = true;
+        if (i == iter->batch_size && !iter->generation_done)
+        {
+            iter->batch_size = 0;
+            i = 0;
+            swapcontext(&iter->geng_user, &iter->geng_worker);
+        }
+
+        memcpy(
+            g,
+            &iter->batch[i * iter->graph_size],
+            sizeof(set) * iter->graph_size
+        );
+        i++;
         return true;
     }
 }
 
+void
+geng_iterator_destroy(struct geng_iterator *iter)
+{
+    free(iter->batch);
+}
+
 // TODO: only on macos
 #pragma GCC diagnostic pop
diff --git a/nauty/geng.c b/nauty/geng.c
index 7ea0085..3066c15 100644
--- a/nauty/geng.c
+++ b/nauty/geng.c
@@ -419,8 +419,6 @@ efficient to use the res/mod feature than to split by numbers of edges.
 #include "geng.h"
 #include <stdint.h>
 
-struct geng_iterator;
-
 /* No need for TLS if not calling from a program. */
 #ifndef GENG_MAIN
 #undef TLS_ATTR
@@ -2279,6 +2277,7 @@ genextend(graph *g, int n, int *deg, int ne, boolean rigid, int xlb, int xub, st
 void
 GENG_MAIN(int argc, uint32_t argv1, uint32_t argv2, uint32_t iter1, uint32_t iter2)
 {
+    // TODO: make macro
     size_t *p_argv = (size_t *) (((size_t) argv1) | (((size_t) argv2) << 32));
     char **argv = (char **) *p_argv;
 
diff --git a/nauty/geng.h b/nauty/geng.h
index 751be5c..42f0149 100644
--- a/nauty/geng.h
+++ b/nauty/geng.h
@@ -11,8 +11,5 @@ struct geng_iterator
     bool iteration_done;
     int batch_size;
     int batch_capacity;
-    graph **batch;
-
-    // TODO: remove
-    graph *cur;
+    set *batch;
 };
diff --git a/src/main.rs b/src/main.rs
index ac9995b..7012c99 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,6 +11,7 @@ extern "C" {
         batch_size: usize,
     );
     fn geng_iterator_next(iter: *const geng_iterator, g: *mut u32) -> bool;
+    fn geng_iterator_destroy(iter: *const geng_iterator);
     fn printgraph(g: *const u32, n: usize);
 }
 
@@ -29,7 +30,7 @@ impl GengIterator {
     fn new(n: usize) -> GengIterator {
         let iter = unsafe {
             let iter: *mut geng_iterator = ptr::null_mut();
-            geng_iterator_create(&iter, n, 100);
+            geng_iterator_create(&iter, n, 10000);
             Box::from_raw(iter)
         };
         GengIterator { size: n, iter }
@@ -54,15 +55,26 @@ impl Iterator for &GengIterator {
     }
 }
 
+impl Drop for GengIterator {
+    fn drop(&mut self) {
+        unsafe {
+            let ptr: *const geng_iterator = &*self.iter;
+            geng_iterator_destroy(ptr);
+        }
+    }
+}
+
 fn main() {
-    let gi = GengIterator::new(2);
+    let gi = GengIterator::new(10);
 
-    // let q = gi.take(2).collect::<Vec<_>>();
+    // for i in &gi {
+    //     print_graph(i, gi.size);
+    // }
+
+    let q = gi.take(2000000).collect::<Vec<_>>();
+    println!("{}", q.len());
     // println!("{:?}", q);
     // for i in q {
     //     print_graph(i, gi.size);
     // }
-    for i in &gi {
-        print_graph(i, gi.size);
-    }
 }