Some cool code gen stuff

cogo/cogo.go

@@ -0,0 +1,12 @@
+package cogo
+
+func Begin() {
+
+}
+
+func Yield() {
+
+}
+
+func End() {
+}

go.mod

@@ -1,3 +1,10 @@
 module github.com/bloeys/cogo
 
 go 1.18
+
+require golang.org/x/tools v0.2.0
+
+require (
+	golang.org/x/mod v0.6.0 // indirect
+	golang.org/x/sys v0.1.0 // indirect
+)

go.sum

@@ -0,0 +1,6 @@
+golang.org/x/mod v0.6.0 h1:b9gGHsz9/HhJ3HF5DHQytPpuwocVTChQJK3AvoLRD5I=
+golang.org/x/mod v0.6.0/go.mod h1:4mET923SAdbXp2ki8ey+zGs1SLqsuM2Y0uvdZR/fUNI=
+golang.org/x/sys v0.1.0 h1:kunALQeHf1/185U1i0GOB/fy1IPRDDpuoOOqRReG57U=
+golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/tools v0.2.0 h1:G6AHpWxTMGY1KyEYoAQ5WTtIekUUvDNjan3ugu60JvE=
+golang.org/x/tools v0.2.0/go.mod h1:y4OqIKeOV/fWJetJ8bXPU1sEVniLMIyDAZWeHdV+NTA=

inliner/main.go

@@ -0,0 +1,122 @@
+package main
+
+import (
+	"fmt"
+	"go/ast"
+	"go/format"
+	"os"
+
+	"golang.org/x/tools/go/ast/astutil"
+	"golang.org/x/tools/go/packages"
+)
+
+func printDebugInfo() {
+
+	fmt.Printf("Running inliner on '%s'\n", os.Getenv("GOFILE"))
+
+	cwd, err := os.Getwd()
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("  cwd = %s\n", cwd)
+	fmt.Printf("  os.Args = %#v\n\n", os.Args)
+}
+
+func main() {
+
+	printDebugInfo()
+
+	cwd, err := os.Getwd()
+	if err != nil {
+		panic(err)
+	}
+
+	// Parse package
+	pkgs, err := packages.Load(&packages.Config{
+		Dir:   cwd,
+		Mode:  packages.NeedName | packages.NeedTypes | packages.NeedTypesInfo | packages.NeedSyntax,
+		Tests: false,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	if len(pkgs) != 1 {
+		panic(fmt.Sprintf("expected to find one package but found %d", len(pkgs)))
+	}
+
+	processPkg(pkgs[0])
+}
+
+func processPkg(pkg *packages.Package) {
+
+	for i, synFile := range pkg.Syntax {
+		pkg.Syntax[i] = astutil.Apply(synFile, processDeclNode, nil).(*ast.File)
+		// for _, dec := range synFile.Decls {
+		// ast.Inspect(dec, processDeclNode)
+		// }
+	}
+
+	// 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())
+	}
+}
+
+func processDeclNode(c *astutil.Cursor) bool {
+
+	n := c.Node()
+	if n == nil {
+		return false
+	}
+
+	funcDecl, ok := n.(*ast.FuncDecl)
+	if !ok || funcDecl.Body == nil {
+		return true
+	}
+
+	if funcDecl.Name.Name != "test" {
+		return false
+	}
+
+	for _, stmt := range funcDecl.Body.List {
+
+		exprStmt, ok := stmt.(*ast.ExprStmt)
+		if !ok {
+			continue
+		}
+
+		callExpr, ok := exprStmt.X.(*ast.CallExpr)
+		if !ok {
+			continue
+		}
+
+		pkgFuncCallExpr, ok := callExpr.Fun.(*ast.SelectorExpr)
+		if !ok {
+			continue
+		}
+
+		pkgIdent, ok := pkgFuncCallExpr.X.(*ast.Ident)
+		if !ok || pkgIdent.Name != "cogo" {
+			continue
+		}
+		fmt.Printf("Found: %+v\n", pkgFuncCallExpr)
+
+		if pkgFuncCallExpr.Sel.Name == "Yield" {
+
+			exprStmt.X = &ast.CallExpr{
+				Fun: &ast.Ident{
+					Name: "Wow",
+				},
+				Args: nil,
+			}
+		}
+	}
+
+	return true
+}

main.go

@@ -1,5 +1,72 @@
+//go:generate go run inliner/main.go
 package main
 
-func main() {
-	println("Hello there, friend.")
+import (
+	"fmt"
+	"runtime"
+	"runtime/debug"
+
+	"github.com/bloeys/cogo/cogo"
+)
+
+type Coroutine[T any] struct {
+	State int32
+	In    *T
+}
+
+func (c *Coroutine[T]) Run(f func(in *T)) {
+	f(c.In)
+}
+
+var state = 0
+
+func Wow() {
+	println("wow")
+}
+
+func test() {
+
+	cogo.Begin()
+
+	println("hi")
+	println("this is from state_0")
+	cogo.Yield()
+	state = 1
+
+	if 1 > 2 {
+		println("gg")
+	}
+
+	println("Bye")
+	println("this is from state_1")
+	cogo.Yield()
+	state = 2
+
+	cogo.End()
+}
+
+func main() {
+
+	test()
+	test()
+	test()
+
+	println("Final state:", state)
+}
+
+func FileLine() int {
+	_, _, lineNum, ok := runtime.Caller(1)
+	if !ok {
+		panic("failed to get line number. Stack trace: " + string(debug.Stack()))
+	}
+	return lineNum
+}
+
+func FileLineString() string {
+	_, _, lineNum, ok := runtime.Caller(1)
+	if !ok {
+		panic("failed to get line number. Stack trace: " + string(debug.Stack()))
+	}
+
+	return fmt.Sprint(lineNum)
 }