From 3d471b31c4ee3110343d3d14d875d20f97f50069 Mon Sep 17 00:00:00 2001
From: Vincent Ambo <tazjin@yandex-team.ru>
Date: Tue, 25 Jul 2023 21:52:36 +0300
Subject: [PATCH] lib/proto: add support for instantiating proto map fields

With this change, it becomes possible to instantiate `map<k,v>` type
fields in Protobuf messages from Starlark.

Maps can be constructed from any Starlark type that implements
starlark.IterableMapping.

Protobuf messages can have most types as keys/values, so the type
conformity is checked individually for each key/value pair (as the
Starlark side of things is dynamically typed).

This has been tested against fairly complex proto messages containing
maps.

Map operations apart from construction are not supported in this CL.
---
 lib/proto/proto.go | 40 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/lib/proto/proto.go b/lib/proto/proto.go
index 3ad37916..36f00233 100644
--- a/lib/proto/proto.go
+++ b/lib/proto/proto.go
@@ -395,7 +395,6 @@ func setField(msg protoreflect.Message, fdesc protoreflect.FieldDescriptor, valu
 		}
 		defer iter.Done()
 
-		// TODO(adonovan): handle maps
 		list := msg.Mutable(fdesc).List()
 		list.Truncate(0)
 		var x starlark.Value
@@ -409,6 +408,45 @@ func setField(msg protoreflect.Message, fdesc protoreflect.FieldDescriptor, valu
 		return nil
 	}
 
+	if fdesc.IsMap() {
+		mapping, ok := value.(starlark.IterableMapping)
+		if !ok {
+			return fmt.Errorf("in map field %s: expected mappable type, but got %s", fdesc.Name(), value.Type())
+		}
+
+		iter := mapping.Iterate()
+		defer iter.Done()
+
+		// Each value is converted using toProto as usual, passing the key/value
+		// field descriptors to check their types.
+		mutMap := msg.Mutable(fdesc).Map()
+		var k starlark.Value
+		for iter.Next(&k) {
+			kproto, err := toProto(fdesc.MapKey(), k)
+			if err != nil {
+				return fmt.Errorf("in key of map field %s: %w", fdesc.Name(), err)
+			}
+
+			// `found` is discarded, as the presence of the key in the
+			// iterator guarantees the presence of some value (even if it is
+			// starlark.None). Mismatching values will be caught in toProto
+			// below.
+			v, _, err := mapping.Get(k)
+			if err != nil {
+				return fmt.Errorf("in map field %s, at key %s: %w", fdesc.Name(), k.String(), err)
+			}
+
+			vproto, err := toProto(fdesc.MapValue(), v)
+			if err != nil {
+				return fmt.Errorf("in map field %s, at key %s: %w", fdesc.Name(), k.String(), err)
+			}
+
+			mutMap.Set(kproto.MapKey(), vproto)
+		}
+
+		return nil
+	}
+
 	v, err := toProto(fdesc, value)
 	if err != nil {
 		return fmt.Errorf("in field %s: %v", fdesc.Name(), err)