Fix custom resolver for map definition

internal/testutil/builder.go

@@ -1218,6 +1218,13 @@ func NewFieldRuleBuilder(value *resolver.Value) *FieldRuleBuilder {
 	return &FieldRuleBuilder{rule: &resolver.FieldRule{Value: value}}
 }
 
+func (b *FieldRuleBuilder) SetAutoBind(field *resolver.Field) *FieldRuleBuilder {
+	b.rule.AutoBindField = &resolver.AutoBindField{
+		Field: field,
+	}
+	return b
+}
+
 func (b *FieldRuleBuilder) SetCustomResolver(v bool) *FieldRuleBuilder {
 	b.rule.CustomResolver = v
 	return b

resolver/message.go

@@ -360,8 +360,13 @@ func (m *Message) CustomResolvers() []*CustomResolver {
 
 func (m *Message) customResolvers(resolver *MessageResolver) []*CustomResolver {
 	var ret []*CustomResolver
-	if def := resolver.VariableDefinition; def != nil && def.Expr.Message != nil {
-		ret = append(ret, def.Expr.Message.Message.CustomResolvers()...)
+	if def := resolver.VariableDefinition; def != nil {
+		for _, expr := range def.MessageExprs() {
+			if expr.Message == nil {
+				continue
+			}
+			ret = append(ret, expr.Message.CustomResolvers()...)
+		}
 	}
 	return ret
 }

resolver/resolver_test.go

@@ -3187,8 +3187,16 @@ func TestMap(t *testing.T) {
 		).
 		AddMessage(
 			testutil.NewMessageBuilder("User").
-				AddFieldWithAutoBind("id", resolver.StringType, ref.Field(t, "org.user", "User", "id")).
-				AddFieldWithAutoBind("name", resolver.StringType, ref.Field(t, "org.user", "User", "name")).
+				AddFieldWithRule(
+					"id",
+					resolver.StringType,
+					testutil.NewFieldRuleBuilder(nil).SetMessageCustomResolver(true).Build(t),
+				).
+				AddFieldWithRule(
+					"name",
+					resolver.StringType,
+					testutil.NewFieldRuleBuilder(nil).SetMessageCustomResolver(true).Build(t),
+				).
 				SetRule(
 					testutil.NewMessageRuleBuilder().
 						AddVariableDefinition(
@@ -3219,6 +3227,7 @@ func TestMap(t *testing.T) {
 								SetBy(testutil.NewCELValueBuilder("res.user", ref.Type(t, "org.user", "User")).Build(t)).
 								Build(t),
 						).
+						SetCustomResolver(true).
 						SetMessageArgument(ref.Message(t, "org.federation", "UserArgument")).
 						SetDependencyGraph(
 							testutil.NewDependencyGraphBuilder().

testdata/map.proto

@@ -83,6 +83,7 @@ message User {
       }
     }
     def { name: "user" by: "res.user" autobind: true }
+    custom_resolver: true
   };
   string id = 1;
   string name = 2;