lattice_db/query/
query_prepared.rs1use std::collections::HashMap;
2
3use bincode::{Decode, Encode};
4
5use crate::{errors::LatticeError, properties::PropertyHandle};
6
7use super::query_builder::*;
8
9type NodeIdx = usize;
10
11#[derive(Hash, PartialEq, Eq, Clone, Debug, Encode, Decode)]
12pub(crate) enum Node {
13 Union(Vec<NodeIdx>),
15 Intersect(Vec<NodeIdx>),
16 Difference(NodeIdx, NodeIdx),
17 Attribute {
19 attr: PropertyHandle,
20 value: u64, },
22 Edge {
24 dir: EdgeDirection,
25 label: PropertyHandle,
26 target: NodeIdx,
27 },
28 SavedQuery(u64),
30}
31
32#[derive(Encode, Decode)]
33pub struct PreparedQuery {
34 pub(crate) nodes: Vec<Node>,
35 pub(crate) root: NodeIdx,
36}
37
38impl QueryBuilder {
39 pub fn compile(&self) -> Result<PreparedQuery, LatticeError> {
41 let root = self.root.ok_or(LatticeError::RootNotFound)?;
42 let build_order = self.get_build_order(root);
43
44 let mut output = vec![]; let mut visited = HashMap::new(); let mut dup_cache = HashMap::new();
51
52 for handle in build_order {
53 let src_node = self
54 .nodes
55 .get(handle.0)
56 .ok_or(LatticeError::VertexNotFound)?;
57
58 let compiled_node = match src_node {
59 QueryNode::Union(handles) => {
60 let mut ids = vec![];
61 for h in handles {
62 ids.push(*visited.get(&h.0.index).unwrap());
63 }
64 ids.sort_unstable(); ids.dedup(); Node::Union(ids)
67 }
68 QueryNode::Intersect(handles) => {
69 let mut ids = vec![];
70 for h in handles {
71 ids.push(*visited.get(&h.0.index).unwrap());
72 }
73 ids.sort_unstable(); ids.dedup(); Node::Intersect(ids)
76 }
77 QueryNode::Difference(a, b) => {
78 let a = *visited.get(&a.0.index).unwrap();
79 let b = *visited.get(&b.0.index).unwrap();
80 Node::Difference(a, b)
81 }
82 QueryNode::Attribute { attr, value } => {
83 let value_hash = value.hash();
84 Node::Attribute {
85 attr: *attr,
86 value: value_hash,
87 }
88 }
89 QueryNode::Edge { dir, label, target } => {
90 let target_id = *visited.get(&target.0.index).unwrap();
91 Node::Edge {
92 dir: *dir,
93 label: *label,
94 target: target_id,
95 }
96 }
97 QueryNode::SavedQuery(id) => Node::SavedQuery(*id),
98 };
99
100 let idx = if let Some(&idx) = dup_cache.get(&compiled_node) {
101 idx
102 } else {
103 let new_idx = output.len();
104 dup_cache.insert(compiled_node.clone(), new_idx);
105 output.push(compiled_node);
106 new_idx
107 };
108
109 visited.insert(handle.0.index, idx);
111 }
112
113 let root_idx = *visited.get(&root.0.index).unwrap();
114 Ok(PreparedQuery {
115 nodes: output,
116 root: root_idx,
117 })
118 }
119
120 fn get_build_order(&self, root: NodeHandle) -> Vec<NodeHandle> {
122 let mut order = vec![];
123 let mut visited = HashMap::new();
124 let mut stack = vec![(root, false)];
125
126 while let Some((handle, children_processed)) = stack.pop() {
127 if visited.contains_key(&handle.0.index) {
129 continue;
130 }
131
132 if children_processed {
133 visited.insert(handle.0.index, true);
134 order.push(handle);
135 } else {
136 stack.push((handle, true)); if let Some(node) = self.nodes.get(handle.0) {
140 match node {
141 QueryNode::Union(children) | QueryNode::Intersect(children) => {
142 for c in children {
143 stack.push((*c, false));
144 }
145 }
146 QueryNode::Difference(a, b) => {
147 stack.extend_from_slice(&[(*a, false), (*b, false)]);
148 }
149 QueryNode::Edge { target, .. } => {
150 stack.push((*target, false));
151 }
152 QueryNode::Attribute { .. } => {}
153 QueryNode::SavedQuery(_) => {}
154 }
155 }
156 }
157 }
158 order
159 }
160}