Skip to main content

hidane_proto/
ordered.rs

1//! The messages whose maps keep their order: `Document`, `MapValue` and `AggregationResult`, and
2//! Datastore mode's `Entity` (#81).
3//!
4//! prost generates a protobuf map as a `HashMap` or a `BTreeMap`. The official emulator returns
5//! a document's fields in the order it holds them, which depends on how they were written
6//! (docs/parity-exceptions.md, #108), so these three messages are written by hand around
7//! [`Fields`], a map that keeps insertion order, and take the generated messages' place through
8//! `extern_path` (build.rs). On the wire they are the same messages, with entries in order.
9
10use std::{collections::BTreeMap, fmt, ops::Index};
11
12use indexmap::IndexMap;
13use prost::{
14    DecodeError, Message,
15    bytes::{Buf, BufMut},
16    encoding::{self, DecodeContext, WireType},
17};
18
19use crate::google::{datastore::v1 as datastore, firestore::v1::Value};
20
21/// Field names and values in insertion order. Inserting a name that is already there keeps its
22/// place; removing one keeps the order of the others. Equality ignores the order, as for
23/// protobuf maps.
24#[derive(Clone, PartialEq)]
25pub struct Ordered<V>(IndexMap<String, V>);
26
27/// A Firestore document's or map's fields.
28pub type Fields = Ordered<Value>;
29
30/// A Datastore entity's properties (#81).
31pub type Properties = Ordered<datastore::Value>;
32
33impl<V> Default for Ordered<V> {
34    fn default() -> Self {
35        Self(IndexMap::new())
36    }
37}
38
39impl<V> Ordered<V> {
40    pub fn new() -> Self {
41        Self::default()
42    }
43
44    pub fn with_capacity(capacity: usize) -> Self {
45        Self(IndexMap::with_capacity(capacity))
46    }
47
48    pub fn len(&self) -> usize {
49        self.0.len()
50    }
51
52    pub fn is_empty(&self) -> bool {
53        self.0.is_empty()
54    }
55
56    pub fn get(&self, key: &str) -> Option<&V> {
57        self.0.get(key)
58    }
59
60    pub fn get_mut(&mut self, key: &str) -> Option<&mut V> {
61        self.0.get_mut(key)
62    }
63
64    pub fn contains_key(&self, key: &str) -> bool {
65        self.0.contains_key(key)
66    }
67
68    /// Sets `key`, in place if it is there, else at the end.
69    pub fn insert(&mut self, key: String, value: V) -> Option<V> {
70        self.0.insert(key, value)
71    }
72
73    /// Removes `key`; the fields after it keep their order.
74    pub fn remove(&mut self, key: &str) -> Option<V> {
75        self.0.shift_remove(key)
76    }
77
78    /// Sets `key` at the end, moving it there if it is already present.
79    pub fn push(&mut self, key: String, value: V) {
80        self.0.shift_remove(&key);
81        self.0.insert(key, value);
82    }
83
84    pub fn entry(&mut self, key: String) -> indexmap::map::Entry<'_, String, V> {
85        self.0.entry(key)
86    }
87
88    pub fn iter(&self) -> indexmap::map::Iter<'_, String, V> {
89        self.0.iter()
90    }
91
92    pub fn iter_mut(&mut self) -> indexmap::map::IterMut<'_, String, V> {
93        self.0.iter_mut()
94    }
95
96    pub fn keys(&self) -> indexmap::map::Keys<'_, String, V> {
97        self.0.keys()
98    }
99
100    pub fn values(&self) -> indexmap::map::Values<'_, String, V> {
101        self.0.values()
102    }
103
104    pub fn values_mut(&mut self) -> indexmap::map::ValuesMut<'_, String, V> {
105        self.0.values_mut()
106    }
107
108    /// The fields in Firestore's canonical order: names by UTF-8 bytes, which is `String`'s
109    /// `Ord`. Comparisons and index keys use this order, never the stored one.
110    pub fn sorted(&self) -> Vec<(&String, &V)> {
111        let mut entries: Vec<_> = self.0.iter().collect();
112        entries.sort_unstable_by(|a, b| a.0.cmp(b.0));
113        entries
114    }
115
116    pub fn clear(&mut self) {
117        self.0.clear();
118    }
119}
120
121impl<V: fmt::Debug> fmt::Debug for Ordered<V> {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        f.debug_map().entries(self.0.iter()).finish()
124    }
125}
126
127impl<V> FromIterator<(String, V)> for Ordered<V> {
128    fn from_iter<I: IntoIterator<Item = (String, V)>>(iter: I) -> Self {
129        Self(iter.into_iter().collect())
130    }
131}
132
133impl<V, const N: usize> From<[(String, V); N]> for Ordered<V> {
134    fn from(entries: [(String, V); N]) -> Self {
135        entries.into_iter().collect()
136    }
137}
138
139/// Panics when `key` is missing, as `BTreeMap` does.
140impl<V> Index<&str> for Ordered<V> {
141    type Output = V;
142
143    fn index(&self, key: &str) -> &V {
144        &self.0[key]
145    }
146}
147
148impl<V> Extend<(String, V)> for Ordered<V> {
149    fn extend<I: IntoIterator<Item = (String, V)>>(&mut self, iter: I) {
150        self.0.extend(iter);
151    }
152}
153
154impl<V> IntoIterator for Ordered<V> {
155    type Item = (String, V);
156    type IntoIter = indexmap::map::IntoIter<String, V>;
157
158    fn into_iter(self) -> Self::IntoIter {
159        self.0.into_iter()
160    }
161}
162
163impl<'a, V> IntoIterator for &'a Ordered<V> {
164    type Item = (&'a String, &'a V);
165    type IntoIter = indexmap::map::Iter<'a, String, V>;
166
167    fn into_iter(self) -> Self::IntoIter {
168        self.0.iter()
169    }
170}
171
172impl<'a, V> IntoIterator for &'a mut Ordered<V> {
173    type Item = (&'a String, &'a mut V);
174    type IntoIter = indexmap::map::IterMut<'a, String, V>;
175
176    fn into_iter(self) -> Self::IntoIter {
177        self.0.iter_mut()
178    }
179}
180
181/// A Firestore document. Same fields and tags as `google.firestore.v1.Document`.
182#[derive(Clone, Default, PartialEq, Debug)]
183pub struct Document {
184    pub name: String,
185    pub fields: Fields,
186    pub create_time: Option<prost_types::Timestamp>,
187    pub update_time: Option<prost_types::Timestamp>,
188}
189
190/// A map value. Same fields and tags as `google.firestore.v1.MapValue`.
191#[derive(Clone, Default, PartialEq, Debug)]
192pub struct MapValue {
193    pub fields: Fields,
194}
195
196/// The result of one aggregation. Same fields and tags as
197/// `google.firestore.v1.AggregationResult`.
198#[derive(Clone, Default, PartialEq, Debug)]
199pub struct AggregationResult {
200    pub aggregate_fields: Fields,
201}
202
203/// One entry per field, in order, each encoded as prost encodes a map entry: the key and the
204/// value are left out when they are empty.
205fn encode_fields<V: Message + Default + PartialEq>(
206    tag: u32,
207    fields: &Ordered<V>,
208    buf: &mut impl BufMut,
209) {
210    for (key, value) in fields {
211        encoding::encode_key(tag, WireType::LengthDelimited, buf);
212        encoding::encode_varint(entry_len(key, value) as u64, buf);
213        if !key.is_empty() {
214            encoding::string::encode(1, key, buf);
215        }
216        if *value != V::default() {
217            encoding::message::encode(2, value, buf);
218        }
219    }
220}
221
222fn entry_len<V: Message + Default + PartialEq>(key: &String, value: &V) -> usize {
223    let key = if key.is_empty() {
224        0
225    } else {
226        encoding::string::encoded_len(1, key)
227    };
228    let value = if *value == V::default() {
229        0
230    } else {
231        encoding::message::encoded_len(2, value)
232    };
233    key + value
234}
235
236fn fields_len<V: Message + Default + PartialEq>(tag: u32, fields: &Ordered<V>) -> usize {
237    fields
238        .iter()
239        .map(|(key, value)| {
240            let len = entry_len(key, value);
241            encoding::key_len(tag) + encoding::encoded_len_varint(len as u64) + len
242        })
243        .sum()
244}
245
246/// One occurrence of a map field is one entry. prost decodes it (framing, unknown fields,
247/// recursion limit); a name seen again keeps its first place and takes the new value, as
248/// protobuf's "last one wins" and Java's `LinkedHashMap` both have it.
249fn merge_fields<V: Message + Default>(
250    wire_type: WireType,
251    fields: &mut Ordered<V>,
252    buf: &mut impl Buf,
253    ctx: DecodeContext,
254) -> Result<(), DecodeError> {
255    encoding::check_wire_type(WireType::LengthDelimited, wire_type)?;
256    let mut entry = BTreeMap::new();
257    encoding::btree_map::merge(
258        encoding::string::merge,
259        encoding::message::merge,
260        &mut entry,
261        buf,
262        ctx,
263    )?;
264    fields.extend(entry);
265    Ok(())
266}
267
268fn context(message: &'static str, field: &'static str) -> impl Fn(DecodeError) -> DecodeError {
269    move |mut error| {
270        error.push(message, field);
271        error
272    }
273}
274
275impl Message for Document {
276    fn encode_raw(&self, buf: &mut impl BufMut) {
277        if !self.name.is_empty() {
278            encoding::string::encode(1, &self.name, buf);
279        }
280        encode_fields(2, &self.fields, buf);
281        if let Some(time) = &self.create_time {
282            encoding::message::encode(3, time, buf);
283        }
284        if let Some(time) = &self.update_time {
285            encoding::message::encode(4, time, buf);
286        }
287    }
288
289    fn merge_field(
290        &mut self,
291        tag: u32,
292        wire_type: WireType,
293        buf: &mut impl Buf,
294        ctx: DecodeContext,
295    ) -> Result<(), DecodeError> {
296        match tag {
297            1 => encoding::string::merge(wire_type, &mut self.name, buf, ctx)
298                .map_err(context("Document", "name")),
299            2 => merge_fields(wire_type, &mut self.fields, buf, ctx)
300                .map_err(context("Document", "fields")),
301            3 => encoding::message::merge(
302                wire_type,
303                self.create_time.get_or_insert_with(Default::default),
304                buf,
305                ctx,
306            )
307            .map_err(context("Document", "create_time")),
308            4 => encoding::message::merge(
309                wire_type,
310                self.update_time.get_or_insert_with(Default::default),
311                buf,
312                ctx,
313            )
314            .map_err(context("Document", "update_time")),
315            _ => encoding::skip_field(wire_type, tag, buf, ctx),
316        }
317    }
318
319    fn encoded_len(&self) -> usize {
320        let name = if self.name.is_empty() {
321            0
322        } else {
323            encoding::string::encoded_len(1, &self.name)
324        };
325        name + fields_len(2, &self.fields)
326            + self
327                .create_time
328                .as_ref()
329                .map_or(0, |time| encoding::message::encoded_len(3, time))
330            + self
331                .update_time
332                .as_ref()
333                .map_or(0, |time| encoding::message::encoded_len(4, time))
334    }
335
336    fn clear(&mut self) {
337        *self = Self::default();
338    }
339}
340
341impl Message for MapValue {
342    fn encode_raw(&self, buf: &mut impl BufMut) {
343        encode_fields(1, &self.fields, buf);
344    }
345
346    fn merge_field(
347        &mut self,
348        tag: u32,
349        wire_type: WireType,
350        buf: &mut impl Buf,
351        ctx: DecodeContext,
352    ) -> Result<(), DecodeError> {
353        match tag {
354            1 => merge_fields(wire_type, &mut self.fields, buf, ctx)
355                .map_err(context("MapValue", "fields")),
356            _ => encoding::skip_field(wire_type, tag, buf, ctx),
357        }
358    }
359
360    fn encoded_len(&self) -> usize {
361        fields_len(1, &self.fields)
362    }
363
364    fn clear(&mut self) {
365        self.fields.clear();
366    }
367}
368
369impl Message for AggregationResult {
370    fn encode_raw(&self, buf: &mut impl BufMut) {
371        encode_fields(2, &self.aggregate_fields, buf);
372    }
373
374    fn merge_field(
375        &mut self,
376        tag: u32,
377        wire_type: WireType,
378        buf: &mut impl Buf,
379        ctx: DecodeContext,
380    ) -> Result<(), DecodeError> {
381        match tag {
382            2 => merge_fields(wire_type, &mut self.aggregate_fields, buf, ctx)
383                .map_err(context("AggregationResult", "aggregate_fields")),
384            _ => encoding::skip_field(wire_type, tag, buf, ctx),
385        }
386    }
387
388    fn encoded_len(&self) -> usize {
389        fields_len(2, &self.aggregate_fields)
390    }
391
392    fn clear(&mut self) {
393        self.aggregate_fields.clear();
394    }
395}
396
397/// A Datastore entity. Same fields and tags as `google.datastore.v1.Entity`: its properties
398/// keep the order they were written in, as the official emulator returns them (#81).
399#[derive(Clone, Default, PartialEq, Debug)]
400pub struct Entity {
401    pub key: Option<datastore::Key>,
402    pub properties: Properties,
403}
404
405impl Message for Entity {
406    fn encode_raw(&self, buf: &mut impl BufMut) {
407        if let Some(key) = &self.key {
408            encoding::message::encode(1, key, buf);
409        }
410        encode_fields(3, &self.properties, buf);
411    }
412
413    fn merge_field(
414        &mut self,
415        tag: u32,
416        wire_type: WireType,
417        buf: &mut impl Buf,
418        ctx: DecodeContext,
419    ) -> Result<(), DecodeError> {
420        match tag {
421            1 => encoding::message::merge(
422                wire_type,
423                self.key.get_or_insert_with(Default::default),
424                buf,
425                ctx,
426            )
427            .map_err(context("Entity", "key")),
428            3 => merge_fields(wire_type, &mut self.properties, buf, ctx)
429                .map_err(context("Entity", "properties")),
430            _ => encoding::skip_field(wire_type, tag, buf, ctx),
431        }
432    }
433
434    fn encoded_len(&self) -> usize {
435        self.key
436            .as_ref()
437            .map_or(0, |key| encoding::message::encoded_len(1, key))
438            + fields_len(3, &self.properties)
439    }
440
441    fn clear(&mut self) {
442        *self = Self::default();
443    }
444}
445
446/// The result of a Datastore aggregation query. Same fields and tags as
447/// `google.datastore.v1.AggregationResult`: its properties come in the aggregations' order (#81).
448#[derive(Clone, Default, PartialEq, Debug)]
449pub struct DatastoreAggregationResult {
450    pub aggregate_properties: Properties,
451}
452
453impl Message for DatastoreAggregationResult {
454    fn encode_raw(&self, buf: &mut impl BufMut) {
455        encode_fields(2, &self.aggregate_properties, buf);
456    }
457
458    fn merge_field(
459        &mut self,
460        tag: u32,
461        wire_type: WireType,
462        buf: &mut impl Buf,
463        ctx: DecodeContext,
464    ) -> Result<(), DecodeError> {
465        match tag {
466            2 => merge_fields(wire_type, &mut self.aggregate_properties, buf, ctx)
467                .map_err(context("AggregationResult", "aggregate_properties")),
468            _ => encoding::skip_field(wire_type, tag, buf, ctx),
469        }
470    }
471
472    fn encoded_len(&self) -> usize {
473        fields_len(2, &self.aggregate_properties)
474    }
475
476    fn clear(&mut self) {
477        self.aggregate_properties.clear();
478    }
479}