zen_types/variable/
mod.rs

1use crate::variable::ref_ser::RefSerializer;
2use ahash::HashMap;
3use rust_decimal::Decimal;
4use rust_decimal::prelude::Zero;
5use serde_json::Value;
6use std::any::Any;
7use std::cell::RefCell;
8use std::collections::hash_map::Entry;
9use std::fmt::{Debug, Display, Formatter};
10use std::ops::Deref;
11use std::rc::Rc;
12
13use crate::rcvalue::RcValue;
14pub use crate::variable::ref_deser::RefDeserializeError;
15use crate::variable::ref_deser::RefDeserializer;
16pub use de::VariableDeserializer;
17pub use impls::ToVariable;
18
19mod conv;
20mod de;
21mod impls;
22mod ref_deser;
23mod ref_ser;
24mod ser;
25
26pub(crate) type RcCell<T> = Rc<RefCell<T>>;
27
28pub type VariableMap = HashMap<Rc<str>, Variable>;
29
30thread_local! {
31    static DOLLAR_KEY: Rc<str> = Rc::from("$");
32    static ROOT_KEY: Rc<str> = Rc::from("$root");
33}
34
35pub enum Variable {
36    Null,
37    Bool(bool),
38    Number(Decimal),
39    String(Rc<str>),
40    Array(RcCell<Vec<Variable>>),
41    Object(RcCell<VariableMap>),
42    Dynamic(Rc<dyn DynamicVariable>),
43}
44
45pub trait DynamicVariable: Display {
46    fn type_name(&self) -> &'static str;
47
48    fn as_any(&self) -> &dyn Any;
49
50    fn to_value(&self) -> Value;
51}
52
53impl Variable {
54    pub fn dollar_key() -> Rc<str> {
55        DOLLAR_KEY.with(Rc::clone)
56    }
57
58    pub fn root_key() -> Rc<str> {
59        ROOT_KEY.with(Rc::clone)
60    }
61
62    pub fn from_array(arr: Vec<Self>) -> Self {
63        Self::Array(Rc::new(RefCell::new(arr)))
64    }
65
66    pub fn serialize_ref(&self) -> RcValue {
67        RefSerializer::new().serialize(self)
68    }
69
70    pub fn deserialize_ref(serialized: RcValue) -> Result<Self, RefDeserializeError> {
71        RefDeserializer::new().deserialize(serialized)
72    }
73
74    pub fn from_object(obj: HashMap<Rc<str>, Self>) -> Self {
75        Self::Object(Rc::new(RefCell::new(obj)))
76    }
77
78    pub fn empty_object() -> Self {
79        Variable::Object(Default::default())
80    }
81
82    pub fn empty_array() -> Self {
83        Variable::Array(Default::default())
84    }
85
86    pub fn as_str(&self) -> Option<&str> {
87        match self {
88            Variable::String(s) => Some(s.as_ref()),
89            _ => None,
90        }
91    }
92
93    pub fn as_rc_str(&self) -> Option<Rc<str>> {
94        match self {
95            Variable::String(s) => Some(s.clone()),
96            _ => None,
97        }
98    }
99
100    pub fn as_array(&self) -> Option<RcCell<Vec<Variable>>> {
101        match self {
102            Variable::Array(arr) => Some(arr.clone()),
103            _ => None,
104        }
105    }
106
107    pub fn is_array(&self) -> bool {
108        match self {
109            Variable::Array(_) => true,
110            _ => false,
111        }
112    }
113
114    pub fn as_object(&self) -> Option<RcCell<HashMap<Rc<str>, Variable>>> {
115        match self {
116            Variable::Object(obj) => Some(obj.clone()),
117            _ => None,
118        }
119    }
120
121    pub fn is_object(&self) -> bool {
122        match self {
123            Variable::Object(_) => true,
124            _ => false,
125        }
126    }
127
128    pub fn as_bool(&self) -> Option<bool> {
129        match self {
130            Variable::Bool(b) => Some(*b),
131            _ => None,
132        }
133    }
134
135    pub fn as_number(&self) -> Option<Decimal> {
136        match self {
137            Variable::Number(n) => Some(*n),
138            _ => None,
139        }
140    }
141
142    pub fn type_name(&self) -> &'static str {
143        match self {
144            Variable::Null => "null",
145            Variable::Bool(_) => "bool",
146            Variable::Number(_) => "number",
147            Variable::String(_) => "string",
148            Variable::Array(_) => "array",
149            Variable::Object(_) => "object",
150            Variable::Dynamic(d) => d.type_name(),
151        }
152    }
153
154    pub fn dynamic<T: DynamicVariable + 'static>(&self) -> Option<&T> {
155        match self {
156            Variable::Dynamic(d) => d.as_any().downcast_ref::<T>(),
157            _ => None,
158        }
159    }
160
161    pub fn to_value(&self) -> Value {
162        Value::from(self.shallow_clone())
163    }
164
165    pub fn dot(&self, key: &str) -> Option<Variable> {
166        key.split('.')
167            .try_fold(self.shallow_clone(), |var, part| match var {
168                Variable::Object(obj) => {
169                    let reference = obj.borrow();
170                    reference.get(part).map(|v| v.shallow_clone())
171                }
172                _ => None,
173            })
174    }
175
176    fn dot_head_detach(&self, key: &str) -> (Variable, Option<Variable>) {
177        let mut parts = Vec::from_iter(key.split('.'));
178        parts.pop();
179
180        let cloned_self = self.depth_clone(1);
181        let head = parts
182            .iter()
183            .try_fold(cloned_self.shallow_clone(), |var, part| match var {
184                Variable::Object(obj) => {
185                    let mut obj_ref = obj.borrow_mut();
186                    Some(match obj_ref.entry(Rc::from(*part)) {
187                        Entry::Occupied(mut occ) => {
188                            let var = occ.get();
189                            let new_obj = match var {
190                                Variable::Object(_) => var.depth_clone(1),
191                                _ => Variable::empty_object(),
192                            };
193
194                            occ.insert(new_obj.shallow_clone());
195                            new_obj
196                        }
197                        Entry::Vacant(vac) => vac.insert(Self::empty_object()).shallow_clone(),
198                    })
199                }
200                _ => None,
201            });
202
203        (cloned_self, head)
204    }
205
206    pub fn dot_remove(&self, key: &str) -> Option<Variable> {
207        let mut parts = key.split('.');
208        let last_part = parts.next_back()?;
209        let head = parts.try_fold(self.shallow_clone(), |var, part| match var {
210            Variable::Object(obj) => {
211                let mut obj_ref = obj.borrow_mut();
212                Some(match obj_ref.entry(Rc::from(part)) {
213                    Entry::Occupied(occ) => occ.get().shallow_clone(),
214                    Entry::Vacant(vac) => vac.insert(Self::empty_object()).shallow_clone(),
215                })
216            }
217            _ => None,
218        })?;
219        let Variable::Object(object_ref) = head else {
220            return None;
221        };
222
223        let mut object = object_ref.borrow_mut();
224        object.remove(last_part)
225    }
226
227    pub fn dot_insert(&self, key: &str, variable: Variable) -> Option<Variable> {
228        let mut parts = key.split('.');
229        let last_part = parts.next_back()?;
230        let head = parts.try_fold(self.shallow_clone(), |var, part| match var {
231            Variable::Object(obj) => {
232                let mut obj_ref = obj.borrow_mut();
233                Some(match obj_ref.entry(Rc::from(part)) {
234                    Entry::Occupied(occ) => occ.get().shallow_clone(),
235                    Entry::Vacant(vac) => vac.insert(Self::empty_object()).shallow_clone(),
236                })
237            }
238            _ => None,
239        })?;
240        let Variable::Object(object_ref) = head else {
241            return None;
242        };
243
244        let mut object = object_ref.borrow_mut();
245        object.insert(Rc::from(last_part), variable)
246    }
247
248    pub fn dot_insert_detached(&self, key: &str, variable: Variable) -> Option<Variable> {
249        let last_part = key.split('.').last()?;
250        let (new_var, head_opt) = self.dot_head_detach(key);
251        let head = head_opt?;
252        let Variable::Object(object_ref) = head else {
253            return None;
254        };
255
256        let mut object = object_ref.borrow_mut();
257        object.insert(Rc::from(last_part), variable);
258        Some(new_var)
259    }
260
261    pub fn merge(&mut self, patch: &Variable) -> Variable {
262        let _ = merge_variables(self, patch, true, MergeStrategy::InPlace);
263
264        self.shallow_clone()
265    }
266
267    pub fn merge_clone(&mut self, patch: &Variable) -> Variable {
268        let mut new_self = self.shallow_clone();
269
270        let _ = merge_variables(&mut new_self, patch, true, MergeStrategy::CloneOnWrite);
271        new_self
272    }
273
274    pub fn shallow_clone(&self) -> Self {
275        match self {
276            Variable::Null => Variable::Null,
277            Variable::Bool(b) => Variable::Bool(*b),
278            Variable::Number(n) => Variable::Number(*n),
279            Variable::String(s) => Variable::String(s.clone()),
280            Variable::Array(a) => Variable::Array(a.clone()),
281            Variable::Object(o) => Variable::Object(o.clone()),
282            Variable::Dynamic(d) => Variable::Dynamic(d.clone()),
283        }
284    }
285
286    pub fn deep_clone(&self) -> Self {
287        match self {
288            Variable::Array(a) => {
289                let arr = a.borrow();
290                Variable::from_array(arr.iter().map(|v| v.deep_clone()).collect())
291            }
292            Variable::Object(o) => {
293                let obj = o.borrow();
294                Variable::from_object(
295                    obj.iter()
296                        .map(|(k, v)| (k.clone(), v.deep_clone()))
297                        .collect(),
298                )
299            }
300            _ => self.shallow_clone(),
301        }
302    }
303
304    pub fn depth_clone(&self, depth: usize) -> Self {
305        match depth.is_zero() {
306            true => self.shallow_clone(),
307            false => match self {
308                Variable::Array(a) => {
309                    let arr = a.borrow();
310                    Variable::from_array(arr.iter().map(|v| v.depth_clone(depth - 1)).collect())
311                }
312                Variable::Object(o) => {
313                    let obj = o.borrow();
314                    Variable::from_object(
315                        obj.iter()
316                            .map(|(k, v)| (k.clone(), v.depth_clone(depth - 1)))
317                            .collect(),
318                    )
319                }
320                _ => self.shallow_clone(),
321            },
322        }
323    }
324}
325
326impl Clone for Variable {
327    fn clone(&self) -> Self {
328        self.shallow_clone()
329    }
330}
331
332#[derive(Copy, Clone)]
333enum MergeStrategy {
334    InPlace,
335    CloneOnWrite,
336}
337
338fn merge_variables(
339    doc: &mut Variable,
340    patch: &Variable,
341    top_level: bool,
342    strategy: MergeStrategy,
343) -> bool {
344    if patch.is_array() && top_level {
345        *doc = patch.shallow_clone();
346        return true;
347    }
348
349    if !patch.is_object() && top_level {
350        return false;
351    }
352
353    if doc.is_object() && patch.is_object() {
354        let doc_ref = doc.as_object().unwrap();
355        let patch_ref = patch.as_object().unwrap();
356        if Rc::ptr_eq(&doc_ref, &patch_ref) {
357            return false;
358        }
359
360        let patch = patch_ref.borrow();
361        match strategy {
362            MergeStrategy::InPlace => {
363                let mut map = doc_ref.borrow_mut();
364                for (key, value) in patch.deref() {
365                    if value == &Variable::Null {
366                        map.remove(key);
367                    } else {
368                        let entry = map.entry(key.clone()).or_insert(Variable::Null);
369                        merge_variables(entry, value, false, strategy);
370                    }
371                }
372
373                return true;
374            }
375            MergeStrategy::CloneOnWrite => {
376                let mut changed = false;
377                let mut new_map = None;
378
379                for (key, value) in patch.deref() {
380                    // Get or create the new map if we haven't yet
381                    let map = if let Some(ref mut m) = new_map {
382                        m
383                    } else {
384                        let m = doc_ref.borrow().clone();
385                        new_map = Some(m);
386                        new_map.as_mut().unwrap()
387                    };
388
389                    if value == &Variable::Null {
390                        // Remove null values
391                        if map.remove(key).is_some() {
392                            changed = true;
393                        }
394                    } else {
395                        // Handle nested merging
396                        let entry = map.entry(key.clone()).or_insert(Variable::Null);
397                        if merge_variables(entry, value, false, strategy) {
398                            changed = true;
399                        }
400                    }
401                }
402
403                // Only update doc if changes were made
404                if changed {
405                    if let Some(new_map) = new_map {
406                        *doc = Variable::Object(Rc::new(RefCell::new(new_map)));
407                    }
408                    return true;
409                }
410
411                return false;
412            }
413        }
414    } else {
415        let new_value = patch.shallow_clone();
416        if *doc != new_value {
417            *doc = new_value;
418            return true;
419        }
420
421        return false;
422    }
423}
424
425impl Default for Variable {
426    fn default() -> Self {
427        Variable::Null
428    }
429}
430
431impl Display for Variable {
432    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
433        match self {
434            Variable::Null => write!(f, "null"),
435            Variable::Bool(b) => match *b {
436                true => write!(f, "true"),
437                false => write!(f, "false"),
438            },
439            Variable::Number(n) => write!(f, "{n}"),
440            Variable::String(s) => write!(f, "\"{s}\""),
441            Variable::Array(arr) => {
442                let arr = arr.borrow();
443                let s = arr
444                    .iter()
445                    .map(|v| v.to_string())
446                    .collect::<Vec<String>>()
447                    .join(",");
448                write!(f, "[{s}]")
449            }
450            Variable::Object(obj) => {
451                let obj = obj.borrow();
452                let s = obj
453                    .iter()
454                    .map(|(k, v)| format!("\"{k}\":{v}"))
455                    .collect::<Vec<String>>()
456                    .join(",");
457
458                write!(f, "{{{s}}}")
459            }
460            Variable::Dynamic(d) => write!(f, "{d}"),
461        }
462    }
463}
464
465impl Debug for Variable {
466    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
467        write!(f, "{}", self)
468    }
469}
470
471impl PartialEq for Variable {
472    fn eq(&self, other: &Self) -> bool {
473        match (&self, &other) {
474            (Variable::Null, Variable::Null) => true,
475            (Variable::Bool(b1), Variable::Bool(b2)) => b1 == b2,
476            (Variable::Number(n1), Variable::Number(n2)) => n1 == n2,
477            (Variable::String(s1), Variable::String(s2)) => s1 == s2,
478            (Variable::Array(a1), Variable::Array(a2)) => a1 == a2,
479            (Variable::Object(obj1), Variable::Object(obj2)) => obj1 == obj2,
480            (Variable::Dynamic(d1), Variable::Dynamic(d2)) => Rc::ptr_eq(d1, d2),
481            _ => false,
482        }
483    }
484}
485
486impl Eq for Variable {}
487
488#[cfg(test)]
489mod tests {
490    use crate::variable::Variable;
491    use rust_decimal_macros::dec;
492    use serde_json::json;
493
494    #[test]
495    fn insert_detached() {
496        let some_data: Variable = json!({ "customer": { "firstName": "John" }}).into();
497
498        let a_a = some_data
499            .dot_insert_detached("a.a", Variable::Number(dec!(1)))
500            .unwrap();
501        let a_b = a_a
502            .dot_insert_detached("a.b", Variable::Number(dec!(2)))
503            .unwrap();
504        let a_c = a_b
505            .dot_insert_detached("a.c", Variable::Number(dec!(3)))
506            .unwrap();
507
508        assert_eq!(a_a.dot("a"), Some(Variable::from(json!({ "a": 1 }))));
509        assert_eq!(
510            a_b.dot("a"),
511            Some(Variable::from(json!({ "a": 1, "b": 2 })))
512        );
513        assert_eq!(
514            a_c.dot("a"),
515            Some(Variable::from(json!({ "a": 1, "b": 2, "c": 3 })))
516        );
517    }
518}
zen_types/variable/mod.rs

zen_types/variable/
mod.rs
