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

hive_router_plan_executor/response/
value.rs
