go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/starlark/starlarkproto/values_map.go (about)

     1  // Copyright 2019 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package starlarkproto
    16  
    17  import (
    18  	"fmt"
    19  	"sort"
    20  
    21  	"go.starlark.net/starlark"
    22  	"go.starlark.net/syntax"
    23  
    24  	"google.golang.org/protobuf/reflect/protoreflect"
    25  
    26  	"go.chromium.org/luci/starlark/typed"
    27  )
    28  
    29  // newStarlarkDict returns a new typed.Dict of an appropriate type.
    30  func newStarlarkDict(l *Loader, fd protoreflect.FieldDescriptor, capacity int) *typed.Dict {
    31  	return typed.NewDict(
    32  		converter(l, fd.MapKey()),
    33  		converter(l, fd.MapValue()),
    34  		capacity)
    35  }
    36  
    37  // toStarlarkDict does protoreflect.Map => typed.Dict conversion.
    38  //
    39  // Panics if type of 'm' (or some of its items) doesn't match 'fd'.
    40  func toStarlarkDict(l *Loader, fd protoreflect.FieldDescriptor, m protoreflect.Map) *typed.Dict {
    41  	keyFD := fd.MapKey()
    42  	valFD := fd.MapValue()
    43  
    44  	// A key of a protoreflect.Map, together with its Starlark copy.
    45  	type key struct {
    46  		pv protoreflect.MapKey // as proto value
    47  		sv starlark.Value      // as starlark value
    48  	}
    49  
    50  	// Protobuf maps are unordered, but Starlark dicts retain order, so collect
    51  	// keys first to sort them before adding to a dict.
    52  	keys := make([]key, 0, m.Len())
    53  	m.Range(func(k protoreflect.MapKey, _ protoreflect.Value) bool {
    54  		keys = append(keys, key{k, toStarlarkSingular(l, keyFD, k.Value())})
    55  		return true
    56  	})
    57  
    58  	// Sort keys as Starlark values. Proto map keys are (u)int(32|64) or bools,
    59  	// they *must* be sortable, so panic on errors.
    60  	sort.Slice(keys, func(i, j int) bool {
    61  		lt, err := starlark.Compare(syntax.LT, keys[i].sv, keys[j].sv)
    62  		if err != nil {
    63  			panic(fmt.Errorf("internal error: when sorting dict keys: %s", err))
    64  		}
    65  		return lt
    66  	})
    67  
    68  	// Construct typed.Dict from sorted keys and values from the proto map.
    69  	d := newStarlarkDict(l, fd, len(keys))
    70  	for _, k := range keys {
    71  		sv := toStarlarkSingular(l, valFD, m.Get(k.pv))
    72  		if err := d.SetKey(k.sv, sv); err != nil {
    73  			panic(fmt.Errorf("internal error: value of key %s: %s", k.pv, err))
    74  		}
    75  	}
    76  
    77  	return d
    78  }
    79  
    80  // assignProtoMap does m.fd = iter, where fd is a map.
    81  //
    82  // Assumes type checks have been done already (this is responsibility of
    83  // prepareRHS). Panics on type mismatch.
    84  func assignProtoMap(m protoreflect.Message, fd protoreflect.FieldDescriptor, dict *typed.Dict) {
    85  	mp := m.Mutable(fd).Map()
    86  	if mp.Len() != 0 {
    87  		panic(fmt.Errorf("internal error: the proto map is not empty"))
    88  	}
    89  
    90  	keyFD := fd.MapKey()
    91  	valFD := fd.MapValue()
    92  
    93  	it := dict.Iterate()
    94  	defer it.Done()
    95  
    96  	var sk starlark.Value
    97  	for it.Next(&sk) {
    98  		switch sv, ok, err := dict.Get(sk); {
    99  		case err != nil:
   100  			panic(fmt.Errorf("internal error: key %s: %s", sk, err))
   101  		case !ok:
   102  			panic(fmt.Errorf("internal error: key %s: suddenly gone while iterating", sk))
   103  		default:
   104  			mp.Set(toProtoSingular(keyFD, sk).MapKey(), toProtoSingular(valFD, sv))
   105  		}
   106  	}
   107  }