1 #include "rust/cpp_kernel/map.h"
2
3 #include <cstddef>
4 #include <cstdint>
5 #include <string>
6 #include <type_traits>
7 #include <utility>
8
9 #include "google/protobuf/map.h"
10 #include "google/protobuf/message.h"
11 #include "google/protobuf/message_lite.h"
12 #include "rust/cpp_kernel/strings.h"
13
14 namespace google {
15 namespace protobuf {
16 namespace rust {
17 namespace {
18
19 template <typename T>
20 struct FromViewType {
21 using type = T;
22 };
23
24 template <>
25 struct FromViewType<PtrAndLen> {
26 using type = std::string;
27 };
28
29 template <typename Key>
30 using KeyMap = internal::KeyMapBase<
31 internal::KeyForBase<typename FromViewType<Key>::type>>;
32
33 template <typename Key>
DestroyMapNode(internal::UntypedMapBase * m,internal::NodeBase * node,internal::MapNodeSizeInfoT size_info)34 void DestroyMapNode(internal::UntypedMapBase* m, internal::NodeBase* node,
35 internal::MapNodeSizeInfoT size_info) {
36 if constexpr (std::is_same<Key, PtrAndLen>::value) {
37 static_cast<std::string*>(node->GetVoidKey())->~basic_string();
38 }
39 internal::RustMapHelper::DestroyMessage(
40 static_cast<MessageLite*>(node->GetVoidValue(size_info)));
41 internal::RustMapHelper::DeallocNode(m, node, size_info);
42 }
43
44 template <typename Key>
Insert(internal::UntypedMapBase * m,internal::MapNodeSizeInfoT size_info,Key key,MessageLite * value)45 bool Insert(internal::UntypedMapBase* m, internal::MapNodeSizeInfoT size_info,
46 Key key, MessageLite* value) {
47 internal::NodeBase* node = internal::RustMapHelper::AllocNode(m, size_info);
48 if constexpr (std::is_same<Key, PtrAndLen>::value) {
49 new (node->GetVoidKey()) std::string(key.ptr, key.len);
50 } else {
51 *static_cast<Key*>(node->GetVoidKey()) = key;
52 }
53
54 MessageLite* new_msg = internal::RustMapHelper::PlacementNew(
55 value, node->GetVoidValue(size_info));
56 auto* full_msg = DynamicCastMessage<Message>(new_msg);
57
58 // If we are working with a full (non-lite) proto, we reflectively swap the
59 // value into place. Otherwise, we have to perform a copy.
60 if (full_msg != nullptr) {
61 full_msg->GetReflection()->Swap(full_msg,
62 DynamicCastMessage<Message>(value));
63 } else {
64 new_msg->CheckTypeAndMergeFrom(*value);
65 }
66
67 node = internal::RustMapHelper::InsertOrReplaceNode(
68 static_cast<KeyMap<Key>*>(m), node);
69 if (node == nullptr) {
70 return true;
71 }
72 DestroyMapNode<Key>(m, node, size_info);
73 return false;
74 }
75
76 template <typename Map, typename Key,
77 typename = typename std::enable_if<
78 !std::is_same<Key, google::protobuf::rust::PtrAndLen>::value>::type>
FindHelper(Map * m,Key key)79 internal::RustMapHelper::NodeAndBucket FindHelper(Map* m, Key key) {
80 return internal::RustMapHelper::FindHelper(
81 m, static_cast<internal::KeyForBase<Key>>(key));
82 }
83
84 template <typename Map>
FindHelper(Map * m,google::protobuf::rust::PtrAndLen key)85 internal::RustMapHelper::NodeAndBucket FindHelper(Map* m,
86 google::protobuf::rust::PtrAndLen key) {
87 return internal::RustMapHelper::FindHelper(
88 m, absl::string_view(key.ptr, key.len));
89 }
90
91 template <typename Key>
Get(internal::UntypedMapBase * m,internal::MapNodeSizeInfoT size_info,Key key,MessageLite ** value)92 bool Get(internal::UntypedMapBase* m, internal::MapNodeSizeInfoT size_info,
93 Key key, MessageLite** value) {
94 auto* map_base = static_cast<KeyMap<Key>*>(m);
95 internal::RustMapHelper::NodeAndBucket result = FindHelper(map_base, key);
96 if (result.node == nullptr) {
97 return false;
98 }
99 *value = static_cast<MessageLite*>(result.node->GetVoidValue(size_info));
100 return true;
101 }
102
103 template <typename Key>
Remove(internal::UntypedMapBase * m,internal::MapNodeSizeInfoT size_info,Key key)104 bool Remove(internal::UntypedMapBase* m, internal::MapNodeSizeInfoT size_info,
105 Key key) {
106 auto* map_base = static_cast<KeyMap<Key>*>(m);
107 internal::RustMapHelper::NodeAndBucket result = FindHelper(map_base, key);
108 if (result.node == nullptr) {
109 return false;
110 }
111 internal::RustMapHelper::EraseNoDestroy(map_base, result.bucket, result.node);
112 DestroyMapNode<Key>(m, result.node, size_info);
113 return true;
114 }
115
116 template <typename Key>
IterGet(const internal::UntypedMapIterator * iter,internal::MapNodeSizeInfoT size_info,Key * key,MessageLite ** value)117 void IterGet(const internal::UntypedMapIterator* iter,
118 internal::MapNodeSizeInfoT size_info, Key* key,
119 MessageLite** value) {
120 internal::NodeBase* node = iter->node_;
121 if constexpr (std::is_same<Key, PtrAndLen>::value) {
122 const std::string* s = static_cast<const std::string*>(node->GetVoidKey());
123 *key = PtrAndLen{s->data(), s->size()};
124 } else {
125 *key = *static_cast<const Key*>(node->GetVoidKey());
126 }
127 *value = static_cast<MessageLite*>(node->GetVoidValue(size_info));
128 }
129
ClearMap(internal::UntypedMapBase * m,internal::MapNodeSizeInfoT size_info,bool key_is_string,bool reset_table)130 void ClearMap(internal::UntypedMapBase* m, internal::MapNodeSizeInfoT size_info,
131 bool key_is_string, bool reset_table) {
132 if (internal::RustMapHelper::IsGlobalEmptyTable(m)) return;
133 uint8_t bits = internal::RustMapHelper::kValueIsProto;
134 if (key_is_string) {
135 bits |= internal::RustMapHelper::kKeyIsString;
136 }
137 internal::RustMapHelper::ClearTable(
138 m, internal::RustMapHelper::ClearInput{size_info, bits, reset_table,
139 /* destroy_node = */ nullptr});
140 }
141
142 } // namespace
143 } // namespace rust
144 } // namespace protobuf
145 } // namespace google
146
147 extern "C" {
148
proto2_rust_thunk_UntypedMapIterator_increment(google::protobuf::internal::UntypedMapIterator * iter)149 void proto2_rust_thunk_UntypedMapIterator_increment(
150 google::protobuf::internal::UntypedMapIterator* iter) {
151 iter->PlusPlus();
152 }
153
proto2_rust_map_new()154 google::protobuf::internal::UntypedMapBase* proto2_rust_map_new() {
155 return new google::protobuf::internal::UntypedMapBase(/* arena = */ nullptr);
156 }
157
proto2_rust_map_free(google::protobuf::internal::UntypedMapBase * m,bool key_is_string,google::protobuf::internal::MapNodeSizeInfoT size_info)158 void proto2_rust_map_free(google::protobuf::internal::UntypedMapBase* m,
159 bool key_is_string,
160 google::protobuf::internal::MapNodeSizeInfoT size_info) {
161 google::protobuf::rust::ClearMap(m, size_info, key_is_string,
162 /* reset_table = */ false);
163 delete m;
164 }
165
proto2_rust_map_clear(google::protobuf::internal::UntypedMapBase * m,bool key_is_string,google::protobuf::internal::MapNodeSizeInfoT size_info)166 void proto2_rust_map_clear(google::protobuf::internal::UntypedMapBase* m,
167 bool key_is_string,
168 google::protobuf::internal::MapNodeSizeInfoT size_info) {
169 google::protobuf::rust::ClearMap(m, size_info, key_is_string, /* reset_table = */ true);
170 }
171
proto2_rust_map_size(google::protobuf::internal::UntypedMapBase * m)172 size_t proto2_rust_map_size(google::protobuf::internal::UntypedMapBase* m) {
173 return m->size();
174 }
175
proto2_rust_map_iter(google::protobuf::internal::UntypedMapBase * m)176 google::protobuf::internal::UntypedMapIterator proto2_rust_map_iter(
177 google::protobuf::internal::UntypedMapBase* m) {
178 return m->begin();
179 }
180
181 #define DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(cpp_type, suffix) \
182 bool proto2_rust_map_insert_##suffix( \
183 google::protobuf::internal::UntypedMapBase* m, \
184 google::protobuf::internal::MapNodeSizeInfoT size_info, cpp_type key, \
185 google::protobuf::MessageLite* value) { \
186 return google::protobuf::rust::Insert(m, size_info, key, value); \
187 } \
188 \
189 bool proto2_rust_map_get_##suffix( \
190 google::protobuf::internal::UntypedMapBase* m, \
191 google::protobuf::internal::MapNodeSizeInfoT size_info, cpp_type key, \
192 google::protobuf::MessageLite** value) { \
193 return google::protobuf::rust::Get(m, size_info, key, value); \
194 } \
195 \
196 bool proto2_rust_map_remove_##suffix( \
197 google::protobuf::internal::UntypedMapBase* m, \
198 google::protobuf::internal::MapNodeSizeInfoT size_info, cpp_type key) { \
199 return google::protobuf::rust::Remove(m, size_info, key); \
200 } \
201 \
202 void proto2_rust_map_iter_get_##suffix( \
203 const google::protobuf::internal::UntypedMapIterator* iter, \
204 google::protobuf::internal::MapNodeSizeInfoT size_info, cpp_type* key, \
205 google::protobuf::MessageLite** value) { \
206 return google::protobuf::rust::IterGet(iter, size_info, key, value); \
207 }
208
209 DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(int32_t, i32)
210 DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(uint32_t, u32)
211 DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(int64_t, i64)
212 DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(uint64_t, u64)
213 DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(bool, bool)
214 DEFINE_KEY_SPECIFIC_MAP_OPERATIONS(google::protobuf::rust::PtrAndLen, ProtoString)
215
216 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(int32_t, i32, int32_t,
217 int32_t, value, cpp_value);
218 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(uint32_t, u32, uint32_t,
219 uint32_t, value, cpp_value);
220 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(float, f32, float, float,
221 value, cpp_value);
222 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(double, f64, double, double,
223 value, cpp_value);
224 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(bool, bool, bool, bool,
225 value, cpp_value);
226 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(uint64_t, u64, uint64_t,
227 uint64_t, value, cpp_value);
228 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(int64_t, i64, int64_t,
229 int64_t, value, cpp_value);
230 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(
231 std::string, ProtoBytes, google::protobuf::rust::PtrAndLen, std::string*,
232 std::move(*value),
233 (google::protobuf::rust::PtrAndLen{cpp_value.data(), cpp_value.size()}));
234 __PB_RUST_EXPOSE_SCALAR_MAP_METHODS_FOR_VALUE_TYPE(
235 std::string, ProtoString, google::protobuf::rust::PtrAndLen, std::string*,
236 std::move(*value),
237 (google::protobuf::rust::PtrAndLen{cpp_value.data(), cpp_value.size()}));
238
239 } // extern "C"
240