[release-branch.go1.17] modfile: in SetRequireSeparateIndirect, arrange requirements consistently

SetRequireSeparateIndirect now makes a stronger attempt to keep
automatically added requirements in two blocks: one containing only
direct requirements and one containing only indirect
requirements. SetRequireSeparateIndirect will find or add these two
blocks (commented blocks are left alone). New requirements are added
to one of these two blocks. Existing requirements may be moved between
these two blocks if their markings change.

SetRequireSeparateIndirect attempts to preserve existing structure in
the file by not adding requirements to or moving requirements from
blocks with block-level comments and blocks other than the last
uncommented direct-only and indirect-only block.

As an exception to aid with migration, if the file contains a single
uncommented block of requirements (as would the be the case if the
user started with a 1.16 go.mod file, changed the go directive to
1.17, then ran 'go mod tidy'), SetRequireSeparateIndirect will split
the block into direct-only and indirect-only blocks.

This is a change in behavior, but it has no semantic effect, and it
should result in cleaner, more stable go.mod files.

For golang/go#47756

Change-Id: Ifa20bb084c6bdaf1e00140600380857de8afa320
Reviewed-on: https://go-review.googlesource.com/c/mod/+/343431
Trust: Jay Conrod <jayconrod@google.com>
Trust: Bryan C. Mills <bcmills@google.com>
Run-TryBot: Jay Conrod <jayconrod@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>
Reviewed-by: Michael Matloob <matloob@golang.org>
(cherry picked from commit 4be982bcc49d69e202fcb02aee4d80ce195aee1b)
Reviewed-on: https://go-review.googlesource.com/c/mod/+/351317
diff --git a/modfile/rule.go b/modfile/rule.go
index 78f83fa..d53d7a8 100644
--- a/modfile/rule.go
+++ b/modfile/rule.go
@@ -956,170 +956,190 @@
 
 // SetRequireSeparateIndirect updates the requirements of f to contain the given
 // requirements. Comment contents (except for 'indirect' markings) are retained
-// from the first existing requirement for each module path, and block structure
-// is maintained as long as the indirect markings match.
+// from the first existing requirement for each module path. Like SetRequire,
+// SetRequireSeparateIndirect adds requirements for new paths in req,
+// updates the version and "// indirect" comment on existing requirements,
+// and deletes requirements on paths not in req. Existing duplicate requirements
+// are deleted.
 //
-// Any requirements on paths not already present in the file are added. Direct
-// requirements are added to the last block containing *any* other direct
-// requirement. Indirect requirements are added to the last block containing
-// *only* other indirect requirements. If no suitable block exists, a new one is
-// added, with the last block containing a direct dependency (if any)
-// immediately before the first block containing only indirect dependencies.
+// As its name suggests, SetRequireSeparateIndirect puts direct and indirect
+// requirements into two separate blocks, one containing only direct
+// requirements, and the other containing only indirect requirements.
+// SetRequireSeparateIndirect may move requirements between these two blocks
+// when their indirect markings change. However, SetRequireSeparateIndirect
+// won't move requirements from other blocks, especially blocks with comments.
 //
