Embed out in coroutine+YieldTo+Yielder interface+Sleeper

README.md

@@ -1 +1,5 @@
-# cogo
+# cogo
+
+Cogo is a **Coroutine** implementation for Go!
+
+## What are Coroutines?

cogo/cogo.go

@@ -2,32 +2,67 @@ package cogo
 
 import "fmt"
 
-type CoroutineFunc[InT, OutT any] func(c *Coroutine[InT, OutT]) (out OutT)
+type CoroutineFunc[InT, OutT any] func(c *Coroutine[InT, OutT])
+
+var _ Yielder = &Coroutine[int, int]{}
 
 type Coroutine[InT, OutT any] struct {
 	State    int32
 	SubState int32
 	In       InT
+	Out      OutT
 	Func     CoroutineFunc[InT, OutT]
+	Yielder  Yielder
 }
 
 func (c *Coroutine[InT, OutT]) Begin() {
 }
 
-func (c *Coroutine[InT, OutT]) Tick() (out OutT, done bool) {
+func (c *Coroutine[InT, OutT]) Tick() (done bool) {
 
 	if c.State == -1 {
-		return out, true
+		return true
 	}
 
-	out = c.Func(c)
-	return out, c.State == -1
+	if c.Yielder != nil {
+		if !c.Yielder.Tick() {
+			return false
+		}
+
+		c.Yielder = nil
+	}
+
+	oldYielder := c.Yielder
+	c.Func(c)
+
+	// On YieldTo() we want to always tick once before returning, so here we check do that.
+	// Also, if the yielder was done after one tick we nil it
+	if c.Yielder != oldYielder {
+		if c.Yielder.Tick() {
+			c.Yielder = nil
+		} else {
+			return false
+		}
+	}
+
+	return c.State == -1
 }
 
 func (c *Coroutine[InT, OutT]) Yield(out OutT) {
 	panic(fmt.Sprintf("Yield got called at runtime, which means the code generator was not run, you used cogo incorrectly, or cogo has a bug. Yield should NOT get called at runtime. coroutine: %+v;;; yield value: %+v;;;", c, out))
 }
 
+func (c *Coroutine[InT, OutT]) YieldTo(y Yielder) {
+	panic(fmt.Sprintf("YieldTo got called at runtime, which means the code generator was not run, you used cogo incorrectly, or cogo has a bug. Yield should NOT get called at runtime. coroutine: %+v;;; yielder value: %+v;;;", c, y))
+}
+
 func HasGen() bool {
 	return true
 }
+
+func New[InT, OutT any](coro CoroutineFunc[InT, OutT], input InT) (c *Coroutine[InT, OutT]) {
+	return &Coroutine[InT, OutT]{
+		Func: coro,
+		In:   input,
+	}
+}

cogo/sleeper.go

@@ -0,0 +1,20 @@
+package cogo
+
+import "time"
+
+var _ Yielder = &Sleeper{}
+
+type Sleeper struct {
+	wakeupTime time.Time
+}
+
+func (s *Sleeper) Tick() bool {
+	return !time.Now().Before(s.wakeupTime)
+}
+
+// NewSleeper returns a sleeper that is done after at least sleepDuration time has passed
+func NewSleeper(sleepDuration time.Duration) *Sleeper {
+	return &Sleeper{
+		wakeupTime: time.Now().Add(sleepDuration),
+	}
+}

cogo/yielder.go

@@ -0,0 +1,5 @@
+package cogo
+
+type Yielder interface {
+	Tick() (done bool)
+}

demo.cogo.go

@@ -0,0 +1,94 @@
+// Code generated by 'cogo'; DO NOT EDIT.
+package main
+
+import (
+	"time"
+
+	"github.com/bloeys/cogo/cogo"
+)
+
+func test_cogo(c *cogo.Coroutine[int, int]) {
+	switch c.State {
+	case 0:
+		switch c.SubState {
+		default:
+		}
+
+		println("test yield:", 1)
+		c.State++
+		c.SubState = -1
+		c.Out = 1
+		return
+	case 1:
+		switch c.SubState {
+		default:
+		}
+		c.State++
+		c.SubState = -1
+		c.Yielder = cogo.NewSleeper(100 * time.Millisecond)
+		return
+	case 2:
+		switch c.SubState {
+		default:
+		}
+		c.State++
+		c.SubState = -1
+		c.Yielder = cogo.New(test2, 0)
+		return
+	case 3:
+		switch c.SubState {
+		default:
+		}
+
+		println("test yield:", 2)
+		c.State++
+		c.SubState = -1
+		c.Out = 2
+		return
+	case 4:
+		switch c.SubState {
+		default:
+		}
+		c.State = -1
+		c.SubState = -1
+	default:
+		c.State = -1
+		c.SubState = -1
+		return
+	}
+}
+
+func test2_cogo(c *cogo.Coroutine[int, int]) {
+	switch c.State {
+	case 0:
+		switch c.SubState {
+		default:
+		}
+
+		println("test2222 yield:", 1)
+		c.State++
+		c.SubState = -1
+		c.Out = 1
+		return
+	case 1:
+		switch c.SubState {
+		default:
+		}
+
+		println("test2222 yield:", 2)
+		c.State++
+		c.SubState = -1
+		c.Out = 2
+		return
+	case 2:
+		switch c.SubState {
+		default:
+		}
+		c.State = -1
+		c.SubState = -1
+	default:
+		c.State = -1
+		c.SubState = -1
+		return
+	}
+}

demo.go

@@ -0,0 +1,54 @@
+//go:generate cogo
+package main
+
+import (
+	"time"
+
+	"github.com/bloeys/cogo/cogo"
+)
+
+func runDemo() {
+
+	c := cogo.New(test, 0)
+
+	ticks := 1
+	for done := c.Tick(); !done; done = c.Tick() {
+		println("Ticks done:", ticks, "; Output:", c.Out, "\n")
+		ticks++
+		time.Sleep(1 * time.Millisecond)
+	}
+}
+
+func test(c *cogo.Coroutine[int, int]) {
+	if cogo.HasGen() {
+		test_cogo(c)
+		return
+	}
+
+	c.Begin()
+
+	println("test yield:", 1)
+	c.Yield(1)
+
+	c.YieldTo(cogo.NewSleeper(100 * time.Millisecond))
+
+	c.YieldTo(cogo.New(test2, 0))
+
+	println("test yield:", 2)
+	c.Yield(2)
+}
+
+func test2(c *cogo.Coroutine[int, int]) {
+	if cogo.HasGen() {
+		test2_cogo(c)
+		return
+	}
+
+	c.Begin()
+
+	println("test2222 yield:", 1)
+	c.Yield(1)
+
+	println("test2222 yield:", 2)
+	c.Yield(2)
+}

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	"go/ast"
 	"go/format"
@@ -15,8 +16,19 @@ import (
 	"golang.org/x/tools/imports"
 )
 
+var (
+	demo = flag.Bool("demo", false, "")
+)
+
 func main() {
 
+	flag.Parse()
+
+	if *demo {
+		runDemo()
+		return
+	}
+
 	cwd, err := os.Getwd()
 	if err != nil {
 		panic(err)
@@ -276,7 +288,7 @@ func (p *processor) genCogoFuncsNodeProcessor(c *astutil.Cursor) bool {
 			// Add everything from the last begin/yield until this yield into a case
 			stmtsSinceLastCogo := funcDecl.Body.List[lastCaseEndBodyListIndex+1 : i]
 
-			caseStmts := make([]ast.Stmt, 0, len(stmtsSinceLastCogo)+2)
+			caseStmts := make([]ast.Stmt, 0, len(stmtsSinceLastCogo)+4)
 
 			caseStmts = append(caseStmts, subSwitchStmt)
 			subSwitchStmt = &ast.SwitchStmt{
@@ -299,9 +311,12 @@ func (p *processor) genCogoFuncsNodeProcessor(c *astutil.Cursor) bool {
 					Tok: token.ASSIGN,
 					Rhs: []ast.Expr{ast.NewIdent("-1")},
 				},
-				&ast.ReturnStmt{
-					Results: callExpr.Args,
+				&ast.AssignStmt{
+					Lhs: []ast.Expr{ast.NewIdent(coroutineParamName + ".Out")},
+					Tok: token.ASSIGN,
+					Rhs: callExpr.Args,
 				},
+				&ast.ReturnStmt{},
 			)
 
 			switchStmt.Body.List = append(switchStmt.Body.List,
@@ -313,6 +328,50 @@ func (p *processor) genCogoFuncsNodeProcessor(c *astutil.Cursor) bool {
 
 			lastCaseEndBodyListIndex = i
 
+		} else if cogoFuncSelExpr.Sel.Name == "YieldTo" {
+
+			// Add everything from the last begin/yield until this yield into a case
+			stmtsSinceLastCogo := funcDecl.Body.List[lastCaseEndBodyListIndex+1 : i]
+
+			caseStmts := make([]ast.Stmt, 0, len(stmtsSinceLastCogo)+5)
+
+			caseStmts = append(caseStmts, subSwitchStmt)
+			subSwitchStmt = &ast.SwitchStmt{
+				Tag: ast.NewIdent(coroutineParamName + ".SubState"),
+				Body: &ast.BlockStmt{
+					List: []ast.Stmt{
+						getCaseWithStmts(nil, []ast.Stmt{}),
+					},
+				},
+			}
+
+			caseStmts = append(caseStmts, stmtsSinceLastCogo...)
+			caseStmts = append(caseStmts,
+				&ast.IncDecStmt{
+					Tok: token.INC,
+					X:   ast.NewIdent(coroutineParamName + ".State"),
+				},
+				&ast.AssignStmt{
+					Lhs: []ast.Expr{ast.NewIdent(coroutineParamName + ".SubState")},
+					Tok: token.ASSIGN,
+					Rhs: []ast.Expr{ast.NewIdent("-1")},
+				},
+				&ast.AssignStmt{
+					Lhs: []ast.Expr{ast.NewIdent(coroutineParamName + ".Yielder")},
+					Tok: token.ASSIGN,
+					Rhs: callExpr.Args,
+				},
+				&ast.ReturnStmt{},
+			)
+
+			switchStmt.Body.List = append(switchStmt.Body.List,
+				getCaseWithStmts(
+					[]ast.Expr{ast.NewIdent(fmt.Sprint(len(switchStmt.Body.List)))},
+					caseStmts,
+				),
+			)
+
+			lastCaseEndBodyListIndex = i
 		}
 	}
 
@@ -435,9 +494,12 @@ func (p *processor) genCogoForStmt(forStmt *ast.ForStmt, coroutineParamName stri
 					Tok: token.ASSIGN,
 					Rhs: []ast.Expr{ast.NewIdent(fmt.Sprint(postInitLblNum))},
 				},
-				&ast.ReturnStmt{
-					Results: selExprArgs,
+				&ast.AssignStmt{
+					Lhs: []ast.Expr{ast.NewIdent(coroutineParamName + ".Out")},
+					Tok: token.ASSIGN,
+					Rhs: selExprArgs,
 				},
+				&ast.ReturnStmt{},
 			},
 		}
 	}
@@ -472,9 +534,12 @@ func (p *processor) genCogoBlockStmt(blockStmt *ast.BlockStmt, coroutineParamNam
 					Tok: token.ASSIGN,
 					Rhs: []ast.Expr{ast.NewIdent(fmt.Sprint(newSubStateNum))},
 				},
-				&ast.ReturnStmt{
-					Results: selExprArgs,
+				&ast.AssignStmt{
+					Lhs: []ast.Expr{ast.NewIdent(coroutineParamName + ".Out")},
+					Tok: token.ASSIGN,
+					Rhs: selExprArgs,
 				},
+				&ast.ReturnStmt{},
 			},
 		}
 	}
@@ -566,12 +631,12 @@ func createHasGenIfStmt(funcDecl *ast.FuncDecl, coroutineParamName string) *ast.
 		Cond: createStmtFromSelFuncCall("cogo", "HasGen").(*ast.ExprStmt).X,
 		Body: &ast.BlockStmt{
 			List: []ast.Stmt{
-				&ast.ReturnStmt{
-					Results: []ast.Expr{&ast.CallExpr{
+				&ast.ExprStmt{
+					X: &ast.CallExpr{
 						Fun:  ast.NewIdent(funcDecl.Name.Name + "_cogo"),
 						Args: []ast.Expr{ast.NewIdent(coroutineParamName)},
 					}},
-				},
+				&ast.ReturnStmt{},
 			},
 		},
 	}