gopls/internal/golang: preserve copyright and build constraint

- This commit updates addTest and extractToFile to preserve
copyright and build constraint comments when creating new files.
- The change introduces utility functions to extract these comments
from an ast.File.

For golang/vscode-go#1594

Change-Id: I2b2f70d6d4de662c8357bca8558c094496c8b2e9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/629756
Reviewed-by: Alan Donovan <[email protected]>
Auto-Submit: Hongxiang Jiang <[email protected]>
Reviewed-by: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: h9jiang

gopls/internal/golang/addtest.go

@@ -330,22 +330,28 @@ func AddTestForFunc(ctx context.Context, snapshot *cache.Snapshot, loc protocol.
 		// package decl based on the originating file.
 		// Search for something that looks like a copyright header, to replicate
 		// in the new file.
-		if groups := pgf.File.Comments; len(groups) > 0 {
-			// Copyright should appear before package decl and must be the first
-			// comment group.
-			// Avoid copying any other comment like package doc or directive comment.
-			if c := groups[0]; c.Pos() < pgf.File.Package && c != pgf.File.Doc &&
-				!isDirective(c.List[0].Text) &&
-				strings.Contains(strings.ToLower(c.List[0].Text), "copyright") {
-				start, end, err := pgf.NodeOffsets(c)
-				if err != nil {
-					return nil, err
-				}
-				header.Write(pgf.Src[start:end])
-				// One empty line between copyright header and package decl.
-				header.WriteString("\n\n")
+		if c := copyrightComment(pgf.File); c != nil {
+			start, end, err := pgf.NodeOffsets(c)
+			if err != nil {
+				return nil, err
 			}
+			header.Write(pgf.Src[start:end])
+			// One empty line between copyright header and following.
+			header.WriteString("\n\n")
 		}
+
+		// If this test file was created by gopls, add build constraints
+		// matching the non-test file.
+		if c := buildConstraintComment(pgf.File); c != nil {
+			start, end, err := pgf.NodeOffsets(c)
+			if err != nil {
+				return nil, err
+			}
+			header.Write(pgf.Src[start:end])
+			// One empty line between build constraint and following.
+			header.WriteString("\n\n")
+		}
+
 		fmt.Fprintf(&header, "package %s_test\n", pkg.Types().Name())
 
 		// Write the copyright and package decl to the beginning of the file.

gopls/internal/golang/extracttofile.go

@@ -133,6 +133,26 @@ func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 	}
 
 	var buf bytes.Buffer
+	if c := copyrightComment(pgf.File); c != nil {
+		start, end, err := pgf.NodeOffsets(c)
+		if err != nil {
+			return nil, err
+		}
+		buf.Write(pgf.Src[start:end])
+		// One empty line between copyright header and following.
+		buf.WriteString("\n\n")
+	}
+
+	if c := buildConstraintComment(pgf.File); c != nil {
+		start, end, err := pgf.NodeOffsets(c)
+		if err != nil {
+			return nil, err
+		}
+		buf.Write(pgf.Src[start:end])
+		// One empty line between build constraint and following.
+		buf.WriteString("\n\n")
+	}
+
 	fmt.Fprintf(&buf, "package %s\n", pgf.File.Name.Name)
 	if len(adds) > 0 {
 		buf.WriteString("import (")
@@ -154,7 +174,6 @@ func ExtractToNewFile(ctx context.Context, snapshot *cache.Snapshot, fh file.Han
 	fileStart := pgf.File.FileStart
 	buf.Write(pgf.Src[start-fileStart : end-fileStart])
 
-	// TODO: attempt to duplicate the copyright header, if any.
 	newFileContent, err := format.Source(buf.Bytes())
 	if err != nil {
 		return nil, err

gopls/internal/golang/util.go

@@ -397,3 +397,45 @@ func AbbreviateVarName(s string) string {
 	}
 	return b.String()
 }
+
+// copyrightComment returns the copyright comment group from the input file, or
+// nil if not found.
+func copyrightComment(file *ast.File) *ast.CommentGroup {
+	if len(file.Comments) == 0 {
+		return nil
+	}
+
+	// Copyright should appear before package decl and must be the first
+	// comment group.
+	if c := file.Comments[0]; c.Pos() < file.Package && c != file.Doc &&
+		!isDirective(c.List[0].Text) &&
+		strings.Contains(strings.ToLower(c.List[0].Text), "copyright") {
+		return c
+	}
+
+	return nil
+}
+
+var buildConstraintRe = regexp.MustCompile(`^//(go:build|\s*\+build).*`)
+
+// buildConstraintComment returns the build constraint comment from the input
+// file.
+// Returns nil if not found.
+func buildConstraintComment(file *ast.File) *ast.Comment {
+	for _, cg := range file.Comments {
+		// In Go files a build constraint must appear before the package clause.
+		// See https://pkg.go.dev/cmd/go#hdr-Build_constraints
+		if cg.Pos() > file.Package {
+			return nil
+		}
+
+		for _, c := range cg.List {
+			// TODO: use ast.ParseDirective when available (#68021).
+			if buildConstraintRe.MatchString(c.Text) {
+				return c
+			}
+		}
+	}
+
+	return nil
+}

gopls/internal/test/marker/testdata/codeaction/addtest.txt

@@ -13,7 +13,7 @@ go 1.18
 	"addTestSourceCodeAction": true
 }
 
--- withcopyright/copyright.go --
+-- copyrightandbuildconstraint/copyrightandbuildconstraint.go --
 // Copyright 2020 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
@@ -23,18 +23,20 @@ go 1.18
 // Package main is for lsp test.
 package main
 
-func Foo(in string) string {return in} //@codeaction("Foo", "source.addTest", edit=with_copyright)
+func Foo(in string) string {return in} //@codeaction("Foo", "source.addTest", edit=with_copyright_build_constraint)
 
--- @with_copyright/withcopyright/copyright_test.go --
-@@ -0,0 +1,30 @@
+-- @with_copyright_build_constraint/copyrightandbuildconstraint/copyrightandbuildconstraint_test.go --
+@@ -0,0 +1,32 @@
 +// Copyright 2020 The Go Authors. All rights reserved.
 +// Use of this source code is governed by a BSD-style
 +// license that can be found in the LICENSE file.
 +
++//go:build go1.18
++
 +package main_test
 +
 +import(
-+	"golang.org/lsptests/addtest/withcopyright"
++	"golang.org/lsptests/addtest/copyrightandbuildconstraint"
 +	"testing"
 +)
 +
@@ -57,20 +59,22 @@ func Foo(in string) string {return in} //@codeaction("Foo", "source.addTest", ed
 +    })
 +  }
 +}
