sonic_rs/pointer/
tree.rs

1use std::collections::HashMap;
2
3use faststr::FastStr;
4
5use crate::index::Index;
6
7/// PointerTree is designed for [`get_many`][`crate::get_many`] and
8/// [`get_many_unchecked`][`crate::get_many_unchecked`].
9///
10/// It is recommended to use `get_many` when you need to get multiple values from json. Instead of
11/// using `get` multiple times.
12///
13/// # Examples
14///
15/// ```
16/// # use sonic_rs::pointer;
17/// # use sonic_rs::PointerTree;
18///
19/// let json = r#"
20/// {"u": 123, "a": {"b" : {"c": [null, "found"]}}}"#;
21///
22/// // build a pointer tree, representing multiple json path
23/// let mut tree = PointerTree::new();
24///
25/// tree.add_path(&["u"]);
26/// tree.add_path(&["unknown_key"]);
27/// tree.add_path(&pointer!["a", "b", "c", 1]);
28///
29/// let nodes = unsafe { sonic_rs::get_many_unchecked(json, &tree) };
30///
31/// match nodes {
32///     Ok(vals) => {
33///         assert_eq!(vals[0].as_ref().unwrap().as_raw_str(), "123");
34///         assert!(vals[1].is_none());
35///         assert_eq!(vals[2].as_ref().unwrap().as_raw_str(), "\"found\"");
36///         for val in vals {
37///             match val {
38///                 Some(_) => println!("{}", val.as_ref().unwrap().as_raw_str()),
39///                 None => println!("None"),
40///             };
41///         }
42///     }
43///     Err(e) => {
44///         println!("err: {:?}", e)
45///     }
46/// }
47/// ```
48
49#[derive(Debug, Default)]
50pub struct PointerTree {
51    // the count of path
52    size: usize,
53    // the root of tree
54    pub(crate) root: PointerTreeNode,
55}
56
57impl PointerTree {
58    /// Creat a empty tree. If `get_many` from empty tree, it will return the whole json.
59    pub fn new() -> Self {
60        Self::default()
61    }
62
63    /// we build tree and return value according by the order of path.
64    /// Allow the repeated path.
65    pub fn add_path<Path: IntoIterator>(&mut self, path: Path)
66    where
67        Path::Item: Index,
68    {
69        self.root.add_path(path, self.size);
70        self.size += 1;
71    }
72
73    /// the count of nodes
74    pub fn size(&self) -> usize {
75        self.size
76    }
77}
78
79#[derive(Debug, Default)]
80pub(crate) enum PointerTreeInner {
81    #[default]
82    Empty,
83    Key(MultiKey),
84    Index(MultiIndex),
85}
86
87// Note: support the repeat path
88#[derive(Debug, Default)]
89pub(crate) struct PointerTreeNode {
90    pub(crate) order: Vec<usize>,
91    pub(crate) children: PointerTreeInner,
92}
93
94impl PointerTreeNode {
95    pub fn add_path<Path: IntoIterator>(&mut self, path: Path, order: usize)
96    where
97        Path::Item: Index,
98    {
99        let mut cur = self;
100        let iter = path.into_iter();
101        for p in iter {
102            if let Some(key) = p.as_key() {
103                if matches!(cur.children, PointerTreeInner::Empty) {
104                    cur.children = PointerTreeInner::Key(HashMap::new());
105                }
106                cur = cur.insert_key(key)
107            } else if let Some(index) = p.as_index() {
108                if matches!(cur.children, PointerTreeInner::Empty) {
109                    cur.children = PointerTreeInner::Index(HashMap::new());
110                }
111                cur = cur.insert_index(index)
112            }
113        }
114        cur.order.push(order);
115    }
116
117    fn insert_key(&mut self, key: &str) -> &mut Self {
118        if let PointerTreeInner::Key(mkey) = &mut self.children {
119            mkey.entry(FastStr::new(key)).or_insert(Self::default())
120        } else {
121            unreachable!()
122        }
123    }
124
125    fn insert_index(&mut self, idx: usize) -> &mut Self {
126        if let PointerTreeInner::Index(midx) = &mut self.children {
127            midx.entry(idx).or_insert(Self::default())
128        } else {
129            unreachable!()
130        }
131    }
132}
133
134#[allow(clippy::mutable_key_type)]
135pub(crate) type MultiKey = HashMap<FastStr, PointerTreeNode>;
136
137pub(crate) type MultiIndex = HashMap<usize, PointerTreeNode>;
138
139#[cfg(test)]
140mod test {
141    use super::*;
142    use crate::pointer;
143
144    #[test]
145    fn test_tree() {
146        let mut tree = PointerTree::default();
147        tree.add_path(["a", "a_b", "a_b_c"].iter());
148        tree.add_path(["a", "a_b"].iter());
149        tree.add_path(pointer!["a", "a_a", 1].iter());
150        tree.add_path(pointer!["a"].iter());
151        tree.add_path(pointer!["a"].iter());
152        tree.add_path(pointer!["b", 2].iter());
153        tree.add_path(pointer![].iter());
154        assert_eq!(tree.size(), 7);
155        println!("tree is {tree:#?}");
156    }
157}