support AssignExpr statements in edit.FindRuleByName()

edit/edit.go

@@ -305,7 +305,17 @@ func IndexOfRuleByName(f *build.File, name string) (int, *build.Rule) {
 	for i, stmt := range f.Stmt {
 		call, ok := stmt.(*build.CallExpr)
 		if !ok {
-			continue
+			// if stmt is not a CallExpr, check if it is an AssignExpr with a CallExpr
+			// on the right-hand side
+			as, ok := stmt.(*build.AssignExpr)
+			if !ok {
+				continue
+			}
+			if as.RHS != nil {
+				if call, ok = as.RHS.(*build.CallExpr); !ok {
+					continue
+				}
+			}
 		}
 		r := f.Rule(call)
 		start, _ := call.X.Span()

edit/edit_test.go

@@ -829,3 +829,35 @@ func TestInterpretLabelForWorkspaceLocation(t *testing.T) {
 	runTestInterpretLabelForWorkspaceLocation(t, "BUILD")
 	runTestInterpretLabelForWorkspaceLocation(t, "BUILD.bazel")
 }
+
+func TestFindRuleByName(t *testing.T) {
+	input := `load("foo.bzl", "bar")
+rule(
+  name="r1",
+	srcs=["a.txt"]
+)
+
+a = rule2(
+	name="r2",
+)
+`
+
+	bld, err := build.Parse("BUILD", []byte(input))
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	for _, name := range []string{"r1", "r2"} {
+		rule := FindRuleByName(bld, name)
+		if rule == nil {
+			t.Errorf("FindRuleByName: could not find rule with name '%s' in BUILD file <%s>", name, input)
+			return
+		}
+	}
+	// expect to not find:
+	rule := FindRuleByName(bld, "r3")
+	if rule != nil {
+		t.Errorf("FindRuleByName: found rule with name 'r3', expected not to find it")
+		return
+	}
+}