Started adding tests for ConcatenateFiles

Author: grantnelson-wf

compiler/astutil/astutil.go

@@ -453,15 +453,24 @@ func FinalizeRemovals(file *ast.File) {
 
 	file.Imports = squeeze(file.Imports)
 
+	rebuildFileComments(file)
+}
+
+// rebuildFileComments will rebuild the top level comments
+// group for the given file using the comments found in the file.
+func rebuildFileComments(file *ast.File) {
 	file.Comments = nil // clear this first so ast.Inspect doesn't walk it.
-	remComments := []*ast.CommentGroup{}
+	comments := []*ast.CommentGroup{}
+	if file.Doc != nil {
+		comments = append(comments, file.Doc)
+	}
 	ast.Inspect(file, func(n ast.Node) bool {
 		if cg, ok := n.(*ast.CommentGroup); ok {
-			remComments = append(remComments, cg)
+			comments = append(comments, cg)
 		}
 		return true
 	})
-	file.Comments = remComments
+	file.Comments = comments
 }
 
 // ConcatenateFiles will concatenate the given tailing files onto the
@@ -488,7 +497,11 @@ func ConcatenateFiles(file *ast.File, tails ...*ast.File) error {
 		if file.Name.Name != tail.Name.Name {
 			return fmt.Errorf("can not concatenate files with different package names: %q != %q", file.Name.Name, tail.Name.Name)
 		}
+		if file.GoVersion != tail.GoVersion {
+			return fmt.Errorf("can not concatenate files with different Go versions: %q != %q", file.GoVersion, tail.GoVersion)
+		}
 
+		// Concatenate the imports.
 		for _, imp := range tail.Imports {
 			path := imp.Path.Value
 			if oldImp, ok := imports[path]; ok {
@@ -503,15 +516,33 @@ func ConcatenateFiles(file *ast.File, tails ...*ast.File) error {
 			imports[imp.Path.Value] = imp
 		}
 
+		// Concatenate the declarations.
 		file.Decls = append(file.Decls, tail.Decls...)
-		file.Comments = append(file.Comments, tail.Comments...)
+
+		// Concatenate the document comments.
+		if tail.Doc != nil {
+			if file.Doc == nil {
+				file.Doc = tail.Doc
+			} else {
+				file.Doc.List = append(file.Doc.List, tail.Doc.List...)
+			}
+		}
+
+		// Concatenate the unresolved identifier and update scope.
+		// Both of these are deprecated. See Object.
+		// We just join them to attempt to keep the file in a valid state.
 		file.Unresolved = append(file.Unresolved, tail.Unresolved...)
 		for name, obj := range tail.Scope.Objects {
 			if _, ok := file.Scope.Objects[name]; ok {
 				return fmt.Errorf("can not concatenate files with duplicate object names: %q", name)
 			}
 			file.Scope.Objects[name] = obj
 		}
+
+		// Update the file end to the new end.
+		file.FileEnd = tail.FileEnd
 	}
+
+	rebuildFileComments(file)
 	return nil
 }

compiler/astutil/astutil_test.go

@@ -770,13 +770,13 @@ func TestFinalizeRemovals(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			st := srctesting.New(t)
 
-			srcFile := st.Parse("testSrc.go", test.src)
+			srcFile := st.Parse(`testSrc.go`, test.src)
 			test.perforator(srcFile)
 			FinalizeRemovals(srcFile)
 			got := srctesting.Format(t, st.FileSet, srcFile)
 
 			// parse and format the expected result so that formatting matches
-			wantFile := st.Parse("testWant.go", test.want)
+			wantFile := st.Parse(`testWant.go`, test.want)
 			want := srctesting.Format(t, st.FileSet, wantFile)
 
 			if got != want {
@@ -785,3 +785,66 @@ func TestFinalizeRemovals(t *testing.T) {
 		})
 	}
 }
+
+func TestConcatenateFiles(t *testing.T) {
+	tests := []struct {
+		name    string
+		srcHead string
+		srcTail string
+		want    string
+		expErr  string
+	}{
+		{
+			name: `add a method`,
+			srcHead: `package testpackage
+					// foo is an original method.
+					func foo() {}`,
+			srcTail: `package testpackage
+					// bar is a concatenated method
+					// from an additional override file.
+					func bar() {}`,
+			want: `package testpackage
+					// foo is an original method.
+					func foo() {}
+					// bar is a concatenated method
+					// from an additional override file.
+					func bar() {}`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			st := srctesting.New(t)
+			if (len(test.want) > 0) == (len(test.expErr) > 0) {
+				t.Fatal(`One and only one of "want" and "expErr" must be set`)
+			}
+
+			headFile := st.Parse(`testHead.go`, test.srcHead)
+			tailFile := st.Parse(`testTail.go`, test.srcTail)
+			err := ConcatenateFiles(headFile, tailFile)
+			if err != nil {
+				if len(test.expErr) == 0 {
+					t.Errorf(`Expected an AST but got an error: %v`, err)
+				} else if err.Error() != test.expErr {
+					t.Errorf("Unexpected error:\n\tgot:  %q\n\twant: %q", err.Error(), test.expErr)
+				}
+				return
+			}
+
+			ast.Print(st.FileSet, headFile) // TODO: REMOVE
+
+			got := srctesting.Format(t, st.FileSet, headFile)
+			if len(test.want) == 0 {
+				t.Errorf("Expected an error but got AST:\n\tgot:  %q\n\twant: %q", got, test.expErr)
+				return
+			}
+
+			// parse and format the expected result so that formatting matches
+			wantFile := st.Parse("testWant.go", test.want)
+			want := srctesting.Format(t, st.FileSet, wantFile)
+			if got != want {
+				t.Errorf("Unexpected resulting AST:\n\tgot:  %q\n\twant: %q", got, want)
+			}
+		})
+	}
+}