modfile: add support for dropping go and toolchain stmts

Also add tests of previous CLs.

For golang/go#57001.

Change-Id: I755429dd07c0e84910108ce9807d607115329b79
Reviewed-on: https://go-review.googlesource.com/c/mod/+/497400
Run-TryBot: Russ Cox <rsc@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Bryan Mills <bcmills@google.com>
diff --git a/modfile/print.go b/modfile/print.go
index 524f930..2a0123d 100644
--- a/modfile/print.go
+++ b/modfile/print.go
@@ -16,7 +16,13 @@
 func Format(f *FileSyntax) []byte {
 	pr := &printer{}
 	pr.file(f)
-	return pr.Bytes()
+
+	// remove trailing blank lines
+	b := pr.Bytes()
+	for len(b) > 0 && b[len(b)-1] == '\n' && (len(b) == 1 || b[len(b)-2] == '\n') {
+		b = b[:len(b)-1]
+	}
+	return b
 }
 
 // A printer collects the state during printing of a file or expression.
@@ -59,7 +65,11 @@
 	}
 
 	p.trim()
-	p.printf("\n")
+	if b := p.Bytes(); len(b) == 0 || (len(b) >= 2 && b[len(b)-1] == '\n' && b[len(b)-2] == '\n') {
+		// skip the blank line at top of file or after a blank line
+	} else {
+		p.printf("\n")
+	}
 	for i := 0; i < p.margin; i++ {
 		p.printf("\t")
 	}
diff --git a/modfile/rule.go b/modfile/rule.go
index 45e7b6a..39f03f2 100644
--- a/modfile/rule.go
+++ b/modfile/rule.go
@@ -986,6 +986,22 @@
 	return nil
 }
 
+// DropGoStmt deletes the go statement from the file.
+func (f *File) DropGoStmt() {
+	if f.Go != nil {
+		f.Go.Syntax.markRemoved()
+		f.Go = nil
+	}
+}
+
+// DropToolchainStmt deletes the toolchain statement from the file.
+func (f *File) DropToolchainStmt() {
+	if f.Toolchain != nil {
+		f.Toolchain.Syntax.markRemoved()
+		f.Toolchain = nil
+	}
+}
+
 func (f *File) AddToolchainStmt(name string) error {
 	if !ToolchainRE.MatchString(name) {
 		return fmt.Errorf("invalid toolchain name %q", name)
@@ -1003,7 +1019,7 @@
 		}
 	} else {
 		f.Toolchain.Name = name
-		f.Syntax.updateLine(f.Go.Syntax, "toolchain", name)
+		f.Syntax.updateLine(f.Toolchain.Syntax, "toolchain", name)
 	}
 	return nil
 }
diff --git a/modfile/rule_test.go b/modfile/rule_test.go
index f8dd174..57c8be6 100644
--- a/modfile/rule_test.go
+++ b/modfile/rule_test.go
@@ -696,6 +696,137 @@
 	},
 }
 
+var dropGoTests = []struct {
+	desc string
+	in   string
+	out  string
+}{
+	{
+		`module_only`,
+		`module m
+		go 1.14
+		`,
+		`module m
+		`,
+	},
+	{
+		`module_before_require`,
+		`module m
+		go 1.14
+		require x.y/a v1.2.3
+		`,
+		`module m
+		require x.y/a v1.2.3
+		`,
+	},
+	{
+		`require_before_module`,
+		`require x.y/a v1.2.3
+		module example.com/inverted
+		go 1.14
+		`,
+		`require x.y/a v1.2.3
+		module example.com/inverted
+		`,
+	},
+	{
+		`require_only`,
+		`require x.y/a v1.2.3
+		go 1.14
+		`,
+		`require x.y/a v1.2.3
+		`,
+	},
+}
+
+var addToolchainTests = []struct {
+	desc    string
+	in      string
+	version string
+	out     string
+}{
+	{
+		`empty`,
+		``,
+		`go1.17`,
+		`toolchain go1.17
+		`,
+	},
+	{
+		`aftergo`,
+		`// this is a comment
+		require x v1.0.0
+
+		go 1.17
+
+		require y v1.0.0
+		`,
+		`go1.17`,
+		`// this is a comment
+		require x v1.0.0
+
+		go 1.17
+
+		toolchain go1.17
+
+		require y v1.0.0
+		`,
+	},
+	{
+		`already_have_toolchain`,
+		`go 1.17
+
+		toolchain go1.18
+		`,
+		`go1.19`,
+		`go 1.17
+
+		toolchain go1.19
+		`,
+	},
+}
+
+var dropToolchainTests = []struct {
+	desc string
+	in   string
+	out  string
+}{
+	{
+		`empty`,
+		`toolchain go1.17
+		`,
+		``,
+	},
+	{
+		`aftergo`,
+		`// this is a comment
+		require x v1.0.0
+
+		go 1.17
+
+		toolchain go1.17
+
+		require y v1.0.0
+		`,
+		`// this is a comment
+		require x v1.0.0
+
+		go 1.17
+
+		require y v1.0.0
+		`,
+	},
+	{
+		`already_have_toolchain`,
+		`go 1.17
+
+		toolchain go1.18
+		`,
+		`go 1.17
+		`,
+	},
+}
+
 var addExcludeTests = []struct {
 	desc    string
 	in      string
@@ -948,7 +1079,7 @@
 		`prefix_one`,
 		`module m
 		//   prefix
-		retract v1.0.0 
+		retract v1.0.0
 		`,
 		`prefix`,
 	},
