helix_db/protocol/
return_values.rs

1use super::{
2    remapping::{Remapping, ResponseRemapping},
3    value::Value,
4};
5use crate::utils::{
6    count::Count,
7    filterable::{Filterable, FilterableType},
8    items::{Edge, Node},
9};
10use crate::helix_engine::graph_core::ops::tr_val::TraversalVal;
11use sonic_rs::{Deserialize, Serialize};
12use std::{cell::RefMut, collections::HashMap};
13
14/// A return value enum that represents different possible outputs from graph operations.
15/// Can contain traversal results, counts, boolean flags, or empty values.
16#[derive(Deserialize, Debug, Clone, PartialEq)]
17pub enum ReturnValue {
18    Array(Vec<ReturnValue>),
19    Object(HashMap<String, ReturnValue>),
20    Value(Value),
21    Empty,
22}
23
24impl ReturnValue {
25    #[inline]
26    #[allow(unused_attributes)]
27    #[ignore = "No use for this function yet, however, I believe it may be useful in the future so I'm keeping it here"]
28    pub fn from_properties(properties: HashMap<String, Value>) -> Self {
29        ReturnValue::Object(
30            properties
31                .into_iter()
32                .map(|(k, v)| (k, ReturnValue::Value(v)))
33                .collect(),
34        )
35    }
36
37    #[inline(always)]
38    fn process_items_with_mixin<T>(
39        item: T,
40        mixin: &mut HashMap<u128, ResponseRemapping>,
41    ) -> ReturnValue
42    where
43        T: Filterable + Clone,
44    {
45        if let Some(m) = mixin.remove(&item.id()) {
46            if m.should_spread {
47                ReturnValue::from(item).mixin_remapping(m.remappings)
48            } else {
49                ReturnValue::default().mixin_remapping(m.remappings)
50            }
51        } else {
52            ReturnValue::from(item)
53        }
54    }
55
56    #[inline]
57    pub fn from_traversal_value_array_with_mixin(
58        traversal_value: Vec<TraversalVal>,
59        mut mixin: RefMut<HashMap<u128, ResponseRemapping>>,
60    ) -> Self {
61        ReturnValue::Array(
62            traversal_value
63                .into_iter()
64                .map(|val| match val {
65                    TraversalVal::Node(node) => {
66                        ReturnValue::process_items_with_mixin(node, &mut mixin)
67                    }
68                    TraversalVal::Edge(edge) => {
69                        ReturnValue::process_items_with_mixin(edge, &mut mixin)
70                    }
71                    TraversalVal::Vector(vector) => {
72                        ReturnValue::process_items_with_mixin(vector, &mut mixin)
73                    }
74                    TraversalVal::Count(count) => ReturnValue::from(count),
75                    TraversalVal::Empty => ReturnValue::Empty,
76                    TraversalVal::Value(value) => ReturnValue::from(value),
77                    TraversalVal::Path((nodes, edges)) => {
78                        let mut properties = HashMap::with_capacity(2);
79                        properties.insert(
80                            "nodes".to_string(),
81                            ReturnValue::Array(
82                                nodes
83                                    .into_iter()
84                                    .map(|node| ReturnValue::from(node))
85                                    .collect(),
86                            ),
87                        );
88                        properties.insert(
89                            "edges".to_string(),
90                            ReturnValue::Array(
91                                edges
92                                    .into_iter()
93                                    .map(|edge| ReturnValue::from(edge))
94                                    .collect(),
95                            ),
96                        );
97                        ReturnValue::Object(properties)
98                    }
99                })
100                .collect(),
101        )
102    }
103
104    #[inline]
105    pub fn from_traversal_value_with_mixin(
106        traversal_value: TraversalVal,
107        mut mixin: RefMut<HashMap<u128, ResponseRemapping>>,
108    ) -> Self {
109        match traversal_value {
110                    TraversalVal::Node(node) => {
111                        println!("node processing");
112                        ReturnValue::process_items_with_mixin(node, &mut mixin)
113                    }
114                    TraversalVal::Edge(edge) => {
115                        ReturnValue::process_items_with_mixin(edge, &mut mixin)
116                    }
117                    TraversalVal::Vector(vector) => {
118                        ReturnValue::process_items_with_mixin(vector, &mut mixin)
119                    }
120                    TraversalVal::Count(count) => ReturnValue::from(count),
121                    TraversalVal::Empty => ReturnValue::Empty,
122                    TraversalVal::Value(value) => ReturnValue::from(value),
123                    TraversalVal::Path((nodes, edges)) => {
124                        let mut properties = HashMap::with_capacity(2);
125                        properties.insert(
126                            "nodes".to_string(),
127                            ReturnValue::Array(
128                                nodes
129                                    .into_iter()
130                                    .map(|node| ReturnValue::from(node))
131                                    .collect(),
132                            ),
133                        );
134                        properties.insert(
135                            "edges".to_string(),
136                            ReturnValue::Array(
137                                edges
138                                    .into_iter()
139                                    .map(|edge| ReturnValue::from(edge))
140                                    .collect(),
141                            ),
142                        );
143                        ReturnValue::Object(properties)
144                    }
145        }
146    }
147
148    #[inline(always)]
149    #[allow(unused_attributes)]
150    #[ignore = "No use for this function yet, however, I believe it may be useful in the future so I'm keeping it here"]
151    pub fn mixin(self, other: ReturnValue) -> Self {
152        match (self, other) {
153            (ReturnValue::Object(mut a), ReturnValue::Object(b)) => {
154                a.extend(b);
155                ReturnValue::Object(a)
156            }
157            _ => unreachable!(),
158        }
159    }
160
161    /// Mixin a remapping to a return value.
162    ///
163    /// This function takes a hashmap of `Remappings` and mixes them into the return value
164    ///
165    /// - If the mapping is an exclude, then the key is removed from the return value
166    /// - If the mapping is a remapping from an old value to a new value, then the key
167    ///     is replaced with the new name and the value is the new value
168    /// - If the mapping is a new mapping, then the key is added to the return value
169    ///     and the value is the new value
170    /// - Otherwise, the key is left unchanged and the value is the original value
171    ///
172    /// Basic usage:
173    ///
174    /// ```rust
175    /// use helix_db::protocol::{ReturnValue, Remapping};
176    /// use std::collections::HashMap;
177    ///
178    /// let remappings = HashMap::new();
179    /// remappings.insert(
180    ///     "old_key".to_string(),
181    ///     Remapping::new(
182    ///         Some("new_key".to_string()),
183    ///         ReturnValue::from("new_value".to_string())
184    ///     )
185    /// );
186    ///
187    /// let return_value = ReturnValue::from("old_value".to_string());
188    /// let return_value = return_value.mixin_remapping(remappings);
189    ///
190    /// assert_eq!(
191    ///     return_value.get("new_key".to_string()),
192    ///     Some(&ReturnValue::from("new_value".to_string()))
193    /// );
194    /// ```
195    #[inline(always)]
196    pub fn mixin_remapping(self, remappings: HashMap<String, Remapping>) -> Self {
197        match self {
198            ReturnValue::Object(mut a) => {
199                remappings.into_iter().for_each(|(k, v)| {
200                    if v.exclude {
201                        let _ = a.remove(&k);
202                    } else if let Some(new_name) = &v.new_name {
203                        if let Some(value) = a.remove(&k) {
204                            a.insert(new_name.clone(), value);
205                        } else {
206                            println!("no value found for key: {:?}", k);
207                            a.insert(k, v.return_value);
208                        }
209                    } else {
210                        println!("inserting value: {:?}", k);
211                        a.insert(k, v.return_value);
212                    }
213                });
214                ReturnValue::Object(a)
215            }
216            _ => unreachable!(),
217        }
218    }
219
220    #[inline(always)]
221    #[allow(unused_attributes)]
222    #[ignore = "No use for this function yet, however, I believe it may be useful in the future so I'm keeping it here"]
223    pub fn mixin_other<I>(&self, item: I, secondary_properties: ResponseRemapping) -> Self
224    where
225        I: Filterable + Clone,
226    {
227        let mut return_val = ReturnValue::default();
228        if !secondary_properties.should_spread {
229            match item.type_name() {
230                FilterableType::Node => {
231                    return_val = ReturnValue::from(item);
232                }
233                FilterableType::Edge => {
234                    return_val = ReturnValue::from(item);
235                }
236                FilterableType::Vector => {
237                    return_val = ReturnValue::from(item);
238                }
239            }
240        }
241        return_val = return_val.mixin_remapping(secondary_properties.remappings);
242        return_val
243    }
244}
245
246impl<I> From<I> for ReturnValue
247where
248    I: Filterable + Clone,
249{
250    #[inline]
251    fn from(item: I) -> Self {
252        let length = match item.properties_ref() {
253            Some(properties) => properties.len(),
254            None => 0,
255        };
256        let mut properties = match item.type_name() {
257            FilterableType::Node => HashMap::with_capacity(Node::NUM_PROPERTIES + length),
258            FilterableType::Edge => {
259                let mut properties = HashMap::with_capacity(Edge::NUM_PROPERTIES + length);
260                properties.insert(
261                    "from_node".to_string(),
262                    ReturnValue::from(item.from_node_uuid()),
263                );
264                properties.insert(
265                    "to_node".to_string(),
266                    ReturnValue::from(item.to_node_uuid()),
267                );
268                properties
269            }
270            FilterableType::Vector => {
271                let data = item.vector_data();
272                let score = item.score();
273
274                let mut properties = HashMap::with_capacity(2 + length);
275                properties.insert("data".to_string(), ReturnValue::from(data));
276                properties.insert("score".to_string(), ReturnValue::from(score));
277                properties
278            }
279        };
280        properties.insert("id".to_string(), ReturnValue::from(item.uuid()));
281        properties.insert(
282            "label".to_string(),
283            ReturnValue::from(item.label().to_string()),
284        );
285        if item.properties_ref().is_some() {
286            properties.extend(
287                item.properties()
288                    .unwrap()
289                    .into_iter()
290                    .map(|(k, v)| (k, ReturnValue::from(v))),
291            );
292        }
293
294        ReturnValue::Object(properties)
295    }
296}
297
298impl From<Value> for ReturnValue {
299    fn from(value: Value) -> Self {
300        ReturnValue::Value(value)
301    }
302}
303
304impl From<&Value> for ReturnValue {
305    fn from(value: &Value) -> Self {
306        ReturnValue::Value(value.clone())
307    }
308}
309
310impl From<Count> for ReturnValue {
311    fn from(count: Count) -> Self {
312        ReturnValue::Value(Value::I32(count.value() as i32))
313    }
314}
315
316impl From<String> for ReturnValue {
317    fn from(string: String) -> Self {
318        ReturnValue::Value(Value::String(string))
319    }
320}
321
322impl From<bool> for ReturnValue {
323    fn from(boolean: bool) -> Self {
324        ReturnValue::Value(Value::Boolean(boolean))
325    }
326}
327
328impl From<&str> for ReturnValue {
329    fn from(string: &str) -> Self {
330        ReturnValue::Value(Value::String(string.to_string()))
331    }
332}
333
334impl From<i32> for ReturnValue {
335    fn from(integer: i32) -> Self {
336        ReturnValue::Value(Value::I32(integer))
337    }
338}
339
340impl From<f64> for ReturnValue {
341    fn from(float: f64) -> Self {
342        ReturnValue::Value(Value::F64(float))
343    }
344}
345
346impl From<u128> for ReturnValue {
347    fn from(integer: u128) -> Self {
348        ReturnValue::Value(Value::U128(integer))
349    }
350}
351
352impl From<Vec<TraversalVal>> for ReturnValue {
353    fn from(array: Vec<TraversalVal>) -> Self {
354        ReturnValue::Array(array.into_iter().map(|val| val.into()).collect())
355    }
356}
357
358impl From<TraversalVal> for ReturnValue {
359    fn from(val: TraversalVal) -> Self {
360        match val {
361            TraversalVal::Node(node) => ReturnValue::from(node),
362            TraversalVal::Edge(edge) => ReturnValue::from(edge),
363            TraversalVal::Vector(vector) => ReturnValue::from(vector),
364            TraversalVal::Count(count) => ReturnValue::from(count),
365            TraversalVal::Value(value) => ReturnValue::from(value),
366            TraversalVal::Empty => ReturnValue::Empty,
367            _ => unreachable!(),
368        }
369    }
370}
371
372impl From<&[f64]> for ReturnValue {
373    fn from(data: &[f64]) -> Self {
374        ReturnValue::Array(data.iter().map(|f| ReturnValue::from(*f)).collect())
375    }
376}
377
378impl Serialize for ReturnValue {
379    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
380    where
381        S: serde::ser::Serializer,
382    {
383        match self {
384            ReturnValue::Value(value) => value.serialize(serializer),
385            ReturnValue::Object(object) => object.serialize(serializer),
386            ReturnValue::Array(array) => array.serialize(serializer),
387            ReturnValue::Empty => serializer.serialize_none(),
388        }
389    }
390}
391
392impl Default for ReturnValue {
393    fn default() -> Self {
394        ReturnValue::Object(HashMap::new())
395    }
396}