hive_router_plan_executor/response/
value.rs

1use core::fmt;
2use hive_router_query_planner::ast::selection_item::SelectionItem;
3use serde::{
4    de::{self, Deserializer, MapAccess, SeqAccess, Visitor},
5    ser::{SerializeMap, SerializeSeq},
6};
7use sonic_rs::{JsonNumberTrait, ValueRef};
8use std::{
9    fmt::Display,
10    hash::{Hash, Hasher},
11};
12use xxhash_rust::xxh3::Xxh3;
13
14use crate::{introspection::schema::PossibleTypes, utils::consts::TYPENAME_FIELD_NAME};
15
16#[derive(Clone)]
17pub enum Value<'a> {
18    Null,
19    F64(f64),
20    I64(i64),
21    U64(u64),
22    Bool(bool),
23    String(&'a str),
24    Array(Vec<Value<'a>>),
25    Object(Vec<(&'a str, Value<'a>)>),
26}
27
28impl Hash for Value<'_> {
29    fn hash<H: Hasher>(&self, state: &mut H) {
30        match self {
31            Value::Null => 0.hash(state),
32            Value::F64(f) => f.to_bits().hash(state),
33            Value::I64(i) => i.hash(state),
34            Value::U64(u) => u.hash(state),
35            Value::Bool(b) => b.hash(state),
36            Value::String(s) => s.hash(state),
37            Value::Array(arr) => arr.hash(state),
38            Value::Object(obj) => obj.hash(state),
39        }
40    }
41}
42
43impl<'a> Value<'a> {
44    pub fn take_entities<'b: 'a>(&'a mut self) -> Option<Vec<Value<'a>>> {
45        match self {
46            Value::Object(data) => {
47                if let Ok(entities_idx) = data.binary_search_by_key(&"_entities", |(k, _)| *k) {
48                    if let Value::Array(arr) = data.remove(entities_idx).1 {
49                        return Some(arr);
50                    }
51                }
52                None
53            }
54            _ => None,
55        }
56    }
57
58    pub fn to_hash(
59        &self,
60        selection_items: &[SelectionItem],
61        possible_types: &PossibleTypes,
62    ) -> u64 {
63        let mut hasher = Xxh3::new();
64        self.hash_with_requires(&mut hasher, selection_items, possible_types);
65        hasher.finish()
66    }
67
68    fn hash_with_requires<H: Hasher>(
69        &self,
70        state: &mut H,
71        selection_items: &[SelectionItem],
72        possible_types: &PossibleTypes,
73    ) {
74        if selection_items.is_empty() {
75            self.hash(state);
76            return;
77        }
78
79        match self {
80            Value::Object(obj) => {
81                Value::hash_object_with_requires(state, obj, selection_items, possible_types);
82            }
83            Value::Array(arr) => {
84                for item in arr {
85                    item.hash_with_requires(state, selection_items, possible_types);
86                }
87            }
88            _ => {
89                self.hash(state);
90            }
91        }
92    }
93
94    fn hash_object_with_requires<H: Hasher>(
95        state: &mut H,
96        obj: &[(&'a str, Value<'a>)],
97        selection_items: &[SelectionItem],
98        possible_types: &PossibleTypes,
99    ) {
100        for item in selection_items {
101            match item {
102                SelectionItem::Field(field_selection) => {
103                    let field_name = &field_selection.name;
104                    if let Ok(idx) = obj.binary_search_by_key(&field_name.as_str(), |(k, _)| k) {
105                        let (key, value) = &obj[idx];
106                        key.hash(state);
107                        value.hash_with_requires(
108                            state,
109                            &field_selection.selections.items,
110                            possible_types,
111                        );
112                    }
113                }
114                SelectionItem::InlineFragment(inline_fragment) => {
115                    let type_condition = &inline_fragment.type_condition;
116                    let type_name = obj
117                        .binary_search_by_key(&TYPENAME_FIELD_NAME, |(k, _)| k)
118                        .ok()
119                        .and_then(|idx| obj[idx].1.as_str())
120                        .unwrap_or(type_condition);
121
122                    if possible_types.entity_satisfies_type_condition(type_name, type_condition) {
123                        Value::hash_object_with_requires(
124                            state,
125                            obj,
126                            &inline_fragment.selections.items,
127                            possible_types,
128                        );
129                    }
130                }
131                SelectionItem::FragmentSpread(_) => {
132                    unreachable!("Fragment spreads should not exist in FetchNode::requires.")
133                }
134            }
135        }
136    }
137
138    pub fn from(json: ValueRef<'a>) -> Value<'a> {
139        match json {
140            ValueRef::Null => Value::Null,
141            ValueRef::Bool(b) => Value::Bool(b),
142            ValueRef::String(s) => Value::String(s),
143            ValueRef::Number(num) => {
144                if let Some(num) = num.as_f64() {
145                    return Value::F64(num);
146                }
147
148                if let Some(num) = num.as_i64() {
149                    return Value::I64(num);
150                }
151
152                if let Some(num) = num.as_u64() {
153                    return Value::U64(num);
154                }
155
156                Value::Null
157            }
158            ValueRef::Array(arr) => {
159                let mut vec = Vec::with_capacity(arr.len());
160                vec.extend(arr.iter().map(|v| Value::from(v.as_ref())));
161                Value::Array(vec)
162            }
163            ValueRef::Object(obj) => {
164                let mut vec = Vec::with_capacity(obj.len());
165                vec.extend(obj.iter().map(|(k, v)| (k, Value::from(v.as_ref()))));
166                vec.sort_unstable_by_key(|(k, _)| *k);
167                Value::Object(vec)
168            }
169        }
170    }
171
172    pub fn as_object(&self) -> Option<&Vec<(&'a str, Value<'a>)>> {
173        match self {
174            Value::Object(obj) => Some(obj),
175            _ => None,
176        }
177    }
178
179    pub fn as_str(&self) -> Option<&str> {
180        match self {
181            Value::String(s) => Some(s),
182            _ => None,
183        }
184    }
185
186    pub fn is_null(&self) -> bool {
187        matches!(self, Value::Null)
188    }
189
190    pub fn is_object(&self) -> bool {
191        matches!(self, Value::Object(_))
192    }
193
194    #[inline]
195    pub fn estimate_size(&self) -> usize {
196        match self {
197            Value::String(s) => key_size(s),
198            Value::Object(obj) => {
199                let mut total = 0usize;
200                for (k, v) in obj.iter() {
201                    total += key_size(k);
202                    total += v.estimate_size();
203                }
204                // braces + commas + colons
205                let len = obj.len();
206                total + 2 + len + len.saturating_sub(1)
207            }
208            Value::Bool(_) => 5, /* false = 5, true = 4 */
209            Value::Null => 4,
210            Value::Array(arr) => {
211                let mut total = 0usize;
212                for v in arr.iter() {
213                    total += v.estimate_size();
214                }
215                // brackets + commas
216                let len = arr.len();
217                total + 2 + len.saturating_sub(1)
218            }
219            Value::F64(_) => 25, /* worst case scenario */
220            Value::I64(n) => i64_len(*n),
221            Value::U64(n) => u64_len(*n),
222        }
223    }
224}
225
226// Our new trait with the desired methods
227pub trait ValueRefExt {
228    fn to_data<'a>(&'a self) -> Option<ValueRef<'a>>;
229    fn to_entities<'a>(&'a self) -> Option<Vec<ValueRef<'a>>>;
230}
231
232// Implement our trait for the foreign type
233impl ValueRefExt for ValueRef<'_> {
234    fn to_data<'a>(&'a self) -> Option<ValueRef<'a>> {
235        match self {
236            ValueRef::Object(obj) => obj.get(&"data").map(|v| v.as_ref()),
237            _ => None,
238        }
239    }
240
241    fn to_entities<'a>(&'a self) -> Option<Vec<ValueRef<'a>>> {
242        match self.to_data().unwrap() {
243            ValueRef::Object(obj) => obj.get(&"_entities").and_then(|v| match v.as_ref() {
244                ValueRef::Array(arr) => Some(arr.iter().map(|v| v.as_ref()).collect()),
245                _ => None,
246            }),
247            _ => None,
248        }
249    }
250}
251
252impl Display for Value<'_> {
253    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254        match self {
255            Value::Null => write!(f, "null"),
256            Value::Bool(b) => write!(f, "{}", b),
257            Value::String(s) => write!(f, "\"{}\"", s),
258            Value::F64(n) => write!(f, "{}", n),
259            Value::U64(n) => write!(f, "{}", n),
260            Value::I64(n) => write!(f, "{}", n),
261            Value::Array(arr) => {
262                write!(f, "[")?;
263                for (i, v) in arr.iter().enumerate() {
264                    if i > 0 {
265                        write!(f, ", ")?;
266                    }
267                    write!(f, "{}", v)?;
268                }
269                write!(f, "]")
270            }
271            Value::Object(obj) => {
272                write!(f, "{{")?;
273                for (i, (k, v)) in obj.iter().enumerate() {
274                    if i > 0 {
275                        write!(f, ", ")?;
276                    }
277                    write!(f, "\"{}\": {}", k, v)?;
278                }
279                write!(f, "}}")
280            }
281        }
282    }
283}
284
285struct ValueVisitor<'a> {
286    // We need a marker to hold the lifetime 'a.
287    _marker: std::marker::PhantomData<&'a ()>,
288}
289
290impl<'de> de::Deserialize<'de> for Value<'de> {
291    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
292    where
293        D: Deserializer<'de>,
294    {
295        deserializer.deserialize_any(ValueVisitor {
296            _marker: std::marker::PhantomData,
297        })
298    }
299}
300
301impl<'de> Visitor<'de> for ValueVisitor<'de> {
302    type Value = Value<'de>;
303
304    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
305        formatter.write_str("any valid JSON value")
306    }
307
308    fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
309        Ok(Value::Bool(value))
310    }
311
312    fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
313        Ok(Value::I64(value))
314    }
315
316    fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
317        Ok(Value::U64(value))
318    }
319
320    fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E> {
321        Ok(Value::F64(value))
322    }
323
324    // This is the zero-copy part. We borrow the string slice directly from the input.
325    fn visit_borrowed_str<E>(self, value: &'de str) -> Result<Self::Value, E>
326    where
327        E: de::Error,
328    {
329        Ok(Value::String(value))
330    }
331
332    fn visit_unit<E>(self) -> Result<Self::Value, E> {
333        Ok(Value::Null)
334    }
335
336    // For arrays, Serde recursively calls `deserialize` for the inner type.
337    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
338    where
339        A: SeqAccess<'de>,
340    {
341        let mut elements = Vec::with_capacity(seq.size_hint().unwrap_or(0));
342        while let Some(elem) = seq.next_element()? {
343            elements.push(elem);
344        }
345        Ok(Value::Array(elements))
346    }
347
348    // For objects, the same happens for keys and values.
349    fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
350    where
351        M: MapAccess<'de>,
352    {
353        let mut entries = Vec::with_capacity(map.size_hint().unwrap_or(0));
354        while let Some((key, value)) = map.next_entry()? {
355            entries.push((key, value));
356        }
357        // IMPORTANT: We keep the sort for binary search compatibility.
358        entries.sort_unstable_by_key(|(k, _)| *k);
359        Ok(Value::Object(entries))
360    }
361}
362
363impl serde::Serialize for Value<'_> {
364    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
365    where
366        S: serde::Serializer,
367    {
368        match self {
369            Value::Null => serializer.serialize_unit(),
370            Value::Bool(b) => serializer.serialize_bool(*b),
371            Value::I64(n) => serializer.serialize_i64(*n),
372            Value::U64(n) => serializer.serialize_u64(*n),
373            Value::F64(n) => serializer.serialize_f64(*n),
374            Value::String(s) => serializer.serialize_str(s),
375            Value::Array(arr) => {
376                let mut seq = serializer.serialize_seq(Some(arr.len()))?;
377                for v in arr {
378                    seq.serialize_element(v)?;
379                }
380                seq.end()
381            }
382            Value::Object(obj) => {
383                let mut map = serializer.serialize_map(Some(obj.len()))?;
384                for (k, v) in obj {
385                    map.serialize_entry(k, v)?;
386                }
387                map.end()
388            }
389        }
390    }
391}
392#[inline]
393fn key_size(key: &str) -> usize {
394    key.len() + 2 /* for "<key>" */
395}
396
397/// Calculates the number of digits in a u64.
398#[inline]
399fn u64_len(n: u64) -> usize {
400    if n == 0 {
401        return 1;
402    }
403    n.ilog10() as usize + 1
404}
405
406/// Calculates the number of characters for an i64, including the sign.
407#[inline]
408fn i64_len(n: i64) -> usize {
409    if n == 0 {
410        return 1;
411    }
412    let len = if n.is_negative() { 1 } else { 0 };
413    len + n.unsigned_abs().ilog10() as usize + 1
414}
hive_router_plan_executor/response/value.rs

hive_router_plan_executor/response/
value.rs
