Skip to main content

linkrs_query/planning/template_extractor/
parameterize.rs

1use super::ParameterizedResult;
2use linkrs_core::types::expr::{ContextualExpression, Expression, FunctionArg};
3
4/// Expression Parameterized Translator
5///
6/// Traverse the expression tree and replace all literals with parameter placeholders.
7pub struct ParameterizingTransformer;
8
9impl ParameterizingTransformer {
10    pub fn new() -> Self {
11        Self
12    }
13
14    /// Parameterizing a single ContextualExpression
15    pub fn parameterize(&mut self, expr: &ContextualExpression) -> ParameterizedResult {
16        let inner_expr = match expr.get_expression() {
17            Some(e) => e,
18            None => {
19                return ParameterizedResult::new();
20            }
21        };
22        let mut result = ParameterizedResult::new();
23        let new_expr = self.transform_with_params(&inner_expr, &mut result);
24        result.expression = new_expr;
25        result
26    }
27
28    /// Parameterizing a single Expression (compatible with older interfaces)
29    pub fn parameterize_expression(&mut self, expr: &Expression) -> ParameterizedResult {
30        let mut result = ParameterizedResult::new();
31        let new_expr = self.transform_with_params(expr, &mut result);
32        result.expression = new_expr;
33        result
34    }
35
36    /// Parameterizing multiple expressions
37    pub fn parameterize_many(
38        &mut self,
39        exprs: &[ContextualExpression],
40    ) -> Vec<ParameterizedResult> {
41        exprs.iter().map(|expr| self.parameterize(expr)).collect()
42    }
43
44    fn transform_with_params(
45        &mut self,
46        expr: &Expression,
47        result: &mut ParameterizedResult,
48    ) -> Expression {
49        match expr {
50            Expression::Literal(value) => {
51                let param_name = result.add_parameter(value.clone());
52                Expression::Variable(param_name)
53            }
54            Expression::Binary { left, op, right } => {
55                let new_left = self.transform_with_params(left, result);
56                let new_right = self.transform_with_params(right, result);
57                Expression::Binary {
58                    left: Box::new(new_left),
59                    op: *op,
60                    right: Box::new(new_right),
61                }
62            }
63            Expression::Unary { op, operand } => {
64                let new_operand = self.transform_with_params(operand, result);
65                Expression::Unary {
66                    op: *op,
67                    operand: Box::new(new_operand),
68                }
69            }
70            Expression::Function { name, args } => {
71                let new_args: Vec<FunctionArg> = args
72                    .iter()
73                    .map(|arg| {
74                        FunctionArg::Positional(self.transform_with_params(arg.as_expr(), result))
75                    })
76                    .collect();
77                Expression::Function {
78                    name: name.clone(),
79                    args: new_args,
80                }
81            }
82            Expression::Aggregate {
83                func,
84                args,
85                distinct,
86                filter,
87            } => {
88                let new_args: Vec<Expression> = args
89                    .iter()
90                    .map(|a| self.transform_with_params(a, result))
91                    .collect();
92                let new_filter = filter
93                    .as_ref()
94                    .map(|f| self.transform_with_params(f, result));
95                Expression::Aggregate {
96                    func: *func,
97                    args: new_args,
98                    distinct: *distinct,
99                    filter: new_filter.map(Box::new),
100                }
101            }
102            Expression::List(items) => {
103                let new_items: Vec<Expression> = items
104                    .iter()
105                    .map(|item| self.transform_with_params(item, result))
106                    .collect();
107                Expression::List(new_items)
108            }
109            Expression::Map(pairs) => {
110                let new_pairs: Vec<(String, Expression)> = pairs
111                    .iter()
112                    .map(|(k, v)| (k.clone(), self.transform_with_params(v, result)))
113                    .collect();
114                Expression::Map(new_pairs)
115            }
116            Expression::Case {
117                test_expr,
118                conditions,
119                default,
120            } => {
121                let new_test_expr = test_expr
122                    .as_ref()
123                    .map(|e| Box::new(self.transform_with_params(e, result)));
124                let new_conditions: Vec<(Expression, Expression)> = conditions
125                    .iter()
126                    .map(|(cond, val)| {
127                        (
128                            self.transform_with_params(cond, result),
129                            self.transform_with_params(val, result),
130                        )
131                    })
132                    .collect();
133                let new_default = default
134                    .as_ref()
135                    .map(|d| Box::new(self.transform_with_params(d, result)));
136                Expression::Case {
137                    test_expr: new_test_expr,
138                    conditions: new_conditions,
139                    default: new_default,
140                }
141            }
142            Expression::TypeCast {
143                expression,
144                target_type,
145            } => {
146                let new_expr = self.transform_with_params(expression, result);
147                Expression::TypeCast {
148                    expression: Box::new(new_expr),
149                    target_type: target_type.clone(),
150                }
151            }
152            Expression::Subscript { collection, index } => {
153                let new_collection = self.transform_with_params(collection, result);
154                let new_index = self.transform_with_params(index, result);
155                Expression::Subscript {
156                    collection: Box::new(new_collection),
157                    index: Box::new(new_index),
158                }
159            }
160            Expression::Range {
161                collection,
162                start,
163                end,
164            } => {
165                let new_collection = self.transform_with_params(collection, result);
166                let new_start = start
167                    .as_ref()
168                    .map(|s| Box::new(self.transform_with_params(s, result)));
169                let new_end = end
170                    .as_ref()
171                    .map(|e| Box::new(self.transform_with_params(e, result)));
172                Expression::Range {
173                    collection: Box::new(new_collection),
174                    start: new_start,
175                    end: new_end,
176                }
177            }
178            Expression::Path(items) => {
179                let new_items: Vec<Expression> = items
180                    .iter()
181                    .map(|item| self.transform_with_params(item, result))
182                    .collect();
183                Expression::Path(new_items)
184            }
185            Expression::Property { object, property } => {
186                let new_object = self.transform_with_params(object, result);
187                Expression::Property {
188                    object: Box::new(new_object),
189                    property: property.clone(),
190                }
191            }
192            Expression::StructField { base, field } => {
193                let new_base = self.transform_with_params(base, result);
194                Expression::StructField {
195                    base: Box::new(new_base),
196                    field: field.clone(),
197                }
198            }
199            Expression::ListComprehension {
200                variable,
201                source,
202                filter,
203                map,
204            } => {
205                let new_source = self.transform_with_params(source, result);
206                let new_filter = filter
207                    .as_ref()
208                    .map(|f| Box::new(self.transform_with_params(f, result)));
209                let new_map = map
210                    .as_ref()
211                    .map(|m| Box::new(self.transform_with_params(m, result)));
212                Expression::ListComprehension {
213                    variable: variable.clone(),
214                    source: Box::new(new_source),
215                    filter: new_filter,
216                    map: new_map,
217                }
218            }
219            Expression::LabelTagProperty { tag, property } => {
220                let new_tag = self.transform_with_params(tag, result);
221                Expression::LabelTagProperty {
222                    tag: Box::new(new_tag),
223                    property: property.clone(),
224                }
225            }
226            Expression::Predicate { func, args } => {
227                let new_args: Vec<Expression> = args
228                    .iter()
229                    .map(|arg| self.transform_with_params(arg, result))
230                    .collect();
231                Expression::Predicate {
232                    func: func.clone(),
233                    args: new_args,
234                }
235            }
236            Expression::Reduce {
237                accumulator,
238                initial,
239                variable,
240                source,
241                mapping,
242            } => {
243                let new_initial = self.transform_with_params(initial, result);
244                let new_source = self.transform_with_params(source, result);
245                let new_mapping = self.transform_with_params(mapping, result);
246                Expression::Reduce {
247                    accumulator: accumulator.clone(),
248                    initial: Box::new(new_initial),
249                    variable: variable.clone(),
250                    source: Box::new(new_source),
251                    mapping: Box::new(new_mapping),
252                }
253            }
254            Expression::PathBuild(exprs) => {
255                let new_exprs: Vec<Expression> = exprs
256                    .iter()
257                    .map(|e| self.transform_with_params(e, result))
258                    .collect();
259                Expression::PathBuild(new_exprs)
260            }
261            Expression::Variable(_)
262            | Expression::Label(_)
263            | Expression::TagProperty { .. }
264            | Expression::EdgeProperty { .. }
265            | Expression::Parameter(_)
266            | Expression::SessionVariable(_)
267            | Expression::Vector(_)
268            | Expression::Exists { .. }
269            | Expression::In { .. }
270            | Expression::WindowFunction { .. } => expr.clone(),
271            Expression::CountSubquery { body } => {
272                Expression::count_subquery(self.transform_subquery_body(body, result))
273            }
274            Expression::ScalarSubquery { body } => {
275                Expression::scalar_subquery(self.transform_subquery_body(body, result))
276            }
277            Expression::Lambda { params, body } => {
278                let transformed_body = self.transform_with_params(body, result);
279                Expression::Lambda {
280                    params: params.clone(),
281                    body: Box::new(transformed_body),
282                }
283            }
284        }
285    }
286
287    fn transform_subquery_body(
288        &mut self,
289        body: &linkrs_core::types::expr::SubqueryBody,
290        result: &mut ParameterizedResult,
291    ) -> linkrs_core::types::expr::SubqueryBody {
292        let where_clause = body.where_clause.as_ref().map(|expr| {
293            let transformed = self.transform_with_params(expr, result);
294            Box::new(transformed)
295        });
296        let return_expr = body.return_expr.as_ref().map(|expr| {
297            let transformed = self.transform_with_params(expr, result);
298            Box::new(transformed)
299        });
300        linkrs_core::types::expr::SubqueryBody {
301            id: body.id,
302            patterns: body.patterns.clone(),
303            where_clause,
304            return_expr,
305        }
306    }
307}
308
309impl Default for ParameterizingTransformer {
310    fn default() -> Self {
311        Self::new()
312    }
313}