1
2 #ifndef UPBC_MESSAGE_LAYOUT_H
3 #define UPBC_MESSAGE_LAYOUT_H
4
5 #include "absl/base/macros.h"
6 #include "absl/container/flat_hash_map.h"
7 #include "google/protobuf/descriptor.h"
8
9 namespace upbc {
10
11 class MessageLayout {
12 public:
13 struct Size {
AddSize14 void Add(const Size& other) {
15 size32 += other.size32;
16 size64 += other.size64;
17 }
18
MaxFromSize19 void MaxFrom(const Size& other) {
20 size32 = std::max(size32, other.size32);
21 size64 = std::max(size64, other.size64);
22 }
23
AlignUpSize24 void AlignUp(const Size& align) {
25 size32 = Align(size32, align.size32);
26 size64 = Align(size64, align.size64);
27 }
28
29 int64_t size32;
30 int64_t size64;
31 };
32
33 struct SizeAndAlign {
34 Size size;
35 Size align;
36
MaxFromSizeAndAlign37 void MaxFrom(const SizeAndAlign& other) {
38 size.MaxFrom(other.size);
39 align.MaxFrom(other.align);
40 }
41 };
42
MessageLayout(const google::protobuf::Descriptor * descriptor)43 MessageLayout(const google::protobuf::Descriptor* descriptor) {
44 ComputeLayout(descriptor);
45 }
46
GetFieldOffset(const google::protobuf::FieldDescriptor * field)47 Size GetFieldOffset(const google::protobuf::FieldDescriptor* field) const {
48 return GetMapValue(field_offsets_, field);
49 }
50
GetOneofCaseOffset(const google::protobuf::OneofDescriptor * oneof)51 Size GetOneofCaseOffset(
52 const google::protobuf::OneofDescriptor* oneof) const {
53 return GetMapValue(oneof_case_offsets_, oneof);
54 }
55
GetHasbitIndex(const google::protobuf::FieldDescriptor * field)56 int GetHasbitIndex(const google::protobuf::FieldDescriptor* field) const {
57 return GetMapValue(hasbit_indexes_, field);
58 }
59
message_size()60 Size message_size() const { return size_; }
61
62 static bool HasHasbit(const google::protobuf::FieldDescriptor* field);
63 static SizeAndAlign SizeOfUnwrapped(
64 const google::protobuf::FieldDescriptor* field);
65
66 private:
67 void ComputeLayout(const google::protobuf::Descriptor* descriptor);
68 void PlaceNonOneofFields(const google::protobuf::Descriptor* descriptor);
69 void PlaceOneofFields(const google::protobuf::Descriptor* descriptor);
70 Size Place(SizeAndAlign size_and_align);
71
72 template <class K, class V>
GetMapValue(const absl::flat_hash_map<K,V> & map,K key)73 static V GetMapValue(const absl::flat_hash_map<K, V>& map, K key) {
74 auto iter = map.find(key);
75 if (iter == map.end()) {
76 fprintf(stderr, "No value for field.\n");
77 abort();
78 }
79 return iter->second;
80 }
81
IsPowerOfTwo(size_t val)82 static bool IsPowerOfTwo(size_t val) {
83 return (val & (val - 1)) == 0;
84 }
85
Align(size_t val,size_t align)86 static size_t Align(size_t val, size_t align) {
87 ABSL_ASSERT(IsPowerOfTwo(align));
88 return (val + align - 1) & ~(align - 1);
89 }
90
91 static SizeAndAlign SizeOf(const google::protobuf::FieldDescriptor* field);
92 static int64_t FieldLayoutRank(
93 const google::protobuf::FieldDescriptor* field);
94
95 absl::flat_hash_map<const google::protobuf::FieldDescriptor*, Size>
96 field_offsets_;
97 absl::flat_hash_map<const google::protobuf::FieldDescriptor*, int>
98 hasbit_indexes_;
99 absl::flat_hash_map<const google::protobuf::OneofDescriptor*, Size>
100 oneof_case_offsets_;
101 Size maxalign_;
102 Size size_;
103 };
104
105 // Returns fields in order of "hotness", eg. how frequently they appear in
106 // serialized payloads. Ideally this will use a profile. When we don't have
107 // that, we assume that fields with smaller numbers are used more frequently.
FieldHotnessOrder(const google::protobuf::Descriptor * message)108 inline std::vector<const google::protobuf::FieldDescriptor*> FieldHotnessOrder(
109 const google::protobuf::Descriptor* message) {
110 std::vector<const google::protobuf::FieldDescriptor*> fields;
111 for (int i = 0; i < message->field_count(); i++) {
112 fields.push_back(message->field(i));
113 }
114 std::sort(fields.begin(), fields.end(),
115 [](const google::protobuf::FieldDescriptor* a,
116 const google::protobuf::FieldDescriptor* b) {
117 return a->number() < b->number();
118 });
119 return fields;
120 }
121
122 } // namespace upbc
123
124 #endif // UPBC_MESSAGE_LAYOUT_H
125