modfile: clean up SetRequire

I started this change by expanding the documentation and tests for
SetRequire. Unfortunately, the tests failed when the existing
contents included duplicates of a module path:

    --- FAIL: TestSetRequire/existing_duplicate (0.00s)
        rule_test.go:1011: after Cleanup, len(Require) = 3; want 1
    --- FAIL: TestSetRequire/existing_duplicate_multi (0.00s)
        rule_test.go:1011: after Cleanup, len(Require) = 3; want 1

So then I fixed the detected bug, by updating the Line entries
(possibly marking them for removal) in the same loop that updates the
Require entries. (We don't need to loop over f.Syntax.Stmt separately
to remove deleted entries because f.Syntax.Cleanup already does that.)

For golang/go#45965

Change-Id: I1b665c0832112de2c4273628f266dc3d966fefdd
Reviewed-on: https://go-review.googlesource.com/c/mod/+/325230
Trust: Bryan C. Mills <bcmills@google.com>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Michael Matloob <matloob@golang.org>
diff --git a/modfile/rule.go b/modfile/rule.go
index 5be5715..71cda84 100644
--- a/modfile/rule.go
+++ b/modfile/rule.go
@@ -870,67 +870,76 @@
 	f.Require = append(f.Require, &Require{module.Version{Path: path, Version: vers}, indirect, line})
 }
 
