generator: Allow generators to take arguments

Using some serious macro magic, this patch extends generators to allow
them to take arbitrary arguments.  The arguments are marshalled into a
structure placed at the far end of the generator's stack when it is
created.  Then, they're unmarshalled back into C parameters when we first
context switch into the generator.
Signed-off-by: David Gibson <david@gibson.dropbear.id.au>

ccan/generator/_info

@@ -59,6 +59,7 @@ int main(int argc, char *argv[])
 		printf("ccan/ptrint\n");
 		printf("ccan/alignof\n");
 		printf("ccan/cppmagic\n");
+		printf("ccan/compiler\n");
 		return 0;
 	}
 

ccan/generator/generator.h

@@ -23,6 +23,7 @@
 #include <ccan/ptrint/ptrint.h>
 #include <ccan/build_assert/build_assert.h>
 #include <ccan/cppmagic/cppmagic.h>
+#include <ccan/compiler/compiler.h>
 
 /*
  * Internals - included just for the use of inlines and macros
@@ -40,6 +41,11 @@ static inline struct generator_ *generator_state_(const void *ret)
 	return (struct generator_ *)ret - 1;
 }
 
+static inline void *generator_argp_(const void *ret)
+{
+	return generator_state_(ret)->base;
+}
+
 struct generator_incomplete_;
 
 #define generator_rtype_(gen_)			\
@@ -77,8 +83,8 @@ void generator_free_(void *ret);
  * Example:
  *	generator_declare(count_to_3, int);
  */
-#define generator_declare(name_, rtype_)	\
-	generator_t(rtype_) name_(void)
+#define generator_declare(name_, rtype_, ...)	\
+	generator_t(rtype_) name_(generator_parms_outer_(__VA_ARGS__))
 
 /**
  * generator_def - define a generator function
@@ -97,11 +103,35 @@ void generator_free_(void *ret);
  *		generator_yield(3);
  *	}
  */
-#define generator_def_(name_, rtype_, storage_)				\
-	static void name_##_generator_(rtype_ *ret_);			\
+#define generator_parm_(t_, n_)			t_ n_
+#define generator_parms_(...)						\
+	CPPMAGIC_2MAP(generator_parm_, __VA_ARGS__)
+#define generator_parms_inner_(...)					\
+	CPPMAGIC_IFELSE(CPPMAGIC_NONEMPTY(__VA_ARGS__))			\
+		(, generator_parms_(__VA_ARGS__))()
+#define generator_parms_outer_(...)					\
+	CPPMAGIC_IFELSE(CPPMAGIC_NONEMPTY(__VA_ARGS__))	\
+		(generator_parms_(__VA_ARGS__))(void)
+#define generator_argfield_(t_, n_)		t_ n_;
+#define generator_argstruct_(...)					\
+	struct {							\
+		CPPMAGIC_JOIN(, CPPMAGIC_2MAP(generator_argfield_,	\
+					      __VA_ARGS__))		\
+	}
+#define generator_arg_unpack_(t_, n_)		args->n_
+#define generator_args_unpack_(...)		\
+	CPPMAGIC_IFELSE(CPPMAGIC_NONEMPTY(__VA_ARGS__))			\
+		(, CPPMAGIC_2MAP(generator_arg_unpack_, __VA_ARGS__))()
+#define generator_arg_pack_(t_, n_)		args->n_ = n_
+#define generator_args_pack_(...)					\
+	CPPMAGIC_JOIN(;, CPPMAGIC_2MAP(generator_arg_pack_, __VA_ARGS__))
+#define generator_def_(name_, rtype_, storage_, ...)			\
+	static void name_##_generator_(rtype_ *ret_			\
+				       generator_parms_inner_(__VA_ARGS__)); \
 	static void name_##_generator__(generator_wrapper_args_())	\
 	{								\
 		struct generator_ *gen;					\
+		UNNEEDED generator_argstruct_(__VA_ARGS__) *args;	\
 		CPPMAGIC_IFELSE(HAVE_POINTER_SAFE_MAKECONTEXT)		\
 			()						\
 			(ptrdiff_t hilo = ((ptrdiff_t)hi << (8*sizeof(int))) \
@@ -110,19 +140,26 @@ void generator_free_(void *ret);
 			BUILD_ASSERT(sizeof(struct generator_ *)	\
 				     <= 2*sizeof(int));)		\
 		gen = generator_state_(ret);				\
-		name_##_generator_(ret);				\
+		args = generator_argp_(ret);				\
+		name_##_generator_(ret generator_args_unpack_(__VA_ARGS__)); \
 		gen->complete = true;					\
 		setcontext(&gen->caller);				\
 		assert(0);						\
 	}								\
