modfile: prefer to add a new 'go' directive after the 'module' directive

The previous behavior put the new 'go' directive at the end of the
file, which is not where anybody actually puts it, and that in turn
complicates the regression tests for golang/go#34822.

Updates golang/go#34822

Change-Id: Ic275eb6999fc0db7e72e2d9636b77c9e3d919539
Reviewed-on: https://go-review.googlesource.com/c/mod/+/205497
Run-TryBot: Bryan C. Mills <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Jay Conrod <[email protected]>

Author: Bryan C. Mills

modfile/read.go

@@ -90,6 +90,19 @@ func (x *FileSyntax) Span() (start, end Position) {
 	return start, end
 }
 
+// addLine adds a line containing the given tokens to the file.
+//
+// If the first token of the hint matches the first token of the
+// line, the new line is added at the end of the block containing hint,
+// extracting hint into a new block if it is not yet in one.
+//
+// If the hint is non-nil buts its first token does not match,
+// the new line is added after the block containing hint
+// (or hint itself, if not in a block).
+//
+// If no hint is provided, addLine appends the line to the end of
+// the last block with a matching first token,
+// or to the end of the file if no such block exists.
 func (x *FileSyntax) addLine(hint Expr, tokens ...string) *Line {
 	if hint == nil {
 		// If no hint given, add to the last statement of the given type.
@@ -111,11 +124,27 @@ func (x *FileSyntax) addLine(hint Expr, tokens ...string) *Line {
 		}
 	}
 
+	newLineAfter := func(i int) *Line {
+		new := &Line{Token: tokens}
+		if i == len(x.Stmt) {
+			x.Stmt = append(x.Stmt, new)
+		} else {
+			x.Stmt = append(x.Stmt, nil)
+			copy(x.Stmt[i+2:], x.Stmt[i+1:])
+			x.Stmt[i+1] = new
+		}
+		return new
+	}
+
 	if hint != nil {
 		for i, stmt := range x.Stmt {
 			switch stmt := stmt.(type) {
 			case *Line:
 				if stmt == hint {
+					if stmt.Token == nil || stmt.Token[0] != tokens[0] {
+						return newLineAfter(i)
+					}
+
 					// Convert line to line block.
 					stmt.InBlock = true
 					block := &LineBlock{Token: stmt.Token[:1], Line: []*Line{stmt}}
@@ -125,15 +154,25 @@ func (x *FileSyntax) addLine(hint Expr, tokens ...string) *Line {
 					block.Line = append(block.Line, new)
 					return new
 				}
+
 			case *LineBlock:
 				if stmt == hint {
+					if stmt.Token[0] != tokens[0] {
+						return newLineAfter(i)
+					}
+
 					new := &Line{Token: tokens[1:], InBlock: true}
 					stmt.Line = append(stmt.Line, new)
 					return new
 				}
+
 				for j, line := range stmt.Line {
 					if line == hint {
-						// Add new line after hint.
+						if stmt.Token[0] != tokens[0] {
+							return newLineAfter(i)
+						}
+
+						// Add new line after hint within the block.
 						stmt.Line = append(stmt.Line, nil)
 						copy(stmt.Line[j+2:], stmt.Line[j+1:])
 						new := &Line{Token: tokens[1:], InBlock: true}

modfile/rule.go

@@ -505,9 +505,13 @@ func (f *File) AddGoStmt(version string) error {
 		return fmt.Errorf("invalid language version string %q", version)
 	}
 	if f.Go == nil {
+		var hint Expr
+		if f.Module != nil && f.Module.Syntax != nil {
+			hint = f.Module.Syntax
+		}
 		f.Go = &Go{
 			Version: version,
-			Syntax:  f.Syntax.addLine(nil, "go", version),
+			Syntax:  f.Syntax.addLine(hint, "go", version),
 		}
 	} else {
 		f.Go.Version = version

modfile/rule_test.go

@@ -96,6 +96,47 @@ var setRequireTests = []struct {
 	},
 }
 
+var addGoTests = []struct {
+	in      string
+	version string
+	out     string
+}{
+	{`module m
+		`,
+		`1.14`,
+		`module m
+		go 1.14
+		`,
+	},
+	{`module m
+		require x.y/a v1.2.3
+		`,
+		`1.14`,
+		`module m
+		go 1.14
+		require x.y/a v1.2.3
+		`,
+	},
+	{
+		`require x.y/a v1.2.3
+		module example.com/inverted
+		`,
+		`1.14`,
+		`require x.y/a v1.2.3
+		module example.com/inverted
+		go 1.14
+		`,
+	},
+	{
+		`require x.y/a v1.2.3
+		`,
+		`1.14`,
+		`require x.y/a v1.2.3
+		go 1.14
+		`,
+	},
+}
+
 func TestAddRequire(t *testing.T) {
 	for i, tt := range addRequireTests {
 		t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {
@@ -167,3 +208,33 @@ func TestSetRequire(t *testing.T) {
 		})
 	}
 }
+
+func TestAddGo(t *testing.T) {
+	for i, tt := range addGoTests {
+		t.Run(fmt.Sprintf("#%d", i), func(t *testing.T) {
+			f, err := Parse("in", []byte(tt.in), nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			g, err := Parse("out", []byte(tt.out), nil)
+			if err != nil {
+				t.Fatal(err)
+			}
+			golden, err := g.Format()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if err := f.AddGoStmt(tt.version); err != nil {
+				t.Fatal(err)
+			}
+			out, err := f.Format()
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !bytes.Equal(out, golden) {
+				t.Errorf("have:\n%s\nwant:\n%s", out, golden)
+			}
+		})
+	}
+}