modfile: add support for go and toolchain lines

As part of the forward compatibility work, a new toolchain line
is being added, and go lines are allowed to specify toolchain
versions like "1.21.0" or "1.21rc1" now. (The lax RE has allowed this for quite
some time; what's new here is allowing it in the main module.)

For golang/go#57001.

Change-Id: I1dc01289381fe080644a7a391b97a65158938f39
Reviewed-on: https://go-review.googlesource.com/c/mod/+/497397
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Bryan Mills <bcmills@google.com>
Run-TryBot: Russ Cox <rsc@golang.org>
diff --git a/modfile/read_test.go b/modfile/read_test.go
index 82c778d..df4f117 100644
--- a/modfile/read_test.go
+++ b/modfile/read_test.go
@@ -434,13 +434,13 @@
 		{desc: "empty", input: "module m\ngo \n", ok: false},
 		{desc: "one", input: "module m\ngo 1\n", ok: false},
 		{desc: "two", input: "module m\ngo 1.22\n", ok: true},
-		{desc: "three", input: "module m\ngo 1.22.333", ok: false},
+		{desc: "three", input: "module m\ngo 1.22.333", ok: true},
 		{desc: "before", input: "module m\ngo v1.2\n", ok: false},
-		{desc: "after", input: "module m\ngo 1.2rc1\n", ok: false},
+		{desc: "after", input: "module m\ngo 1.2rc1\n", ok: true},
 		{desc: "space", input: "module m\ngo 1.2 3.4\n", ok: false},
-		{desc: "alt1", input: "module m\ngo 1.2.3\n", ok: false, laxOK: true},
-		{desc: "alt2", input: "module m\ngo 1.2rc1\n", ok: false, laxOK: true},
-		{desc: "alt3", input: "module m\ngo 1.2beta1\n", ok: false, laxOK: true},
+		{desc: "alt1", input: "module m\ngo 1.2.3\n", ok: true, laxOK: true},
+		{desc: "alt2", input: "module m\ngo 1.2rc1\n", ok: true, laxOK: true},
+		{desc: "alt3", input: "module m\ngo 1.2beta1\n", ok: true, laxOK: true},
 		{desc: "alt4", input: "module m\ngo 1.2.beta1\n", ok: false, laxOK: true},
 		{desc: "alt1", input: "module m\ngo v1.2.3\n", ok: false, laxOK: true},
 		{desc: "alt2", input: "module m\ngo v1.2rc1\n", ok: false, laxOK: true},
diff --git a/modfile/rule.go b/modfile/rule.go
index c20aef1..3306f6f 100644
--- a/modfile/rule.go
+++ b/modfile/rule.go
@@ -35,12 +35,13 @@
 
 // A File is the parsed, interpreted form of a go.mod file.
 type File struct {
-	Module  *Module
-	Go      *Go
-	Require []*Require
-	Exclude []*Exclude
-	Replace []*Replace
-	Retract []*Retract
+	Module    *Module
+	Go        *Go
+	Toolchain *Toolchain
+	Require   []*Require
+	Exclude   []*Exclude
+	Replace   []*Replace
+	Retract   []*Retract
 
 	Syntax *FileSyntax
 }
@@ -58,6 +59,12 @@
 	Syntax  *Line
 }
 
+// A Toolchain is the toolchain statement.
+type Toolchain struct {
+	Name   string // "go1.21rc1"
+	Syntax *Line
+}
+
 // An Exclude is a single exclude statement.
 type Exclude struct {
 	Mod    module.Version
@@ -296,9 +303,13 @@
 	return f, nil
 }
 
