Skip to main content

hidane_proto/
ordered.rs

1//! The messages whose maps keep their order: `Document`, `MapValue` and `AggregationResult`.
2//!
3//! prost generates a protobuf map as a `HashMap` or a `BTreeMap`. The official emulator returns
4//! a document's fields in the order it holds them, which depends on how they were written
5//! (docs/parity-exceptions.md, #108), so these three messages are written by hand around
6//! [`Fields`], a map that keeps insertion order, and take the generated messages' place through
7//! `extern_path` (build.rs). On the wire they are the same messages, with entries in order.
8
9use 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/// Field names and values in insertion order. Inserting a name that is already there keeps its
21/// place; removing one keeps the order of the others. Equality ignores the order, as for
22/// protobuf maps.
23#[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    /// Sets `key`, in place if it is there, else at the end.
56    pub fn insert(&mut self, key: String, value: Value) -> Option<Value> {
57        self.0.insert(key, value)
58    }
59
60    /// Removes `key`; the fields after it keep their order.
61    pub fn remove(&mut self, key: &str) -> Option<Value> {
62        self.0.shift_remove(key)
63    }
64
65    /// Sets `key` at the end, moving it there if it is already present.
66    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    /// The fields in Firestore's canonical order: names by UTF-8 bytes, which is `String`'s
96    /// `Ord`. Comparisons and index keys use this order, never the stored one.
97    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
126/// Panics when `key` is missing, as `BTreeMap` does.
127impl 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/// A Firestore document. Same fields and tags as `google.firestore.v1.Document`.
169#[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/// A map value. Same fields and tags as `google.firestore.v1.MapValue`.
178#[derive(Clone, Default, PartialEq, Debug)]
179pub struct MapValue {
180    pub fields: Fields,
181}
182
183/// The result of one aggregation. Same fields and tags as
184/// `google.firestore.v1.AggregationResult`.
185#[derive(Clone, Default, PartialEq, Debug)]
186pub struct AggregationResult {
187    pub aggregate_fields: Fields,
188}
189
190/// One entry per field, in order, each encoded as prost encodes a map entry: the key and the
191/// value are left out when they are empty.
192fn 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
229/// One occurrence of a map field is one entry. prost decodes it (framing, unknown fields,
230/// recursion limit); a name seen again keeps its first place and takes the new value, as
231/// protobuf's "last one wins" and Java's `LinkedHashMap` both have it.
232fn 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}