From 69888d4d90a0f399f08e1e735e75b3bfc3cb6725 Mon Sep 17 00:00:00 2001
From: bloeys <bloeys@outlook.com>
Date: Thu, 27 Oct 2022 05:06:30 +0400
Subject: [PATCH] Stuff

---
 inliner/main.go | 57 +++++++++++++++++++++++++++++++++++++------------
 main.go         | 22 ++++++++++++-------
 2 files changed, 57 insertions(+), 22 deletions(-)

diff --git a/inliner/main.go b/inliner/main.go
index 3cc7c1e..c8cac04 100755
--- a/inliner/main.go
+++ b/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,20 +215,10 @@ 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{
-			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"),
-						}},
-					},
-				},
-			},
-		},
-	)
+
+	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)
@@ -228,6 +226,37 @@ func (p *processor) processDeclNode(c *astutil.Cursor) bool {
 	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)},
+					}},
+				},
+			},
+		},
+	}
+}
+
+func ifStmtIsHasGen(stmt *ast.IfStmt) bool {
+
+	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 {
 	return &ast.ExprStmt{
 		X: &ast.CallExpr{
diff --git a/main.go b/main.go
index 24b6a7d..ba881c6 100755
--- a/main.go
+++ b/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
 }