-// SetRequire sets the requirements of f to contain exactly req, preserving
-// any existing line comments contents (except for 'indirect' markings)
-// for the module versions named in req.
+// SetRequire updates the requirements of f to contain exactly req, preserving
+// the existing block structure and line comment contents (except for 'indirect'
+// markings) for the first requirement on each named module path.
+//
+// The Syntax field is ignored for the requirements in req.
+//
+// Any requirements not already present in the file are added to the block
+// containing the last require line.
+//
+// The requirements in req must specify at most one distinct version for each
+// module path.
+//
+// If any existing requirements may be removed, the caller should call Cleanup
+// after all edits are complete.
 func (f *File) SetRequire(req []*Require) {
 	need := make(map[string]string)
 	indirect := make(map[string]bool)
 	for _, r := range req {
+		if prev, dup := need[r.Mod.Path]; dup && prev != r.Mod.Version {
+			panic(fmt.Errorf("SetRequire called with conflicting versions for path %s (%s and %s)", r.Mod.Path, prev, r.Mod.Version))
+		}
 		need[r.Mod.Path] = r.Mod.Version
 		indirect[r.Mod.Path] = r.Indirect
 	}
 
+	// Update or delete the existing Require entries to preserve
+	// only the first for each module path in req.
 	for _, r := range f.Require {
-		if v, ok := need[r.Mod.Path]; ok {
-			r.Mod.Version = v
-			r.Indirect = indirect[r.Mod.Path]
-		} else {
+		v, ok := need[r.Mod.Path]
+		if !ok {
+			// This line is redundant or its path is no longer required at all.
+			// Mark the requirement for deletion in Cleanup.
+			r.Syntax.Token = nil
 			*r = Require{}
 		}
-	}
 
-	var newStmts []Expr
-	for _, stmt := range f.Syntax.Stmt {
-		switch stmt := stmt.(type) {
-		case *LineBlock:
-			if len(stmt.Token) > 0 && stmt.Token[0] == "require" {
-				var newLines []*Line
-				for _, line := range stmt.Line {
-					if p, err := parseString(&line.Token[0]); err == nil && need[p] != "" {
-						if len(line.Comments.Before) == 1 && len(line.Comments.Before[0].Token) == 0 {
-							line.Comments.Before = line.Comments.Before[:0]
-						}
-						line.Token[1] = need[p]
-						delete(need, p)
-						setIndirect(line, indirect[p])
-						newLines = append(newLines, line)
-					}
+		r.Mod.Version = v
+		r.Indirect = indirect[r.Mod.Path]
+
+		if line := r.Syntax; line != nil && len(line.Token) > 0 {
+			if line.InBlock {
+				// If the line is preceded by an empty line, remove it; see
+				// https://golang.org/issue/33779.
+				if len(line.Comments.Before) == 1 && len(line.Comments.Before[0].Token) == 0 {
+					line.Comments.Before = line.Comments.Before[:0]
 				}
-				if len(newLines) == 0 {
-					continue // drop stmt
+				if len(line.Token) >= 2 { // example.com v1.2.3
+					line.Token[1] = v
 				}
-				stmt.Line = newLines
+			} else {
+				if len(line.Token) >= 3 { // require example.com v1.2.3
+					line.Token[2] = v
+				}
 			}
 
-		case *Line:
-			if len(stmt.Token) > 0 && stmt.Token[0] == "require" {
-				if p, err := parseString(&stmt.Token[1]); err == nil && need[p] != "" {
-					stmt.Token[2] = need[p]
-					delete(need, p)
-					setIndirect(stmt, indirect[p])
-				} else {
-					continue // drop stmt
-				}
-			}
+			setIndirect(line, r.Indirect)
 		}
-		newStmts = append(newStmts, stmt)
-	}
-	f.Syntax.Stmt = newStmts
 
+		delete(need, r.Mod.Path)
+	}
+
+	// Add new entries in the last block of the file for any paths that weren't
+	// already present.
+	//
+	// This step is nondeterministic, but the final result will be deterministic
+	// because we will sort the block.
 	for path, vers := range need {
 		f.AddNewRequire(path, vers, indirect[path])
 	}
+
 	f.SortBlocks()
 }
 
diff --git a/modfile/rule_test.go b/modfile/rule_test.go
index 2ec24ea..7c6f0e8 100644
--- a/modfile/rule_test.go
+++ b/modfile/rule_test.go
@@ -92,15 +92,16 @@
 	},
 }
 
+type require struct {
+	path, vers string
+	indirect   bool
+}
+
 var setRequireTests = []struct {
 	desc string
 	in   string
-	mods []struct {
-		path     string
-		vers     string
-		indirect bool
-	}
-	out string
+	mods []require
+	out  string
 }{
 	{
 		`https://golang.org/issue/45932`,
@@ -111,11 +112,7 @@
 			x.y/c v1.2.3
 		)
 		`,
-		[]struct {
-			path     string
-			vers     string
-			indirect bool
-		}{
+		[]require{
 			{"x.y/a", "v1.2.3", false},
 			{"x.y/b", "v1.2.3", false},
 			{"x.y/c", "v1.2.3", false},
@@ -138,11 +135,7 @@
 			x.y/d v1.2.3
 		)
 		`,
-		[]struct {
-			path     string
-			vers     string
-			indirect bool
-		}{
+		[]require{
 			{"x.y/a", "v1.2.3", false},
 			{"x.y/b", "v1.2.3", false},
 			{"x.y/c", "v1.2.3", false},
@@ -168,11 +161,7 @@
 			x.y/g v1.2.3 //	indirect
 		)
 		`,
-		[]struct {
-			path     string
-			vers     string
-			indirect bool
-		}{
+		[]require{
 			{"x.y/a", "v1.2.3", true},
 			{"x.y/b", "v1.2.3", true},
 			{"x.y/c", "v1.2.3", true},
@@ -193,6 +182,76 @@
 		)
 		`,
 	},
+	{
+		`existing_multi`,
+		`module m
+		require x.y/a v1.2.3
+		require x.y/b v1.2.3
+		require x.y/c v1.0.0 // not v1.2.3!
+		require x.y/d v1.2.3 // comment kept
+		require x.y/e v1.2.3 // comment kept
+		require x.y/f v1.2.3 // indirect
+		require x.y/g v1.2.3 // indirect
+		`,
+		[]require{
+			{"x.y/h", "v1.2.3", false},
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", false},
+			{"x.y/c", "v1.2.3", false},
+			{"x.y/d", "v1.2.3", false},
+			{"x.y/e", "v1.2.3", true},
+			{"x.y/f", "v1.2.3", false},
+			{"x.y/g", "v1.2.3", false},
+		},
+		`module m
+		require x.y/a v1.2.3
+
+		require x.y/b v1.2.3
+
+		require x.y/c v1.2.3 // not v1.2.3!
+
+		require x.y/d v1.2.3 // comment kept
+
+		require x.y/e v1.2.3 // indirect; comment kept
+
+		require x.y/f v1.2.3
+
+		require (
+			x.y/g v1.2.3
+			x.y/h v1.2.3
+		)
+		`,
+	},
+	{
+		`existing_duplicate`,
+		`module m
+		require (
+			x.y/a v1.0.0 // zero
+			x.y/a v1.1.0 // one
+			x.y/a v1.2.3 // two
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", true},
+		},
+		`module m
+		require x.y/a v1.2.3 // indirect; zero
+		`,
+	},
+	{
+		`existing_duplicate_multi`,
+		`module m
+		require x.y/a v1.0.0 // zero
+		require x.y/a v1.1.0 // one
+		require x.y/a v1.2.3 // two
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", true},
+		},
+		`module m
+		require x.y/a v1.2.3 // indirect; zero
+		`,
+	},
 }
 
 var addGoTests = []struct {
@@ -942,10 +1001,10 @@
 
 			f := testEdit(t, tt.in, tt.out, true, func(f *File) error {
 				f.SetRequire(mods)
+				f.Cleanup()
 				return nil
 			})
 
-			f.Cleanup()
 			if len(f.Require) != len(mods) {
 				t.Errorf("after Cleanup, len(Require) = %v; want %v", len(f.Require), len(mods))
 			}