generalize

Author: Simn

src/hxcoro/generators/SyncGenerator.hx

@@ -1,5 +1,6 @@
 package hxcoro.generators;
 
+import haxe.Unit;
 import haxe.Exception;
 import haxe.coro.Coroutine;
 import haxe.coro.IContinuation;
@@ -10,21 +11,21 @@ import haxe.exceptions.CoroutineException;
 import hxcoro.Coro.*;
 
 @:coroutine.restrictedSuspension
-typedef Yield<T> = Coroutine<T -> Void>;
+typedef Yield<T, R> = Coroutine<T -> R>;
 
-private class GeneratorImpl<T> extends Dispatcher implements IContinuation<Any> {
+private class GeneratorImpl<T, R> extends Dispatcher implements IContinuation<Unit> {
 	public var context(get, null):Context;
 
-	final f:Coroutine<Yield<T> -> Void>;
+	final f:Coroutine<Yield<T, R> -> Void>;
 	var nextValue:Null<T>;
-	var nextStep:Null<IContinuation<T>>;
+	var nextStep:Null<IContinuation<R>>;
 	var raisedException:Null<Exception>;
 	var resumed:Bool;
 
-	public function new(f:Coroutine<Yield<T> -> Void>) {
+	public function new(f:Coroutine<Yield<T, R> -> Void>) {
 		this.context = Context.create(this);
 		this.f = f;
-		resumed = true;
+		resumed = false;
 	}
 
 	function get_context() {
@@ -38,21 +39,21 @@ private class GeneratorImpl<T> extends Dispatcher implements IContinuation<Any>
 	public function hasNext() {
 		if (nextStep == null) {
 			f(this, yield);
-		} else if (!resumed) {
-			nextStep.resume(null, null);
 		}
 		return !resumed;
 	}
 
-	public function next() {
+	public function next(value:R) {
 		if (raisedException != null) {
 			resumed = true;
 			throw raisedException;
 		}
-		return nextValue;
+		var current = nextValue;
+		nextStep.resume(value, null);
+		return current;
 	}
 
-	public function resume(result:Null<Any>, error:Null<Exception>) {
+	public function resume(result:Null<Unit>, error:Null<Exception>) {
 		if (error != null) {
 			raisedException = error;
 		} else {
@@ -64,10 +65,10 @@ private class GeneratorImpl<T> extends Dispatcher implements IContinuation<Any>
 		obj.onDispatch();
 	}
 
-	@:coroutine function yield(value:T) {
+	@:coroutine function yield(value:T):R {
 		resumed = false;
 		nextValue = value;
-		suspend(cont -> {
+		return suspend(cont -> {
 			nextStep = cont;
 		});
 	}
@@ -76,19 +77,19 @@ private class GeneratorImpl<T> extends Dispatcher implements IContinuation<Any>
 /**
 	A synchronous generator that can be used as an `Iterator`.
 **/
-abstract SyncGenerator<T>(GeneratorImpl<T>) {
+abstract SyncGenerator<T>(GeneratorImpl<T, Unit>) from GeneratorImpl<T, Unit> {
 	/**
 		@see `Iterator.hasNext`
 	**/
-	public inline function hasNext() {
+	public inline function hasNext():Bool {
 		return this.hasNext();
 	}
 
 	/**
 		@see `Iterator.next`
 	**/
-	public inline function next() {
-		return this.next();
+	public inline function next():Null<T> {
+		return this.next(Unit);
 	}
 
 	/**
@@ -97,7 +98,21 @@ abstract SyncGenerator<T>(GeneratorImpl<T>) {
 		The coroutine `f` is executed in a restricted suspension scope, which means
 		that it cannot call arbitrary coroutines that might suspend.
 	**/
-	static public function create<T>(f:Coroutine<Yield<T> -> Void>) {
+	static public function create<T>(f:Coroutine<Yield<T, Unit> -> Void>):SyncGenerator<T> {
+		return new GeneratorImpl(f);
+	}
+}
+
+abstract SyncValueGenerator<T, R>(GeneratorImpl<T, R>) from GeneratorImpl<T, R> {
+	public inline function hasNext() {
+		return this.hasNext();
+	}
+
+	public inline function next(value:R) {
+		return this.next(value);
+	}
+
+	static public function create<T, R>(f:Coroutine<Yield<T, R> -> Void>):SyncValueGenerator<T, R> {
 		return new GeneratorImpl(f);
 	}
 }

tests/src/TestGenerator.hx

@@ -1,3 +1,4 @@
+import haxe.Unit;
 import hxcoro.dispatchers.TrampolineDispatcher;
 import hxcoro.task.CoroTask;
 import haxe.coro.context.Context;
@@ -15,13 +16,13 @@ class TestGenerator extends utest.Test {
 	}
 
 	function testTreeIter() {
-		@:coroutine function iterTreeRec<T>(yield:Yield<T>, tree:Tree<T>) {
+		@:coroutine function iterTreeRec<T>(yield:Yield<T, Unit>, tree:Tree<T>) {
 			yield(tree.leaf);
 			if (tree.left != null) iterTreeRec(yield, tree.left);
 			if (tree.right != null) iterTreeRec(yield, tree.right);
 		}
 
-		function iterTree<T>(tree:Tree<T>):Iterator<T> {
+		function iterTree<T>(tree:Tree<T>) {
 			return SyncGenerator.create(yield -> iterTreeRec(yield, tree));
 		}