Skip to main content

lattice_db/query/
query_builder.rs

1use bincode::{Decode, Encode};
2
3use crate::{
4    errors::LatticeError,
5    properties::PropertyHandle,
6    utils::{
7        generational_vector::{GenVec, Handle},
8        values::{Primitive, Value},
9    },
10};
11
12#[derive(Clone, Copy, Hash, PartialEq, Eq, Debug, Encode, Decode)]
13pub enum EdgeDirection {
14    Outgoing,
15    Incoming,
16}
17
18#[derive(Clone, Copy)]
19pub struct NodeHandle(pub(crate) Handle);
20
21pub enum QueryNode {
22    // set logic
23    Union(Vec<NodeHandle>),
24    Intersect(Vec<NodeHandle>),
25    Difference(NodeHandle, NodeHandle),
26    // search for values
27    Attribute {
28        attr: PropertyHandle,
29        value: Primitive,
30    },
31    // search for linked nodes
32    Edge {
33        dir: EdgeDirection,
34        label: PropertyHandle,
35        target: NodeHandle,
36    },
37    // saved query
38    SavedQuery(u64),
39}
40
41impl QueryNode {
42    pub fn attribute<V: Value>(attr: PropertyHandle, value: V) -> Self {
43        Self::Attribute {
44            attr,
45            value: value.to_primitive(),
46        }
47    }
48}
49
50pub struct QueryBuilder {
51    pub(crate) nodes: GenVec<QueryNode>,
52    pub(crate) root: Option<NodeHandle>,
53}
54
55impl QueryBuilder {
56    /// Find vertices by attribute value.
57    pub fn match_attr<V: Value>(
58        &mut self,
59        attr: PropertyHandle,
60        value: V,
61    ) -> Result<NodeHandle, LatticeError> {
62        let value = value.to_primitive();
63        value.verify()?;
64        let handle = self.nodes.add(QueryNode::Attribute { attr, value });
65        Ok(NodeHandle(handle))
66    }
67
68    /// All vertices that are pointed to by `subject` via label.
69    pub fn match_outgoing(
70        &mut self,
71        label: PropertyHandle,
72        subject: NodeHandle,
73    ) -> Result<NodeHandle, LatticeError> {
74        self.nodes
75            .get(subject.0)
76            .ok_or(LatticeError::EdgeNotFound)?;
77        let handle = self.nodes.add(QueryNode::Edge {
78            dir: EdgeDirection::Outgoing,
79            label,
80            target: subject,
81        });
82        Ok(NodeHandle(handle))
83    }
84
85    /// All vertices that point to `target` via label.
86    pub fn match_incoming(
87        &mut self,
88        label: PropertyHandle,
89        target: NodeHandle,
90    ) -> Result<NodeHandle, LatticeError> {
91        self.nodes.get(target.0).ok_or(LatticeError::EdgeNotFound)?;
92        let handle = self.nodes.add(QueryNode::Edge {
93            dir: EdgeDirection::Incoming,
94            label,
95            target,
96        });
97        Ok(NodeHandle(handle))
98    }
99
100    /// Find a vertex that satisfies multiple features within children.
101    pub fn group_and(&mut self, children: Vec<NodeHandle>) -> Result<NodeHandle, LatticeError> {
102        for c in &children {
103            self.nodes.get(c.0).ok_or(LatticeError::EdgeNotFound)?;
104        }
105        let handle = self.nodes.add(QueryNode::Intersect(children));
106        Ok(NodeHandle(handle))
107    }
108
109    /// Find a vertex that satisfies one of multiple features within children.
110    pub fn group_or(&mut self, children: Vec<NodeHandle>) -> Result<NodeHandle, LatticeError> {
111        for c in &children {
112            self.nodes.get(c.0).ok_or(LatticeError::EdgeNotFound)?;
113        }
114        let handle = self.nodes.add(QueryNode::Union(children));
115        Ok(NodeHandle(handle))
116    }
117
118    /// Find a vertex that satisfies include, but does not satisfy exclude.
119    pub fn difference(
120        &mut self,
121        include: NodeHandle,
122        exclude: NodeHandle,
123    ) -> Result<NodeHandle, LatticeError> {
124        self.nodes
125            .get(include.0)
126            .ok_or(LatticeError::EdgeNotFound)?;
127        self.nodes
128            .get(exclude.0)
129            .ok_or(LatticeError::EdgeNotFound)?;
130        let handle = self.nodes.add(QueryNode::Difference(include, exclude));
131        Ok(NodeHandle(handle))
132    }
133}
134
135impl Default for QueryBuilder {
136    fn default() -> Self {
137        Self::new()
138    }
139}
140
141impl QueryBuilder {
142    pub fn new() -> Self {
143        Self {
144            nodes: GenVec::new(),
145            root: None,
146        }
147    }
148
149    pub fn set_root(&mut self, handle: NodeHandle) {
150        self.root = Some(handle);
151    }
152
153    pub fn get_root(&self) -> Option<NodeHandle> {
154        self.root
155    }
156}