1 files changed, 34 insertions, 19 deletions

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