go/ssa: support custom TestMain functions in test packages

Supporting user-defined TestMain functions requires that we generate a
"testmain" package for each testable package, rather than a single one
for the entire program.  This entails these API changes:

1. (*ssa.Program).{CreateTestMainPackage,FindTests} both now
   accept only a single package.  Existing clients that pass them
   multiple packages must call them from a loop.

2. (*ssa.Program).FindTests returns an additional result, the the
   optional TestMain *ssa.Function.  Existing clients may discard it.

Also:
- Test the new logic using the SSA interpreter
- add ssautil.MainPackages helper
- callgraph: allow multiple main packages, and analyze them all
- ssadump -run: allow multiple main/test packages, and run each in a new interpreter
- minor simplifications to some callers (e.g. guru)

Fixes golang/go#9553

Change-Id: Ia7de9bd27448fb08b8d172ba5cdbcf37a762b7a0
Reviewed-on: https://go-review.googlesource.com/25102
Reviewed-by: Robert Griesemer <gri@golang.org>

cmd/callgraph/main.go

@@ -221,12 +221,12 @@ func doCallgraph(ctxt *build.Context, algo, format string, tests bool, args []st
 			}
 		}
 
-		main, err := mainPackage(prog, tests)
+		mains, err := mainPackages(prog, tests)
 		if err != nil {
 			return err
 		}
 		config := &pointer.Config{
-			Mains:          []*ssa.Package{main},
+			Mains:          mains,
 			BuildCallGraph: true,
 			Log:            ptalog,
 		}
