Skip to main content

raphtory_storage/graph/
locked.rs

1use raphtory_api::core::{
2    entities::{GidRef, VID},
3    storage::dict_mapper::MaybeNew,
4};
5use raphtory_core::{
6    entities::graph::{logical_to_physical::InvalidNodeId, tgraph::TemporalGraph},
7    storage::{
8        raw_edges::{LockedEdges, WriteLockedEdges},
9        ReadLockedStorage, WriteLockedNodes,
10    },
11};
12use std::sync::Arc;
13
14#[derive(Debug)]
15pub struct LockedGraph {
16    pub(crate) nodes: Arc<ReadLockedStorage>,
17    pub(crate) edges: Arc<LockedEdges>,
18    pub graph: Arc<TemporalGraph>,
19}
20
21impl<'de> serde::Deserialize<'de> for LockedGraph {
22    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
23    where
24        D: serde::Deserializer<'de>,
25    {
26        TemporalGraph::deserialize(deserializer).map(|graph| LockedGraph::new(Arc::new(graph)))
27    }
28}
29
30impl serde::Serialize for LockedGraph {
31    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
32    where
33        S: serde::Serializer,
34    {
35        self.graph.serialize(serializer)
36    }
37}
38
39impl LockedGraph {
40    pub fn new(graph: Arc<TemporalGraph>) -> Self {
41        let nodes = Arc::new(graph.storage.nodes_read_lock());
42        let edges = Arc::new(graph.storage.edges_read_lock());
43        Self {
44            nodes,
45            edges,
46            graph,
47        }
48    }
49}
50
51impl Clone for LockedGraph {
52    fn clone(&self) -> Self {
53        LockedGraph {
54            nodes: self.nodes.clone(),
55            edges: self.edges.clone(),
56            graph: self.graph.clone(),
57        }
58    }
59}
60
61pub struct WriteLockedGraph<'a> {
62    pub nodes: WriteLockedNodes<'a>,
63    pub edges: WriteLockedEdges<'a>,
64    pub graph: &'a TemporalGraph,
65}
66
67impl<'a> WriteLockedGraph<'a> {
68    pub(crate) fn new(graph: &'a TemporalGraph) -> Self {
69        let nodes = graph.storage.nodes.write_lock();
70        let edges = graph.storage.edges.write_lock();
71        Self {
72            nodes,
73            edges,
74            graph,
75        }
76    }
77
78    pub fn num_nodes(&self) -> usize {
79        self.graph.storage.nodes.len()
80    }
81    pub fn resolve_node(&self, gid: GidRef) -> Result<MaybeNew<VID>, InvalidNodeId> {
82        self.graph
83            .logical_to_physical
84            .get_or_init(gid, || self.graph.storage.nodes.next_id())
85    }
86
87    pub fn resolve_node_type(&self, node_type: Option<&str>) -> MaybeNew<usize> {
88        node_type
89            .map(|node_type| self.graph.node_meta.get_or_create_node_type_id(node_type))
90            .unwrap_or_else(|| MaybeNew::Existing(0))
91    }
92
93    pub fn num_shards(&self) -> usize {
94        self.nodes.num_shards().max(self.edges.num_shards())
95    }
96
97    pub fn edges_mut(&mut self) -> &mut WriteLockedEdges<'a> {
98        &mut self.edges
99    }
100
101    pub fn graph(&self) -> &TemporalGraph {
102        self.graph
103    }
104}