goprepare() function added

Signed-off-by: Martin Sustrik <sustrik@250bpm.com>

chan.c

@@ -292,10 +292,11 @@ int mill_choose_wait(void) {
     return mill_suspend();
 }
 
-void *mill_choose_val(void) {
-    /* We can ask for valbuf of size 0 here because we are sure it was
-       preallocated to the correct size beforehand. */
-    return mill_valbuf(mill_running, 0);
+void *mill_choose_val(size_t sz) {
+    /* The assumption here is that by supplying the same size as before
+       we are going to get the same buffer which already has the data
+       written into it. */
+    return mill_valbuf(mill_running, sz);
 }
 
 void mill_chs(chan ch, void *val, size_t sz, const char *current) {
@@ -318,7 +319,7 @@ void *mill_chr(chan ch, size_t sz, const char *current) {
     struct mill_clause cl;
     mill_choose_in(&cl, ch, sz, 0);
     mill_choose_wait();
-    return mill_choose_val();
+    return mill_choose_val(sz);
 }
 
 void mill_chdone(chan ch, void *val, size_t sz, const char *current) {

cr.c

@@ -38,15 +38,29 @@
    allocated at the moment. */
 size_t mill_valbuf_size = 128;
 
+/* Valbuf for tha main coroutine. */
+char mill_main_valbuf[128];
+
 volatile int mill_unoptimisable1 = 1;
 volatile void *mill_unoptimisable2 = NULL;
 
 struct mill_cr mill_main = {0};
+
 struct mill_cr *mill_running = &mill_main;
 
 /* Queue of coroutines scheduled for execution. */
 static struct mill_slist mill_ready = {0};
 
+int goprepare(int count, size_t stack_size, size_t val_size) {
+    if(mill_slow(mill_hascrs()))
+        mill_panic("goprepare called while coroutines are running");
+    /* If needed, make val size slightly bigger to align properly. */
+    mill_valbuf_size = (val_size + 15) & ~((size_t)0xf);
+    /* Allocate the stacks. */
+    return mill_preparestacks(count, stack_size + mill_valbuf_size +
+        sizeof(struct mill_cr));
+}
+
 int mill_suspend(void) {
     /* Store the context of the current coroutine, if any. */
     if(mill_running && mill_setjmp(&mill_running->ctx))
@@ -121,12 +135,23 @@ void mill_yield(const char *current) {
 }
 
 void *mill_valbuf(struct mill_cr *cr, size_t size) {
-    if(mill_fast(size <= mill_valbuf_size))
-        return (void*)(((char*)cr) - mill_valbuf_size);
+    /* Small valbufs don't require dynamic allocation. Also note that main
+       coroutine doesn't have a stack allocated on the heap like other
+       coroutines, so we have to handle valbuf in a special way. */
+    if(mill_fast(cr != &mill_main)) {
+        if(mill_fast(size <= mill_valbuf_size))
+            return (void*)(((char*)cr) - mill_valbuf_size);
+    }
+    else {
+        if(mill_fast(size <= sizeof(mill_main_valbuf)))
+            return (void*)mill_main_valbuf;
+    }
+    /* Large valbufs are simply allocated on heap. */
     if(mill_fast(cr->valbuf && cr->valbuf_sz <= size))
         return cr->valbuf;
     cr->valbuf = realloc(cr->valbuf, size);
-    mill_assert(cr->valbuf);
+    if(!cr->valbuf)
+        mill_panic("not enough memory to receive from channel");
     cr->valbuf_sz = size;
     return cr->valbuf;
 }

debug.c

@@ -235,3 +235,8 @@ void mill_preserve_debug(void) {
     gotrace(0);
 }
 
+int mill_hascrs(void) {
+    return (mill_all_crs.first == &mill_main.debug.item &&
+        mill_all_crs.last == &mill_main.debug.item) ? 0 : 1;
+}
+

debug.h

@@ -67,4 +67,7 @@ extern int mill_tracelevel;
 #define mill_trace if(mill_slow(mill_tracelevel)) mill_trace_
 void mill_trace_(const char *location, const char *format, ...);
 
+/* Returns 1 if there are any coroutines running, 0 otherwise. */
+int mill_hascrs(void);
+
 #endif

libmill.h

@@ -82,10 +82,7 @@ MILL_EXPORT int64_t now(void);
 /*  Coroutines                                                                */
 /******************************************************************************/
 
-MILL_EXPORT int gotune(
-    int max_coroutine_count,
-    size_t max_stack_size,
-    size_t max_chan_item_size);
+MILL_EXPORT int goprepare(int count, size_t stack_size, size_t val_size);
 
 MILL_EXPORT extern volatile int mill_unoptimisable1;
 MILL_EXPORT extern volatile void *mill_unoptimisable2;
@@ -200,7 +197,7 @@ MILL_EXPORT void mill_chclose(chan ch, const char *current);
                 type name;\
                 mill_concat(mill_label, idx):\
                 if(mill_idx == idx) {\
-                    name = *(type*)mill_choose_val();\
+                    name = *(type*)mill_choose_val(sizeof(type));\
                     goto mill_concat(mill_dummylabel, idx);\
                     mill_concat(mill_dummylabel, idx)
 
@@ -254,7 +251,7 @@ MILL_EXPORT void mill_choose_out(void *clause, chan ch, void *val, size_t sz,
     int idx);
 MILL_EXPORT void mill_choose_otherwise(void);
 MILL_EXPORT int mill_choose_wait(void);
-MILL_EXPORT void *mill_choose_val(void);
+MILL_EXPORT void *mill_choose_val(size_t sz);
 
 /******************************************************************************/
 /*  IP address library                                                        */

stack.c

@@ -30,7 +30,7 @@
 #include <stdint.h>
 #include <stdlib.h>
 
-/* Size of stack for new coroutines. In bytes. Default size is slightly smaller
+/* Size of stack for all coroutines. In bytes. Default size is slightly smaller
    than page size to account for malloc's chunk header. */
 size_t mill_stack_size = 256 * 1024 - 256;
 
@@ -46,13 +46,41 @@ int mill_max_cached_stacks = 64;
 static int mill_num_cached_stacks = 0;
 static struct mill_slist mill_cached_stacks = {0};
 
+int mill_preparestacks(int count, size_t stack_size) {
+    /* Purge the cached stacks. */
+    while(1) {
+        struct mill_slist_item *item = mill_slist_pop(&mill_cached_stacks);
+        if(!item)
+            break;
+        free(((char*)(item + 1)) - mill_stack_size);
+    }
+    /* Now that there are no stacks allocated, we can adjust the stack size. */
+    mill_stack_size = stack_size;
+    /* Make sure that the stacks won't get deallocated even if they aren't used
+       at the moment. */
+    if(count > mill_max_cached_stacks)
+        mill_max_cached_stacks = count;
+    /* Allocate the new stacks. */
+    int i;
+    for(i = 0; i != count; ++i) {
+        char *ptr = malloc(mill_stack_size);
+        if(!ptr)
+            break;
+        ptr += mill_stack_size;
+        struct mill_slist_item *item = ((struct mill_slist_item*)ptr) - 1;
+        mill_slist_push_back(&mill_cached_stacks, item);
+    }
+    return i;
+}
+
 void *mill_allocstack(void) {
     if(!mill_slist_empty(&mill_cached_stacks)) {
         --mill_num_cached_stacks;
         return (void*)(mill_slist_pop(&mill_cached_stacks) + 1);
     }
     char *ptr = malloc(mill_stack_size);
-    mill_assert(ptr);
+    if(!ptr)
+        mill_panic("not enough memory to allocate coroutine stack");
     return ptr + mill_stack_size;
 }
 

stack.h

@@ -25,6 +25,12 @@
 #ifndef MILL_STACK_INCLUDED
 #define MILL_STACK_INCLUDED
 
+#include <stddef.h>
+
+/* Purges all the existing cached stacks and preallocated 'count' new stacks
+   of size 'stack_size'. Returns number of actually allocated stacks. */
+int mill_preparestacks(int count, size_t stack_size);
+
 /* Allocates new stack. Returns pointer to the *top* of the stack.
    For now we assume that the stack grows downwards. */
 void *mill_allocstack(void);