mutable_bson/
parsed_document.rs

1use std::{borrow::Cow, ops::Index};
2
3use bson::{Document, RawDocument};
4use bytes::BufMut;
5use indexmap::IndexMap;
6
7use crate::{MutableValue, put_raw_cstr, raw_cstr_len};
8
9#[derive(Default, Clone, Debug)]
10pub struct ParsedDocument<'a>(IndexMap<Cow<'a, str>, MutableValue<'a>>);
11
12impl<'a> ParsedDocument<'a> {
13    pub fn new() -> Self {
14        Self::default()
15    }
16
17    pub fn insert<V: Into<MutableValue<'static>>>(
18        &mut self,
19        key: impl Into<String>,
20        value: V,
21    ) -> Option<MutableValue<'_>> {
22        self.0.insert(Cow::from(key.into()), value.into())
23    }
24
25    /// Remove key and return the value for that key if present.
26    ///
27    /// Runs in _O(n)_ time.
28    pub fn remove(&mut self, key: impl AsRef<str>) -> Option<MutableValue<'a>> {
29        self.0.shift_remove(key.as_ref())
30    }
31
32    pub fn clear(&mut self) {
33        self.0.clear()
34    }
35
36    pub fn contains_key(&self, key: impl AsRef<str>) -> bool {
37        self.0.contains_key(key.as_ref())
38    }
39
40    pub fn get(&self, key: impl AsRef<str>) -> Option<&MutableValue<'a>> {
41        self.0.get(key.as_ref())
42    }
43
44    pub fn get_mut(&mut self, key: impl AsRef<str>) -> Option<&mut MutableValue<'a>> {
45        self.0.get_mut(key.as_ref())
46    }
47
48    pub fn iter(&self) -> impl Iterator<Item = (&str, &MutableValue<'a>)> {
49        self.0.iter().map(|(k, v)| (k.as_ref(), v))
50    }
51
52    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&str, &mut MutableValue<'a>)> {
53        self.0.iter_mut().map(|(k, v)| (k.as_ref(), v))
54    }
55
56    pub fn len(&self) -> usize {
57        self.0.len()
58    }
59
60    pub fn is_empty(&self) -> bool {
61        self.0.is_empty()
62    }
63
64    pub(super) fn raw_len(&self) -> usize {
65        self.0
66            .iter()
67            // 1 byte for type, key, 1 byte null terminator, value.
68            .map(|(k, v)| 1 + raw_cstr_len(k.as_ref()) + v.raw_len())
69            .sum::<usize>()
70            // 4 bytes for doc length, 1 byte for null terminator.
71            + 4usize + 1usize
72    }
73
74    pub(super) fn put(&self, buf: &mut impl BufMut) -> Result<(), bson::ser::Error> {
75        buf.put_i32_le(
76            self.raw_len()
77                .try_into()
78                .expect("message len checked before put"),
79        );
80        for (k, v) in self.0.iter() {
81            buf.put_u8(v.element_type() as u8);
82            put_raw_cstr(k.as_ref(), buf)?;
83            v.put(buf)?;
84        }
85        buf.put_u8(0);
86        Ok(())
87    }
88
89    // TODO: keys()
90    // TODO: values()
91    // TODO: values_mut()
92    // TODO: hash/btree map style entry()
93}
94
95impl<'a> TryFrom<&'a RawDocument> for ParsedDocument<'a> {
96    type Error = bson::raw::Error;
97
98    fn try_from(value: &'a RawDocument) -> Result<Self, Self::Error> {
99        let mut fields = IndexMap::new();
100        for e in value.iter() {
101            let (k, v) = e?;
102            fields.insert(k.into(), v.into());
103        }
104        Ok(Self(fields))
105    }
106}
107
108impl From<Document> for ParsedDocument<'_> {
109    fn from(value: Document) -> Self {
110        Self(
111            value
112                .into_iter()
113                .map(|(k, v)| (k.into(), v.into()))
114                .collect(),
115        )
116    }
117}
118
119impl<'a, S: AsRef<str>> Index<S> for ParsedDocument<'a> {
120    type Output = MutableValue<'a>;
121
122    fn index(&self, index: S) -> &Self::Output {
123        &self.0[index.as_ref()]
124    }
125}
126
127#[cfg(test)]
128mod test {
129    use bson::{
130        Binary, Bson, DateTime, Decimal128, Document, JavaScriptCodeWithScope, RawDocumentBuf,
131        Regex, Timestamp, bson, doc, oid::ObjectId, rawdoc, to_raw_document_buf,
132    };
133
134    use crate::MutableValue;
135
136    use super::ParsedDocument;
137
138    fn doc_to_vec(doc: &ParsedDocument<'_>) -> Vec<u8> {
139        let mut out = vec![];
140        doc.put(&mut out).unwrap();
141        out
142    }
143
144    fn doc_all_types_owned() -> Document {
145        let mut doc = Document::new();
146        doc.insert("a", 1.0f64);
147        doc.insert("b", "str");
148        doc.insert("c", doc! { "text": "the quick brown fox" });
149        doc.insert("d", bson!([1, "value"]));
150        doc.insert(
151            "e",
152            Binary {
153                subtype: bson::spec::BinarySubtype::Generic,
154                bytes: vec![1u8, 2, 3],
155            },
156        );
157        doc.insert("f", Bson::Undefined);
158        doc.insert("g", ObjectId::from_bytes([0xae; 12]));
159        doc.insert("h", true);
160        doc.insert("i", DateTime::from_millis(1234567890));
161        doc.insert("j", Bson::Null);
162        doc.insert(
163            "k",
164            Regex {
165                pattern: "foo.*".into(),
166                options: "i".into(),
167            },
168        );
169        doc.insert("m", Bson::JavaScriptCode("some code".into()));
170        doc.insert("n", Bson::Symbol("symbol".into()));
171        doc.insert(
172            "o",
173            JavaScriptCodeWithScope {
174                code: "more code".into(),
175                scope: doc! {"text": "jumped over the lazy dog"},
176            },
177        );
178        doc.insert("p", 7i32);
179        doc.insert(
180            "q",
181            Timestamp {
182                time: 1234567890,
183                increment: 2,
184            },
185        );
186        doc.insert("r", 11i64);
187        doc.insert("s", Decimal128::from_bytes([0xaf; 16]));
188        doc.insert("t", Bson::MinKey);
189        doc.insert("u", Bson::MaxKey);
190        doc
191    }
192
193    fn doc_all_types_unowned() -> RawDocumentBuf {
194        to_raw_document_buf(&doc_all_types_owned()).unwrap()
195    }
196
197    #[test]
198    fn empty() {
199        let doc = ParsedDocument::new();
200        assert_eq!(doc.len(), 0);
201        assert!(doc.is_empty());
202        assert!(doc.iter().next().is_none());
203        assert_eq!(doc.raw_len(), 5);
204        assert_eq!(doc_to_vec(&doc), vec![5, 0, 0, 0, 0]);
205    }
206
207    #[test]
208    fn all_types_owned() {
209        let doc = ParsedDocument::from(doc_all_types_owned());
210        assert_eq!(doc_to_vec(&doc), doc_all_types_unowned().as_bytes());
211    }
212
213    #[test]
214    fn all_types_unowned() {
215        let raw_doc = doc_all_types_unowned();
216        let doc = ParsedDocument::try_from(raw_doc.as_ref()).unwrap();
217        assert_eq!(doc_to_vec(&doc), doc_all_types_unowned().as_bytes());
218    }
219
220    #[test]
221    fn contains_key() {
222        let doc = ParsedDocument::from(doc_all_types_owned());
223        assert!(doc.contains_key("e"));
224        assert!(!doc.contains_key("z"));
225    }
226
227    #[test]
228    fn get() {
229        let doc = ParsedDocument::from(doc_all_types_owned());
230        assert_eq!(doc.get("p").and_then(MutableValue::as_i32).unwrap(), 7);
231        assert!(doc.get("z").is_none());
232    }
233
234    #[test]
235    fn index() {
236        let doc = ParsedDocument::from(doc_all_types_owned());
237        assert_eq!(doc["p"].as_i32().unwrap(), 7);
238        assert!(doc.get("z").is_none());
239    }
240
241    #[test]
242    #[should_panic]
243    fn index_invalid() {
244        let doc = ParsedDocument::from(doc_all_types_owned());
245        let _ = doc["z"];
246    }
247
248    #[test]
249    fn insert_and_replace() {
250        let mut doc = ParsedDocument::new();
251        assert!(doc.insert("foo", 5).is_none());
252        assert!(doc.insert("bar", "bat").is_none());
253        assert_eq!(
254            doc_to_vec(&doc),
255            rawdoc! { "foo": 5, "bar": "bat" }.as_bytes()
256        );
257
258        assert!(doc.insert("foo", 16384).is_some());
259        assert_eq!(
260            doc_to_vec(&doc),
261            rawdoc! { "foo": 16384, "bar": "bat" }.as_bytes()
262        );
263    }
264
265    #[test]
266    fn remove() {
267        let mut doc = ParsedDocument::new();
268        assert!(doc.insert("foo", 5).is_none());
269        assert!(doc.insert("bar", "bat").is_none());
270        assert_eq!(
271            doc_to_vec(&doc),
272            rawdoc! { "foo": 5, "bar": "bat" }.as_bytes()
273        );
274
275        assert!(doc.remove("foo").is_some());
276        assert_eq!(doc_to_vec(&doc), rawdoc! { "bar": "bat" }.as_bytes());
277
278        assert!(doc.remove("foo").is_none());
279    }
280
281    #[test]
282    fn get_mut() {
283        let mut doc = ParsedDocument::new();
284        assert!(doc.insert("foo", 5).is_none());
285        assert!(doc.get_mut("bar").is_none());
286
287        let v = doc.get_mut("foo").expect("foo is present");
288        if let MutableValue::Int32(i) = v {
289            assert_eq!(*i, 5);
290        } else {
291            panic!("not an int32");
292        }
293
294        *v = "bar".into();
295        assert_eq!(doc_to_vec(&doc), rawdoc! { "foo": "bar" }.as_bytes());
296    }
297
298    #[test]
299    fn clear() {
300        let mut doc = ParsedDocument::from(doc_all_types_owned());
301        assert_eq!(doc_to_vec(&doc), doc_all_types_unowned().as_bytes());
302
303        doc.clear();
304        assert_eq!(doc_to_vec(&doc), vec![5, 0, 0, 0, 0]);
305    }
306
307    #[test]
308    fn mutate_unowned() {
309        let raw_doc = rawdoc! { "foo": "bar", "bat": 5 };
310        let mut doc = ParsedDocument::try_from(raw_doc.as_ref()).unwrap();
311        assert!(doc.insert("bat", true).is_some());
312        assert!(doc.insert("quux", 7).is_none());
313        assert_eq!(
314            doc_to_vec(&doc),
315            rawdoc! { "foo": "bar", "bat": true, "quux": 7 }.as_bytes()
316        );
317    }
318
319    #[test]
320    fn mutate_unowned_embedded_doc() {
321        let raw_doc = rawdoc! { "foo": 5, "bar": { "id": 0, "score": 1.1 }};
322        let mut doc = ParsedDocument::try_from(raw_doc.as_ref()).unwrap();
323        let emb_doc = doc
324            .get_mut("bar")
325            .and_then(MutableValue::as_doc_mut)
326            .unwrap()
327            .to_parsed()
328            .unwrap();
329        assert!(emb_doc.insert("score", 2.1).is_some());
330        assert!(emb_doc.insert("price", 7.17).is_none());
331        assert_eq!(
332            doc_to_vec(&doc),
333            rawdoc! { "foo": 5, "bar": { "id": 0, "score": 2.1, "price": 7.17 }}.as_bytes()
334        );
335    }
336
337    #[test]
338    fn mutate_embedded_array() {
339        let raw_doc = rawdoc! { "foo": 5, "vec": [0, "foo", 2, 3]};
340        let mut doc = ParsedDocument::try_from(raw_doc.as_ref()).unwrap();
341        let emb_array = doc
342            .get_mut("vec")
343            .and_then(MutableValue::as_array_mut)
344            .unwrap()
345            .to_parsed()
346            .unwrap();
347        emb_array[2] = "bar".into();
348        emb_array.push(4.into());
349        assert_eq!(
350            doc_to_vec(&doc),
351            rawdoc! { "foo": 5, "vec": [0, "foo", "bar", 3, 4]}.as_bytes()
352        );
353    }
354}
mutable_bson/parsed_document.rs

mutable_bson/
parsed_document.rs