@@ -959,7 +1090,7 @@
 		//
 		//     two
 		//
-		// three  
+		// three
 		retract v1.0.0`,
 		`one
 
@@ -1083,7 +1214,7 @@
 	{
 		`deprecated_paragraph_space`,
 		`// Deprecated: the next line has a space
-		// 
+		//
 		// c
 		module m`,
 		"the next line has a space",
@@ -1501,6 +1632,38 @@
 	}
 }
 
+func TestDropGo(t *testing.T) {
+	for _, tt := range dropGoTests {
+		t.Run(tt.desc, func(t *testing.T) {
+			testEdit(t, tt.in, tt.out, true, func(f *File) error {
+				f.DropGoStmt()
+				return nil
+			})
+		})
+	}
+}
+
+func TestAddToolchain(t *testing.T) {
+	for _, tt := range addToolchainTests {
+		t.Run(tt.desc, func(t *testing.T) {
+			testEdit(t, tt.in, tt.out, true, func(f *File) error {
+				return f.AddToolchainStmt(tt.version)
+			})
+		})
+	}
+}
+
+func TestDropToolchain(t *testing.T) {
+	for _, tt := range dropToolchainTests {
+		t.Run(tt.desc, func(t *testing.T) {
+			testEdit(t, tt.in, tt.out, true, func(f *File) error {
+				f.DropToolchainStmt()
+				return nil
+			})
+		})
+	}
+}
+
 func TestAddExclude(t *testing.T) {
 	for _, tt := range addExcludeTests {
 		t.Run(tt.desc, func(t *testing.T) {
diff --git a/modfile/work.go b/modfile/work.go
index 827a01d..75dc1c5 100644
--- a/modfile/work.go
+++ b/modfile/work.go
@@ -168,6 +168,22 @@
 	return nil
 }
 
+// DropGoStmt deletes the go statement from the file.
+func (f *WorkFile) DropGoStmt() {
+	if f.Go != nil {
+		f.Go.Syntax.markRemoved()
+		f.Go = nil
+	}
+}
+
+// DropToolchainStmt deletes the toolchain statement from the file.
+func (f *WorkFile) DropToolchainStmt() {
+	if f.Toolchain != nil {
+		f.Toolchain.Syntax.markRemoved()
+		f.Toolchain = nil
+	}
+}
+
 func (f *WorkFile) AddUse(diskPath, modulePath string) error {
 	need := true
 	for _, d := range f.Use {
diff --git a/modfile/work_test.go b/modfile/work_test.go
index 096ed5c..dcc0810 100644
--- a/modfile/work_test.go
+++ b/modfile/work_test.go
@@ -212,6 +212,53 @@
 	},
 }
 
+var workAddToolchainTests = []struct {
+	desc    string
+	in      string
+	version string
+	out     string
+}{
+	{
+		`empty`,
+		``,
+		`go1.17`,
+		`toolchain go1.17
+		`,
+	},
+	{
+		`aftergo`,
+		`// this is a comment
+		use foo
+
+		go 1.17
+
+		use bar
+		`,
+		`go1.17`,
+		`// this is a comment
+		use foo
+
+		go 1.17
+
+		toolchain go1.17
+
+		use bar
+		`,
+	},
+	{
+		`already_have_toolchain`,
+		`go 1.17
+
+		toolchain go1.18
+		`,
+		`go1.19`,
+		`go 1.17
+
+		toolchain go1.19
+		`,
+	},
+}
+
 var workSortBlocksTests = []struct {
 	desc, in, out string
 }{
@@ -284,6 +331,16 @@
 	}
 }
 
+func TestWorkAddToolchain(t *testing.T) {
+	for _, tt := range workAddToolchainTests {
+		t.Run(tt.desc, func(t *testing.T) {
+			testWorkEdit(t, tt.in, tt.out, func(f *WorkFile) error {
+				return f.AddToolchainStmt(tt.version)
+			})
+		})
+	}
+}
+
 func TestWorkSortBlocks(t *testing.T) {
 	for _, tt := range workSortBlocksTests {
 		t.Run(tt.desc, func(t *testing.T) {