Skip to main content

core_storage/
edge_props.rs

1use crate::types::Value;
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5#[derive(Debug, Default, Clone, Serialize, Deserialize)]
6pub struct EdgeProps {
7    map: BTreeMap<(u32, u32, u32), BTreeMap<String, Value>>,
8}
9
10impl EdgeProps {
11    pub fn new() -> Self {
12        Self::default()
13    }
14
15    pub fn set(&mut self, etype: u32, src: u32, dst: u32, field: &str, value: Value) {
16        self.map
17            .entry((etype, src, dst))
18            .or_default()
19            .insert(field.to_owned(), value);
20    }
21
22    pub fn get(&self, etype: u32, src: u32, dst: u32, field: &str) -> Option<&Value> {
23        self.map.get(&(etype, src, dst))?.get(field)
24    }
25
26    pub fn remove_edge(&mut self, etype: u32, src: u32, dst: u32) {
27        self.map.remove(&(etype, src, dst));
28    }
29}
30
31#[cfg(test)]
32mod tests {
33    use super::*;
34
35    #[test]
36    fn set_get_remove_edge_props() {
37        let mut e = EdgeProps::new();
38        e.set(0, 1, 2, "score", Value::Float(0.5));
39        e.set(0, 1, 2, "score", Value::Float(0.7)); // overwrite
40        assert_eq!(e.get(0, 1, 2, "score"), Some(&Value::Float(0.7)));
41        assert_eq!(e.get(0, 2, 1, "score"), None); // directed
42        e.remove_edge(0, 1, 2);
43        assert_eq!(e.get(0, 1, 2, "score"), None);
44    }
45}