setapi: add/remove go_features.proto import where needed

fixes golang/open2opaque#5

PiperOrigin-RevId: 707487761
Change-Id: I5423f3a236eb08e66260e293ed75fe31a09c3e7a
Reviewed-on: https://go-review.googlesource.com/c/open2opaque/+/636616
Auto-Submit: Michael Stapelberg <stapelberg@google.com>
Reviewed-by: Michael Stapelberg <stapelberg@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Chressie Himpel <chressie@google.com>
diff --git a/internal/o2o/setapi/setapi.go b/internal/o2o/setapi/setapi.go
index 1143a8b..49e8c2c 100644
--- a/internal/o2o/setapi/setapi.go
+++ b/internal/o2o/setapi/setapi.go
@@ -433,15 +433,11 @@
 		if err != nil {
 			return nil, fmt.Errorf("fileOptionLineNumber: %v", err)
 		}
-		lines := bytes.Split(content, []byte{'\n'})
-		insertLine := fmt.Sprintf("option go_api_flag = %q;", fromFeatureToOld(targetAPI))
+		levelLine := fmt.Sprintf("option go_api_flag = %q;", fromFeatureToOld(targetAPI))
 		if fopt.Syntax == "editions" {
-			insertLine = fmt.Sprintf("option features.(pb.go).api_level = %s;", targetAPI)
+			levelLine = fmt.Sprintf("option features.(pb.go).api_level = %s;", targetAPI)
 		}
-		result := slices.Clone(lines[:ln])
-		result = append(result, []byte(insertLine))
-		result = append(result, lines[ln:]...)
-		return bytes.Join(result, []byte{'\n'}), nil
+		return insertLine(content, levelLine, ln), nil
 	}
 
 	// File is currently explicitly set to API value and target API isn't the
@@ -630,7 +626,7 @@
 	// case of editions proto that use the new edition feature, the parent of
 	// a nested message is its parent message.
 	//
-	// Parse again after the previous cleanup might have modified conten.
+	// Parse again after the previous cleanup might have modified content.
 	fopt, err = parse(path, content, false)
 	if err != nil {
 		return nil, err
@@ -670,22 +666,117 @@
 			return nil, err
 		}
 	}
