more tests for serializable closures (#45)

I was hoping to exercise more closure memory layouts but it appears that
when a closure escapes to the heap, all its captured variables have to
escape and therefore the closure captures them as pointers. In our
current model, local variables will always escape so we will never have
closure types that contain non-pointer fields.

compiler/compile.go

@@ -337,7 +337,7 @@ func (c *compiler) compileFunction(p *packages.Package, fn *ast.FuncDecl, color
 	// declarations to the function prologue. We downgrade inline var decls and
 	// assignments that use := to assignments that use =. Constant decls are
 	// hoisted and also have their value assigned in the function prologue.
-	decls := extractDecls(fn, p.TypesInfo)
+	decls := extractDecls(fn.Body, p.TypesInfo)
 	renameObjects(fn, p.TypesInfo, decls)
 	for _, decl := range decls {
 		gen.Body.List = append(gen.Body.List, &ast.DeclStmt{Decl: decl})

compiler/coroutine_test.go

@@ -121,8 +121,26 @@ func TestCoroutineYield(t *testing.T) {
 		},
 
 		{
-			name:   "range over closure",
-			coro:   func() { Range10Closure() },
+			name:   "range over closure capturing values",
+			coro:   Range10ClosureCapturingValues,
+			yields: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
+		},
+
+		{
+			name:   "range over closure capturing pointers",
+			coro:   Range10ClosureCapturingPointers,
+			yields: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
+		},
+
+		{
+			name:   "range over closure capturing heterogenous values",
+			coro:   Range10ClosureHeterogenousCapture,
+			yields: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
+		},
+
+		{
+			name:   "range with heterogenous values",
+			coro:   Range10Heterogenous,
 			yields: []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
 		},
 

compiler/decls.go

@@ -21,9 +21,13 @@ import (
 // Note that declarations are extracted from all nested scopes within the
 // function body, so there may be duplicate identifiers. Identifiers can be
 // disambiguated using (*types.Info).ObjectOf(ident).
-func extractDecls(fn *ast.FuncDecl, info *types.Info) (decls []*ast.GenDecl) {
-	ast.Inspect(fn.Body, func(node ast.Node) bool {
+func extractDecls(tree ast.Node, info *types.Info) (decls []*ast.GenDecl) {
+	ast.Inspect(tree, func(node ast.Node) bool {
 		switch n := node.(type) {
+		case *ast.FuncLit:
+			// Stop when we encounter a function listeral so we don't hoist its
+			// local variables into the scope of its parent function.
+			return false
 		case *ast.GenDecl: // const, var, type
 			if n.Tok == token.TYPE || n.Tok == token.CONST {
 				decls = append(decls, n)
@@ -135,6 +139,8 @@ func renameObjects(tree ast.Node, info *types.Info, decls []*ast.GenDecl) {
 func removeDecls(tree ast.Node) {
 	astutil.Apply(tree, func(cursor *astutil.Cursor) bool {
 		switch n := cursor.Node().(type) {
+		case *ast.FuncLit:
+			return false
 		case *ast.AssignStmt:
 			if n.Tok == token.DEFINE {
 				if _, ok := cursor.Parent().(*ast.TypeSwitchStmt); ok {

compiler/function.go

@@ -114,16 +114,19 @@ func generateFunctypesInit(pkg *ssa.Package, fn *ssa.Function, init *ast.BlockSt
 		return cmp.Compare(f1.Name(), f2.Name())
 	})
 
-	for index, anonFunc := range anonFuncs {
-		_, colored := colors[anonFunc]
-		if colored {
-			// Colored functions (those rewritten into coroutines) have a
-			// deferred anonymous function injected at the beginning to perform
-			// stack unwinding, which takes the ".func1" name.
-			index++
-		}
-		name = anonFuncLinkName(name, index)
-		generateFunctypesInit(pkg, anonFunc, init, name, colors)
+	index := 0
+	// Colored functions (those rewritten into coroutines) have a
+	// deferred anonymous function injected at the beginning to perform
+	// stack unwinding, which takes the ".func1" name.
+	_, colored := colors[fn]
+	if colored {
+		index++
+	}
+
+	for _, anonFunc := range anonFuncs {
+		index++
+		anonFuncName := anonFuncLinkName(name, index)
+		generateFunctypesInit(pkg, anonFunc, init, anonFuncName, colors)
 	}
 }
 
@@ -133,5 +136,5 @@ func generateFunctypesInit(pkg *ssa.Package, fn *ssa.Function, init *ast.BlockSt
 // The function works with multiple levels of nesting as each level adds another
 // ".func<index>" suffix, with the index being local to the parent scope.
 func anonFuncLinkName(base string, index int) string {
-	return fmt.Sprintf("%s.func%d", base, index+1)
+	return fmt.Sprintf("%s.func%d", base, index)
 }

compiler/testdata/coroutine.go

@@ -279,7 +279,7 @@ func RangeTripleFuncValue(n int) {
 	Range(n, f)
 }
 
-func Range10Closure() {
+func Range10ClosureCapturingValues() {
 	i := 0
 	n := 10
 	f := func() bool {
@@ -295,6 +295,112 @@ func Range10Closure() {
 	}
 }
 
+func Range10ClosureCapturingPointers() {
+	i, n := 0, 10
+	p := &i
+	q := &n
+	f := func() bool {
+		if *p < *q {
+			coroutine.Yield[int, any](*p)
+			(*p)++
+			return true
+		}
+		return false
+	}
+
+	for f() {
+	}
+}
+
+func Range10ClosureHeterogenousCapture() {
+	var (
+		a int8    = 0
+		b int16   = 1
+		c int32   = 2
+		d int64   = 3
+		e uint8   = 4
+		f uint16  = 5
+		g uint32  = 6
+		h uint64  = 7
+		i uintptr = 8
+		j         = func() int { return int(i) + 1 }
+	)
+
+	n := 0
+	x := func() bool {
+		var v int
+		switch n {
+		case 0:
+			v = int(a)
+		case 1:
+			v = int(b)
+		case 2:
+			v = int(c)
+		case 3:
+			v = int(d)
+		case 4:
+			v = int(e)
+		case 5:
+			v = int(f)
+		case 6:
+			v = int(g)
+		case 7:
+			v = int(h)
+		case 8:
+			v = int(i)
+		case 9:
+			v = j()
+		}
+		coroutine.Yield[int, any](v)
+		n++
+		return n < 10
+	}
+
+	for x() {
+	}
+}
+
+func Range10Heterogenous() {
+	var (
+		a int8    = 0
+		b int16   = 1
+		c int32   = 2
+		d int64   = 3
+		e uint8   = 4
+		f uint16  = 5
+		g uint32  = 6
+		h uint64  = 7
+		i uintptr = 8
+	)
+
+	for n := 0; n < 10; n++ {
+		var v int
+		switch n {
+		case 0:
+			v = int(a)
+		case 1:
+			v = int(b)
+		case 2:
+			v = int(c)
+		case 3:
+			v = int(d)
+		case 4:
+			v = int(e)
+		case 5:
+			v = int(f)
+		case 6:
+			v = int(g)
+		case 7:
+			v = int(h)
+		case 8:
+			v = int(i)
+		case 9:
+			v = int(n)
+		}
+		coroutine.Yield[int, any](v)
+	}
+}
+
 func Select(n int) {
 	select {
 	default: