modfile: fix AddNewRequire block ordering

modfile/read.go

@@ -15,6 +15,9 @@ import (
 	"unicode/utf8"
 )
 
+// requireToken is the token used for require statements.
+const requireToken = "require"
+
 // A Position describes an arbitrary source position in a file, including the
 // file, line, column, and byte offset.
 type Position struct {
@@ -187,6 +190,108 @@ func (x *FileSyntax) addLine(hint Expr, tokens ...string) *Line {
 	return new
 }
 
+// requireHint returns a hint for adding a new require line.
+// It tries to maintain the standard order of require blocks, with
+// the first block being the direct requires and the second block
+// being the indirect requires.
+func (x *FileSyntax) requireHint(indirect bool) Expr {
+	if indirect {
+		// Indirect block requested return the first require
+		// block with indirect dependencies.
+		for _, stmt := range x.Stmt {
+			switch stmt := stmt.(type) {
+			case *Line:
+				if stmt.Token != nil && stmt.Token[0] == requireToken && isIndirect(stmt) {
+					return stmt
+				}
+			case *LineBlock:
+				if stmt.Token[0] == requireToken {
+					for _, line := range stmt.Line {
+						if isIndirect(line) {
+							return line
+						}
+					}
+				}
+			}
+		}
+
+		// No indirect require found, append a new block and return
+		// is as the hint. This prevents adding the indirect to an
+		// existing direct block.
+		block := &LineBlock{
+			Token: []string{requireToken},
+		}
+		x.Stmt = append(x.Stmt, block)
+
+		return block
+	}
+
+	// Direct block requested.
+	var last Expr
+	for _, stmt := range x.Stmt {
+		switch stmt := stmt.(type) {
+		case *Line:
+			if stmt.Token != nil && stmt.Token[0] == requireToken {
+				if !isIndirect(stmt) {
+					// Direct require line found, return it as hint to
+					// combine with it.
+					return stmt
+				}
+
+				// Indirect line first, which is unexpected. We return
+				// the last stmt as hint to create a new require block
+				// before it, to try to maintain the standard order.
+				if last != nil {
+					return last
+				}
+
+				// Indirect line at the beginning of the file, prepend
+				// a new require block and return it as hint, to try to
+				// maintain the standard order.
+				block := &LineBlock{
+					Token: []string{requireToken},
+				}
+				x.Stmt = append([]Expr{block}, x.Stmt...)
+
+				return block
+			}
+		case *LineBlock:
+			if stmt.Token[0] == requireToken {
+				for _, line := range stmt.Line {
+					if !isIndirect(line) {
+						// Direct require block found, return it as hint to
+						// combine with it.
+						return line
+					}
+				}
+
+				// Indirect block before first, which is unexpected.
+				// Return the last stmt as hint to create a new require
+				// block before it to try to maintain the standard order.
+				if last != nil {
+					return last
+				}
+
+				// Indirect block at the beginning of the file, which is
+				// unexpected. Prepend a new require block and return it
+				// as hint to try to maintain the standard order.
+				block := &LineBlock{
+					Token: []string{requireToken},
+				}
+				x.Stmt = append([]Expr{block}, x.Stmt...)
+
+				return block
+			}
+		}
+
+		last = stmt
+		continue
+	}
+
+	// No requires found, addLine will create one.
+	return nil
+}
+
 func (x *FileSyntax) updateLine(line *Line, tokens ...string) {
 	if line.InBlock {
 		tokens = tokens[1:]

modfile/rule.go

@@ -1144,7 +1144,7 @@ func (f *File) addNewGodebug(key, value string) {
 // other lines for path.
 //
 // If no line currently exists for path, AddRequire adds a new line
-// at the end of the last require block.
+// using AddNewRequire.
 func (f *File) AddRequire(path, vers string) error {
 	need := true
 	for _, r := range f.Require {
@@ -1166,10 +1166,14 @@ func (f *File) AddRequire(path, vers string) error {
 	return nil
 }
 
-// AddNewRequire adds a new require line for path at version vers at the end of
-// the last require block, regardless of any existing require lines for path.
+// AddNewRequire adds a new require line for path at version vers, regardless
+// of any existing require lines for path.
+// Similar to SetRequireSeparateIndirect it attempts to maintain the standard
+// of having direct requires in the first block and indirect requires in the
+// second block.
 func (f *File) AddNewRequire(path, vers string, indirect bool) {
-	line := f.Syntax.addLine(nil, "require", AutoQuote(path), vers)
+	hint := f.Syntax.requireHint(indirect)
+	line := f.Syntax.addLine(hint, "require", AutoQuote(path), vers)
 	r := &Require{
 		Mod:    module.Version{Path: path, Version: vers},
 		Syntax: line,
@@ -1248,6 +1252,12 @@ func (f *File) SetRequire(req []*Require) {
 // 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.
+//
+// New entries in req must not be added to the file.Require slice before calling
+// otherwise this function will panic. So the calling convention should be along
+// the lines of:
+//
+//	File.SetRequireSeparateIndirect(append(File.Require, newReqs...))
 func (f *File) SetRequireSeparateIndirect(req []*Require) {
 	// hasComments returns whether a line or block has comments
 	// other than "indirect".