so many hacks

cogo/cogo.go

@@ -1,10 +1,12 @@
 package cogo
 
-func Begin() {
+// func Tick(c any) {
+// }
 
+func Begin() {
 }
 
-func Yield() {
+func Yield[T any](out T) {
 
 }
 

inliner/main.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"go/ast"
 	"go/format"
+	"go/token"
 	"os"
 
 	"golang.org/x/tools/go/ast/astutil"
@@ -50,8 +51,11 @@ func main() {
 
 func processPkg(pkg *packages.Package) {
 
+	p := processor{
+		fset: pkg.Fset,
+	}
 	for i, synFile := range pkg.Syntax {
-		pkg.Syntax[i] = astutil.Apply(synFile, processDeclNode, nil).(*ast.File)
+		pkg.Syntax[i] = astutil.Apply(synFile, p.processDeclNode, nil).(*ast.File)
 	}
 
 	// f, err := os.Create("main_gen.go")
@@ -65,7 +69,11 @@ func processPkg(pkg *packages.Package) {
 	}
 }
 
-func processDeclNode(c *astutil.Cursor) bool {
+type processor struct {
+	fset *token.FileSet
+}
+
+func (p *processor) processDeclNode(c *astutil.Cursor) bool {
 
 	n := c.Node()
 	if n == nil {
@@ -81,67 +89,70 @@ func processDeclNode(c *astutil.Cursor) bool {
 		return false
 	}
 
+	beginBodyListIndex := -1
+	lastCaseEndBodyListIndex := -1
+	switchStmt := &ast.SwitchStmt{
+		Tag: ast.NewIdent("c.state"),
+		Body: &ast.BlockStmt{
+			List: []ast.Stmt{},
+		},
+	}
+
 	for i, stmt := range funcDecl.Body.List {
 
+		var cogoFuncCallExpr *ast.SelectorExpr
+
+		// ifStmt, ifStmtOk := stmt.(*ast.IfStmt)
+		// if ifStmtOk {
+		// 	handleNestedCogo(ifStmt.Body)
+		// 	continue
+		// }
+
 		// Find functions calls in the style of 'cogo.ABC123()'
-		exprStmt, ok := stmt.(*ast.ExprStmt)
+		exprStmt, exprStmtOk := stmt.(*ast.ExprStmt)
-		if !ok {
+		if !exprStmtOk {
 			continue
 		}
 
-		callExpr, ok := exprStmt.X.(*ast.CallExpr)
+		callExpr, exprStmtOk := exprStmt.X.(*ast.CallExpr)
-		if !ok {
+		if !exprStmtOk {
 			continue
 		}
 
-		pkgFuncCallExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+		cogoFuncCallExpr, exprStmtOk = callExpr.Fun.(*ast.SelectorExpr)
-		if !ok {
+		if !exprStmtOk {
 			continue
 		}
 
-		pkgIdent, ok := pkgFuncCallExpr.X.(*ast.Ident)
+		if !funcCallHasPkgName(cogoFuncCallExpr, "cogo") {
-		if !ok || pkgIdent.Name != "cogo" {
 			continue
 		}
-		fmt.Printf("Found: %+v\n", pkgFuncCallExpr)
+
+		cogoFuncCallLineNum := p.fset.File(cogoFuncCallExpr.Pos()).Line(cogoFuncCallExpr.Pos())
+		fmt.Printf("Found: '%+v' at line %d\n", cogoFuncCallExpr, cogoFuncCallLineNum)
 
 		// Now that we found a call to cogo decide what to do
-		if pkgFuncCallExpr.Sel.Name == "Begin" {
+		if cogoFuncCallExpr.Sel.Name == "Begin" {
 
-			beginStmt := &ast.SwitchStmt{
+			beginBodyListIndex = i
-				Tag: ast.NewIdent("state"),
+			lastCaseEndBodyListIndex = i
-				Body: &ast.BlockStmt{
+			continue
-					List: []ast.Stmt{
+		} else if cogoFuncCallExpr.Sel.Name == "Yield" || cogoFuncCallExpr.Sel.Name == "End" {
-						&ast.CaseClause{
-							List: nil,
-							Body: []ast.Stmt{
-								&ast.ExprStmt{
-									X: &ast.CallExpr{
-										Fun: &ast.Ident{
-											Name: "Wow",
-										},
-										Args: nil,
-									},
-								},
-							},
-						},
-					},
-				},
-			}
 
-			funcDecl.Body.List[i] = beginStmt
+			// Add everything from the last begin/yield until this yield into a case
 
-		} else if pkgFuncCallExpr.Sel.Name == "Yield" {
+			stmtsSinceLastCogo := funcDecl.Body.List[lastCaseEndBodyListIndex+1 : i]
+			switchStmt.Body.List = append(switchStmt.Body.List, getCaseWithStmts(
+				stmtsSinceLastCogo,
+				[]ast.Expr{ast.NewIdent(fmt.Sprint(cogoFuncCallLineNum))},
+			))
 
-			exprStmt.X = &ast.CallExpr{
+			lastCaseEndBodyListIndex = i
-				Fun: &ast.Ident{
-					Name: "Wow",
-				},
-				Args: nil,
-			}
 		}
 	}
 
+	funcDecl.Body.List = funcDecl.Body.List[:beginBodyListIndex]
+	funcDecl.Body.List = append(funcDecl.Body.List, switchStmt)
 	return true
 }
 
@@ -169,13 +180,17 @@ func funcDeclCallsCogo(fd *ast.FuncDecl) bool {
 			continue
 		}
 
-		pkgIdent, ok := pkgFuncCallExpr.X.(*ast.Ident)
+		return funcCallHasPkgName(pkgFuncCallExpr, "cogo")
-		return ok && pkgIdent.Name == "cogo"
 	}
 
 	return false
 }
 
+func funcCallHasPkgName(selExpr *ast.SelectorExpr, pkgName string) bool {
+	pkgIdent, ok := selExpr.X.(*ast.Ident)
+	return ok && pkgIdent.Name == pkgName
+}
+
 func filter[T any](arr []T, where func(x T) bool) []T {
 
 	out := []T{}
@@ -190,3 +205,10 @@ func filter[T any](arr []T, where func(x T) bool) []T {
 
 	return out
 }
+
+func getCaseWithStmts(stmts []ast.Stmt, conditions []ast.Expr) *ast.CaseClause {
+	return &ast.CaseClause{
+		List: conditions,
+		Body: stmts,
+	}
+}

main.go

@@ -9,70 +9,120 @@ import (
 	"github.com/bloeys/cogo/cogo"
 )
 
-type Coroutine[T any] struct {
+type CoroutineFunc[InT, OutT any] func(c *Coroutine[InT, OutT]) (out OutT)
+
+type Coroutine[InT, OutT any] struct {
 	State int32
-	In    *T
+	In    InT
+	Func  CoroutineFunc[InT, OutT]
 }
 
-func (c *Coroutine[T]) Run(f func(in *T)) {
+func (c *Coroutine[InT, OutT]) Tick() (out OutT, done bool) {
-	f(c.In)
+
+	if c.State == -1 {
+		return out, true
 	}
 
-var state = 0
+	out = c.Func(c)
+	return out, c.State == -1
+}
+
+// func (c *Coroutine[InT, OutT]) Yield(out OutT) {
+// }
+
+func (c *Coroutine[InT, OutT]) Break() {
+}
 
 func Wow() {
 	println("wow")
 }
 
-// func test() {
+func test(c *Coroutine[int, int]) (out int) {
 
-// 	cogo.Begin()
+	cogo.Begin()
 
-// 	println("hi")
+	println("Tick 1")
-// 	println("this is from state_0")
+	cogo.Yield(1)
-// 	cogo.Yield()
-// 	state = 1
 
-// 	if 1 > 2 {
+	println("Tick 2")
-// 		println("gg")
+	cogo.Yield(2)
+
+	println("Tick 3")
+	cogo.Yield(3)
+
+	println("Tick 4")
+	cogo.Yield(4)
+
+	cogo.End()
+
+	// switch c.State {
+	// case 0:
+	// 	println("Tick 0")
+	// 	c.State++
+	// 	return 1, false
+	// case 1:
+	// 	println("Tick 1")
+	// 	c.State++
+	// 	return 2, false
+	// case 2:
+	// 	println("Tick 2")
+	// 	c.State++
+	// 	return 3, false
+	// case 3:
+	// 	println("Tick 3")
+	// 	c.State++
+	// 	return 4, false
+	// default:
+	// 	return out, true
 	// }
 
-// 	println("Bye")
+	return out
-// 	println("this is from state_1")
+}
+
+// func test2() {
+
+// 	// cogo.Begin()
+
+// 	println("Hey")
+// 	cogo.Yield()
+
+// 	println("How you?")
+// 	cogo.Yield()
+
+// 	println("Bye")
 // 	cogo.Yield()
-// 	state = 2
 
 // 	cogo.End()
 // }
 
-func test2() {
-
-	cogo.Begin()
-
-	println("Hey")
-	cogo.Yield()
-
-	println("How you?")
-	cogo.Yield()
-
-	println("Bye")
-	cogo.Yield()
-
-	cogo.End()
-}
-
 func main() {
 
-	// test()
+	x := 1
-	// test()
+switch_start:
-	// test()
+	switch x {
+	case 1:
+		println(1)
+		x = 3
+		goto switch_start
+	case 2:
+		println(2)
+	case 3:
+		println(3)
+	}
+	return
+	c := &Coroutine[int, int]{
+		Func: test,
+		In:   0,
+	}
 
-	test2()
+	for out, done := c.Tick(); !done; out, done = c.Tick() {
-	test2()
+		println(out)
-	test2()
+	}
-	test2()
 
-	println("Final state:", state)
+	// test2()
+	// test2()
+	// test2()
+	// test2()
 }
 
 func FileLine() int {