range.go 2.7 KB

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