--- withoutcopyright/copyright.go --
+-- buildconstraint/buildconstraint.go --
 //go:build go1.18
 
 // Package copyright is for lsp test.
 package copyright
 
-func Foo(in string) string {return in} //@codeaction("Foo", "source.addTest", edit=without_copyright)
+func Foo(in string) string {return in} //@codeaction("Foo", "source.addTest", edit=with_build_constraint)
 
--- @without_copyright/withoutcopyright/copyright_test.go --
-@@ -0,0 +1,26 @@
+-- @with_build_constraint/buildconstraint/buildconstraint_test.go --
+@@ -0,0 +1,28 @@
++//go:build go1.18
++
 +package copyright_test
 +
 +import(
-+	"golang.org/lsptests/addtest/withoutcopyright"
++	"golang.org/lsptests/addtest/buildconstraint"
 +	"testing"
 +)
 +

gopls/internal/test/marker/testdata/codeaction/extracttofile.txt

@@ -275,3 +275,77 @@ func F() {} //@codeaction("func", "refactor.extract.toNewFile", edit=blank_impor
 -    return time1.Now().string()
 -}
 -
+-- copyright.go --
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+// docs
+const C = "" //@codeaction("const", "refactor.extract.toNewFile", edit=copyright)
+
+-- @copyright/c.go --
+@@ -0,0 +1,8 @@
++// Copyright 2020 The Go Authors. All rights reserved.
++// Use of this source code is governed by a BSD-style
++// license that can be found in the LICENSE file.
++
++package main
++
++// docs
++const C = ""
+-- @copyright/copyright.go --
+@@ -7,2 +7 @@
+-// docs
+-const C = "" //@codeaction("const", "refactor.extract.toNewFile", edit=copyright)
++//@codeaction("const", "refactor.extract.toNewFile", edit=copyright)
+-- buildconstraint.go --
+//go:build go1.18
+
+package main
+
+// docs
+const C = "" //@codeaction("const", "refactor.extract.toNewFile", edit=buildconstraint)
+
+-- @buildconstraint/buildconstraint.go --
+@@ -5,2 +5 @@
+-// docs
+-const C = "" //@codeaction("const", "refactor.extract.toNewFile", edit=buildconstraint)
++//@codeaction("const", "refactor.extract.toNewFile", edit=buildconstraint)
+-- @buildconstraint/c.go --
+@@ -0,0 +1,6 @@
++//go:build go1.18
++
++package main
++
++// docs
++const C = ""
+-- copyrightandbuildconstraint.go --
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.18
+
+package main
+
+// docs
+const C = "" //@codeaction("const", "refactor.extract.toNewFile", edit=copyrightandbuildconstraint)
+-- @copyrightandbuildconstraint/c.go --
+@@ -0,0 +1,10 @@
++// Copyright 2020 The Go Authors. All rights reserved.
++// Use of this source code is governed by a BSD-style
++// license that can be found in the LICENSE file.
++
++//go:build go1.18
++
++package main
++
++// docs
++const C = ""
+-- @copyrightandbuildconstraint/copyrightandbuildconstraint.go --
+@@ -9,2 +9 @@
+-// docs
+-const C = "" //@codeaction("const", "refactor.extract.toNewFile", edit=copyrightandbuildconstraint)
++//@codeaction("const", "refactor.extract.toNewFile", edit=copyrightandbuildconstraint)