Skip to main content

lattice_db/query/
query_prepared.rs

1use 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    // set logic
14    Union(Vec<NodeIdx>),
15    Intersect(Vec<NodeIdx>),
16    Difference(NodeIdx, NodeIdx),
17    // search for values
18    Attribute {
19        attr: PropertyHandle,
20        value: u64, // hashed value
21    },
22    // search for linked nodes
23    Edge {
24        dir: EdgeDirection,
25        label: PropertyHandle,
26        target: NodeIdx,
27    },
28    // saved query
29    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    // assumes query has no dangling node links
40    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![]; // ordered nodes to query for
45        let mut visited = HashMap::new(); // maps self.nodes index to output index
46
47        // extremely primitive node de-duplication (exact matches)
48        // overhead may not be worth it if query is optimized
49        // removes nodes with different indices but exact match on children or attributes
50        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(); // these lines help dup_cache
65                    ids.dedup(); //
66                    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(); // these lines help dup_cache
74                    ids.dedup(); //
75                    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            // dedup visiting twice (diamond case)
110            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    // children before parents
121    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            // nodes only visited once (prevent dup visits on diamond formations)
128            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)); // children will be processed
137
138                // add children to be processed
139                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}