yrs/json_path/
iter_txn.rs

1use crate::any::AnyArrayIter;
2use crate::json_path::JsonPathToken;
3use crate::{Any, Array, JsonPath, JsonPathEval, Map, Out, ReadTxn, Xml, XmlFragment};
4
5impl<T> JsonPathEval for T
6where
7    T: ReadTxn,
8{
9    type Iter<'a>
10        = JsonPathIter<'a, T>
11    where
12        T: 'a;
13
14    /// Evaluate JSON path on the current transaction, starting from current transaction [Doc] as
15    /// its root.
16    ///
17    /// # Example
18    /// ```rust
19    /// use yrs::{any, Array, ArrayPrelim, Doc, In, JsonPath, JsonPathEval, Map, MapPrelim, Out, Transact, WriteTxn};
20    ///
21    /// let doc = Doc::new();
22    /// let mut txn = doc.transact_mut();
23    /// let users = txn.get_or_insert_array("users");
24    ///
25    /// // populate the document with some data to query
26    /// users.insert(&mut txn, 0, MapPrelim::from([
27    ///     ("name".to_string(), In::Any(any!("Alice"))),
28    ///     ("surname".into(), In::Any(any!("Smith"))),
29    ///     ("age".into(), In::Any(any!(25))),
30    ///     (
31    ///         "friends".into(),
32    ///         In::from(ArrayPrelim::from([
33    ///             any!({ "name": "Bob", "nick": "boreas" }),
34    ///             any!({ "nick": "crocodile91" }),
35    ///         ])),
36    ///     ),
37    /// ]));
38    ///
39    /// let query = JsonPath::parse("$.users..friends.*.nick").unwrap();
40    /// let values: Vec<Out> = txn.json_path(&query).collect();
41    /// assert_eq!(values, vec![Out::Any(any!("boreas")), Out::Any(any!("crocodile91"))]);
42    /// ```
43    fn json_path<'a>(&'a self, path: &'a JsonPath<'a>) -> Self::Iter<'a> {
44        JsonPathIter::new(self, path.as_ref())
45    }
46}
47
48fn slice_iter<'a, T: ReadTxn>(
49    txn: &'a T,
50    value: Option<Out>,
51    from: usize,
52    to: usize,
53    by: usize,
54) -> Option<Box<dyn Iterator<Item = Out> + 'a>> {
55    match value {
56        Some(Out::Any(Any::Array(array))) => {
57            let iter = AnyArrayIter::new(array);
58            Some(Box::new(
59                iter.skip(from).take(to - from).step_by(by).map(Out::Any),
60            ))
61        }
62        Some(Out::YArray(array)) => {
63            let iter = array.iter(txn);
64            Some(Box::new(iter.skip(from).take(to - from).step_by(by)))
65        }
66        Some(Out::YXmlElement(xml)) => {
67            let iter = xml.children(txn);
68            Some(Box::new(
69                iter.skip(from).take(to - from).step_by(by).map(Out::from),
70            ))
71        }
72        Some(Out::YXmlFragment(xml)) => {
73            let iter = xml.children(txn);
74            Some(Box::new(
75                iter.skip(from).take(to - from).step_by(by).map(Out::from),
76            ))
77        }
78        _ => None,
79    }
80}
81
82fn any_iter<'a, T: ReadTxn>(
83    txn: &'a T,
84    out: Option<Out>,
85) -> Option<Box<dyn Iterator<Item = Out> + 'a>> {
86    #[inline]
87    fn dyn_iter<'a, I: Iterator<Item = Out> + 'a>(iter: I) -> Box<dyn Iterator<Item = Out> + 'a> {
88        Box::new(iter)
89    }
90
91    match out {
92        None => Some(dyn_iter(txn.root_refs().map(|(_, out)| out))),
93        Some(Out::Any(any)) => {
94            let iter = any.try_into_iter();
95            iter.map(|iter| dyn_iter(iter.map(|(_, v)| Out::Any(v))))
96        }
97        Some(Out::YArray(array)) => Some(dyn_iter(array.iter(txn))),
98        Some(Out::YXmlElement(elem)) => Some(dyn_iter(elem.children(txn).map(Out::from))),
99        Some(Out::YXmlFragment(elem)) => Some(dyn_iter(elem.children(txn).map(Out::from))),
100        Some(Out::YMap(map)) => Some(dyn_iter(map.into_iter(txn).map(|(_, value)| value))),
101        Some(Out::UndefinedRef(branch)) => {
102            // undefined ref only happens for root level types. These can be: ArrayRef, MapRef,
103            // XmlFragmentRef, TextRef. For JsonPath, we only care about ArrayRef and MapRef.
104            if branch.start.is_some() {
105                // a list component is not empty: we assume it's an Array
106                let array = crate::ArrayRef::from(branch);
107                Some(dyn_iter(array.iter(txn)))
108            } else {
109                // we assume it's a YMap
110                let map = crate::MapRef::from(branch);
111                Some(dyn_iter(map.into_iter(txn).map(|(_, value)| value)))
112            }
113        }
114        _ => None,
115    }
116}
117
118fn member_union_iter<'a, T: ReadTxn>(
119    txn: &'a T,
120    value: Option<Out>,
121    members: &'a [&'a str],
122) -> Option<Box<dyn Iterator<Item = Out> + 'a>> {
123    match value {
124        Some(Out::Any(Any::Map(map))) => {
125            let iter = members
126                .into_iter()
127                .flat_map(move |key| map.get(*key).cloned().map(Out::Any));
128            Some(Box::new(iter))
129        }
130        Some(Out::YMap(map)) => {
131            let iter = members.into_iter().flat_map(move |key| map.get(txn, *key));
132            Some(Box::new(iter))
133        }
134        None => {
135            let iter = members.into_iter().flat_map(move |key| txn.get(key));
136            Some(Box::new(iter))
137        }
138        _ => None,
139    }
140}
141
142fn index_union_iter<'a, T: ReadTxn>(
143    txn: &'a T,
144    value: Option<Out>,
145    indices: &'a [i32],
146) -> Option<Box<dyn Iterator<Item = Out> + 'a>> {
147    match value {
148        Some(Out::Any(Any::Array(array))) => {
149            let iter = indices.into_iter().flat_map(move |i| {
150                let i = if *i < 0 {
151                    (array.len() as i32 + *i) as usize
152                } else {
153                    *i as usize
154                };
155                array.get(i).cloned().map(Out::Any)
156            });
157            Some(Box::new(iter))
158        }
159        Some(Out::YArray(array)) => {
160            let len = array.len(txn);
161            let iter = indices.into_iter().flat_map(move |i| {
162                let i = if *i < 0 {
163                    (len as i32 + *i) as u32
164                } else {
165                    *i as u32
166                };
167                array.get(txn, i)
168            });
169            Some(Box::new(iter))
170        }
171        Some(Out::YXmlFragment(xml)) => {
172            let len = xml.len(txn);
173            let iter = indices.into_iter().flat_map(move |i| {
174                let i = if *i < 0 {
175                    (len as i32 + *i) as u32
176                } else {
177                    *i as u32
178                };
179                xml.get(txn, i).map(Out::from)
180            });
181            Some(Box::new(iter))
182        }
183        Some(Out::YXmlElement(xml)) => {
184            let len = xml.len(txn);
185            let iter = indices.into_iter().flat_map(move |i| {
186                let i = if *i < 0 {
187                    (len as i32 + *i) as u32
188                } else {
189                    *i as u32
190                };
191                xml.get(txn, i).map(Out::from)
192            });
193            Some(Box::new(iter))
194        }
195        _ => None,
196    }
197}
198
199pub struct JsonPathIter<'a, T> {
200    txn: &'a T,
201    pattern: &'a [JsonPathToken<'a>],
202    frame: ExecutionFrame<'a>,
203}
204
205impl<'a, T> JsonPathIter<'a, T>
206where
207    T: ReadTxn,
208{
209    fn new(txn: &'a T, path: &'a [JsonPathToken<'a>]) -> Self {
210        Self {
211            txn,
212            pattern: path.as_ref(),
213            frame: ExecutionFrame::new(None, 0, None),
214        }
215    }
216}
217
218impl<'a, T> Iterator for JsonPathIter<'a, T>
219where
220    T: ReadTxn,
221{
222    type Item = Out;
223
224    fn next(&mut self) -> Option<Self::Item> {
225        if let Some(iter) = &mut self.frame.iter {
226            match iter.next() {
227                None => {
228                    self.frame.iter = None;
229                    return if self.frame.ascend() {
230                        self.next()
231                    } else {
232                        None
233                    };
234                }
235                Some(curr) => {
236                    self.frame.descend(curr);
237                }
238            }
239        } else if self.frame.index == self.pattern.len() {
240            return if self.frame.ascend() {
241                self.next()
242            } else {
243                None
244            };
245        }
246        // early return only works if the loop after has at least one iteration
247        let mut early_return = false;
248        while self.frame.index < self.pattern.len() {
249            let segment = &self.pattern[self.frame.index];
250            self.frame.index += 1;
251            match segment {
252                JsonPathToken::Root => self.frame.current = None,
253                JsonPathToken::Current => { /* do nothing */ }
254                JsonPathToken::Member(key) => {
255                    if let Some(value) = get_member(self.txn, self.frame.current.as_ref(), *key) {
256                        self.frame.current = Some(value);
257                    } else {
258                        early_return = true;
259                    }
260                }
261                JsonPathToken::Index(idx) => {
262                    if let Some(value) = get_index(self.txn, self.frame.current.as_ref(), *idx) {
263                        self.frame.current = Some(value);
264                    } else {
265                        early_return = true;
266                    }
267                }
268                JsonPathToken::Wildcard => {
269                    if let Some(iter) = any_iter(self.txn, self.frame.current.clone()) {
270                        self.frame.iter = Some(iter);
271                        return self.next();
272                    }
273                    early_return = true;
274                }
275                JsonPathToken::RecursiveDescend => {
276                    if let Some(iter) = any_iter(self.txn, self.frame.current.clone()) {
277                        self.frame.iter = Some(iter);
278                        self.frame.is_descending = true;
279                        return self.next();
280                    }
281                    early_return = true;
282                }
283                JsonPathToken::Slice(from, to, by) => {
284                    if let Some(iter) = slice_iter(
285                        self.txn,
286                        self.frame.current.clone(),
287                        *from as usize,
288                        *to as usize,
289                        *by as usize,
290                    ) {
291                        self.frame.iter = Some(iter);
292                        return self.next();
293                    }
294                    early_return = true;
295                }
296                JsonPathToken::MemberUnion(members) => {
297                    if let Some(iter) =
298                        member_union_iter(self.txn, self.frame.current.clone(), &members)
299                    {
300                        self.frame.iter = Some(iter);
301                        return self.next();
302                    }
303                    early_return = true;
304                }
305                JsonPathToken::IndexUnion(indices) => {
306                    if let Some(iter) =
307                        index_union_iter(self.txn, self.frame.current.clone(), &indices)
308                    {
309                        self.frame.iter = Some(iter);
310                        return self.next();
311                    }
312                    early_return = true;
313                }
314            }
315
316            if early_return {
317                break;
318            }
319        }
320
321        if !early_return {
322            return self.frame.current.clone();
323        } else if self.frame.is_descending {
324            if let Some(iter) = any_iter(self.txn, self.frame.current.clone()) {
325                self.frame.iter = Some(iter);
326                self.frame.is_descending = true;
327                self.frame.index -= 1; // '..' means we're not consuming the segment in this iteration
328                return self.next();
329            }
330        }
331
332        if self.frame.iter.is_none() && !self.frame.ascend() {
333            None // we got to the end of the evaluator
334        } else {
335            self.next()
336        }
337    }
338}
339
340fn get_member<T: ReadTxn>(txn: &T, out: Option<&Out>, key: &str) -> Option<Out> {
341    match out {
342        None => txn.get(key),
343        Some(Out::YMap(map)) => map.get(txn, key),
344        Some(Out::Any(Any::Map(map))) => map.get(key).map(|any| Out::Any(any.clone())),
345        Some(Out::YXmlElement(elem)) => elem.get_attribute(txn, key),
346        Some(Out::YXmlText(elem)) => elem.get_attribute(txn, key),
347        Some(Out::UndefinedRef(branch)) => {
348            // we assume it's a YMap
349            let map = crate::MapRef::from(*branch);
350            map.get(txn, key)
351        }
352        _ => None,
353    }
354}
355
356fn get_index<T: ReadTxn>(txn: &T, out: Option<&Out>, idx: i32) -> Option<Out> {
357    match out {
358        Some(Out::YArray(array)) => {
359            let idx = if idx < 0 {
360                array.len(txn) as i32 + idx
361            } else {
362                idx
363            } as u32;
364            array.get(txn, idx)
365        }
366        Some(Out::Any(Any::Array(array))) => {
367            let idx = if idx < 0 {
368                array.len() as i32 + idx
369            } else {
370                idx
371            } as usize;
372            array.get(idx).cloned().map(Out::Any)
373        }
374        Some(Out::YXmlFragment(elem)) => {
375            let idx = if idx < 0 {
376                elem.len(txn) as i32 + idx
377            } else {
378                idx
379            } as u32;
380            elem.get(txn, idx).map(Out::from)
381        }
382        Some(Out::YXmlElement(elem)) => {
383            let idx = if idx < 0 {
384                elem.len(txn) as i32 + idx
385            } else {
386                idx
387            } as u32;
388            elem.get(txn, idx).map(Out::from)
389        }
390        Some(Out::UndefinedRef(array)) => {
391            // we assume it's a YMap
392            let array = crate::ArrayRef::from(*array);
393            let idx = if idx < 0 {
394                array.len(txn) as i32 + idx
395            } else {
396                idx
397            } as u32;
398            array.get(txn, idx)
399        }
400        _ => None,
401    }
402}
403
404/// Scope used for recursive iteration, i.e. wildcard, descent or slice.
405struct ExecutionFrame<'a> {
406    /// Offset to tokens array, where the current scope starts.
407    index: usize,
408    /// Whether we're in recursive descent scope.
409    is_descending: bool,
410    /// Current object this scope is iterating over.
411    current: Option<Out>,
412    /// Iterator used by this scope.
413    iter: Option<ScopeIterator<'a>>,
414    /// Scopes can be nested in each other i.e. `$.people.*.friends[*]name`. In such case they
415    /// are organized in a linked list, with the first elements being the innermost scopes.
416    next: Option<Box<ExecutionFrame<'a>>>,
417}
418
419impl<'a> ExecutionFrame<'a> {
420    fn new(current: Option<Out>, index: usize, iter: Option<ScopeIterator<'a>>) -> Self {
421        Self {
422            index,
423            is_descending: false,
424            current,
425            iter,
426            next: None,
427        }
428    }
429
430    /// Descent into given iterator context, moving current frame to the stack and replacing it with
431    /// a new one executing in a context of that iterator.
432    fn descend(&mut self, current: Out) {
433        let new_self = ExecutionFrame {
434            index: self.index,
435            is_descending: self.is_descending,
436            current: Some(current),
437            iter: None,
438            next: None,
439        };
440        let old_frame = std::mem::replace(self, new_self);
441        self.next = Some(Box::new(old_frame));
442    }
443
444    /// Return from current scope, restoring previous iter frame from the stack.
445    /// If there are no more iter frames, return false.
446    fn ascend(&mut self) -> bool {
447        match self.next.take() {
448            None => false,
449            Some(next) => {
450                *self = *next;
451                true
452            }
453        }
454    }
455}
456
457type ScopeIterator<'a> = Box<dyn Iterator<Item = Out> + 'a>;
458
459#[cfg(test)]
460mod test {
461    use crate::updates::decoder::Decode;
462    use crate::{
463        any, Array, ArrayPrelim, Doc, In, JsonPath, JsonPathEval, MapPrelim, Out, ReadTxn,
464        Transact, Update, WriteTxn,
465    };
466
467    fn mixed_sample() -> Doc {
468        let doc = Doc::new();
469        let mut tx = doc.transact_mut();
470        let users = tx.get_or_insert_array("users");
471        users.insert(
472            &mut tx,
473            0,
474            MapPrelim::from([
475                ("name".to_string(), In::Any(any!("Alice"))),
476                ("surname".into(), In::Any(any!("Smith"))),
477                ("age".into(), In::Any(any!(25))),
478                (
479                    "friends".into(),
480                    In::from(ArrayPrelim::from([
481                        any!({ "name": "Bob", "nick": "boreas" }),
482                        any!({ "nick": "crocodile91" }),
483                    ])),
484                ),
485            ]),
486        );
487        users.insert(
488            &mut tx,
489            1,
490            MapPrelim::from([
491                ("name".to_string(), In::Any(any!("Bob"))),
492                ("nick".into(), In::Any(any!("boreas"))),
493                ("age".into(), In::Any(any!(30))),
494            ]),
495        );
496        users.insert(
497            &mut tx,
498            2,
499            MapPrelim::from([
500                ("nick".to_string(), In::Any(any!("crocodile91"))),
501                ("age".into(), In::Any(any!(35))),
502            ]),
503        );
504        users.insert(
505            &mut tx,
506            3,
507            MapPrelim::from([
508                ("name".to_string(), In::Any(any!("Damian"))),
509                ("surname".into(), In::Any(any!("Smith"))),
510                ("age".into(), In::Any(any!(30))),
511            ]),
512        );
513        users.insert(
514            &mut tx,
515            4,
516            MapPrelim::from([
517                ("name".to_string(), In::Any(any!("Elise"))),
518                ("age".into(), In::Any(any!(35))),
519            ]),
520        );
521        drop(tx);
522        doc
523    }
524
525    #[test]
526    fn eval_member_partial() {
527        let doc = mixed_sample();
528        let path = JsonPath::parse("$.users").unwrap();
529        let tx = doc.transact();
530        let values: Vec<_> = tx.json_path(&path).collect();
531        let expected = tx.get("users").unwrap();
532        assert_eq!(values, vec![expected]);
533    }
534
535    #[test]
536    fn eval_member_full() {
537        let doc = mixed_sample();
538        let path = JsonPath::parse("$.users[0].name").unwrap();
539        let tx = doc.transact();
540        let values: Vec<_> = tx.json_path(&path).collect();
541        assert_eq!(values, vec![Out::Any(any!("Alice"))]);
542    }
543
544    #[test]
545    fn eval_member_negative_index() {
546        let doc = mixed_sample();
547        let path = JsonPath::parse("$.users[-1].name").unwrap();
548        let tx = doc.transact();
549        let values: Vec<_> = tx.json_path(&path).collect();
550        assert_eq!(values, vec![Out::Any(any!("Elise"))]);
551    }
552
553    #[test]
554    fn eval_member_wildcard_array() {
555        let doc = mixed_sample();
556        let path = JsonPath::parse("$.users[*].name").unwrap();
557        let tx = doc.transact();
558        let values: Vec<_> = tx.json_path(&path).collect();
559        assert_eq!(
560            values,
561            vec![
562                Out::Any(any!("Alice")),
563                Out::Any(any!("Bob")),
564                Out::Any(any!("Damian")),
565                Out::Any(any!("Elise"))
566            ]
567        );
568    }
569
570    #[test]
571    fn eval_member_slice() {
572        let doc = mixed_sample();
573        let path = JsonPath::parse("$.users[1:3].nick").unwrap();
574        let tx = doc.transact();
575        let values: Vec<_> = tx.json_path(&path).collect();
576        assert_eq!(
577            values,
578            vec![Out::Any(any!("boreas")), Out::Any(any!("crocodile91"))]
579        );
580    }
581
582    #[test]
583    fn eval_index_union() {
584        let doc = mixed_sample();
585        let path = JsonPath::parse("$.users[1,3].name").unwrap();
586        let tx = doc.transact();
587        let values: Vec<_> = tx.json_path(&path).collect();
588        assert_eq!(
589            values,
590            vec![Out::Any(any!("Bob")), Out::Any(any!("Damian"))]
591        );
592    }
593
594    #[test]
595    fn eval_member_union() {
596        let doc = mixed_sample();
597        let path = JsonPath::parse("$.users[0]['name','surname']").unwrap();
598        let tx = doc.transact();
599        let values: Vec<_> = tx.json_path(&path).collect();
600        assert_eq!(
601            values,
602            vec![Out::Any(any!("Alice")), Out::Any(any!("Smith"))]
603        );
604    }
605
606    #[test]
607    fn eval_descent_flat() {
608        let doc = mixed_sample();
609        let path = JsonPath::parse("$.users..name").unwrap();
610        let tx = doc.transact();
611        let values: Vec<_> = tx.json_path(&path).collect();
612        assert_eq!(
613            values,
614            vec![
615                Out::Any(any!("Alice")),
616                Out::Any(any!("Bob")),
617                Out::Any(any!("Damian")),
618                Out::Any(any!("Elise"))
619            ]
620        );
621    }
622    #[test]
623    fn eval_on_fresh_document() {
624        let doc_state = mixed_sample()
625            .transact()
626            .encode_state_as_update_v1(&Default::default());
627        let doc = Doc::new();
628        doc.transact_mut()
629            .apply_update(Update::decode_v1(&doc_state).unwrap())
630            .unwrap();
631        let path = JsonPath::parse("$.users[*].name").unwrap();
632        let tx = doc.transact();
633        let values: Vec<_> = tx.json_path(&path).collect();
634        assert_eq!(
635            values,
636            vec![
637                Out::Any(any!("Alice")),
638                Out::Any(any!("Bob")),
639                Out::Any(any!("Damian")),
640                Out::Any(any!("Elise"))
641            ]
642        );
643    }
644}
yrs/json_path/iter_txn.rs

yrs/json_path/
iter_txn.rs
