linkrs_query/optimizer/heuristic/
context.rs

1//! Rework the definition of the context.
2//!
3//! Define the RewriteContext structure to manage the state during the rewriting process.
4//! This is a simplified version that has been separated from the optimizer layer, focusing on the requirements for heuristic rewriting rules.
5
6use std::cell::RefCell;
7use std::collections::HashMap;
8use std::rc::Rc;
9use std::sync::Arc;
10
11use crate::planning::plan::PlanNodeEnum;
12use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
13
14/// Rewrite the context
15///
16/// Manage the status and node information during the rewriting process.
17/// Compared to the OptContext of the optimizer, this is a lightweight version.
18/// It does not include features specific to optimizers, such as caches for statistical information or cost calculations.
19#[derive(Debug)]
20pub struct RewriteContext {
21    /// Node ID Counter – Generates unique node IDs
22    node_id_counter: usize,
23    /// Mapping of plan nodes to IDs
24    plan_node_to_id: RefCell<HashMap<usize, usize>>,
25    /// Mapping from IDs to planned nodes
26    nodes_by_id: RefCell<HashMap<usize, Rc<RefCell<PlanNodeWrapper>>>>,
27    /// Expression context
28    expr_context: Arc<ExpressionAnalysisContext>,
29    /// ID of the node whose rules are currently being applied.  The plan root
30    /// always receives the first allocated id (0), so whole-plan rules can
31    /// detect they are running at the root.
32    current_node_id: usize,
33}
34
35/// Plan Node Wrapper
36///
37/// Wrap the PlanNodeEnum and add the necessary metadata for the override.
38#[derive(Debug, Clone)]
39pub struct PlanNodeWrapper {
40    pub id: usize,
41    pub plan_node: PlanNodeEnum,
42    pub dependencies: Vec<usize>,
43}
44
45impl PlanNodeWrapper {
46    pub fn new(id: usize, plan_node: PlanNodeEnum) -> Self {
47        Self {
48            id,
49            plan_node,
50            dependencies: Vec::new(),
51        }
52    }
53}
54
55impl Default for RewriteContext {
56    fn default() -> Self {
57        Self {
58            node_id_counter: 0,
59            plan_node_to_id: RefCell::new(HashMap::new()),
60            nodes_by_id: RefCell::new(HashMap::new()),
61            expr_context: Arc::new(ExpressionAnalysisContext::new()),
62            current_node_id: usize::MAX,
63        }
64    }
65}
66
67impl RewriteContext {
68    /// Create a new rewriting context (using the default ExpressionContext).
69    pub fn new() -> Self {
70        Self::default()
71    }
72
73    /// Create a new rewriting context (using the specified ExpressionContext).
74    pub fn with_expr_context(expr_context: Arc<ExpressionAnalysisContext>) -> Self {
75        Self {
76            node_id_counter: 0,
77            plan_node_to_id: RefCell::new(HashMap::new()),
78            nodes_by_id: RefCell::new(HashMap::new()),
79            expr_context,
80            current_node_id: usize::MAX,
81        }
82    }
83
84    /// Assign a new node ID
85    pub fn allocate_node_id(&mut self) -> usize {
86        let id = self.node_id_counter;
87        self.node_id_counter += 1;
88        id
89    }
90
91    /// Register a plan node
92    pub fn register_node(
93        &mut self,
94        node_id: usize,
95        plan_node: PlanNodeEnum,
96    ) -> Rc<RefCell<PlanNodeWrapper>> {
97        let wrapper = Rc::new(RefCell::new(PlanNodeWrapper::new(node_id, plan_node)));
98        self.nodes_by_id
99            .borrow_mut()
100            .insert(node_id, wrapper.clone());
101        wrapper
102    }
103
104    /// Set the id of the node whose rules are currently being applied.
105    pub fn set_current_node_id(&mut self, node_id: usize) {
106        self.current_node_id = node_id;
107    }
108
109    /// ID of the node whose rules are currently being applied.  The plan root
110    /// is registered with id 0.
111    pub fn current_node_id(&self) -> usize {
112        self.current_node_id
113    }
114
115    /// Find a node by its ID
116    pub fn find_node_by_id(&self, id: usize) -> Option<Rc<RefCell<PlanNodeWrapper>>> {
117        self.nodes_by_id.borrow().get(&id).cloned()
118    }
119
120    /// Add node mapping
121    pub fn add_plan_node_mapping(&self, plan_node_id: usize, rewrite_node_id: usize) {
122        self.plan_node_to_id
123            .borrow_mut()
124            .insert(plan_node_id, rewrite_node_id);
125    }
126
127    /// Find the rewrite node ID by using the planned node ID.
128    pub fn find_rewrite_id_by_plan_id(&self, plan_node_id: usize) -> Option<usize> {
129        self.plan_node_to_id.borrow().get(&plan_node_id).copied()
130    }
131
132    /// Get the current node count.
133    pub fn node_count(&self) -> usize {
134        self.node_id_counter
135    }
136
137    /// Obtain the context of the expression.
138    pub fn expr_context(&self) -> Arc<ExpressionAnalysisContext> {
139        self.expr_context.clone()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146    use crate::planning::plan::core::nodes::access::graph_scan_node::ScanVerticesNode;
147
148    #[test]
149    fn test_context_creation() {
150        let ctx = RewriteContext::new();
151        assert_eq!(ctx.node_count(), 0);
152    }
153
154    #[test]
155    fn test_allocate_node_id() {
156        let mut ctx = RewriteContext::new();
157        assert_eq!(ctx.allocate_node_id(), 0);
158        assert_eq!(ctx.allocate_node_id(), 1);
159        assert_eq!(ctx.allocate_node_id(), 2);
160    }
161
162    #[test]
163    fn test_register_and_find_node() {
164        let mut ctx = RewriteContext::new();
165        let node_id = ctx.allocate_node_id();
166        let plan_node = PlanNodeEnum::ScanVertices(ScanVerticesNode::new(1, "default"));
167
168        let wrapper = ctx.register_node(node_id, plan_node);
169        assert_eq!(wrapper.borrow().id, node_id);
170
171        let found = ctx.find_node_by_id(node_id);
172        assert!(found.is_some());
173        assert_eq!(found.expect("Failed to find node").borrow().id, node_id);
174    }
175}
linkrs_query/optimizer/heuristic/context.rs

linkrs_query/optimizer/heuristic/
context.rs