-// The Syntax field is ignored for requirements in the given blocks.
+// If the file initially has one uncommented block of requirements,
+// SetRequireSeparateIndirect will split it into a direct-only and indirect-only
+// block. This aids in the transition to separate blocks.
 func (f *File) SetRequireSeparateIndirect(req []*Require) {
-	type modKey struct {
-		path     string
-		indirect bool
-	}
-	need := make(map[modKey]string)
-	for _, r := range req {
-		need[modKey{r.Mod.Path, r.Indirect}] = r.Mod.Version
+	// hasComments returns whether a line or block has comments
+	// other than "indirect".
+	hasComments := func(c Comments) bool {
+		return len(c.Before) > 0 || len(c.After) > 0 || len(c.Suffix) > 1 ||
+			(len(c.Suffix) == 1 &&
+				strings.TrimSpace(strings.TrimPrefix(c.Suffix[0].Token, string(slashSlash))) != "indirect")
 	}
 
-	comments := make(map[string]Comments)
-	for _, r := range f.Require {
-		v, ok := need[modKey{r.Mod.Path, r.Indirect}]
-		if !ok {
-			if _, ok := need[modKey{r.Mod.Path, !r.Indirect}]; ok {
-				if _, dup := comments[r.Mod.Path]; !dup {
-					comments[r.Mod.Path] = r.Syntax.Comments
-				}
+	// moveReq adds r to block. If r was in another block, moveReq deletes
+	// it from that block and transfers its comments.
+	moveReq := func(r *Require, block *LineBlock) {
+		var line *Line
+		if r.Syntax == nil {
+			line = &Line{Token: []string{AutoQuote(r.Mod.Path), r.Mod.Version}}
+			r.Syntax = line
+			if r.Indirect {
+				r.setIndirect(true)
 			}
-			r.markRemoved()
-			continue
+		} else {
+			line = new(Line)
+			*line = *r.Syntax
+			if !line.InBlock && len(line.Token) > 0 && line.Token[0] == "require" {
+				line.Token = line.Token[1:]
+			}
+			r.Syntax.Token = nil // Cleanup will delete the old line.
+			r.Syntax = line
 		}
-		r.setVersion(v)
-		delete(need, modKey{r.Mod.Path, r.Indirect})
+		line.InBlock = true
+		block.Line = append(block.Line, line)
 	}
 
+	// Examine existing require lines and blocks.
 	var (
-		lastDirectOrMixedBlock Expr
-		firstIndirectOnlyBlock Expr
-		lastIndirectOnlyBlock  Expr
+		// We may insert new requirements into the last uncommented
+		// direct-only and indirect-only blocks. We may also move requirements
+		// to the opposite block if their indirect markings change.
+		lastDirectBlock, lastIndirectBlock *LineBlock
+		lastDirectBlockIndex               = -1
+		lastIndirectBlockIndex             = -1
+
+		// If there are no direct-only or indirect-only blocks, a new block may
+		// be inserted after the last require line or block.
+		lastRequireIndex = -1
+
+		// If there's only one require line or block, and it's uncommented,
+		// we'll move its requirements to the direct-only or indirect-only blocks.
+		requireLineOrBlockCount = 0
+
+		// Track the block each requirement belongs to (if any) so we can
+		// move them later.
+		lineToBlock = make(map[*Line]*LineBlock)
 	)
-	for _, stmt := range f.Syntax.Stmt {
+	for i, stmt := range f.Syntax.Stmt {
 		switch stmt := stmt.(type) {
 		case *Line:
 			if len(stmt.Token) == 0 || stmt.Token[0] != "require" {
 				continue
 			}
-			if isIndirect(stmt) {
-				lastIndirectOnlyBlock = stmt
-			} else {
-				lastDirectOrMixedBlock = stmt
-			}
+			lastRequireIndex = i
+			requireLineOrBlockCount++
+
 		case *LineBlock:
 			if len(stmt.Token) == 0 || stmt.Token[0] != "require" {
 				continue
 			}
-			indirectOnly := true
+			lastRequireIndex = i
+			requireLineOrBlockCount++
+			allDirect := len(stmt.Line) > 0 && !hasComments(stmt.Comments)
+			allIndirect := len(stmt.Line) > 0 && !hasComments(stmt.Comments)
 			for _, line := range stmt.Line {
-				if len(line.Token) == 0 {
-					continue
-				}
-				if !isIndirect(line) {
-					indirectOnly = false
-					break
-				}
-			}
-			if indirectOnly {
-				lastIndirectOnlyBlock = stmt
-				if firstIndirectOnlyBlock == nil {
-					firstIndirectOnlyBlock = stmt
-				}
-			} else {
-				lastDirectOrMixedBlock = stmt
-			}
-		}
-	}
-
-	isOrContainsStmt := func(stmt Expr, target Expr) bool {
-		if stmt == target {
-			return true
-		}
-		if stmt, ok := stmt.(*LineBlock); ok {
-			if target, ok := target.(*Line); ok {
-				for _, line := range stmt.Line {
-					if line == target {
-						return true
-					}
-				}
-			}
-		}
-		return false
-	}
-
-	addRequire := func(path, vers string, indirect bool, comments Comments) {
-		var line *Line
-		if indirect {
-			if lastIndirectOnlyBlock != nil {
-				line = f.Syntax.addLine(lastIndirectOnlyBlock, "require", path, vers)
-			} else {
-				// Add a new require block after the last direct-only or mixed "require"
-				// block (if any).
-				//
-				// (f.Syntax.addLine would add the line to an existing "require" block if
-				// present, but here the existing "require" blocks are all direct-only, so
-				// we know we need to add a new block instead.)
-				line = &Line{Token: []string{"require", path, vers}}
-				lastIndirectOnlyBlock = line
-				firstIndirectOnlyBlock = line // only block implies first block
-				if lastDirectOrMixedBlock == nil {
-					f.Syntax.Stmt = append(f.Syntax.Stmt, line)
+				lineToBlock[line] = stmt
+				if hasComments(line.Comments) {
+					allDirect = false
+					allIndirect = false
+				} else if isIndirect(line) {
+					allDirect = false
 				} else {
-					for i, stmt := range f.Syntax.Stmt {
-						if isOrContainsStmt(stmt, lastDirectOrMixedBlock) {
-							f.Syntax.Stmt = append(f.Syntax.Stmt, nil)     // increase size
-							copy(f.Syntax.Stmt[i+2:], f.Syntax.Stmt[i+1:]) // shuffle elements up
-							f.Syntax.Stmt[i+1] = line
-							break
-						}
-					}
+					allIndirect = false
 				}
 			}
-		} else {
-			if lastDirectOrMixedBlock != nil {
-				line = f.Syntax.addLine(lastDirectOrMixedBlock, "require", path, vers)
+			if allDirect {
+				lastDirectBlock = stmt
+				lastDirectBlockIndex = i
+			}
+			if allIndirect {
+				lastIndirectBlock = stmt
+				lastIndirectBlockIndex = i
+			}
+		}
+	}
+
+	oneFlatUncommentedBlock := requireLineOrBlockCount == 1 &&
+		!hasComments(*f.Syntax.Stmt[lastRequireIndex].Comment())
+
+	// Create direct and indirect blocks if needed. It's okay if they're empty.
+	// Cleanup will remove them and convert one-line blocks to lines.
+	if lastDirectBlock == nil {
+		lastDirectBlock = &LineBlock{Token: []string{"require"}}
+		i := len(f.Syntax.Stmt)
+		if lastIndirectBlockIndex >= 0 {
+			i = lastIndirectBlockIndex
+			lastIndirectBlockIndex++
+		} else if lastRequireIndex >= 0 {
+			i = lastRequireIndex + 1
+		}
+		f.Syntax.Stmt = append(f.Syntax.Stmt, nil)
+		copy(f.Syntax.Stmt[i+1:], f.Syntax.Stmt[i:])
+		f.Syntax.Stmt[i] = lastDirectBlock
+		lastDirectBlockIndex = i
+	}
+
+	if lastIndirectBlock == nil {
+		lastIndirectBlock = &LineBlock{Token: []string{"require"}}
+		i := len(f.Syntax.Stmt)
+		if lastDirectBlockIndex >= 0 {
+			i = lastDirectBlockIndex + 1
+		}
+		f.Syntax.Stmt = append(f.Syntax.Stmt, nil)
+		copy(f.Syntax.Stmt[i+1:], f.Syntax.Stmt[i:])
+		f.Syntax.Stmt[i] = lastIndirectBlock
+		lastIndirectBlockIndex = i
+	}
+
+	// Delete requirements we don't want anymore.
+	// Update versions and indirect comments on requirements we want to keep.
+	// If a requirement is in last{Direct,Indirect}Block with the wrong
+	// indirect marking after this, or if the requirement is in an single
+	// uncommented mixed block (oneFlatUncommentedBlock), move it to the
+	// correct block.
+	//
+	// Some blocks may be empty after this. Cleanup will remove them.
+	need := make(map[string]*Require)
+	for _, r := range req {
+		need[r.Mod.Path] = r
+	}
+	have := make(map[string]*Require)
+	for _, r := range f.Require {
+		path := r.Mod.Path
+		if need[path] == nil || have[path] != nil {
+			// Requirement not needed, or duplicate requirement. Delete.
+			r.markRemoved()
+			continue
+		}
+		have[r.Mod.Path] = r
+		r.setVersion(need[path].Mod.Version)
+		r.setIndirect(need[path].Indirect)
+		if need[path].Indirect &&
+			(oneFlatUncommentedBlock || lineToBlock[r.Syntax] == lastDirectBlock) {
+			moveReq(r, lastIndirectBlock)
+		} else if !need[path].Indirect &&
+			(oneFlatUncommentedBlock || lineToBlock[r.Syntax] == lastIndirectBlock) {
+			moveReq(r, lastDirectBlock)
+		}
+	}
+
+	// Add new requirements.
+	for path, r := range need {
+		if have[path] == nil {
+			if r.Indirect {
+				moveReq(r, lastIndirectBlock)
 			} else {
-				// Add a new require block before the first indirect block (if any).
-				//
-				// That way if the file initially contains only indirect lines,
-				// the direct lines still appear before it: we preserve existing
-				// structure, but only to the extent that that structure already
-				// reflects the direct/indirect split.
-				line = &Line{Token: []string{"require", path, vers}}
-				lastDirectOrMixedBlock = line
-				if firstIndirectOnlyBlock == nil {
-					f.Syntax.Stmt = append(f.Syntax.Stmt, line)
-				} else {
-					for i, stmt := range f.Syntax.Stmt {
-						if isOrContainsStmt(stmt, firstIndirectOnlyBlock) {
-							f.Syntax.Stmt = append(f.Syntax.Stmt, nil)   // increase size
-							copy(f.Syntax.Stmt[i+1:], f.Syntax.Stmt[i:]) // shuffle elements up
-							f.Syntax.Stmt[i] = line
-							break
-						}
-					}
-				}
+				moveReq(r, lastDirectBlock)
 			}
+			f.Require = append(f.Require, r)
 		}
-
-		line.Comments.Before = commentsAdd(line.Comments.Before, comments.Before)
-		line.Comments.Suffix = commentsAdd(line.Comments.Suffix, comments.Suffix)
-
-		r := &Require{
-			Mod:      module.Version{Path: path, Version: vers},
-			Indirect: indirect,
-			Syntax:   line,
-		}
-		r.setIndirect(indirect)
-		f.Require = append(f.Require, r)
 	}
 
-	for k, vers := range need {
-		addRequire(k.path, vers, k.indirect, comments[k.path])
-	}
 	f.SortBlocks()
 }
 
diff --git a/modfile/rule_test.go b/modfile/rule_test.go
index 68d74c5..93dfae0 100644
--- a/modfile/rule_test.go
+++ b/modfile/rule_test.go
@@ -367,24 +367,16 @@
 		},
 		`module m
 		require x.y/a v1.2.3
-
+		require x.y/b v1.2.3 // indirect; demoted to indirect
 		require x.y/c v1.2.3 // not v1.2.3!
-
-		require (
-			x.y/d v1.2.3 // comment kept
-			x.y/f v1.2.3 // promoted to direct
-			// promoted to direct
-			x.y/g v1.2.3
-			x.y/h 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 // promoted to direct
+		// promoted to direct
+		require x.y/g v1.2.3
 		require x.y/i v1.2.3 // indirect
-
-		require (
-			x.y/b v1.2.3 // indirect; demoted to indirect
-			x.y/e v1.2.3 // indirect; comment kept
-			x.y/j v1.2.3 // indirect
-		)
+		require x.y/j v1.2.3 // indirect
+		require x.y/h v1.2.3
 		`,
 	},
 	{
@@ -430,7 +422,7 @@
 			{"x.y/a", "v1.2.3", true},
 		},
 		`module m
-		require x.y/a v1.2.3 // indirect; one
+		require x.y/a v1.2.3 // indirect; zero
 		`,
 	},
 	{
@@ -446,7 +438,7 @@
 			{"x.y/a", "v1.2.3", false},
 		},
 		`module m
-		require x.y/a v1.2.3 // one
+		require x.y/a v1.2.3 // zero
 		`,
 	},
 	{
@@ -508,10 +500,6 @@
 		},
 		`module m
 		require (
-			x.y/a v1.2.3
-			x.y/d v1.2.3
-		)
-		require (
 			x.y/b v1.2.3 // indirect; comment b preserved
 			x.y/e v1.2.3 // indirect; comment d preserved
 		)
@@ -519,6 +507,10 @@
 			x.y/c v1.2.3 // indirect; comment c preserved
 			x.y/f v1.2.3 // indirect; comment e preserved
 		)
+		require (
+			x.y/a v1.2.3
+			x.y/d v1.2.3
+		)
 		`,
 	},
 	{
@@ -531,23 +523,105 @@
 		`,
 		[]require{
 			{"x.y/a", "v1.2.3", false},
-			{"x.y/b", "v1.2.3", true},  // should remain in the existing mixed block
-			{"x.y/c", "v1.2.3", true},  // should be added in an indirect-only block
-			{"x.y/d", "v1.2.3", false}, // should appear in the existing mixed block
-			{"x.y/e", "v1.2.3", true},  // should be added in the same indirect-only block
+			{"x.y/b", "v1.2.3", true},
+			{"x.y/c", "v1.2.3", true},
+			{"x.y/d", "v1.2.3", false},
+			{"x.y/e", "v1.2.3", true},
 		},
 		`module m
 		require (
 			x.y/a v1.2.3
-			x.y/b v1.2.3 // indirect
 			x.y/d v1.2.3
 		)
 		require (
+			x.y/b v1.2.3 // indirect
 			x.y/c v1.2.3 // indirect
 			x.y/e v1.2.3 // indirect
 		)
 		`,
 	},
+	{
+		`preserve_block_comment_indirect_to_direct`,
+		`module m
+		// save
+		require (
+			x.y/a v1.2.3 // indirect
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+		},
+		`module m
+
+		// save
+		require x.y/a v1.2.3
+		`,
+	},
+	{
+		`preserve_block_comment_direct_to_indirect`,
+		`module m
+		// save
+		require (
+			x.y/a v1.2.3
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", true},
+		},
+		`module m
+
+		// save
+		require x.y/a v1.2.3 // indirect
+		`,
+	},
+	{
+		`regroup_flat_uncommented_block`,
+		`module m
+		require (
+			x.y/a v1.0.0 // a
+			x.y/b v1.0.0 // indirect; b
+			x.y/c v1.0.0 // indirect
+		)`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", true},
+			{"x.y/c", "v1.2.3", true},
+			{"x.y/d", "v1.2.3", false},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3 // a
+			x.y/d v1.2.3
+		)
+		require (
+			x.y/b v1.2.3 // indirect; b
+			x.y/c v1.2.3 // indirect
+		)`,
+	},
+	{
+		`dont_regroup_flat_commented_block`,
+		`module m
+		// dont regroup
+		require (
+			x.y/a v1.0.0
+			x.y/b v1.0.0 // indirect
+			x.y/c v1.0.0 // indirect
+		)`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", true},
+			{"x.y/c", "v1.2.3", true},
+			{"x.y/d", "v1.2.3", false},
+		},
+		`module m
+		// dont regroup
+		require (
+			x.y/a v1.2.3
+			x.y/b v1.2.3 // indirect
+			x.y/c v1.2.3 // indirect
+		)
+		require x.y/d v1.2.3`,
+	},
 }
 
 var addGoTests = []struct {