zen_expression/variable/
mod.rs

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