-	if len(removeByteRanges) == 0 {
-		return content, nil
-	}
+	if len(removeByteRanges) > 0 {
+		slices.SortFunc(removeByteRanges, func(a, b byteRange) int {
+			return cmp.Compare(a.from, b.from)
+		})
+		for i := 1; i < len(removeByteRanges); i++ {
+			if removeByteRanges[i].from < removeByteRanges[i-1].to {
+				return nil, fmt.Errorf("text ranges overlap")
+			}
+		}
 
-	slices.SortFunc(removeByteRanges, func(a, b byteRange) int {
-		return cmp.Compare(a.from, b.from)
-	})
-	for i := 1; i < len(removeByteRanges); i++ {
-		if removeByteRanges[i].from < removeByteRanges[i-1].to {
-			return nil, fmt.Errorf("text ranges overlap")
+		for _, br := range slices.Backward(removeByteRanges) {
+			content = slices.Delete(content, br.from, br.to)
 		}
 	}
 
-	for _, br := range slices.Backward(removeByteRanges) {
-		content = slices.Delete(content, br.from, br.to)
+	// Cleanup 3: ensure the go_features.proto import is present if and only if
+	// any parts of the file set the features.(pb.go).api_level option.
+	return cleanupGoFeaturesImport(path, content)
+}
+
+func insertLine(content []byte, insertLine string, ln int32) []byte {
+	lines := bytes.Split(content, []byte{'\n'})
+	result := slices.Clone(lines[:ln])
+	result = append(result, []byte(insertLine))
+	result = append(result, lines[ln:]...)
+	return bytes.Join(result, []byte{'\n'})
+}
+
+func messageUsesFeature(mopt *protoparse.MessageOpt) bool {
+	if mopt.IsExplicit {
+		return true // editions feature set at message level
 	}
+	for _, mo := range mopt.Children {
+		if messageUsesFeature(mo) {
+			return true
+		}
+	}
+	return false
+}
+
+func usesFeature(fopt *protoparse.FileOpt) bool {
+	if fopt.Syntax != "editions" {
+		return false // only files on editions can use editions features
+	}
+	if fopt.IsExplicit {
+		return true // editions feature set at file level
+	}
+	for _, mo := range fopt.MessageOpts {
+		if messageUsesFeature(mo) {
+			return true
+		}
+	}
+	return false
+}
+
+func cleanupGoFeaturesImport(path string, content []byte) ([]byte, error) {
+	const featuresProto = "google/protobuf/go_features.proto"
+	parser := protoparse.NewParserWithAccessor(func(string) (io.ReadCloser, error) {
+		return io.NopCloser(bytes.NewReader(content)), nil
+	})
+	fopt, err := parser.ParseFile(path, false)
+	if err != nil {
+		return nil, fmt.Errorf("error reading file %s: %w", path, err)
+	}
+	desc := fopt.Desc
+	goFeaturesImported := slices.ContainsFunc(
+		desc.GetDependency(),
+		func(path string) bool {
+			return path == featuresProto
+		})
+	usesFeature := usesFeature(fopt)
+
+	if usesFeature && !goFeaturesImported {
+		ln, err := featuresImportLineNumber(desc.GetSourceCodeInfo())
+		if err != nil {
+			return nil, err
+		}
+		importLine := fmt.Sprintf("import %q;", featuresProto)
+		return insertLine(content, importLine, ln), nil
+	}
+
+	if !usesFeature && goFeaturesImported {
+		importNum := 0
+		for _, loc := range desc.GetSourceCodeInfo().GetLocation() {
+			path := loc.GetPath()
+			span := loc.GetSpan()
+			if len(path) < 1 {
+				continue
+			}
+			if path[0] != importFieldNum {
+				continue
+			}
+			// Found an import. Is it the go_features.proto import?
+			if max := len(desc.GetDependency()); importNum >= max {
+				return nil, fmt.Errorf("BUG: too many imports in SourceCodeInfo: got index %d, want at most %d", importNum, max)
+			}
+			imported := desc.GetDependency()[importNum]
+			importNum++
+			if imported != featuresProto {
+				continue
+			}
+			ln := spanEndLine(span)
+			lines := bytes.Split(content, []byte{'\n'})
+			result := slices.Clone(lines[:ln])
+			// lines[ln] deleted by omitting it.
+			result = append(result, lines[ln+1:]...)
+			return bytes.Join(result, []byte{'\n'}), nil
+		}
+		return nil, fmt.Errorf("BUG: could not locate import line in SourceCodeInfo")
+	}
+
 	return content, nil
 }
 
@@ -697,26 +788,17 @@
 	return span[0]
 }
 
-// fileOptionLineNumber returns the line number to insert the file-level
-// go_api_flag option. It uses the following heuristics to determine the
-// line number:
-// If there are any file option settings, return the line number after the last
-// one.
-// If there are any import lines, return the line number after the last one.
-// If there is a package statement line, return the line number after it.
-// If there is a syntax statement line, return the line number after it.
-// Note that this func assumes that an earlier protoparser.Parser.ParseFile
-// invocation will return error already if there is no syntax statement.
-func fileOptionLineNumber(info *descpb.SourceCodeInfo) (int32, error) {
-	const (
-		// https://github.com/protocolbuffers/protobuf/blob/v29.1/src/google/protobuf/descriptor.proto#L105-L137
-		syntaxFieldNum      = 12
-		editionFieldNum     = 14
-		editionDeprFieldNum = 13
-		packageFieldNum     = 2
-		importFieldNum      = 3
-		optionFieldNum      = 8
-	)
+const (
+	// https://github.com/protocolbuffers/protobuf/blob/v29.1/src/google/protobuf/descriptor.proto#L105-L137
+	syntaxFieldNum      = 12
+	editionFieldNum     = 14
+	editionDeprFieldNum = 13
+	packageFieldNum     = 2
+	importFieldNum      = 3
+	optionFieldNum      = 8
+)
+
+func fieldNumToLine(info *descpb.SourceCodeInfo) map[int32]int32 {
 	// pos contains a mapping of proto field numbers to line numbers. For syntax
 	// and package constructs, the line numbers represent where they are declared.
 	// For import and options, the line number represents the last import/option
@@ -737,6 +819,51 @@
 			}
 		}
 	}
