Skip to main content

zen_types/variable_type/
util.rs

1use crate::variable_type::VariableType;
2use ahash::{HashMap, HashMapExt};
3use rust_decimal::prelude::Zero;
4use std::cell::RefCell;
5use std::collections::hash_map::Entry;
6use std::rc::Rc;
7
8impl VariableType {
9    pub fn unwrap_nullable(&self) -> (&VariableType, bool) {
10        match self {
11            VariableType::Nullable(inner) => (inner.as_ref(), true),
12            other => (other, false),
13        }
14    }
15
16    pub fn iterator(&self) -> Option<Rc<VariableType>> {
17        match self {
18            VariableType::Array(item) => Some(item.clone()),
19            VariableType::Interval => Some(Rc::new(VariableType::Number)),
20            VariableType::Nullable(inner) => inner.iterator(),
21            _ => None,
22        }
23    }
24
25    pub fn as_const_str(&self) -> Option<Rc<str>> {
26        match self {
27            VariableType::Const(s) => Some(s.clone()),
28            VariableType::Nullable(inner) => inner.as_const_str(),
29            _ => None,
30        }
31    }
32
33    pub fn get(&self, key: &str) -> VariableType {
34        match self {
35            VariableType::Object(obj) => {
36                let obj = obj.borrow();
37                obj.get(key).cloned().unwrap_or(VariableType::Any)
38            }
39            VariableType::Nullable(inner) => inner.get(key),
40            _ => VariableType::Null,
41        }
42    }
43
44    pub fn satisfies(&self, constraint: &Self) -> bool {
45        match (self, constraint) {
46            (VariableType::Any, _) | (_, VariableType::Any) => true,
47            (VariableType::Nullable(a), VariableType::Nullable(b)) => a.satisfies(b),
48            (VariableType::Nullable(_), _) => false,
49            (other, VariableType::Nullable(inner)) => other.satisfies(inner),
50
51            (VariableType::Null, VariableType::Null) => true,
52            (VariableType::Bool, VariableType::Bool) => true,
53            (VariableType::String, VariableType::String) => true,
54            (VariableType::Number, VariableType::Number) => true,
55            (VariableType::Date, VariableType::Date) => true,
56            (VariableType::Number, VariableType::Date) => true,
57            (_, VariableType::Date) if self.widen().is_string() => true,
58            (VariableType::Interval, VariableType::Interval) => true,
59            (VariableType::Array(a1), VariableType::Array(a2)) => a1.satisfies(a2),
60            (VariableType::Object(o1), VariableType::Object(o2)) => {
61                let o1 = o1.borrow();
62                let o2 = o2.borrow();
63
64                o2.iter().all(|(k, v)| match o1.get(k) {
65                    Some(tv) => tv.satisfies(v),
66                    None => matches!(
67                        v,
68                        VariableType::Any | VariableType::Null | VariableType::Nullable(_)
69                    ),
70                })
71            }
72
73            (VariableType::Const(c1), VariableType::Const(c2)) => c1 == c2,
74            (VariableType::Const(c), VariableType::Enum(_, e)) => e.iter().any(|e| e == c),
75            (VariableType::Const(_), VariableType::String) => true,
76            (VariableType::String, VariableType::Const(_)) => true,
77
78            (VariableType::Enum(_, e1), VariableType::Enum(_, e2)) => {
79                e1.iter().all(|c| e2.contains(c))
80            }
81            (VariableType::Enum(_, e), VariableType::Const(c)) => e.iter().all(|i| i == c),
82            (VariableType::Enum(_, _), VariableType::String) => true,
83            (VariableType::String, VariableType::Enum(_, _)) => true,
84
85            (_, _) => false,
86        }
87    }
88
89    pub fn is_array(&self) -> bool {
90        match self {
91            VariableType::Any | VariableType::Array(_) => true,
92            VariableType::Nullable(inner) => inner.is_array(),
93            _ => false,
94        }
95    }
96
97    pub fn is_iterable(&self) -> bool {
98        match self {
99            VariableType::Any | VariableType::Interval | VariableType::Array(_) => true,
100            VariableType::Nullable(inner) => inner.is_iterable(),
101            _ => false,
102        }
103    }
104
105    pub fn is_string(&self) -> bool {
106        match self {
107            VariableType::String => true,
108            VariableType::Nullable(inner) => inner.is_string(),
109            _ => false,
110        }
111    }
112
113    pub fn is_object(&self) -> bool {
114        match self {
115            VariableType::Any | VariableType::Object(_) => true,
116            VariableType::Nullable(inner) => inner.is_object(),
117            _ => false,
118        }
119    }
120
121    pub fn is_null(&self) -> bool {
122        match self {
123            VariableType::Null => true,
124            _ => false,
125        }
126    }
127
128    pub fn is_nullable(&self) -> bool {
129        matches!(self, VariableType::Nullable(_) | VariableType::Null)
130    }
131
132    pub fn widen(&self) -> Self {
133        match self {
134            VariableType::Const(_) | VariableType::Enum(_, _) => VariableType::String,
135            VariableType::Nullable(inner) => {
136                let widened = inner.widen();
137                VariableType::Nullable(Rc::new(widened))
138            }
139            _ => self.clone(),
140        }
141    }
142
143    pub fn merge(&self, other: &Self) -> Self {
144        let (left, left_nullable) = self.unwrap_nullable();
145        let (right, right_nullable) = other.unwrap_nullable();
146        let nullable = left_nullable || right_nullable;
147
148        let result = match (left, right) {
149            (VariableType::Any, _) | (_, VariableType::Any) => VariableType::Any,
150            (VariableType::Null, VariableType::Null) => VariableType::Null,
151            (VariableType::Null, other) | (other, VariableType::Null) => {
152                VariableType::Nullable(Rc::new(other.clone()))
153            }
154            (VariableType::Bool, VariableType::Bool) => VariableType::Bool,
155            (VariableType::String, VariableType::String) => VariableType::String,
156            (VariableType::Number, VariableType::Number) => VariableType::Number,
157            (VariableType::Date, VariableType::Date) => VariableType::Date,
158            (VariableType::Interval, VariableType::Interval) => VariableType::Interval,
159            (VariableType::Array(a1), VariableType::Array(a2)) => {
160                if Rc::ptr_eq(a1, a2) {
161                    VariableType::Array(a1.clone())
162                } else {
163                    match (a1.as_ref(), a2.as_ref()) {
164                        (VariableType::Any, other) | (other, VariableType::Any) => {
165                            VariableType::Array(Rc::new(other.clone()))
166                        }
167                        (l, r) => VariableType::Array(Rc::new(l.merge(r))),
168                    }
169                }
170            }
171            (VariableType::Object(o1), VariableType::Object(o2)) => {
172                let o1 = o1.borrow();
173                let o2 = o2.borrow();
174
175                let mut merged = HashMap::with_capacity(o1.len().max(o2.len()));
176                for (k, v) in o1.iter() {
177                    merged.insert(k.clone(), v.clone());
178                }
179
180                for (k, v) in o2.iter() {
181                    match merged.entry(k.clone()) {
182                        Entry::Occupied(mut entry) => {
183                            let current = entry.get();
184                            entry.insert(current.merge(v));
185                        }
186                        Entry::Vacant(entry) => {
187                            entry.insert(v.clone());
188                        }
189                    }
190                }
191
192                VariableType::Object(Rc::new(RefCell::new(merged)))
193            }
194            (VariableType::Const(c), VariableType::Enum(_, values)) => {
195                let mut merged = values.clone();
196                if !merged.contains(c) {
197                    merged.push(c.clone());
198                }
199                VariableType::Enum(None, merged)
200            }
201            (VariableType::Const(c1), VariableType::Const(c2)) => {
202                if Rc::ptr_eq(c1, c2) || c1 == c2 {
203                    VariableType::Const(c1.clone())
204                } else {
205                    VariableType::Enum(None, vec![c1.clone(), c2.clone()])
206                }
207            }
208            (VariableType::Const(_), VariableType::String)
209            | (VariableType::String, VariableType::Const(_)) => VariableType::String,
210            (VariableType::Enum(n1, a), VariableType::Enum(n2, b)) => {
211                let mut merged = a.clone();
212                for val in b {
213                    if !merged.contains(val) {
214                        merged.push(val.clone());
215                    }
216                }
217
218                let name = match (n1, n2) {
219                    (Some(n1), Some(n2)) => Some(Rc::<str>::from(format!("{} | {}", n1, n2))),
220                    _ => None,
221                };
222
223                VariableType::Enum(name, merged)
224            }
225            (VariableType::Enum(_, values), VariableType::Const(c)) => {
226                let mut merged = values.clone();
227                if !merged.contains(c) {
228                    merged.push(c.clone());
229                }
230                VariableType::Enum(None, merged)
231            }
232            (VariableType::Enum(_, _), VariableType::String)
233            | (VariableType::String, VariableType::Enum(_, _)) => VariableType::String,
234            (_, _) => VariableType::Any,
235        };
236
237        if nullable
238            && !matches!(
239                &result,
240                VariableType::Any | VariableType::Null | VariableType::Nullable(_)
241            )
242        {
243            VariableType::Nullable(Rc::new(result))
244        } else {
245            result
246        }
247    }
248
249    pub fn shallow_clone(&self) -> Self {
250        match self {
251            VariableType::Any => VariableType::Any,
252            VariableType::Null => VariableType::Null,
253            VariableType::Bool => VariableType::Bool,
254            VariableType::String => VariableType::String,
255            VariableType::Number => VariableType::Number,
256            VariableType::Date => VariableType::Date,
257            VariableType::Interval => VariableType::Interval,
258            VariableType::Array(arr) => VariableType::Array(arr.clone()),
259            VariableType::Object(obj) => VariableType::Object(obj.clone()),
260            VariableType::Const(c) => VariableType::Const(c.clone()),
261            VariableType::Enum(name, options) => VariableType::Enum(name.clone(), options.clone()),
262            VariableType::Nullable(inner) => VariableType::Nullable(Rc::new(inner.shallow_clone())),
263        }
264    }
265
266    pub fn dot_head(&self, key: &str) -> Option<Self> {
267        let mut parts = Vec::from_iter(key.split('.'));
268        parts.pop();
269
270        parts
271            .iter()
272            .try_fold(self.shallow_clone(), |var, part| match var {
273                VariableType::Object(obj) => {
274                    let mut obj_ref = obj.borrow_mut();
275                    Some(match obj_ref.entry(Rc::from(*part)) {
276                        Entry::Occupied(occ) => occ.get().shallow_clone(),
277                        Entry::Vacant(vac) => vac.insert(Self::empty_object()).shallow_clone(),
278                    })
279                }
280                _ => None,
281            })
282    }
283
284    pub fn dot_head_detach(&self, key: &str) -> (Self, Option<Self>) {
285        let mut parts = Vec::from_iter(key.split('.'));
286        parts.pop();
287
288        let cloned_self = self.depth_clone(1);
289        let head = parts
290            .iter()
291            .try_fold(cloned_self.shallow_clone(), |var, part| match var {
292                VariableType::Object(obj) => {
293                    let mut obj_ref = obj.borrow_mut();
294                    Some(match obj_ref.entry(Rc::from(*part)) {
295                        Entry::Occupied(mut occ) => {
296                            let var = occ.get();
297                            let new_obj = match var {
298                                VariableType::Object(_) => var.depth_clone(1),
299                                _ => VariableType::empty_object(),
300                            };
301
302                            occ.insert(new_obj.shallow_clone());
303                            new_obj
304                        }
305                        Entry::Vacant(vac) => vac.insert(Self::empty_object()).shallow_clone(),
306                    })
307                }
308                _ => None,
309            });
310
311        (cloned_self, head)
312    }
313
314    pub fn depth_clone(&self, depth: usize) -> Self {
315        match depth.is_zero() {
316            true => self.shallow_clone(),
317            false => match self {
318                VariableType::Object(o) => {
319                    let obj = o.borrow();
320                    VariableType::Object(Rc::new(RefCell::new(
321                        obj.iter()
322                            .map(|(k, v)| (k.clone(), v.depth_clone(depth - 1)))
323                            .collect(),
324                    )))
325                }
326                VariableType::Nullable(inner) => {
327                    VariableType::Nullable(Rc::new(inner.depth_clone(depth)))
328                }
329                _ => self.shallow_clone(),
330            },
331        }
332    }
333
334    pub fn empty_object() -> Self {
335        VariableType::Object(Rc::new(RefCell::new(HashMap::new())))
336    }
337
338    pub fn dot_insert_detached(&self, key: &str, variable: Self) -> Option<Self> {
339        let last_part = key.split('.').last()?;
340        let (new_var, head_opt) = self.dot_head_detach(key);
341        let head = head_opt?;
342        let VariableType::Object(object_ref) = head else {
343            return None;
344        };
345
346        let mut object = object_ref.borrow_mut();
347        object.insert(Rc::from(last_part), variable);
348        Some(new_var)
349    }
350
351    pub fn dot_insert(&self, key: &str, variable: Self) -> Option<Self> {
352        let last_part = key.split('.').last()?;
353        let head = self.dot_head(key)?;
354        let Self::Object(object_ref) = head else {
355            return None;
356        };
357
358        let mut object = object_ref.borrow_mut();
359        object.insert(Rc::from(last_part), variable)
360    }
361
362    pub fn dot(&self, key: &str) -> Option<Self> {
363        key.split('.')
364            .try_fold(self.shallow_clone(), |var, part| match var {
365                Self::Object(obj) => {
366                    let reference = obj.borrow();
367                    reference.get(part).map(|v| v.shallow_clone())
368                }
369                _ => None,
370            })
371    }
372}