From 94090dc9ab6a669df2825752ace9221fd8e741a7 Mon Sep 17 00:00:00 2001
From: bloeys <bloeys@outlook.com>
Date: Thu, 27 Oct 2022 05:58:58 +0400
Subject: [PATCH] Masha Allah it works!

---
 inliner/main.go | 86 ++++++++++++++++++++++++++-----------------------
 main.go         | 19 ++++++-----
 2 files changed, 56 insertions(+), 49 deletions(-)

diff --git a/inliner/main.go b/inliner/main.go
index 12931ae..5366e4b 100755
--- a/inliner/main.go
+++ b/inliner/main.go
@@ -83,7 +83,7 @@ func genCogoFuncs(cwd string) {
 
 				origFName := pkg.Fset.File(synFile.Pos()).Name()
 				newFName := strings.Split(origFName, ".")[0] + ".cogo.go"
-				writeAst(newFName, pkg.Fset, root)
+				writeAst(newFName, "// Code generated by 'cogo'; DO NOT EDIT.\n", pkg.Fset, root)
 
 				p.funcDeclsToWrite = p.funcDeclsToWrite[:0]
 			}
@@ -91,31 +91,6 @@ func genCogoFuncs(cwd string) {
 	}
 }
 
-func writeAst(fName string, fset *token.FileSet, node any) {
-
-	f, err := os.Create(fName)
-	if err != nil {
-		panic("Failed to create file to write new AST. Err: " + err.Error())
-	}
-	defer f.Close()
-
-	err = format.Node(f, fset, node)
-	if err != nil {
-		panic(err.Error())
-	}
-
-	b, err := imports.Process(fName, nil, nil)
-	if err != nil {
-		panic("Failed to process imports on file " + fName + ". Err: " + err.Error())
-	}
-
-	f.Seek(0, io.SeekStart)
-	_, err = f.Write(b)
-	if err != nil {
-		panic(err.Error())
-	}
-}
-
 func genHasGenChecksOnOriginalFuncs(cwd string) {
 
 	pkgs, err := packages.Load(&packages.Config{
@@ -144,7 +119,7 @@ func genHasGenChecksOnOriginalFuncs(cwd string) {
 
 			if len(p.funcDeclsToWrite) > 0 {
 				origFName := pkg.Fset.File(synFile.Pos()).Name()
-				writeAst(origFName, pkg.Fset, pkg.Syntax[i])
+				writeAst(origFName, "", pkg.Fset, pkg.Syntax[i])
 
 				p.funcDeclsToWrite = p.funcDeclsToWrite[:0]
 			}
@@ -184,15 +159,13 @@ func (p *processor) genCogoFuncsNodeProcessor(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
+			funcDecl.Body.List[i] = &ast.EmptyStmt{}
 			continue
 		}
 
@@ -271,6 +244,19 @@ func (p *processor) genCogoFuncsNodeProcessor(c *astutil.Cursor) bool {
 			[]ast.Expr{ast.NewIdent(fmt.Sprint(len(switchStmt.Body.List)))},
 			caseStmts,
 		),
+
+		// default case
+		getCaseWithStmts(
+			nil,
+			[]ast.Stmt{
+				&ast.AssignStmt{
+					Lhs: []ast.Expr{ast.NewIdent(coroutineParamName + ".State")},
+					Tok: token.ASSIGN,
+					Rhs: []ast.Expr{ast.NewIdent("-1")},
+				},
+				&ast.ReturnStmt{},
+			},
+		),
 	)
 
 	// Apply changes
@@ -282,9 +268,6 @@ func (p *processor) genCogoFuncsNodeProcessor(c *astutil.Cursor) bool {
 	originalList := funcDecl.Body.List
 	funcDecl.Body.List = make([]ast.Stmt, 0, len(funcDecl.Body.List)+1)
 
-	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)
@@ -362,14 +345,6 @@ func ifStmtIsHasGen(stmt *ast.IfStmt) bool {
 	return selExprIs(selExpr, "cogo", "HasGen")
 }
 
-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{
@@ -483,3 +458,32 @@ func getCaseWithStmts(caseConditions []ast.Expr, stmts []ast.Stmt) *ast.CaseClau
 		Body: stmts,
 	}
 }
+
+func writeAst(fName, topComment string, fset *token.FileSet, node any) {
+
+	f, err := os.Create(fName)
+	if err != nil {
+		panic("Failed to create file to write new AST. Err: " + err.Error())
+	}
+	defer f.Close()
+
+	if topComment != "" {
+		f.WriteString(topComment)
+	}
+
+	err = format.Node(f, fset, node)
+	if err != nil {
+		panic(err.Error())
+	}
+
+	b, err := imports.Process(fName, nil, nil)
+	if err != nil {
+		panic("Failed to process imports on file " + fName + ". Err: " + err.Error())
+	}
+
+	f.Seek(0, io.SeekStart)
+	_, err = f.Write(b)
+	if err != nil {
+		panic(err.Error())
+	}
+}
diff --git a/main.go b/main.go
index 4851e62..c873901 100755
--- a/main.go
+++ b/main.go
@@ -10,22 +10,25 @@ import (
 )
 
 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
 }