Merge branch 'main' into defers

compiler/color.go

@@ -27,6 +27,13 @@ func colorFunctions(cg *callgraph.Graph, yieldInstances functionColors) (functio
 type functionColors map[*ssa.Function]*types.Signature
 
 func colorFunctions0(cg *callgraph.Graph, colors functionColors, fn *ssa.Function, color *types.Signature) error {
+	if origin := fn.Origin(); origin != nil && origin.Pkg != nil {
+		// Don't follow edges into and through the coroutine package.
+		if pkgPath := origin.Pkg.Pkg.Path(); pkgPath == coroutinePackage {
+			return nil
+		}
+	}
+
 	existing, ok := colors[fn]
 	if ok {
 		if !types.Identical(existing, color) {

compiler/compile.go

@@ -14,6 +14,7 @@ import (
 
 	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/go/callgraph/cha"
+	"golang.org/x/tools/go/callgraph/vta"
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/go/ssa"
 	"golang.org/x/tools/go/ssa/ssautil"
@@ -109,28 +110,33 @@ func (c *compiler) compile(path string) error {
 			return fmt.Errorf("pattern more than one module (%s + %s)", moduleDir, p.Module.Dir)
 		}
 	}
-	flatpkgs := flattenPackages(pkgs)
-	for _, p := range flatpkgs {
-		for _, err := range p.Errors {
-			return err
+	err = nil
+	packages.Visit(pkgs, func(p *packages.Package) bool {
+		for _, e := range p.Errors {
+			err = e
+			break
 		}
+		return err == nil
+	}, nil)
+	if err != nil {
+		return err
 	}
 
 	log.Printf("building SSA program")
-	prog, _ := ssautil.Packages(pkgs, ssa.InstantiateGenerics|ssa.GlobalDebug)
+	prog, _ := ssautil.AllPackages(pkgs, ssa.InstantiateGenerics|ssa.GlobalDebug)
 	prog.Build()
 
 	log.Printf("building call graph")
-	cg := cha.CallGraph(prog)
+	cg := vta.CallGraph(ssautil.AllFunctions(prog), cha.CallGraph(prog))
 
 	log.Printf("finding generic yield instantiations")
 	var coroutinePkg *packages.Package
-	for _, p := range flatpkgs {
+	packages.Visit(pkgs, func(p *packages.Package) bool {
 		if p.PkgPath == coroutinePackage {
 			coroutinePkg = p
-			break
 		}
-	}
+		return coroutinePkg == nil
+	}, nil)
 	if coroutinePkg == nil {
 		log.Printf("%s not imported by the module. Nothing to do", coroutinePackage)
 		return nil
@@ -149,9 +155,10 @@ func (c *compiler) compile(path string) error {
 		return err
 	}
 	pkgsByTypes := map[*types.Package]*packages.Package{}
-	for _, p := range flatpkgs {
+	packages.Visit(pkgs, func(p *packages.Package) bool {
 		pkgsByTypes[p.Types] = p
-	}
+		return true
+	}, nil)
 	colorsByPkg := map[*packages.Package]functionColors{}
 	for fn, color := range colors {
 		if fn.Pkg == nil {
@@ -167,10 +174,22 @@ func (c *compiler) compile(path string) error {
 		pkgColors[fn] = color
 	}
 
-	for p, colors := range colorsByPkg {
+	var needVendoring []*packages.Package
+	for p := range colorsByPkg {
 		if p.Module == nil || p.Module.Dir != moduleDir {
-			return fmt.Errorf("not implemented: compilation for packages outside module (need to compile %s)", p.PkgPath)
+			needVendoring = append(needVendoring, p)
+			break
+		}
+	}
+	if len(needVendoring) > 0 {
+		log.Printf("vendoring packages")
+		newRoot := filepath.Join(moduleDir, "goroot")
+		if err := vendor(newRoot, needVendoring); err != nil {
+			return err
 		}
+	}
+
+	for p, colors := range colorsByPkg {
 		if err := c.compilePackage(p, colors, prog); err != nil {
 			return err
 		}

compiler/desugar.go

@@ -706,13 +706,20 @@ func (d *desugarer) decomposeExpression(expr ast.Expr, flags exprFlags) (ast.Exp
 				// Need to hoist the CallExpr out into a temporary variable in
 				// this case, so that the relative order of calls (and their
 				// prerequisites) is preserved.
-				queue[i] = decompose(e)
-			} else {
-				e.Fun = decompose(e.Fun)
-				for i, arg := range e.Args {
-					e.Args[i] = decompose(arg)
+				switch d.info.TypeOf(e).(type) {
+				case *types.Tuple:
+					// TODO: can't hoist like this when it's a function
+					//  that returns multiple values
+				default:
+					queue[i] = decompose(e)
+					continue
 				}
 			}
+			e.Fun = decompose(e.Fun)
+			for i, arg := range e.Args {
+				e.Args[i] = decompose(arg)
+			}
+
 		case *ast.CompositeLit:
 			for i, elt := range e.Elts {
 				e.Elts[i] = decompose(elt)

compiler/desugar_test.go

@@ -1147,6 +1147,26 @@ _l0:
 	_v2 := c(_v3)
 	_v0 <- _v2
 }
+`,
+		},
+		{
+			name: "FIXME: don't hoist function calls when there are multiple results",
+			body: "a, b, c = d(e())",
+			info: func(stmts []ast.Stmt, info *types.Info) {
+				callExpr := stmts[0].(*ast.AssignStmt).Rhs[0]
+				info.Types[callExpr] = types.TypeAndValue{
+					Type: types.NewTuple(
+						types.NewVar(0, nil, "a", intType),
+						types.NewVar(0, nil, "b", intType),
+						types.NewVar(0, nil, "c", intType),
+					),
+				}
+			},
+			expect: `
+{
+	_v0 := e()
+	a, b, c = d(_v0)
+}
 `,
 		},
 	} {

compiler/go.mod

@@ -5,6 +5,7 @@ go 1.21.0
 require (
 	github.com/google/go-cmp v0.5.9
 	github.com/stealthrocket/coroutine v0.0.0-20230906012022-7474cda88ddc
+	golang.org/x/sync v0.3.0
 	golang.org/x/tools v0.13.0
 )
 

compiler/testdata/http/.gitignore

@@ -0,0 +1 @@
+/goroot