x/mod: allow for aggressively conslidating requires

This change adds a new function `SetRequireSimplify` that will
aggressively consolidate all go.mod requirements into at most two blocks
(one direct, one indirect).

For golang/go#56471.

Change-Id: If71e124725cbd9b1d9e6fc5ad0a4cf33f34cb3ae
Reviewed-on: https://go-review.googlesource.com/c/mod/+/739221
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Pratt <mpratt@google.com>
diff --git a/modfile/rule.go b/modfile/rule.go
index c5b8305..65eb72b 100644
--- a/modfile/rule.go
+++ b/modfile/rule.go
@@ -327,6 +327,7 @@
 }
 
 var GoVersionRE = lazyregexp.New(`^([1-9][0-9]*)\.(0|[1-9][0-9]*)(\.(0|[1-9][0-9]*))?([a-z]+[0-9]+)?$`)
+
 var laxGoVersionRE = lazyregexp.New(`^v?(([1-9][0-9]*)\.(0|[1-9][0-9]*))([^0-9].*)$`)
 
 // Toolchains must be named beginning with `go1`,
@@ -1272,6 +1273,17 @@
 // 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) {
+	f.setRequireSeparateIndirect(req, false)
+}
+
+// SetRequireAtMostTwo is like SetRequireSeparateIndirect but it aggressively
+// consolidates all requirements into at most two blocks (one direct, one indirect).
+// It ignores existing blocks and comments when deciding where to place requirements.
+func (f *File) SetRequireAtMostTwo(req []*Require) {
+	f.setRequireSeparateIndirect(req, true)
+}
+
+func (f *File) setRequireSeparateIndirect(req []*Require, simplify bool) {
 	// hasComments returns whether a line or block has comments
 	// other than "indirect".
 	hasComments := func(c Comments) bool {
@@ -1304,6 +1316,17 @@
 	}
 
 	// Examine existing require lines and blocks.
+	need := make(map[string]*Require)
+	for _, r := range req {
+		need[r.Mod.Path] = r
+	}
+	lineIndirect := make(map[*Line]bool)
+	for _, r := range f.Require {
+		if n := need[r.Mod.Path]; n != nil {
+			lineIndirect[r.Syntax] = n.Indirect
+		}
+	}
+
 	var (
 		// We may insert new requirements into the last uncommented
 		// direct-only and indirect-only blocks. We may also move requirements
@@ -1321,7 +1344,9 @@
 
 		// Track the block each requirement belongs to (if any) so we can
 		// move them later.
-		lineToBlock = make(map[*Line]*LineBlock)
+		lineToBlock           = make(map[*Line]*LineBlock)
+		directBlockComments   []Comment
+		indirectBlockComments []Comment
 	)
 	for i, stmt := range f.Syntax.Stmt {
 		switch stmt := stmt.(type) {
@@ -1364,6 +1389,24 @@
 			if allIndirect {
 				lastIndirectIndex = i
 			}
+			if simplify {
+				anyDirect := false
+				for _, line := range stmt.Line {
+					if ind, ok := lineIndirect[line]; ok && !ind {
+						anyDirect = true
+						break
+					}
+				}
+				target := &directBlockComments
+				if !anyDirect && len(stmt.Line) > 0 {
+					target = &indirectBlockComments
+				}
+				if len(*target) > 0 && len(stmt.Comments.Before) > 0 {
+					*target = append(*target, Comment{Token: "//"})
+				}
+				*target = append(*target, stmt.Comments.Before...)
+				stmt.Comments.Before = nil
+			}
 		}
 	}
 
@@ -1422,6 +1465,15 @@
 		lastIndirectBlock = ensureBlock(lastIndirectIndex)
 	}
 
+	if simplify {
+		if len(directBlockComments) > 0 {
+			lastDirectBlock.Comments.Before = append(lastDirectBlock.Comments.Before, directBlockComments...)
+		}
+		if len(indirectBlockComments) > 0 {
+			lastIndirectBlock.Comments.Before = append(lastIndirectBlock.Comments.Before, indirectBlockComments...)
+		}
+	}
+
 	// 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
@@ -1430,10 +1482,6 @@
 	// 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
@@ -1446,10 +1494,10 @@
 		r.setVersion(need[path].Mod.Version)
 		r.setIndirect(need[path].Indirect)
 		if need[path].Indirect &&