+	return pos
+}
+
+// featuresImportLineNumber returns the line number to insert the file-level
+// go_features.proto import. It uses the following heuristics to determine the
+// line number:
+// If there are any import lines, return the line number after the last one.
+// If there is a package statement line, return the line number after it.
+// If there is a syntax statement line, return the line number after it.
+// Note that this func assumes that an earlier protoparser.Parser.ParseFile
+// invocation will return error already if there is no syntax statement.
+func featuresImportLineNumber(info *descpb.SourceCodeInfo) (int32, error) {
+	pos := fieldNumToLine(info)
+
+	if lnum, ok := pos[importFieldNum]; ok {
+		return lnum + 1, nil
+	}
+	if lnum, ok := pos[packageFieldNum]; ok {
+		return lnum + 1, nil
+	}
+	if lnum, ok := pos[syntaxFieldNum]; ok {
+		return lnum + 1, nil
+	}
+	if lnum, ok := pos[editionFieldNum]; ok {
+		return lnum + 1, nil
+	}
+	if lnum, ok := pos[editionDeprFieldNum]; ok {
+		return lnum + 1, nil
+	}
+
+	return 0, fmt.Errorf("cannot determine line number for go_features.proto import")
+}
+
+// fileOptionLineNumber returns the line number to insert the file-level
+// go_api_flag option. It uses the following heuristics to determine the
+// line number:
+// If there are any file option settings, return the line number after the last
+// one.
+// If there are any import lines, return the line number after the last one.
+// If there is a package statement line, return the line number after it.
+// If there is a syntax statement line, return the line number after it.
+// Note that this func assumes that an earlier protoparser.Parser.ParseFile
+// invocation will return error already if there is no syntax statement.
+func fileOptionLineNumber(info *descpb.SourceCodeInfo) (int32, error) {
+	pos := fieldNumToLine(info)
 
 	if lnum, ok := pos[optionFieldNum]; ok {
 		return lnum + 1, nil
diff --git a/internal/o2o/setapi/setapi_test.go b/internal/o2o/setapi/setapi_test.go
index b70ad49..598d4b0 100644
--- a/internal/o2o/setapi/setapi_test.go
+++ b/internal/o2o/setapi/setapi_test.go
@@ -508,6 +508,8 @@
 			protoIn: `
 edition = "2023";
 package pkg;
+import "some/unrelated/import.proto";
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -515,6 +517,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "some/unrelated/import.proto";
 
 message M{}
 `,
@@ -531,6 +534,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -545,6 +549,7 @@
 			targetAPI: gofeaturespb.GoFeatures_API_HYBRID,
 			protoWant: `
 edition = "2023";
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -562,6 +567,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 option java_package = "foo";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
@@ -581,6 +587,7 @@
 edition = "2023";
 package pkg;
 import "foo.proto";
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -597,6 +604,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -611,6 +619,7 @@
 			targetAPI: gofeaturespb.GoFeatures_API_HYBRID,
 			protoWant: `
 edition = "2023";
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -628,6 +637,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -645,6 +655,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
 `,
@@ -662,6 +673,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 option features.(pb.go).api_level = API_OPEN;
 message M{}
 `,
@@ -694,6 +706,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 // leading comment
 option features.(pb.go).api_level = API_HYBRID;
 message M{}
@@ -727,6 +740,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 // leading comment
 option features.(pb.go).api_level = API_OPAQUE;
 message M{}
@@ -881,6 +895,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
   message A1{
     
@@ -915,6 +930,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
   message A1{
     option features.(pb.go).api_level = API_HYBRID; // end-of-line comment
@@ -949,6 +965,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
   message A1{
     option features.(pb.go).api_level = API_OPEN; // end-of-line comment
@@ -987,6 +1004,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
   message A1{
     
@@ -1030,6 +1048,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
   message A1{
     option features.(pb.go).api_level = API_OPEN;
@@ -1151,6 +1170,7 @@
 			protoIn: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
 option features.(pb.go).api_level = API_OPAQUE; // end-of-line comment
 message A1{
@@ -1168,6 +1188,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
 
 message A1{
@@ -1206,6 +1227,7 @@
 			protoWant: `
 edition = "2023";
 package pkg;
+import "google/protobuf/go_features.proto";
 message A{
 
 message A1{
diff --git a/internal/protoparse/protoparse.go b/internal/protoparse/protoparse.go
index 892a0b9..5663d01 100644
--- a/internal/protoparse/protoparse.go
+++ b/internal/protoparse/protoparse.go
@@ -129,6 +129,8 @@
 	// Options of messages defined at the file level. Nested messages are stored
 	// as their children.
 	MessageOpts []*MessageOpt
+	// Desc is the descriptor proto of the parsed file.
+	Desc *descpb.FileDescriptorProto
 	// SourceCodeInfo is set if parsed results includes it.
 	SourceCodeInfo *descpb.SourceCodeInfo
 	// Proto syntax: "proto2", "proto3", "editions", or "editions_go_api_flag".
@@ -308,6 +310,7 @@
 		IsExplicit:     explicit,
 		APIInfo:        info,
 		MessageOpts:    mopts,
+		Desc:           desc,
 		SourceCodeInfo: desc.GetSourceCodeInfo(),
 		Syntax:         syntax,
 	}, nil
diff --git a/internal/protoparse/protoparse_test.go b/internal/protoparse/protoparse_test.go
index c77a84c..19245e1 100644
--- a/internal/protoparse/protoparse_test.go
+++ b/internal/protoparse/protoparse_test.go
@@ -132,6 +132,7 @@
 				cmpopts.IgnoreFields(protoparse.FileOpt{}, "APIInfo"),
 				cmpopts.IgnoreFields(protoparse.FileOpt{}, "SourceCodeInfo"),
 				cmpopts.IgnoreFields(protoparse.FileOpt{}, "MessageOpts"),
+				cmpopts.IgnoreFields(protoparse.FileOpt{}, "Desc"),
 			}
 			if diff := cmp.Diff(tc.wantOpt, got, ingores...); diff != "" {
 				t.Errorf("parser.ParseFile(): diff (-want +got):\n%v\n", diff)