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_f64() {
146                    return Value::F64(num);
147                }
148
149                if let Some(num) = num.as_i64() {
150                    return Value::I64(num);
151                }
152
153                if let Some(num) = num.as_u64() {
154                    return Value::U64(num);
155                }
156
157                Value::Null
158            }
159            ValueRef::Array(arr) => {
160                let mut vec = Vec::with_capacity(arr.len());
161                vec.extend(arr.iter().map(|v| Value::from(v.as_ref())));
162                Value::Array(vec)
163            }
164            ValueRef::Object(obj) => {
165                let mut vec = Vec::with_capacity(obj.len());
166                vec.extend(obj.iter().map(|(k, v)| (k, Value::from(v.as_ref()))));
167                vec.sort_unstable_by_key(|(k, _)| *k);
168                Value::Object(vec)
169            }
170        }
171    }
172
173    pub fn as_object(&self) -> Option<&Vec<(&'a str, Value<'a>)>> {
174        match self {
175            Value::Object(obj) => Some(obj),
176            _ => None,
177        }
178    }
179
180    pub fn as_str(&self) -> Option<&str> {
181        match self {
182            Value::String(s) => Some(s),
183            _ => None,
184        }
185    }
186
187    pub fn is_null(&self) -> bool {
188        matches!(self, Value::Null)
189    }
190
191    pub fn is_object(&self) -> bool {
192        matches!(self, Value::Object(_))
193    }
194}
195
196// Our new trait with the desired methods
197pub trait ValueRefExt {
198    fn to_data<'a>(&'a self) -> Option<ValueRef<'a>>;
199    fn to_entities<'a>(&'a self) -> Option<Vec<ValueRef<'a>>>;
200}
201
202// Implement our trait for the foreign type
203impl ValueRefExt for ValueRef<'_> {
204    fn to_data<'a>(&'a self) -> Option<ValueRef<'a>> {
205        match self {
206            ValueRef::Object(obj) => obj.get(&"data").map(|v| v.as_ref()),
207            _ => None,
208        }
209    }
210
211    fn to_entities<'a>(&'a self) -> Option<Vec<ValueRef<'a>>> {
212        match self.to_data().unwrap() {
213            ValueRef::Object(obj) => obj.get(&"_entities").and_then(|v| match v.as_ref() {
214                ValueRef::Array(arr) => Some(arr.iter().map(|v| v.as_ref()).collect()),
215                _ => None,
216            }),
217            _ => None,
218        }
219    }
220}
221
222impl Display for Value<'_> {
223    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224        match self {
225            Value::Null => write!(f, "null"),
226            Value::Bool(b) => write!(f, "{}", b),
227            Value::String(s) => write!(f, "\"{}\"", s),
228            Value::F64(n) => write!(f, "{}", n),
229            Value::U64(n) => write!(f, "{}", n),
230            Value::I64(n) => write!(f, "{}", n),
231            Value::Array(arr) => {
232                write!(f, "[")?;
233                for (i, v) in arr.iter().enumerate() {
234                    if i > 0 {
235                        write!(f, ", ")?;
236                    }
237                    write!(f, "{}", v)?;
238                }
239                write!(f, "]")
240            }
241            Value::Object(obj) => {
242                write!(f, "{{")?;
243                for (i, (k, v)) in obj.iter().enumerate() {
244                    if i > 0 {
245                        write!(f, ", ")?;
246                    }
247                    write!(f, "\"{}\": {}", k, v)?;
248                }
249                write!(f, "}}")
250            }
251        }
252    }
253}
254
255struct ValueVisitor<'a> {
256    // We need a marker to hold the lifetime 'a.
257    _marker: std::marker::PhantomData<&'a ()>,
258}
259
260impl<'de> de::Deserialize<'de> for Value<'de> {
261    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
262    where
263        D: Deserializer<'de>,
264    {
265        deserializer.deserialize_any(ValueVisitor {
266            _marker: std::marker::PhantomData,
267        })
268    }
269}
270
271impl<'de> Visitor<'de> for ValueVisitor<'de> {
272    type Value = Value<'de>;
273
274    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
275        formatter.write_str("any valid JSON value")
276    }
277
278    fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
279        Ok(Value::Bool(value))
280    }
281
282    fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
283        Ok(Value::I64(value))
284    }
285
286    fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
287        Ok(Value::U64(value))
288    }
289
290    fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E> {
291        Ok(Value::F64(value))
292    }
293
294    // This is the zero-copy part. We borrow the string slice directly from the input.
295    fn visit_borrowed_str<E>(self, value: &'de str) -> Result<Self::Value, E>
296    where
297        E: de::Error,
298    {
299        Ok(Value::String(value.into()))
300    }
301
302    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
303    where
304        E: de::Error,
305    {
306        Ok(Value::String(v.to_owned().into()))
307    }
308
309    fn visit_string<E>(self, value: String) -> Result<Self::Value, E>
310    where
311        E: de::Error,
312    {
313        Ok(Value::String(value.into()))
314    }
315
316    fn visit_unit<E>(self) -> Result<Self::Value, E> {
317        Ok(Value::Null)
318    }
319
320    // For arrays, Serde recursively calls `deserialize` for the inner type.
321    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
322    where
323        A: SeqAccess<'de>,
324    {
325        let mut elements = Vec::with_capacity(seq.size_hint().unwrap_or(0));
326        while let Some(elem) = seq.next_element()? {
327            elements.push(elem);
328        }
329        Ok(Value::Array(elements))
330    }
331
332    // For objects, the same happens for keys and values.
333    fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
334    where
335        M: MapAccess<'de>,
336    {
337        let mut entries = Vec::with_capacity(map.size_hint().unwrap_or(0));
338        while let Some((key, value)) = map.next_entry()? {
339            entries.push((key, value));
340        }
341        // IMPORTANT: We keep the sort for binary search compatibility.
342        entries.sort_unstable_by_key(|(k, _)| *k);
343        Ok(Value::Object(entries))
344    }
345}
346
347impl serde::Serialize for Value<'_> {
348    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
349    where
350        S: serde::Serializer,
351    {
352        match self {
353            Value::Null => serializer.serialize_unit(),
354            Value::Bool(b) => serializer.serialize_bool(*b),
355            Value::I64(n) => serializer.serialize_i64(*n),
356            Value::U64(n) => serializer.serialize_u64(*n),
357            Value::F64(n) => serializer.serialize_f64(*n),
358            Value::String(s) => serializer.serialize_str(s),
359            Value::Array(arr) => {
360                let mut seq = serializer.serialize_seq(Some(arr.len()))?;
361                for v in arr {
362                    seq.serialize_element(v)?;
363                }
364                seq.end()
365            }
366            Value::Object(obj) => {
367                let mut map = serializer.serialize_map(Some(obj.len()))?;
368                for (k, v) in obj {
369                    map.serialize_entry(k, v)?;
370                }
371                map.end()
372            }
373        }
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::Value;
380    use serde::Deserialize;
381    use std::borrow::Cow;
382
383    #[test]
384    fn deserializes_escaped_string_as_owned() {
385        let bytes = br#"{"message": "hello\nworld"}"#;
386        let mut deserializer = sonic_rs::Deserializer::from_slice(bytes);
387
388        let value = Value::deserialize(&mut deserializer).unwrap();
389
390        let obj = match value {
391            Value::Object(obj) => obj,
392            _ => panic!("Expected Value::Object"),
393        };
394
395        let message_value = &obj.iter().find(|(k, _)| *k == "message").unwrap().1;
396
397        match message_value {
398            Value::String(value) => {
399                assert_eq!(value, "hello\nworld");
400                assert!(
401                    matches!(value, Cow::Owned(_)),
402                    "Expected Cow::Owned for escaped string"
403                );
404            }
405            _ => panic!("Expected Value::String"),
406        }
407    }
408
409    #[test]
410    fn deserializes_simple_string_as_borrowed() {
411        let bytes = br#"{"message": "hello world"}"#;
412        let mut deserializer = sonic_rs::Deserializer::from_slice(bytes);
413        let value = Value::deserialize(&mut deserializer).unwrap();
414
415        let obj = match value {
416            Value::Object(obj) => obj,
417            _ => panic!("Expected Value::Object"),
418        };
419
420        let message_value = &obj.iter().find(|(k, _)| *k == "message").unwrap().1;
421
422        match message_value {
423            Value::String(value) => {
424                assert_eq!(value, "hello world");
425                assert!(
426                    matches!(value, Cow::Borrowed(_)),
427                    "Expected Cow::Borrowed for simple string"
428                );
429            }
430            _ => panic!("Expected Value::String"),
431        }
432    }
433}
hive_router_plan_executor/response/value.rs

hive_router_plan_executor/response/
value.rs