-	storage_ generator_t(rtype_) name_(void)			\
+	storage_ generator_t(rtype_)					\
+	name_(generator_parms_outer_(__VA_ARGS__))			\
 	{								\
-		return generator_new_(name_##_generator__,		\
+		generator_t(rtype_) gen = generator_new_(name_##_generator__, \
 							 sizeof(rtype_)); \
+		UNNEEDED generator_argstruct_(__VA_ARGS__) *args =	\
+			generator_argp_(gen);				\
+		generator_args_pack_(__VA_ARGS__);			\
+		return gen;						\
 	}								\
-	static void name_##_generator_(rtype_ *ret_)
-#define generator_def(name_, rtype_)		\
-	generator_def_(name_, rtype_, )
+	static void name_##_generator_(rtype_ *ret_			\
+				       generator_parms_inner_(__VA_ARGS__))
+#define generator_def(name_, rtype_, ...)	\
+	generator_def_(name_, rtype_, , __VA_ARGS__)
 
 /**
  * generator_def_static - define a private / local generator function
@@ -132,8 +169,8 @@ void generator_free_(void *ret);
  * As generator_def, but the resulting generator function will be
  * local to this module.
  */
-#define generator_def_static(name_, rtype_)	\
-	generator_def_(name_, rtype_, static)
+#define generator_def_static(name_, rtype_, ...)	\
+	generator_def_(name_, rtype_, static, __VA_ARGS__)
 
 /**
  * generator_yield - yield (return) a value from a generator

ccan/generator/test/api.c

@@ -32,6 +32,40 @@ static void test1(void)
 	generator_free(state1);
 }
 
+static void test2(void)
+{
+	generator_t(int) state2 = gen2(100);
+	int *ret;
+
+	ok1((ret = generator_next(state2)) != NULL);
+	ok1(*ret == 101);
+	ok1((ret = generator_next(state2)) != NULL);
+	ok1(*ret == 103);
+	ok1((ret = generator_next(state2)) != NULL);
+	ok1(*ret == 117);
+	ok1((ret = generator_next(state2)) == NULL);
+
+	generator_free(state2);
+}
+
+static void test3(void)
+{
+	int i;
+
+	for (i = 0; i < 4; i++) {
+		generator_t(const char *) state3 = gen3("test", i);
+		const char *s;
+		int j;
+
+		for (j = 0; j < i; j++) {
+			ok1(generator_next_val(s, state3));
+			ok1(streq(s, "test"));
+		}
+		ok1(!generator_next_val(s, state3));
+		generator_free(state3);
+	}
+}
+
 static void testx(void)
 {
 	generator_t(const char *) statex = genx();
@@ -52,9 +86,11 @@ static void testx(void)
 int main(void)
 {
 	/* This is how many tests you plan to run */
-	plan_tests(8 + 9);
+	plan_tests(8 + 7 + 16 + 9);
 
 	test1();
+	test2();
+	test3();
 	testx();
 
 	/* This exits depending on whether all tests passed */

ccan/generator/test/example-gens.c

+#include <stdio.h>
+
 #include <ccan/generator/generator.h>
 
 #include "example-gens.h"
@@ -8,3 +10,19 @@ generator_def(gen1, int)
 	generator_yield(3);
 	generator_yield(17);
 }
+
+generator_def(gen2, int, int, base)
+{
+	generator_yield(base + 1);
+	generator_yield(base + 3);
+	generator_yield(base + 17);
+}
+
+generator_def(gen3, const char *, const char *, str, int, count)
+{
+	int i;
+
+	for (i = 0; i < count; i++)
+		generator_yield(str);
+}
+

ccan/generator/test/example-gens.h

@@ -4,5 +4,7 @@
 #include <ccan/generator/generator.h>
 
 generator_declare(gen1, int);
+generator_declare(gen2, int, int, base);
+generator_declare(gen3, const char *, const char *, str, int, count);
 
 #endif /* _EXAMPLE_GENS_H */