From 4da1db7989714dd43bf9c2e25f391d5c7ddf325d Mon Sep 17 00:00:00 2001
From: Maxim Tschumak <tschumak@google.com>
Date: Thu, 11 Nov 2021 14:15:24 +0100
Subject: [PATCH] feat: add MethodOption location function

---
 locations/method_locations.go      |  7 +++++++
 locations/method_locations_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/locations/method_locations.go b/locations/method_locations.go
index 116f7f38b..bcf2b094f 100644
--- a/locations/method_locations.go
+++ b/locations/method_locations.go
@@ -19,6 +19,7 @@ import (
 	"github.com/jhump/protoreflect/desc"
 	apb "google.golang.org/genproto/googleapis/api/annotations"
 	lrpb "google.golang.org/genproto/googleapis/longrunning"
+	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 // MethodRequestType returns the precise location of the method's input type.
@@ -48,3 +49,9 @@ func MethodOperationInfo(m *desc.MethodDescriptor) *dpb.SourceCodeInfo_Location
 func MethodSignature(m *desc.MethodDescriptor, index int) *dpb.SourceCodeInfo_Location {
 	return pathLocation(m, 4, int(apb.E_MethodSignature.TypeDescriptor().Number()), index) // MethodDescriptor.options == 4
 }
+
+// MethodOption returns the precise location of the method's option with the given name, if any.
+func MethodOption(m *desc.MethodDescriptor, option string) *dpb.SourceCodeInfo_Location {
+	fd := m.GetMethodOptions().ProtoReflect().Descriptor().Fields().ByName(protoreflect.Name(option))
+	return pathLocation(m, 4, int(fd.Number())) // MethodDescriptor.options == 4
+}
diff --git a/locations/method_locations_test.go b/locations/method_locations_test.go
index 48b40e579..3b52a899d 100644
--- a/locations/method_locations_test.go
+++ b/locations/method_locations_test.go
@@ -117,3 +117,33 @@ func TestMethodSignature(t *testing.T) {
 		}
 	}
 }
+
+func TestMethodOption(t *testing.T) {
+	f := parse(t, `
+	  import "google/api/client.proto";
+		service Library {
+		  rpc GetBook(GetBookRequest) returns (Book) {
+		    option deadline = 60.0;
+		  }
+		  rpc WriteBook(NewBook) returns (Book) {}
+		}
+		message GetBookRequest{}
+		message Book {}
+		message NewBook {}
+	`)
+
+	for _, test := range []struct {
+		name      string
+		methodIdx int
+		option    string
+		want      []int32
+	}{
+		{"First", 0, "deadline", []int32{5, 20, 43}},
+		{"Second", 1, "deadline", nil},
+	} {
+		loc := MethodOption(f.GetServices()[0].GetMethods()[test.methodIdx], test.option)
+		if diff := cmp.Diff(loc.GetSpan(), test.want); diff != "" {
+			t.Errorf("%s: %s", test.name, diff)
+		}
+	}
+}