-var GoVersionRE = lazyregexp.New(`^([1-9][0-9]*)\.(0|[1-9][0-9]*)$`)
+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` or containing `-go1` as a substring,
+// like "go1.20.3" or "gccgo-go1.20.3". As a special case, "local" is also permitted.
+var ToolchainRE = lazyregexp.New(`^local$|(^|-)go1`)
+
 func (f *File) add(errs *ErrorList, block *LineBlock, line *Line, verb string, args []string, fix VersionFixer, strict bool) {
 	// If strict is false, this module is a dependency.
 	// We ignore all unknown directives as well as main-module-only
@@ -926,7 +937,7 @@
 
 func (f *File) AddGoStmt(version string) error {
 	if !GoVersionRE.MatchString(version) {
-		return fmt.Errorf("invalid language version string %q", version)
+		return fmt.Errorf("invalid language version %q", version)
 	}
 	if f.Go == nil {
 		var hint Expr
@@ -944,6 +955,28 @@
 	return nil
 }
 
+func (f *File) AddToolchainStmt(name string) error {
+	if !ToolchainRE.MatchString(name) {
+		return fmt.Errorf("invalid toolchain name %q", name)
+	}
+	if f.Toolchain == nil {
+		var hint Expr
+		if f.Go != nil && f.Go.Syntax != nil {
+			hint = f.Go.Syntax
+		} else if f.Module != nil && f.Module.Syntax != nil {
+			hint = f.Module.Syntax
+		}
+		f.Toolchain = &Toolchain{
+			Name:   name,
+			Syntax: f.Syntax.addLine(hint, "toolchain", name),
+		}
+	} else {
+		f.Toolchain.Name = name
+		f.Syntax.updateLine(f.Go.Syntax, "toolchain", name)
+	}
+	return nil
+}
+
 // AddRequire sets the first require line for path to version vers,
 // preserving any existing comments for that line and removing all
 // other lines for path.
diff --git a/modfile/work.go b/modfile/work.go
index 0c0e521..827a01d 100644
--- a/modfile/work.go
+++ b/modfile/work.go
@@ -12,9 +12,10 @@
 
 // A WorkFile is the parsed, interpreted form of a go.work file.
 type WorkFile struct {
-	Go      *Go
-	Use     []*Use
-	Replace []*Replace
+	Go        *Go
+	Toolchain *Toolchain
+	Use       []*Use
+	Replace   []*Replace
 
 	Syntax *FileSyntax
 }
@@ -109,7 +110,7 @@
 
 func (f *WorkFile) AddGoStmt(version string) error {
 	if !GoVersionRE.MatchString(version) {
-		return fmt.Errorf("invalid language version string %q", version)
+		return fmt.Errorf("invalid language version %q", version)
 	}
 	if f.Go == nil {
 		stmt := &Line{Token: []string{"go", version}}
@@ -117,7 +118,7 @@
 			Version: version,
 			Syntax:  stmt,
 		}
-		// Find the first non-comment-only block that's and add
+		// Find the first non-comment-only block and add
 		// the go statement before it. That will keep file comments at the top.
 		i := 0
 		for i = 0; i < len(f.Syntax.Stmt); i++ {
@@ -133,6 +134,40 @@
 	return nil
 }
 
+func (f *WorkFile) AddToolchainStmt(name string) error {
+	if !ToolchainRE.MatchString(name) {
+		return fmt.Errorf("invalid toolchain name %q", name)
+	}
+	if f.Toolchain == nil {
+		stmt := &Line{Token: []string{"toolchain", name}}
+		f.Toolchain = &Toolchain{
+			Name:   name,
+			Syntax: stmt,
+		}
+		// Find the go line and add the toolchain line after it.
+		// Or else find the first non-comment-only block and add
+		// the toolchain line before it. That will keep file comments at the top.
+		i := 0
+		for i = 0; i < len(f.Syntax.Stmt); i++ {
+			if line, ok := f.Syntax.Stmt[i].(*Line); ok && len(line.Token) > 0 && line.Token[0] == "go" {
+				i++
+				goto Found
+			}
+		}
+		for i = 0; i < len(f.Syntax.Stmt); i++ {
+			if _, ok := f.Syntax.Stmt[i].(*CommentBlock); !ok {
+				break
+			}
+		}
+	Found:
+		f.Syntax.Stmt = append(append(f.Syntax.Stmt[:i:i], stmt), f.Syntax.Stmt[i:]...)
+	} else {
+		f.Toolchain.Name = name
+		f.Syntax.updateLine(f.Toolchain.Syntax, "toolchain", name)
+	}
+	return nil
+}
+
 func (f *WorkFile) AddUse(diskPath, modulePath string) error {
 	need := true
 	for _, d := range f.Use {