Skip to main content

lattice_db/graph/
graph_builder.rs

1use std::collections::HashMap;
2
3use crate::{
4    errors::LatticeError,
5    graph::graph_prepared::PreparedGraph,
6    properties::PropertyHandle,
7    utils::{
8        generational_vector::{GenVec, Handle},
9        values::{Primitive, Value},
10    },
11};
12
13pub(crate) type GlobalVertexId = u64;
14pub(crate) type GraphId = u64;
15
16#[derive(Clone, Copy, PartialEq, Eq, Hash)]
17pub struct VertexHandle(pub(crate) Handle);
18
19#[derive(Clone, Copy, PartialEq, Eq)]
20pub struct EdgeHandle(Handle);
21
22pub struct VertexData {
23    pub global_id: Option<GlobalVertexId>,
24    pub attributes: Vec<(PropertyHandle, Primitive)>,
25    pub incoming_edges: Vec<EdgeHandle>,
26    pub outgoing_edges: Vec<EdgeHandle>,
27}
28
29pub struct EdgeData {
30    pub from: VertexHandle,
31    pub to: VertexHandle,
32    pub label: PropertyHandle,
33}
34
35pub(crate) struct OldGraphData {
36    pub(crate) id: GraphId,
37    pub(crate) graph: PreparedGraph,
38}
39
40pub struct GraphBuilder {
41    new_vertex_count: u64,
42    pub(crate) old_graph_data: Option<OldGraphData>,
43    pub(crate) vertices: GenVec<VertexData>,
44    pub(crate) edges: GenVec<EdgeData>,
45}
46
47impl Default for GraphBuilder {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53impl GraphBuilder {
54    pub fn new_vertex<'a>(&'a mut self) -> VertexBuilder<'a> {
55        self.new_vertex_count += 1;
56        let handle = self.vertices.add(VertexData {
57            global_id: None,
58            attributes: vec![],
59            incoming_edges: vec![],
60            outgoing_edges: vec![],
61        });
62        VertexBuilder::new(self, handle)
63    }
64
65    /// Removes a vertex and all attatched edges.
66    pub fn remove_vertex(&mut self, handle: VertexHandle) -> Result<(), LatticeError> {
67        let removed = self
68            .vertices
69            .remove(handle.0)
70            .ok_or(LatticeError::VertexNotFound)?;
71        if removed.global_id.is_none() {
72            self.new_vertex_count -= 1;
73        } // removing a new vertex decreases new vertex count
74
75        // remove all edges, ignore errors from self loop edges
76        let edges = [removed.incoming_edges, removed.outgoing_edges].concat();
77        for edge in edges {
78            let _ = self.remove_edge(edge);
79        }
80        Ok(())
81    }
82
83    /// Return a VertexBuilder to edit the vertex.
84    pub fn edit_vertex<'a>(
85        &'a mut self,
86        handle: VertexHandle,
87    ) -> Result<VertexBuilder<'a>, LatticeError> {
88        self.vertices
89            .get(handle.0)
90            .ok_or(LatticeError::VertexNotFound)?;
91        Ok(VertexBuilder::new(self, handle.0))
92    }
93
94    /// Get vertex data.
95    pub fn get_vertex(&self, handle: VertexHandle) -> Option<&VertexData> {
96        self.vertices.get(handle.0)
97    }
98
99    /// Create a new edge to link vertices.
100    pub fn new_edge(
101        &mut self,
102        from: VertexHandle,
103        label: PropertyHandle,
104        to: VertexHandle,
105    ) -> Result<&mut Self, LatticeError> {
106        self.vertices
107            .get(from.0)
108            .ok_or(LatticeError::VertexNotFound)?; // confirm 'from' exists
109        let to_vertex = self
110            .vertices
111            .get_mut(to.0)
112            .ok_or(LatticeError::VertexNotFound)?;
113
114        // create edge
115        let edge_handle = EdgeHandle(self.edges.add(EdgeData { from, label, to }));
116
117        // add edge
118        to_vertex.incoming_edges.push(edge_handle);
119        let from_vertex = self.vertices.get_mut(from.0).unwrap();
120        from_vertex.outgoing_edges.push(edge_handle);
121        Ok(self)
122    }
123
124    /// Removes an edge from the graph.
125    pub fn remove_edge(&mut self, handle: EdgeHandle) -> Result<(), LatticeError> {
126        let edge = self
127            .edges
128            .remove(handle.0)
129            .ok_or(LatticeError::EdgeNotFound)?;
130
131        // detatch from source
132        if let Some(from) = self.vertices.get_mut(edge.from.0)
133            && let Some(idx) = from.outgoing_edges.iter().position(|p| *p == handle)
134        {
135            from.outgoing_edges.swap_remove(idx);
136        }
137
138        // detatch from destination
139        if let Some(to) = self.vertices.get_mut(edge.to.0)
140            && let Some(idx) = to.incoming_edges.iter().position(|p| *p == handle)
141        {
142            to.incoming_edges.swap_remove(idx);
143        }
144
145        Ok(())
146    }
147
148    /// Return an EdgeBuilder to edit the edge.
149    pub fn edit_edge<'a>(
150        &'a mut self,
151        handle: EdgeHandle,
152    ) -> Result<EdgeBuilder<'a>, LatticeError> {
153        self.edges.get(handle.0).ok_or(LatticeError::EdgeNotFound)?;
154        Ok(EdgeBuilder::new(self, handle.0))
155    }
156
157    /// Get edge data.
158    pub fn get_edge(&self, handle: EdgeHandle) -> Option<&EdgeData> {
159        self.edges.get(handle.0)
160    }
161}
162
163pub struct EdgeBuilder<'a> {
164    graph: &'a mut GraphBuilder,
165    handle: Handle,
166}
167
168impl<'a> EdgeBuilder<'a> {
169    fn new(graph: &'a mut GraphBuilder, handle: Handle) -> Self {
170        Self { graph, handle }
171    }
172
173    fn get_self(&mut self) -> &mut EdgeData {
174        self.graph.edges.get_mut(self.handle).unwrap() // must exist to be using the fn
175    }
176
177    /// Change the source (from) of the edge.
178    pub fn set_source(&mut self, handle: VertexHandle) -> Result<&mut Self, LatticeError> {
179        self.graph
180            .edges
181            .get(handle.0)
182            .ok_or(LatticeError::VertexNotFound)?;
183        self.get_self().from = handle;
184        Ok(self)
185    }
186
187    /// Change the destination (to) of the edge.
188    pub fn set_destination(&mut self, handle: VertexHandle) -> Result<&mut Self, LatticeError> {
189        self.graph
190            .edges
191            .get(handle.0)
192            .ok_or(LatticeError::VertexNotFound)?;
193        self.get_self().to = handle;
194        Ok(self)
195    }
196
197    /// Change the edge label.
198    pub fn set_label(&mut self, label: PropertyHandle) -> Result<&mut Self, LatticeError> {
199        self.get_self().label = label;
200        Ok(self)
201    }
202}
203
204pub struct VertexBuilder<'a> {
205    graph: &'a mut GraphBuilder,
206    handle: Handle,
207}
208
209impl<'a> VertexBuilder<'a> {
210    fn new(graph: &'a mut GraphBuilder, handle: Handle) -> Self {
211        Self { graph, handle }
212    }
213
214    fn get_self(&mut self) -> &mut VertexData {
215        self.graph.vertices.get_mut(self.handle).unwrap() // must exist to be using the fn
216    }
217
218    /// Create a new attribute with a value.
219    pub fn new_attribute<V: Value>(
220        &mut self,
221        attr: PropertyHandle,
222        value: V,
223    ) -> Result<&mut Self, LatticeError> {
224        let v = value.to_primitive();
225        v.verify()?;
226        self.get_self().attributes.push((attr, v));
227        Ok(self)
228    }
229
230    /// Create a new edge to link vertices.
231    pub fn new_edge(
232        &mut self,
233        label: PropertyHandle,
234        handle: VertexHandle,
235    ) -> Result<&mut Self, LatticeError> {
236        let to = self
237            .graph
238            .vertices
239            .get_mut(handle.0)
240            .ok_or(LatticeError::VertexNotFound)?;
241
242        // create edge
243        let edge_handle = EdgeHandle(self.graph.edges.add(EdgeData {
244            from: VertexHandle(self.handle),
245            label,
246            to: handle,
247        }));
248
249        // add edge
250        to.incoming_edges.push(edge_handle);
251        self.get_self().outgoing_edges.push(edge_handle);
252        Ok(self)
253    }
254
255    /// Returns the vertex handle.
256    pub fn handle(&self) -> VertexHandle {
257        VertexHandle(self.handle)
258    }
259}
260
261impl GraphBuilder {
262    pub fn get_vertex_global_id(
263        &self,
264        handle: VertexHandle,
265    ) -> Result<Option<GlobalVertexId>, LatticeError> {
266        Ok(self
267            .vertices
268            .get(handle.0)
269            .ok_or(LatticeError::VertexNotFound)?
270            .global_id)
271    }
272
273    pub fn get_mut_attributes(
274        &mut self,
275        handle: VertexHandle,
276    ) -> Result<&mut Vec<(PropertyHandle, Primitive)>, LatticeError> {
277        Ok(&mut self
278            .vertices
279            .get_mut(handle.0)
280            .ok_or(LatticeError::VertexNotFound)?
281            .attributes)
282    }
283
284    pub fn get_incoming_edges(
285        &self,
286        handle: VertexHandle,
287    ) -> Result<&Vec<EdgeHandle>, LatticeError> {
288        Ok(&self
289            .vertices
290            .get(handle.0)
291            .ok_or(LatticeError::VertexNotFound)?
292            .incoming_edges)
293    }
294
295    pub fn get_outgoing_edges(
296        &self,
297        handle: VertexHandle,
298    ) -> Result<&Vec<EdgeHandle>, LatticeError> {
299        Ok(&self
300            .vertices
301            .get(handle.0)
302            .ok_or(LatticeError::VertexNotFound)?
303            .outgoing_edges)
304    }
305
306    /// Iterate through all graph vertices.
307    pub fn iter_vertices(&self) -> impl Iterator<Item = (Handle, &VertexData)> {
308        self.vertices.iter()
309    }
310
311    /// Iterate through all graph edges.
312    pub fn iter_edges(&self) -> impl Iterator<Item = (Handle, &EdgeData)> {
313        self.edges.iter()
314    }
315}
316
317impl GraphBuilder {
318    pub fn new() -> Self {
319        Self {
320            new_vertex_count: 0,
321            old_graph_data: None,
322            vertices: GenVec::new(),
323            edges: GenVec::new(),
324        }
325    }
326
327    pub(crate) fn count_new_vertices(&self) -> u64 {
328        self.new_vertex_count
329    }
330
331    pub(crate) fn from_prepared(graph: PreparedGraph) -> Self {
332        let mut builder = Self::new();
333        let mut id_map = HashMap::new(); // maybe an actual use for Rapid
334
335        // populate vertices
336        for v in &graph.vertices {
337            let handle = builder.vertices.add(VertexData {
338                global_id: Some(v.id),
339                attributes: v.attrs.clone(),
340                incoming_edges: vec![],
341                outgoing_edges: vec![],
342            });
343            id_map.insert(v.id, VertexHandle(handle));
344        }
345
346        // populate edges
347        for e in &graph.edges {
348            let from = *id_map.get(&e.from).expect("Edge source vertex missing");
349            let to = *id_map.get(&e.to).expect("Edge destination vertex missing");
350            let handle = builder.edges.add(EdgeData {
351                from,
352                to,
353                label: e.label,
354            });
355            builder
356                .vertices
357                .get_mut(from.0)
358                .unwrap()
359                .outgoing_edges
360                .push(EdgeHandle(handle));
361            builder
362                .vertices
363                .get_mut(to.0)
364                .unwrap()
365                .incoming_edges
366                .push(EdgeHandle(handle));
367        }
368
369        builder.old_graph_data = Some(OldGraphData {
370            id: graph.id,
371            graph,
372        });
373        builder
374    }
375}