@@ -237,13 +237,13 @@ func doCallgraph(ctxt *build.Context, algo, format string, tests bool, args []st
 		cg = ptares.CallGraph
 
 	case "rta":
-		main, err := mainPackage(prog, tests)
+		mains, err := mainPackages(prog, tests)
 		if err != nil {
 			return err
 		}
-		roots := []*ssa.Function{
-			main.Func("init"),
-			main.Func("main"),
+		var roots []*ssa.Function
+		for _, main := range mains {
+			roots = append(roots, main.Func("init"), main.Func("main"))
 		}
 		rtares := rta.Analyze(roots, true)
 		cg = rtares.CallGraph
@@ -303,35 +303,31 @@ func doCallgraph(ctxt *build.Context, algo, format string, tests bool, args []st
 	return nil
 }
 
-// mainPackage returns the main package to analyze.
-// The resulting package has a main() function.
-func mainPackage(prog *ssa.Program, tests bool) (*ssa.Package, error) {
-	pkgs := prog.AllPackages()
-
-	// TODO(adonovan): allow independent control over tests, mains and libraries.
-	// TODO(adonovan): put this logic in a library; we keep reinventing it.
+// mainPackages returns the main packages to analyze.
+// Each resulting package is named "main" and has a main function.
+func mainPackages(prog *ssa.Program, tests bool) ([]*ssa.Package, error) {
+	pkgs := prog.AllPackages() // TODO(adonovan): use only initial packages
 
+	// If tests, create a "testmain" package for each test.
+	var mains []*ssa.Package
 	if tests {
-		// If -test, use all packages' tests.
-		if len(pkgs) > 0 {
-			if main := prog.CreateTestMainPackage(pkgs...); main != nil {
-				return main, nil
+		for _, pkg := range pkgs {
+			if main := prog.CreateTestMainPackage(pkg); main != nil {
+				mains = append(mains, main)
 			}
 		}
-		return nil, fmt.Errorf("no tests")
-	}
-
-	// Otherwise, use the first package named main.
-	for _, pkg := range pkgs {
-		if pkg.Pkg.Name() == "main" {
-			if pkg.Func("main") == nil {
-				return nil, fmt.Errorf("no func main() in main package")
-			}
-			return pkg, nil
+		if mains == nil {
+			return nil, fmt.Errorf("no tests")
 		}
+		return mains, nil
 	}
 
-	return nil, fmt.Errorf("no main package")
+	// Otherwise, use the main packages.
+	mains = append(mains, ssautil.MainPackages(pkgs)...)
+	if len(mains) == 0 {
+		return nil, fmt.Errorf("no main packages")
+	}
+	return mains, nil
 }
 
 type Edge struct {

cmd/callgraph/main_test.go

@@ -47,14 +47,14 @@ func TestCallgraph(t *testing.T) {
 		}},
 		// tests: main is not called.
 		{"rta", format, true, []string{
+			`pkg$testmain.init --> pkg.init`,
 			`pkg.Example --> (pkg.C).f`,
-			`test$main.init --> pkg.init`,
 		}},
 		{"pta", format, true, []string{
+			`<root> --> pkg$testmain.init`,
 			`<root> --> pkg.Example`,
-			`<root> --> test$main.init`,
+			`pkg$testmain.init --> pkg.init`,
 			`pkg.Example --> (pkg.C).f`,
-			`test$main.init --> pkg.init`,
 		}},
 	} {
 		stdout = new(bytes.Buffer)

cmd/guru/guru.go

@@ -129,26 +129,18 @@ func setPTAScope(lconf *loader.Config, scope []string) error {
 // Create a pointer.Config whose scope is the initial packages of lprog
 // and their dependencies.
 func setupPTA(prog *ssa.Program, lprog *loader.Program, ptaLog io.Writer, reflection bool) (*pointer.Config, error) {
-	// TODO(adonovan): the body of this function is essentially
-	// duplicated in all go/pointer clients.  Refactor.
-
 	// For each initial package (specified on the command line),
 	// if it has a main function, analyze that,
 	// otherwise analyze its tests, if any.
-	var testPkgs, mains []*ssa.Package
+	var mains []*ssa.Package
 	for _, info := range lprog.InitialPackages() {
-		initialPkg := prog.Package(info.Pkg)
+		p := prog.Package(info.Pkg)
 
 		// Add package to the pointer analysis scope.
-		if initialPkg.Func("main") != nil {
-			mains = append(mains, initialPkg)
-		} else {
-			testPkgs = append(testPkgs, initialPkg)
-		}
-	}
-	if testPkgs != nil {
-		if p := prog.CreateTestMainPackage(testPkgs...); p != nil {
+		if p.Pkg.Name() == "main" && p.Func("main") != nil {
 			mains = append(mains, p)
+		} else if main := prog.CreateTestMainPackage(p); main != nil {
+			mains = append(mains, main)
 		}
 	}
 	if mains == nil {

cmd/ssadump/main.go

@@ -125,46 +125,44 @@ func doMain() error {
 	}
 
 	// Load, parse and type-check the whole program.
-	iprog, err := conf.Load()
+	lprog, err := conf.Load()
 	if err != nil {
 		return err
 	}
 
 	// Create and build SSA-form program representation.
-	prog := ssautil.CreateProgram(iprog, mode)
+	prog := ssautil.CreateProgram(lprog, mode)
 
 	// Build and display only the initial packages
 	// (and synthetic wrappers), unless -run is specified.
-	for _, info := range iprog.InitialPackages() {
-		prog.Package(info.Pkg).Build()
+	var initpkgs []*ssa.Package
+	for _, info := range lprog.InitialPackages() {
+		ssapkg := prog.Package(info.Pkg)
+		ssapkg.Build()
+		if info.Pkg.Path() != "runtime" {
+			initpkgs = append(initpkgs, ssapkg)
+		}
 	}
 
 	// Run the interpreter.
 	if *runFlag {
 		prog.Build()
 
-		var main *ssa.Package
-		pkgs := prog.AllPackages()
+		var mains []*ssa.Package
 		if *testFlag {
-			// If -test, run all packages' tests.
-			if len(pkgs) > 0 {
-				main = prog.CreateTestMainPackage(pkgs...)
+			// If -test, run the tests.
+			for _, pkg := range initpkgs {
+				if main := prog.CreateTestMainPackage(pkg); main != nil {
+					mains = append(mains, main)
+				}
 			}
-			if main == nil {
+			if mains == nil {
 				return fmt.Errorf("no tests")
 			}
 		} else {
-			// Otherwise, run main.main.
-			for _, pkg := range pkgs {
-				if pkg.Pkg.Name() == "main" {
-					main = pkg
-					if main.Func("main") == nil {
-						return fmt.Errorf("no func main() in main package")
-					}
-					break
-				}
-			}
-			if main == nil {
+			// Otherwise, run the main packages.
+			mains := ssautil.MainPackages(initpkgs)
+			if len(mains) == 0 {
 				return fmt.Errorf("no main package")
 			}
 		}
@@ -174,7 +172,12 @@ func doMain() error {
 				build.Default.GOARCH, runtime.GOARCH)
 		}
 
-		interp.Interpret(main, interpMode, conf.TypeChecker.Sizes, main.Pkg.Path(), args)
+		for _, main := range mains {
+			if len(mains) > 1 {
+				fmt.Fprintf(os.Stderr, "Running: %s\n", main.Pkg.Path())
+			}
+			interp.Interpret(main, interpMode, conf.TypeChecker.Sizes, main.Pkg.Path(), args)
+		}
 	}
 	return nil
 }

go/pointer/stdlib_test.go

@@ -60,20 +60,22 @@ func TestStdlib(t *testing.T) {
 	}
 
 	// Determine the set of packages/tests to analyze.
-	var testPkgs []*ssa.Package
+	var mains []*ssa.Package
 	for _, info := range iprog.InitialPackages() {
-		testPkgs = append(testPkgs, prog.Package(info.Pkg))
+		ssapkg := prog.Package(info.Pkg)
+		if main := prog.CreateTestMainPackage(ssapkg); main != nil {
+			mains = append(mains, main)
+		}
 	}
-	testmain := prog.CreateTestMainPackage(testPkgs...)
-	if testmain == nil {
-		t.Fatal("analysis scope has tests")
+	if mains == nil {
+		t.Fatal("no tests found in analysis scope")
 	}
 
 	// Run the analysis.
 	config := &Config{
 		Reflection:     false, // TODO(adonovan): fix remaining bug in rVCallConstraint, then enable.
 		BuildCallGraph: true,
-		Mains:          []*ssa.Package{testmain},
+		Mains:          mains,
 	}
 	// TODO(adonovan): add some query values (affects track bits).
 

go/ssa/interp/interp_test.go

@@ -222,29 +222,25 @@ func run(t *testing.T, dir, input string, success successPredicate) bool {
 	prog := ssautil.CreateProgram(iprog, ssa.SanityCheckFunctions)
 	prog.Build()
 
+	// Find first main or test package among the initial packages.
 	var mainPkg *ssa.Package
-	var initialPkgs []*ssa.Package
 	for _, info := range iprog.InitialPackages() {
 		if info.Pkg.Path() == "runtime" {
 			continue // not an initial package
 		}
 		p := prog.Package(info.Pkg)
-		initialPkgs = append(initialPkgs, p)
-		if mainPkg == nil && p.Func("main") != nil {
+		if p.Pkg.Name() == "main" && p.Func("main") != nil {
 			mainPkg = p
+			break
+		}
+
+		mainPkg = prog.CreateTestMainPackage(p)
+		if mainPkg != nil {
+			break
 		}
 	}
 	if mainPkg == nil {
-		testmainPkg := prog.CreateTestMainPackage(initialPkgs...)
-		if testmainPkg == nil {
-			t.Errorf("CreateTestMainPackage(%s) returned nil", mainPkg)
-			return false
-		}
-		if testmainPkg.Func("main") == nil {
-			t.Errorf("synthetic testmain package has no main")
-			return false
-		}
-		mainPkg = testmainPkg
+		t.Fatalf("no main or test packages among initial packages: %s", inputs)
 	}
 
 	var out bytes.Buffer
@@ -346,6 +342,23 @@ func TestTestmainPackage(t *testing.T) {
 		return nil
 	}
 	run(t, "testdata"+slash, "a_test.go", success)
+
+	// Run a test with a custom TestMain function and ensure that it
+	// is executed, and that m.Run runs the tests.
+	success = func(exitcode int, output string) error {
+		if exitcode != 0 {
+			return fmt.Errorf("unexpected failure; output=%s", output)
+		}
+		if want := `TestMain start
+TestC
+PASS
+TestMain end
+`; output != want {
+			return fmt.Errorf("output was %q, want %q", output, want)
+		}
+		return nil
+	}
+	run(t, "testdata"+slash, "c_test.go", success)
 }
 
 // CreateTestMainPackage should return nil if there were no tests.

go/ssa/interp/testdata/c_test.go

@@ -0,0 +1,17 @@
+package c_test
+
+import (
+	"os"
+	"testing"
+)
+
+func TestC(t *testing.T) {
+	println("TestC")
+}
+
+func TestMain(m *testing.M) {
+	println("TestMain start")
+	code := m.Run()
+	println("TestMain end")
+	os.Exit(code)
+}

go/ssa/ssautil/visit.go

@@ -64,3 +64,16 @@ func (visit *visitor) function(fn *ssa.Function) {
 		}
 	}
 }
+
+// MainPackages returns the subset of the specified packages
+// named "main" that define a main function.
+// The result may include synthetic "testmain" packages.
+func MainPackages(pkgs []*ssa.Package) []*ssa.Package {
+	var mains []*ssa.Package
+	for _, pkg := range pkgs {
+		if pkg.Pkg.Name() == "main" && pkg.Func("main") != nil {
+			mains = append(mains, pkg)
+		}
+	}
+	return mains
+}