From ee92efc193e3a66ef1fbd962fbd343d699e1f280 Mon Sep 17 00:00:00 2001
From: Alan Donovan <adonovan@google.com>
Date: Tue, 11 Feb 2014 15:59:20 -0500
Subject: [PATCH] go.tools/ssa: tweaks to CreateTestMainPackage.

Move mutation of program to the final step.
+ 2 assertions;
+ switches.go: comment fix.

LGTM=gri
R=gri
CC=golang-codereviews
https://golang.org/cl/61510044
---
 go/ssa/ssautil/testdata/switches.go |  4 ++--
 go/ssa/testmain.go                  | 10 +++++++++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/go/ssa/ssautil/testdata/switches.go b/go/ssa/ssautil/testdata/switches.go
index 8db7e5ab..8ab4c118 100644
--- a/go/ssa/ssautil/testdata/switches.go
+++ b/go/ssa/ssautil/testdata/switches.go
@@ -2,9 +2,9 @@
 
 package main
 
-// This file is the input to TestFindSwitches in switch_test.go.
+// This file is the input to TestSwitches in switch_test.go.
 // Each multiway conditional with constant or type cases (Switch)
-// discovered by FindSwitches is printed, and compared with the
+// discovered by Switches is printed, and compared with the
 // comments.
 //
 // The body of each case is printed as the value of its first
diff --git a/go/ssa/testmain.go b/go/ssa/testmain.go
index fb330acc..dcae06da 100644
--- a/go/ssa/testmain.go
+++ b/go/ssa/testmain.go
@@ -25,13 +25,15 @@ import (
 // It returns nil if the program contains no tests.
 //
 func (prog *Program) CreateTestMainPackage(pkgs ...*Package) *Package {
+	if len(pkgs) == 0 {
+		return nil
+	}
 	testmain := &Package{
 		Prog:    prog,
 		Members: make(map[string]Member),
 		values:  make(map[types.Object]Value),
 		Object:  types.NewPackage("testmain", "testmain", nil),
 	}
-	prog.packages[testmain.Object] = testmain
 
 	// Build package's init function.
 	init := &Function{
@@ -42,8 +44,12 @@ func (prog *Program) CreateTestMainPackage(pkgs ...*Package) *Package {
 		Prog:      prog,
 	}
 	init.startBody()
+	// TODO(adonovan): use lexical order.
 	var expfuncs []*Function // all exported functions of *_test.go in pkgs, unordered
 	for _, pkg := range pkgs {
+		if pkg.Prog != prog {
+			panic("wrong Program")
+		}
 		// Initialize package to test.
 		var v Call
 		v.Call.Value = pkg.init
@@ -135,6 +141,8 @@ func (prog *Program) CreateTestMainPackage(pkgs ...*Package) *Package {
 		sanityCheckPackage(testmain)
 	}
 
+	prog.packages[testmain.Object] = testmain
+
 	return testmain
 }