-			(oneFlatUncommentedBlock || lineToBlock[r.Syntax] == lastDirectBlock) {
+			(simplify || oneFlatUncommentedBlock || lineToBlock[r.Syntax] == lastDirectBlock) {
 			moveReq(r, lastIndirectBlock)
 		} else if !need[path].Indirect &&
-			(oneFlatUncommentedBlock || lineToBlock[r.Syntax] == lastIndirectBlock) {
+			(simplify || oneFlatUncommentedBlock || lineToBlock[r.Syntax] == lastIndirectBlock) {
 			moveReq(r, lastDirectBlock)
 		}
 	}
diff --git a/modfile/rule_test.go b/modfile/rule_test.go
index 5ee8e3e..ea9dd5a 100644
--- a/modfile/rule_test.go
+++ b/modfile/rule_test.go
@@ -666,6 +666,404 @@
 	},
 }
 
+var setRequireSimplifyTests = []struct {
+	desc string
+	in   string
+	mods []require
+	out  string
+}{
+	{
+		`https://golang.org/issue/45932`,
+		`module m
+		require (
+			x.y/a v1.2.3 //indirect
+			x.y/b v1.2.3
+			x.y/c v1.2.3
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", false},
+			{"x.y/c", "v1.2.3", false},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3
+			x.y/b v1.2.3
+			x.y/c v1.2.3
+		)
+		`,
+	},
+	{
+		`existing`,
+		`module m
+		require (
+			x.y/b v1.2.3
+
+			x.y/a v1.2.3
+			x.y/d v1.2.3
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", false},
+			{"x.y/c", "v1.2.3", false},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3
+			x.y/b v1.2.3
+			x.y/c v1.2.3
+		)
+		`,
+	},
+	{
+		`existing_indirect`,
+		`module m
+		require (
+			x.y/a v1.2.3
+			x.y/b v1.2.3 //
+			x.y/c v1.2.3 //c
+		)
+		require (
+			x.y/d v1.2.3 //   c
+			x.y/e v1.2.3 // indirect
+		)
+		require (
+			x.y/f v1.2.3 //indirect
+			x.y/g v1.2.3 //	indirect
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", true},
+			{"x.y/b", "v1.2.3", true},
+			{"x.y/c", "v1.2.3", true},
+			{"x.y/d", "v1.2.3", true},
+			{"x.y/e", "v1.2.3", true},
+			{"x.y/f", "v1.2.3", true},
+			{"x.y/g", "v1.2.3", true},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3 // indirect
+			x.y/b v1.2.3 // indirect
+			x.y/c v1.2.3 // indirect; c
+			x.y/d v1.2.3 // indirect; c
+			x.y/e v1.2.3 // indirect
+			x.y/f v1.2.3 //indirect
+			x.y/g v1.2.3 //	indirect
+		)
+		`,
+	},
+	{
+		`existing_line`,
+		`module m
+		require x.y/a v1.0.0
+		require x.y/c v1.0.0 // indirect
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", false},
+			{"x.y/c", "v1.2.3", true},
+			{"x.y/d", "v1.2.3", true},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3
+			x.y/b v1.2.3
+		)
+		require (
+			x.y/c v1.2.3 // indirect
+			x.y/d v1.2.3 // indirect
+		)`,
+	},
+	{
+		`existing_multi`,
+		`module m
+		require x.y/a v1.2.3
+		require x.y/b v1.2.3 // demoted to indirect
+		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; promoted to direct
+		// promoted to direct
+		require x.y/g v1.2.3 // indirect
+		require x.y/i v1.2.3 // indirect
+		require x.y/j v1.2.3 // indirect
+		`,
+		[]require{
+			{"x.y/h", "v1.2.3", false}, // out of alphabetical order
+			{"x.y/i", "v1.2.3", true},
+			{"x.y/j", "v1.2.3", true},
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", true},
+			{"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
+			x.y/c v1.2.3 // not v1.2.3!
+			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/b v1.2.3 // indirect; demoted to indirect
+			x.y/e v1.2.3 // indirect; comment kept
+			x.y/i v1.2.3 // indirect
+			x.y/j v1.2.3 // indirect
+		)
+		`,
+	},
+	{
+		`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
+		`,
+	},
+	{
+		`existing_duplicate_mix_indirect`,
+		`module m
+		require (
+			x.y/a v1.0.0 // zero
+			x.y/a v1.1.0 // indirect; one
+			x.y/a v1.2.3 // indirect; two
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", true},
+		},
+		`module m
+		require x.y/a v1.2.3 // indirect; zero
+		`,
+	},
+	{
+		`existing_duplicate_mix_direct`,
+		`module m
+		require (
+			x.y/a v1.0.0 // indirect; zero
+			x.y/a v1.1.0 // one
+			x.y/a v1.2.3 // two
+		)
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+		},
+		`module m
+		require x.y/a v1.2.3 // zero
+		`,
+	},
+	{
+		`add_indirect_after_last_direct`,
+		`module m
+		require (
+			x.y/a v1.0.0 // comment a preserved
+			x.y/d v1.0.0 // comment d preserved
+		)
+		require (
+			x.y/b v1.0.0 // comment b preserved
+			x.y/e v1.0.0 // comment e preserved
+		)
+		go 1.17
+		`,
+		[]require{
+			{"x.y/a", "v1.2.3", false},
+			{"x.y/b", "v1.2.3", false},
+			{"x.y/c", "v1.2.3", true},
+			{"x.y/d", "v1.2.3", false},
+			{"x.y/e", "v1.2.3", false},
+			{"x.y/f", "v1.2.3", true},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3 // comment a preserved
+			x.y/b v1.2.3 // comment b preserved
+			x.y/d v1.2.3 // comment d preserved
+			x.y/e v1.2.3 // comment e preserved
+		)
+		require (
+			x.y/c v1.2.3 // indirect
+			x.y/f v1.2.3 // indirect
+		)
+		go 1.17
+		`,
+	},
+	{
+		`add_direct_before_first_indirect`,
+		`module m
+		require (
+			x.y/b v1.0.0 // indirect; comment b preserved
+			x.y/e v1.0.0 // indirect; comment d preserved
+		)
+		require (
+			x.y/c v1.0.0 // indirect; comment c preserved
+			x.y/f v1.0.0 // indirect; comment e preserved
+		)
+		`,
+		[]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},
+			{"x.y/e", "v1.2.3", true},
+			{"x.y/f", "v1.2.3", true},
+		},
+		`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/c v1.2.3 // indirect; comment c preserved
+			x.y/e v1.2.3 // indirect; comment d preserved
+			x.y/f v1.2.3 // indirect; comment e preserved
+		)
+		`,
+	},
+	{
+		`add_indirect_after_mixed`,
+		`module m
+		require (
+			x.y/a v1.0.0
+			x.y/b 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},
+			{"x.y/e", "v1.2.3", true},
+		},
+		`module m
+		require (
+			x.y/a v1.2.3
+			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/d v1.2.3
+		)
+		require (
+			x.y/b v1.2.3 // indirect
+			x.y/c v1.2.3 // indirect
+		)`,
+	},
+}
+
 var addGoTests = []struct {
 	desc    string
 	in      string
@@ -1931,6 +2329,33 @@
 	}
 }
 
+func TestSetRequireAtMostTwo(t *testing.T) {
+	for _, tt := range setRequireSimplifyTests {
+		t.Run(tt.desc, func(t *testing.T) {
+			var mods []*Require
+			for _, mod := range tt.mods {
+				mods = append(mods, &Require{
+					Mod: module.Version{
+						Path:    mod.path,
+						Version: mod.vers,
+					},
+					Indirect: mod.indirect,
+				})
+			}
+
+			f := testEdit(t, tt.in, tt.out, true, func(f *File) error {
+				f.SetRequireAtMostTwo(mods)
+				f.Cleanup()
+				return nil
+			})
+
+			if len(f.Require) != len(mods) {
+				t.Errorf("after Cleanup, len(Require) = %v; want %v", len(f.Require), len(mods))
+			}
+		})
+	}
+}
+
 func TestAddGo(t *testing.T) {
 	for _, tt := range addGoTests {
 		t.Run(tt.desc, func(t *testing.T) {
@@ -2183,7 +2608,6 @@
 				t.Fatal(err)
 			}
 			got, err := f.Format()
-
 			if err != nil {
 				t.Fatal(err)
 			}