Stuff

inliner/main.go

@@ -118,10 +118,18 @@ func (p *processor) processDeclNode(c *astutil.Cursor) bool {
 		},
 	}
 
+	hasGenCheckExists := false
 	for i, stmt := range funcDecl.Body.List {
 
 		var cogoFuncSelExpr *ast.SelectorExpr
 
+		ifStmt, ok := stmt.(*ast.IfStmt)
+		if ok && ifStmtIsHasGen(ifStmt) {
+			funcDecl.Body.List[i] = createHasGenIfStmt(funcDecl, coroutineParamName)
+			hasGenCheckExists = true
+			continue
+		}
+
 		// Find functions calls in the style of 'xyz.ABC123()'
 		exprStmt, exprStmtOk := stmt.(*ast.ExprStmt)
 		if !exprStmtOk {
@@ -207,25 +215,46 @@ func (p *processor) processDeclNode(c *astutil.Cursor) bool {
 
 	originalList := funcDecl.Body.List
 	funcDecl.Body.List = make([]ast.Stmt, 0, len(funcDecl.Body.List)+1)
-	funcDecl.Body.List = append(funcDecl.Body.List,
-		&ast.IfStmt{
+
+	if !hasGenCheckExists {
+		funcDecl.Body.List = append(funcDecl.Body.List, createHasGenIfStmt(funcDecl, coroutineParamName))
+	}
+	funcDecl.Body.List = append(funcDecl.Body.List, originalList...)
+
+	p.funcDeclsToWrite = append(p.funcDeclsToWrite, funcDecl)
+
+	return true
+}
+
+func createHasGenIfStmt(funcDecl *ast.FuncDecl, coroutineParamName string) *ast.IfStmt {
+	return &ast.IfStmt{
 		Cond: createStmtFromSelFuncCall("cogo", "HasGen").(*ast.ExprStmt).X,
 		Body: &ast.BlockStmt{
 			List: []ast.Stmt{
 				&ast.ReturnStmt{
 					Results: []ast.Expr{&ast.CallExpr{
 						Fun:  ast.NewIdent(funcDecl.Name.Name + "_cogo"),
+						Args: []ast.Expr{ast.NewIdent(coroutineParamName)},
 					}},
 				},
 			},
 		},
-		},
-	)
-	funcDecl.Body.List = append(funcDecl.Body.List, originalList...)
+	}
+}
 
-	p.funcDeclsToWrite = append(p.funcDeclsToWrite, funcDecl)
+func ifStmtIsHasGen(stmt *ast.IfStmt) bool {
 
-	return true
+	callExpr, ok := stmt.Cond.(*ast.CallExpr)
+	if !ok {
+		return false
+	}
+
+	selExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+	if !ok {
+		return false
+	}
+
+	return selExprIs(selExpr, "cogo", "HasGen")
 }
 
 func createStmtFromFuncCall(funcName string) ast.Stmt {

main.go

@@ -9,23 +9,29 @@ import (
 	"github.com/bloeys/cogo/cogo"
 )
 
+func test_cogo(c *cogo.Coroutine[int, int]) (out int) { return 0 }
+
 func test(c *cogo.Coroutine[int, int]) (out int) {
 
+	if cogo.HasGen() {
+		return test_cogo(c)
+	}
+
 	c.Begin()
 
 	println("Tick 1")
-	c.Yield(1)
+	// c.Yield(1)
 
-	println("Tick 2")
-	c.Yield(2)
+	// println("Tick 2")
+	// c.Yield(2)
 
-	println("Tick 3")
-	c.Yield(3)
+	// println("Tick 3")
+	// c.Yield(3)
 
-	println("Tick 4")
-	c.Yield(4)
+	// println("Tick 4")
+	// c.Yield(4)
 
-	println("Tick before end")
+	// println("Tick before end")
 
 	return out
 }