1use std::{collections::BTreeMap, fmt, ops::Index};
10
11use indexmap::IndexMap;
12use prost::{
13 DecodeError, Message,
14 bytes::{Buf, BufMut},
15 encoding::{self, DecodeContext, WireType},
16};
17
18use crate::google::firestore::v1::Value;
19
20#[derive(Clone, Default, PartialEq)]
24pub struct Fields(IndexMap<String, Value>);
25
26impl Fields {
27 pub fn new() -> Self {
28 Self::default()
29 }
30
31 pub fn with_capacity(capacity: usize) -> Self {
32 Self(IndexMap::with_capacity(capacity))
33 }
34
35 pub fn len(&self) -> usize {
36 self.0.len()
37 }
38
39 pub fn is_empty(&self) -> bool {
40 self.0.is_empty()
41 }
42
43 pub fn get(&self, key: &str) -> Option<&Value> {
44 self.0.get(key)
45 }
46
47 pub fn get_mut(&mut self, key: &str) -> Option<&mut Value> {
48 self.0.get_mut(key)
49 }
50
51 pub fn contains_key(&self, key: &str) -> bool {
52 self.0.contains_key(key)
53 }
54
55 pub fn insert(&mut self, key: String, value: Value) -> Option<Value> {
57 self.0.insert(key, value)
58 }
59
60 pub fn remove(&mut self, key: &str) -> Option<Value> {
62 self.0.shift_remove(key)
63 }
64
65 pub fn push(&mut self, key: String, value: Value) {
67 self.0.shift_remove(&key);
68 self.0.insert(key, value);
69 }
70
71 pub fn entry(&mut self, key: String) -> indexmap::map::Entry<'_, String, Value> {
72 self.0.entry(key)
73 }
74
75 pub fn iter(&self) -> indexmap::map::Iter<'_, String, Value> {
76 self.0.iter()
77 }
78
79 pub fn iter_mut(&mut self) -> indexmap::map::IterMut<'_, String, Value> {
80 self.0.iter_mut()
81 }
82
83 pub fn keys(&self) -> indexmap::map::Keys<'_, String, Value> {
84 self.0.keys()
85 }
86
87 pub fn values(&self) -> indexmap::map::Values<'_, String, Value> {
88 self.0.values()
89 }
90
91 pub fn values_mut(&mut self) -> indexmap::map::ValuesMut<'_, String, Value> {
92 self.0.values_mut()
93 }
94
95 pub fn sorted(&self) -> Vec<(&String, &Value)> {
98 let mut entries: Vec<_> = self.0.iter().collect();
99 entries.sort_unstable_by(|a, b| a.0.cmp(b.0));
100 entries
101 }
102
103 pub fn clear(&mut self) {
104 self.0.clear();
105 }
106}
107
108impl fmt::Debug for Fields {
109 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110 f.debug_map().entries(self.0.iter()).finish()
111 }
112}
113
114impl FromIterator<(String, Value)> for Fields {
115 fn from_iter<I: IntoIterator<Item = (String, Value)>>(iter: I) -> Self {
116 Self(iter.into_iter().collect())
117 }
118}
119
120impl<const N: usize> From<[(String, Value); N]> for Fields {
121 fn from(entries: [(String, Value); N]) -> Self {
122 entries.into_iter().collect()
123 }
124}
125
126impl Index<&str> for Fields {
128 type Output = Value;
129
130 fn index(&self, key: &str) -> &Value {
131 &self.0[key]
132 }
133}
134
135impl Extend<(String, Value)> for Fields {
136 fn extend<I: IntoIterator<Item = (String, Value)>>(&mut self, iter: I) {
137 self.0.extend(iter);
138 }
139}
140
141impl IntoIterator for Fields {
142 type Item = (String, Value);
143 type IntoIter = indexmap::map::IntoIter<String, Value>;
144
145 fn into_iter(self) -> Self::IntoIter {
146 self.0.into_iter()
147 }
148}
149
150impl<'a> IntoIterator for &'a Fields {
151 type Item = (&'a String, &'a Value);
152 type IntoIter = indexmap::map::Iter<'a, String, Value>;
153
154 fn into_iter(self) -> Self::IntoIter {
155 self.0.iter()
156 }
157}
158
159impl<'a> IntoIterator for &'a mut Fields {
160 type Item = (&'a String, &'a mut Value);
161 type IntoIter = indexmap::map::IterMut<'a, String, Value>;
162
163 fn into_iter(self) -> Self::IntoIter {
164 self.0.iter_mut()
165 }
166}
167
168#[derive(Clone, Default, PartialEq, Debug)]
170pub struct Document {
171 pub name: String,
172 pub fields: Fields,
173 pub create_time: Option<prost_types::Timestamp>,
174 pub update_time: Option<prost_types::Timestamp>,
175}
176
177#[derive(Clone, Default, PartialEq, Debug)]
179pub struct MapValue {
180 pub fields: Fields,
181}
182
183#[derive(Clone, Default, PartialEq, Debug)]
186pub struct AggregationResult {
187 pub aggregate_fields: Fields,
188}
189
190fn encode_fields(tag: u32, fields: &Fields, buf: &mut impl BufMut) {
193 for (key, value) in fields {
194 encoding::encode_key(tag, WireType::LengthDelimited, buf);
195 encoding::encode_varint(entry_len(key, value) as u64, buf);
196 if !key.is_empty() {
197 encoding::string::encode(1, key, buf);
198 }
199 if value.value_type.is_some() {
200 encoding::message::encode(2, value, buf);
201 }
202 }
203}
204
205fn entry_len(key: &String, value: &Value) -> usize {
206 let key = if key.is_empty() {
207 0
208 } else {
209 encoding::string::encoded_len(1, key)
210 };
211 let value = if value.value_type.is_some() {
212 encoding::message::encoded_len(2, value)
213 } else {
214 0
215 };
216 key + value
217}
218
219fn fields_len(tag: u32, fields: &Fields) -> usize {
220 fields
221 .iter()
222 .map(|(key, value)| {
223 let len = entry_len(key, value);
224 encoding::key_len(tag) + encoding::encoded_len_varint(len as u64) + len
225 })
226 .sum()
227}
228
229fn merge_fields(
233 wire_type: WireType,
234 fields: &mut Fields,
235 buf: &mut impl Buf,
236 ctx: DecodeContext,
237) -> Result<(), DecodeError> {
238 encoding::check_wire_type(WireType::LengthDelimited, wire_type)?;
239 let mut entry = BTreeMap::new();
240 encoding::btree_map::merge(
241 encoding::string::merge,
242 encoding::message::merge,
243 &mut entry,
244 buf,
245 ctx,
246 )?;
247 fields.extend(entry);
248 Ok(())
249}
250
251fn context(message: &'static str, field: &'static str) -> impl Fn(DecodeError) -> DecodeError {
252 move |mut error| {
253 error.push(message, field);
254 error
255 }
256}
257
258impl Message for Document {
259 fn encode_raw(&self, buf: &mut impl BufMut) {
260 if !self.name.is_empty() {
261 encoding::string::encode(1, &self.name, buf);
262 }
263 encode_fields(2, &self.fields, buf);
264 if let Some(time) = &self.create_time {
265 encoding::message::encode(3, time, buf);
266 }
267 if let Some(time) = &self.update_time {
268 encoding::message::encode(4, time, buf);
269 }
270 }
271
272 fn merge_field(
273 &mut self,
274 tag: u32,
275 wire_type: WireType,
276 buf: &mut impl Buf,
277 ctx: DecodeContext,
278 ) -> Result<(), DecodeError> {
279 match tag {
280 1 => encoding::string::merge(wire_type, &mut self.name, buf, ctx)
281 .map_err(context("Document", "name")),
282 2 => merge_fields(wire_type, &mut self.fields, buf, ctx)
283 .map_err(context("Document", "fields")),
284 3 => encoding::message::merge(
285 wire_type,
286 self.create_time.get_or_insert_with(Default::default),
287 buf,
288 ctx,
289 )
290 .map_err(context("Document", "create_time")),
291 4 => encoding::message::merge(
292 wire_type,
293 self.update_time.get_or_insert_with(Default::default),
294 buf,
295 ctx,
296 )
297 .map_err(context("Document", "update_time")),
298 _ => encoding::skip_field(wire_type, tag, buf, ctx),
299 }
300 }
301
302 fn encoded_len(&self) -> usize {
303 let name = if self.name.is_empty() {
304 0
305 } else {
306 encoding::string::encoded_len(1, &self.name)
307 };
308 name + fields_len(2, &self.fields)
309 + self
310 .create_time
311 .as_ref()
312 .map_or(0, |time| encoding::message::encoded_len(3, time))
313 + self
314 .update_time
315 .as_ref()
316 .map_or(0, |time| encoding::message::encoded_len(4, time))
317 }
318
319 fn clear(&mut self) {
320 *self = Self::default();
321 }
322}
323
324impl Message for MapValue {
325 fn encode_raw(&self, buf: &mut impl BufMut) {
326 encode_fields(1, &self.fields, buf);
327 }
328
329 fn merge_field(
330 &mut self,
331 tag: u32,
332 wire_type: WireType,
333 buf: &mut impl Buf,
334 ctx: DecodeContext,
335 ) -> Result<(), DecodeError> {
336 match tag {
337 1 => merge_fields(wire_type, &mut self.fields, buf, ctx)
338 .map_err(context("MapValue", "fields")),
339 _ => encoding::skip_field(wire_type, tag, buf, ctx),
340 }
341 }
342
343 fn encoded_len(&self) -> usize {
344 fields_len(1, &self.fields)
345 }
346
347 fn clear(&mut self) {
348 self.fields.clear();
349 }
350}
351
352impl Message for AggregationResult {
353 fn encode_raw(&self, buf: &mut impl BufMut) {
354 encode_fields(2, &self.aggregate_fields, buf);
355 }
356
357 fn merge_field(
358 &mut self,
359 tag: u32,
360 wire_type: WireType,
361 buf: &mut impl Buf,
362 ctx: DecodeContext,
363 ) -> Result<(), DecodeError> {
364 match tag {
365 2 => merge_fields(wire_type, &mut self.aggregate_fields, buf, ctx)
366 .map_err(context("AggregationResult", "aggregate_fields")),
367 _ => encoding::skip_field(wire_type, tag, buf, ctx),
368 }
369 }
370
371 fn encoded_len(&self) -> usize {
372 fields_len(2, &self.aggregate_fields)
373 }
374
375 fn clear(&mut self) {
376 self.aggregate_fields.clear();
377 }
378}