go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/proto/protowalk/plan.go (about)

     1  // Copyright 2022 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 protowalk
    16  
    17  import (
    18  	"reflect"
    19  	"sort"
    20  
    21  	"google.golang.org/protobuf/reflect/protoreflect"
    22  
    23  	"go.chromium.org/luci/common/data/stringset"
    24  )
    25  
    26  type planItem struct {
    27  	processorIdx int
    28  	ProcessAttr
    29  	recurseAttr
    30  }
    31  
    32  type plan struct {
    33  	desc       protoreflect.MessageDescriptor
    34  	fieldOrder []protoreflect.FieldNumber
    35  	items      map[protoreflect.FieldNumber][]planItem
    36  }
    37  
    38  func (p plan) each(cb func(protoreflect.FieldDescriptor, []planItem)) {
    39  	fields := p.desc.Fields()
    40  	for _, fieldNum := range p.fieldOrder {
    41  		cb(fields.ByNumber(fieldNum), p.items[fieldNum])
    42  	}
    43  }
    44  
    45  type procBundle struct {
    46  	proc reflect.Type
    47  	sel  FieldSelector
    48  	fp   FieldProcessor
    49  }
    50  
    51  // makePlan generates a plan from msg+processors.
    52  //
    53  // The returned plan is traversable in a deterministic order (using the sorted
    54  // order of the affected Fields' tag numbers in the message).
    55  //
    56  // For each affected field, one or more processors will either need to directly
    57  // process the field, or will need to recurse deeper.
    58  func makePlan(desc protoreflect.MessageDescriptor, processors []*procBundle) (ret plan) {
    59  	size := desc.Fields().Len()
    60  	ret.fieldOrder = make([]protoreflect.FieldNumber, 0, size)
    61  	ret.desc = desc
    62  	ret.items = make(map[protoreflect.FieldNumber][]planItem, size)
    63  
    64  	// Now for each processor, go over its cached (msg+processor) data, and merge
    65  	// it into the overall plan.
    66  	for procIdx, proc := range processors {
    67  		for _, entry := range setCacheEntry(desc, proc, stringset.New(0), map[string]*cacheEntryBuilder{}) {
    68  			curPlan, ok := ret.items[entry.fieldNum]
    69  			if !ok { // need to add entry.field to ret.fieldOrder
    70  				foIdx := sort.Search(
    71  					len(ret.fieldOrder),
    72  					func(i int) bool {
    73  						return ret.fieldOrder[i] >= entry.fieldNum
    74  					})
    75  				// -1 is invalid, but gets overwritten; if we have a bug, -1 will cause an
    76  				// explosion later.
    77  				ret.fieldOrder = append(ret.fieldOrder, -1)
    78  				copy(ret.fieldOrder[foIdx+1:], ret.fieldOrder[foIdx:])
    79  				ret.fieldOrder[foIdx] = entry.fieldNum
    80  			}
    81  
    82  			// and fold the entry into the overall plan.
    83  			ret.items[entry.fieldNum] = append(curPlan, planItem{
    84  				processorIdx: procIdx,
    85  				ProcessAttr:  entry.ProcessAttr,
    86  				recurseAttr:  entry.recurseAttr,
    87  			})
    88  		}
    89  	}
    90  
    91  	return
    92  }