1 // Copyright 2020 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
4 5 // Package order provides ordered access to messages and maps.
6 package order
7 8 import (
9 "sort"
10 "sync"
11 12 "google.golang.org/protobuf/reflect/protoreflect"
13 )
14 15 type messageField struct {
16 fd protoreflect.FieldDescriptor
17 v protoreflect.Value
18 }
19 20 var messageFieldPool = sync.Pool{
21 New: func() any { return new([]messageField) },
22 }
23 24 type (
25 // FieldRnger is an interface for visiting all fields in a message.
26 // The protoreflect.Message type implements this interface.
27 FieldRanger interface{ Range(VisitField) }
28 // VisitField is called every time a message field is visited.
29 VisitField = func(protoreflect.FieldDescriptor, protoreflect.Value) bool
30 )
31 32 // RangeFields iterates over the fields of fs according to the specified order.
33 func RangeFields(fs FieldRanger, less FieldOrder, fn VisitField) {
34 if less == nil {
35 fs.Range(fn)
36 return
37 }
38 39 // Obtain a pre-allocated scratch buffer.
40 p := messageFieldPool.Get().(*[]messageField)
41 fields := (*p)[:0]
42 defer func() {
43 if cap(fields) < 1024 {
44 *p = fields
45 messageFieldPool.Put(p)
46 }
47 }()
48 49 // Collect all fields in the message and sort them.
50 fs.Range(func(fd protoreflect.FieldDescriptor, v protoreflect.Value) bool {
51 fields = append(fields, messageField{fd, v})
52 return true
53 })
54 sort.Slice(fields, func(i, j int) bool {
55 return less(fields[i].fd, fields[j].fd)
56 })
57 58 // Visit the fields in the specified ordering.
59 for _, f := range fields {
60 if !fn(f.fd, f.v) {
61 return
62 }
63 }
64 }
65 66 type mapEntry struct {
67 k protoreflect.MapKey
68 v protoreflect.Value
69 }
70 71 var mapEntryPool = sync.Pool{
72 New: func() any { return new([]mapEntry) },
73 }
74 75 type (
76 // EntryRanger is an interface for visiting all fields in a message.
77 // The protoreflect.Map type implements this interface.
78 EntryRanger interface{ Range(VisitEntry) }
79 // VisitEntry is called every time a map entry is visited.
80 VisitEntry = func(protoreflect.MapKey, protoreflect.Value) bool
81 )
82 83 // RangeEntries iterates over the entries of es according to the specified order.
84 func RangeEntries(es EntryRanger, less KeyOrder, fn VisitEntry) {
85 if less == nil {
86 es.Range(fn)
87 return
88 }
89 90 // Obtain a pre-allocated scratch buffer.
91 p := mapEntryPool.Get().(*[]mapEntry)
92 entries := (*p)[:0]
93 defer func() {
94 if cap(entries) < 1024 {
95 *p = entries
96 mapEntryPool.Put(p)
97 }
98 }()
99 100 // Collect all entries in the map and sort them.
101 es.Range(func(k protoreflect.MapKey, v protoreflect.Value) bool {
102 entries = append(entries, mapEntry{k, v})
103 return true
104 })
105 sort.Slice(entries, func(i, j int) bool {
106 return less(entries[i].k, entries[j].k)
107 })
108 109 // Visit the entries in the specified ordering.
110 for _, e := range entries {
111 if !fn(e.k, e.v) {
112 return
113 }
114 }
115 }
116