From 5b1ed9ed71c0e0c4768a5ed96d99d2ff0e909848 Mon Sep 17 00:00:00 2001
From: bloeys <bloeys@outlook.com>
Date: Thu, 27 Oct 2022 04:49:52 +0400
Subject: [PATCH] Moving

---
 cogo/cogo.go    |  4 +++
 inliner/main.go | 91 ++++++++++++++++++++++++++++++++++++++-----------
 main.go         | 24 -------------
 3 files changed, 75 insertions(+), 44 deletions(-)

diff --git a/cogo/cogo.go b/cogo/cogo.go
index 60e5a3f..2d5eebb 100755
--- a/cogo/cogo.go
+++ b/cogo/cogo.go
@@ -23,3 +23,7 @@ func (c *Coroutine[InT, OutT]) Tick() (out OutT, done bool) {
 
 func (c *Coroutine[InT, OutT]) Yield(out OutT) {
 }
+
+func HasGen() bool {
+	return true
+}
diff --git a/inliner/main.go b/inliner/main.go
index 22cd9e7..3cc7c1e 100755
--- a/inliner/main.go
+++ b/inliner/main.go
@@ -11,8 +11,6 @@ import (
 	"golang.org/x/tools/go/packages"
 )
 
-const cogoSwitchLbl = "cogoSwitchLbl"
-
 func printDebugInfo() {
 
 	fmt.Printf("Running inliner on '%s'\n", os.Getenv("GOFILE"))
@@ -54,25 +52,44 @@ func main() {
 func processPkg(pkg *packages.Package) {
 
 	p := processor{
-		fset: pkg.Fset,
+		fset:             pkg.Fset,
+		funcDeclsToWrite: []*ast.FuncDecl{},
 	}
 	for i, synFile := range pkg.Syntax {
+
 		pkg.Syntax[i] = astutil.Apply(synFile, p.processDeclNode, nil).(*ast.File)
+
+		if len(p.funcDeclsToWrite) > 0 {
+
+			root := &ast.File{
+				Name:    synFile.Name,
+				Imports: synFile.Imports,
+				Decls:   []ast.Decl{},
+			}
+
+			for _, v := range p.funcDeclsToWrite {
+				root.Decls = append(root.Decls, &ast.FuncDecl{
+					Name: ast.NewIdent(v.Name.Name + "_cogo"),
+					Type: v.Type,
+					Body: v.Body,
+				})
+			}
+
+			// imports.Process()
+			err := format.Node(os.Stdout, pkg.Fset, root)
+			if err != nil {
+				panic(err.Error())
+			}
+
+			p.funcDeclsToWrite = p.funcDeclsToWrite[:0]
+		}
 	}
 
-	// f, err := os.Create("main_gen.go")
-	// if err != nil {
-	// 	panic(err.Error())
-	// }
-
-	err := format.Node(os.Stdout, pkg.Fset, pkg.Syntax[0])
-	if err != nil {
-		panic(err.Error())
-	}
 }
 
 type processor struct {
-	fset *token.FileSet
+	fset             *token.FileSet
+	funcDeclsToWrite []*ast.FuncDecl
 }
 
 func (p *processor) processDeclNode(c *astutil.Cursor) bool {
@@ -105,13 +122,7 @@ func (p *processor) processDeclNode(c *astutil.Cursor) bool {
 
 		var cogoFuncSelExpr *ast.SelectorExpr
 
-		// ifStmt, ifStmtOk := stmt.(*ast.IfStmt)
-		// if ifStmtOk {
-		// 	handleNestedCogo(ifStmt.Body)
-		// 	continue
-		// }
-
-		// Find functions calls in the style of 'cogo.ABC123()'
+		// Find functions calls in the style of 'xyz.ABC123()'
 		exprStmt, exprStmtOk := stmt.(*ast.ExprStmt)
 		if !exprStmtOk {
 			continue
@@ -193,9 +204,49 @@ func (p *processor) processDeclNode(c *astutil.Cursor) bool {
 	funcDecl.Body.List = append(funcDecl.Body.List,
 		switchStmt,
 	)
+
+	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"),
+						}},
+					},
+				},
+			},
+		},
+	)
+	funcDecl.Body.List = append(funcDecl.Body.List, originalList...)
+
+	p.funcDeclsToWrite = append(p.funcDeclsToWrite, funcDecl)
+
 	return true
 }
 
+func createStmtFromFuncCall(funcName string) ast.Stmt {
+	return &ast.ExprStmt{
+		X: &ast.CallExpr{
+			Fun: ast.NewIdent(funcName),
+		},
+	}
+}
+
+func createStmtFromSelFuncCall(lhs, rhs string) ast.Stmt {
+	return &ast.ExprStmt{
+		X: &ast.CallExpr{
+			Fun: &ast.SelectorExpr{
+				X:   ast.NewIdent(lhs),
+				Sel: ast.NewIdent(rhs),
+			},
+		},
+	}
+}
+
 func funcDeclCallsCoroutineBegin(fd *ast.FuncDecl, coroutineParamName string) bool {
 
 	if fd.Body == nil || len(fd.Body.List) == 0 {
diff --git a/main.go b/main.go
index 76c3f86..24b6a7d 100755
--- a/main.go
+++ b/main.go
@@ -9,10 +9,6 @@ import (
 	"github.com/bloeys/cogo/cogo"
 )
 
-func Wow() {
-	println("wow")
-}
-
 func test(c *cogo.Coroutine[int, int]) (out int) {
 
 	c.Begin()
@@ -34,22 +30,6 @@ func test(c *cogo.Coroutine[int, int]) (out int) {
 	return out
 }
 
-// func test2() {
-
-// 	// cogo.Begin()
-
-// 	println("Hey")
-// 	cogo.Yield()
-
-// 	println("How you?")
-// 	cogo.Yield()
-
-// 	println("Bye")
-// 	cogo.Yield()
-
-// 	cogo.End()
-// }
-
 func main() {
 
 	c := &cogo.Coroutine[int, int]{
@@ -61,10 +41,6 @@ func main() {
 		println(out)
 	}
 
-	// test2()
-	// test2()
-	// test2()
-	// test2()
 }
 
 func FileLine() int {