1use crate::optimizer::cost::config::CostModelConfig;
10use linkrs_core::types::{BinaryOperator, Expression, UnaryOperator};
11use linkrs_core::value::Value;
12
13#[derive(Debug, Clone)]
17pub struct ExpressionParser {
18 config: CostModelConfig,
20}
21
22impl ExpressionParser {
23 pub fn config(&self) -> &CostModelConfig {
25 &self.config
26 }
27}
28
29impl ExpressionParser {
30 pub fn new(config: CostModelConfig) -> Self {
32 Self { config }
33 }
34
35 pub fn parse_list_size(&self, expr: &str) -> Option<f64> {
45 let expr = expr.trim();
46
47 if expr.starts_with('[') && expr.ends_with(']') {
49 return self.parse_array_literal(expr);
50 }
51
52 if expr.starts_with("range(") && expr.ends_with(')') {
54 return self.parse_range_function(expr);
55 }
56
57 if expr.contains("..") {
59 return self.parse_range_expression(expr);
60 }
61
62 if let Some(size) = self.parse_collection_function(expr) {
64 return Some(size);
65 }
66
67 if let Some(size) = self.parse_split_function(expr) {
69 return Some(size);
70 }
71
72 let expr_lower = expr.to_lowercase();
74 if expr_lower.starts_with("collect(") || expr_lower.contains(".collect()") {
75 return Some(self.config.default_unwind_list_size * 2.0);
77 }
78
79 None
80 }
81
82 pub fn fold_constants(&self, expr: &Expression) -> Expression {
95 match expr {
96 Expression::Binary { left, op, right } => {
97 let folded_left = self.fold_constants(left);
98 let folded_right = self.fold_constants(right);
99
100 if let (Expression::Literal(l), Expression::Literal(r)) =
102 (&folded_left, &folded_right)
103 {
104 if let Some(result) = self.evaluate_binary_op(op, l, r) {
105 return Expression::Literal(result);
106 }
107 }
108
109 Expression::Binary {
110 left: Box::new(folded_left),
111 op: *op,
112 right: Box::new(folded_right),
113 }
114 }
115 Expression::Unary { op, operand } => {
116 let folded_operand = self.fold_constants(operand);
117
118 if let Expression::Literal(v) = &folded_operand {
120 if let Some(result) = self.evaluate_unary_op(op, v) {
121 return Expression::Literal(result);
122 }
123 }
124
125 Expression::Unary {
126 op: *op,
127 operand: Box::new(folded_operand),
128 }
129 }
130 Expression::Function { name, args } => {
131 let folded_args: Vec<Expression> = args
132 .iter()
133 .map(|arg| self.fold_constants(arg.as_expr()))
134 .collect();
135
136 if folded_args
138 .iter()
139 .all(|arg| matches!(arg, Expression::Literal(_)))
140 {
141 let arg_values: Vec<&Value> = folded_args
142 .iter()
143 .filter_map(|arg| match arg {
144 Expression::Literal(v) => Some(v),
145 _ => None,
146 })
147 .collect();
148
149 if let Some(result) = self.evaluate_function(name, &arg_values) {
150 return Expression::Literal(result);
151 }
152 }
153
154 Expression::function(name.clone(), folded_args)
155 }
156 Expression::List(items) => {
157 let folded_items: Vec<Expression> =
158 items.iter().map(|item| self.fold_constants(item)).collect();
159 Expression::List(folded_items)
160 }
161 Expression::Map(entries) => {
162 let folded_entries: Vec<(String, Expression)> = entries
163 .iter()
164 .map(|(k, v)| (k.clone(), self.fold_constants(v)))
165 .collect();
166 Expression::Map(folded_entries)
167 }
168 _ => expr.clone(),
169 }
170 }
171
172 fn evaluate_binary_op(
174 &self,
175 op: &BinaryOperator,
176 left: &Value,
177 right: &Value,
178 ) -> Option<Value> {
179 match op {
180 BinaryOperator::Add => self.add_values(left, right),
181 BinaryOperator::Subtract => self.subtract_values(left, right),
182 BinaryOperator::Multiply => self.multiply_values(left, right),
183 BinaryOperator::Divide => self.divide_values(left, right),
184 BinaryOperator::Modulo => self.modulo_values(left, right),
185 BinaryOperator::Equal => Some(Value::Bool(
186 self.compare_values(left, right) == Some(std::cmp::Ordering::Equal),
187 )),
188 BinaryOperator::NotEqual => Some(Value::Bool(
189 self.compare_values(left, right) != Some(std::cmp::Ordering::Equal),
190 )),
191 BinaryOperator::LessThan => Some(Value::Bool(
192 self.compare_values(left, right) == Some(std::cmp::Ordering::Less),
193 )),
194 BinaryOperator::GreaterThan => Some(Value::Bool(
195 self.compare_values(left, right) == Some(std::cmp::Ordering::Greater),
196 )),
197 BinaryOperator::LessThanOrEqual => {
198 let cmp = self.compare_values(left, right);
199 Some(Value::Bool(
200 cmp == Some(std::cmp::Ordering::Less) || cmp == Some(std::cmp::Ordering::Equal),
201 ))
202 }
203 BinaryOperator::GreaterThanOrEqual => {
204 let cmp = self.compare_values(left, right);
205 Some(Value::Bool(
206 cmp == Some(std::cmp::Ordering::Greater)
207 || cmp == Some(std::cmp::Ordering::Equal),
208 ))
209 }
210 BinaryOperator::And => self.logical_and(left, right),
211 BinaryOperator::Or => self.logical_or(left, right),
212 BinaryOperator::StringConcat => self.concat_values(left, right),
213 _ => None,
214 }
215 }
216
217 fn evaluate_unary_op(&self, op: &UnaryOperator, operand: &Value) -> Option<Value> {
219 match op {
220 UnaryOperator::Not => match operand {
221 Value::Bool(b) => Some(Value::Bool(!b)),
222 _ => None,
223 },
224 UnaryOperator::Minus => match operand {
225 Value::Int(i) => Some(Value::Int(-i)),
226 Value::Float(f) => Some(Value::Float(-f)),
227 _ => None,
228 },
229 UnaryOperator::IsNull => Some(Value::Bool(operand.is_null())),
230 UnaryOperator::IsNotNull => Some(Value::Bool(!operand.is_null())),
231 _ => None,
232 }
233 }
234
235 fn evaluate_function(&self, name: &str, args: &[&Value]) -> Option<Value> {
237 let name_lower = name.to_lowercase();
238
239 match name_lower.as_str() {
240 "abs" if args.len() == 1 => match args[0] {
241 Value::SmallInt(i) => Some(Value::SmallInt(i.abs())),
242 Value::Int(i) => Some(Value::Int(i.abs())),
243 Value::BigInt(i) => Some(Value::BigInt(i.abs())),
244 Value::Float(f) => Some(Value::Float(f.abs())),
245 Value::Double(f) => Some(Value::Double(f.abs())),
246 _ => None,
247 },
248 "length" | "size" if args.len() == 1 => match args[0] {
249 Value::String(s) => Some(Value::BigInt(s.len() as i64)),
250 Value::List(list) => Some(Value::BigInt(list.len() as i64)),
251 _ => None,
252 },
253 "upper" | "toupper" if args.len() == 1 => match args[0] {
254 Value::String(s) => Some(Value::string(s.to_uppercase())),
255 _ => None,
256 },
257 "lower" | "tolower" if args.len() == 1 => match args[0] {
258 Value::String(s) => Some(Value::string(s.to_lowercase())),
259 _ => None,
260 },
261 "substring" | "substr" if args.len() >= 2 => match (args[0], args.get(1).copied()) {
262 (Value::String(s), Some(Value::BigInt(start))) => {
263 let start_idx = if *start >= 0 { *start as usize } else { 0 };
264 let end_idx = if args.len() >= 3 {
265 match args[2] {
266 Value::BigInt(len) => start_idx + (*len as usize),
267 _ => s.len(),
268 }
269 } else {
270 s.len()
271 };
272 Some(Value::string(
273 s.chars()
274 .skip(start_idx)
275 .take(end_idx - start_idx)
276 .collect::<String>(),
277 ))
278 }
279 _ => None,
280 },
281 "trim" if args.len() == 1 => match args[0] {
282 Value::String(s) => Some(Value::string(s.trim())),
283 _ => None,
284 },
285 _ => None,
286 }
287 }
288
289 fn add_values(&self, left: &Value, right: &Value) -> Option<Value> {
292 match (left, right) {
293 (Value::SmallInt(l), Value::SmallInt(r)) => Some(Value::SmallInt(l + r)),
294 (Value::Int(l), Value::Int(r)) => Some(Value::Int(l + r)),
295 (Value::BigInt(l), Value::BigInt(r)) => Some(Value::BigInt(l + r)),
296 (Value::Float(l), Value::Float(r)) => Some(Value::Float(l + r)),
297 (Value::Double(l), Value::Double(r)) => Some(Value::Double(l + r)),
298 (Value::String(l), Value::String(r)) => Some(Value::string(format!("{}{}", l, r))),
299 _ => None,
300 }
301 }
302
303 fn subtract_values(&self, left: &Value, right: &Value) -> Option<Value> {
304 match (left, right) {
305 (Value::SmallInt(l), Value::SmallInt(r)) => Some(Value::SmallInt(l - r)),
306 (Value::Int(l), Value::Int(r)) => Some(Value::Int(l - r)),
307 (Value::BigInt(l), Value::BigInt(r)) => Some(Value::BigInt(l - r)),
308 (Value::Float(l), Value::Float(r)) => Some(Value::Float(l - r)),
309 (Value::Double(l), Value::Double(r)) => Some(Value::Double(l - r)),
310 _ => None,
311 }
312 }
313
314 fn multiply_values(&self, left: &Value, right: &Value) -> Option<Value> {
315 match (left, right) {
316 (Value::SmallInt(l), Value::SmallInt(r)) => Some(Value::SmallInt(l * r)),
317 (Value::Int(l), Value::Int(r)) => Some(Value::Int(l * r)),
318 (Value::BigInt(l), Value::BigInt(r)) => Some(Value::BigInt(l * r)),
319 (Value::Float(l), Value::Float(r)) => Some(Value::Float(l * r)),
320 (Value::Double(l), Value::Double(r)) => Some(Value::Double(l * r)),
321 _ => None,
322 }
323 }
324
325 fn divide_values(&self, left: &Value, right: &Value) -> Option<Value> {
326 match (left, right) {
327 (Value::SmallInt(l), Value::SmallInt(r)) if *r != 0 => Some(Value::SmallInt(l / r)),
328 (Value::Int(l), Value::Int(r)) if *r != 0 => Some(Value::Int(l / r)),
329 (Value::BigInt(l), Value::BigInt(r)) if *r != 0 => Some(Value::BigInt(l / r)),
330 (Value::Float(l), Value::Float(r)) if *r != 0.0 => Some(Value::Float(l / r)),
331 (Value::Double(l), Value::Double(r)) if *r != 0.0 => Some(Value::Double(l / r)),
332 _ => None,
333 }
334 }
335
336 fn modulo_values(&self, left: &Value, right: &Value) -> Option<Value> {
337 match (left, right) {
338 (Value::SmallInt(l), Value::SmallInt(r)) if *r != 0 => Some(Value::SmallInt(l % r)),
339 (Value::Int(l), Value::Int(r)) if *r != 0 => Some(Value::Int(l % r)),
340 (Value::BigInt(l), Value::BigInt(r)) if *r != 0 => Some(Value::BigInt(l % r)),
341 _ => None,
342 }
343 }
344
345 fn compare_values(&self, left: &Value, right: &Value) -> Option<std::cmp::Ordering> {
346 match (left, right) {
347 (Value::SmallInt(l), Value::SmallInt(r)) => Some(l.cmp(r)),
348 (Value::Int(l), Value::Int(r)) => Some(l.cmp(r)),
349 (Value::BigInt(l), Value::BigInt(r)) => Some(l.cmp(r)),
350 (Value::Float(l), Value::Float(r)) => l.partial_cmp(r),
351 (Value::Double(l), Value::Double(r)) => l.partial_cmp(r),
352 (Value::String(l), Value::String(r)) => Some(l.cmp(r)),
353 (Value::Bool(l), Value::Bool(r)) => Some(l.cmp(r)),
354 _ => None,
355 }
356 }
357
358 fn logical_and(&self, left: &Value, right: &Value) -> Option<Value> {
359 match (left, right) {
360 (Value::Bool(l), Value::Bool(r)) => Some(Value::Bool(*l && *r)),
361 _ => None,
362 }
363 }
364
365 fn logical_or(&self, left: &Value, right: &Value) -> Option<Value> {
366 match (left, right) {
367 (Value::Bool(l), Value::Bool(r)) => Some(Value::Bool(*l || *r)),
368 _ => None,
369 }
370 }
371
372 fn concat_values(&self, left: &Value, right: &Value) -> Option<Value> {
373 match (left, right) {
374 (Value::String(l), Value::String(r)) => Some(Value::string(format!("{}{}", l, r))),
375 _ => None,
376 }
377 }
378
379 pub fn parse_loop_iterations(&self, condition: &str) -> Option<u32> {
389 let condition = condition.trim();
390
391 if let Ok(num) = condition.parse::<u32>() {
393 return Some(num.max(1));
394 }
395
396 if condition.contains("..") {
398 return self.parse_range_expression_u32(condition);
399 }
400
401 if condition.starts_with("range(") && condition.ends_with(")") {
403 return self.parse_range_function(condition).map(|v| v as u32);
404 }
405
406 if let Some(iterations) = self.parse_comparison_expression(condition) {
408 return Some(iterations as u32);
409 }
410
411 if condition.starts_with('[') && condition.ends_with(']') {
413 return Some(self.parse_collection_size(condition));
414 }
415
416 None
417 }
418
419 fn parse_array_literal(&self, expr: &str) -> Option<f64> {
421 let inner = &expr[1..expr.len() - 1];
422 if inner.trim().is_empty() {
423 return Some(0.0);
424 }
425 let count = self.count_top_level_commas(inner) as f64;
428 Some(count)
429 }
430
431 fn count_top_level_commas(&self, s: &str) -> usize {
433 let mut count = 0;
434 let mut depth = 0;
435 let chars = s.chars().peekable();
436
437 for c in chars {
438 match c {
439 '[' | '(' | '{' => depth += 1,
440 ']' | ')' | '}' => depth -= 1,
441 ',' if depth == 0 => count += 1,
442 _ => {}
443 }
444 }
445
446 count + 1 }
448
449 fn parse_range_function(&self, expr: &str) -> Option<f64> {
451 let args_str = &expr[6..expr.len() - 1];
452 let args: Vec<&str> = args_str.split(',').map(|s| s.trim()).collect();
453
454 if args.len() >= 2 {
455 let start: i64 = args[0].parse().ok()?;
456 let end: i64 = args[1].parse().ok()?;
457 let step: i64 = if args.len() >= 3 {
458 args[2].parse().ok()?
459 } else {
460 1
461 };
462
463 if step != 0 {
464 let count = ((end - start) / step).abs() as f64;
465 return Some(count.max(0.0));
466 }
467 }
468 None
469 }
470
471 fn parse_range_expression(&self, expr: &str) -> Option<f64> {
473 if let Some(pos) = expr.find("..") {
474 let start_str = expr[..pos].trim();
475 let end_part = &expr[pos + 2..];
476
477 let (end_str, inclusive) = if let Some(stripped) = end_part.strip_prefix('=') {
478 (stripped, true)
479 } else {
480 (end_part, false)
481 };
482 if let (Ok(start), Ok(end)) = (start_str.parse::<i64>(), end_str.trim().parse::<i64>())
483 {
484 if end > start {
485 let count = if inclusive {
486 end - start + 1
487 } else {
488 end - start
489 };
490 return Some(count.max(0) as f64);
491 }
492 }
493 }
494 None
495 }
496
497 fn parse_range_expression_u32(&self, expr: &str) -> Option<u32> {
499 if let Some(pos) = expr.find("..") {
501 let start_str = expr[..pos].trim();
502 let end_part = &expr[pos + 2..];
503
504 let (end_str, inclusive) = if let Some(stripped) = end_part.strip_prefix('=') {
506 (stripped, true)
507 } else {
508 (end_part, false)
509 };
510
511 if let (Ok(start), Ok(end)) = (start_str.parse::<i64>(), end_str.trim().parse::<i64>())
512 {
513 if end > start {
514 let count = if inclusive {
515 end - start + 1
516 } else {
517 end - start
518 };
519 return Some(count.max(1) as u32);
520 }
521 }
522 }
523 None
524 }
525
526 fn parse_comparison_expression(&self, expr: &str) -> Option<f64> {
528 let operators = [("<", 0u32), ("<=", 0u32), (">", 0u32), (">=", 0u32)];
530
531 for (op, _) in &operators {
532 if let Some(pos) = expr.find(op) {
533 let right_side = &expr[pos + op.len()..];
534 if let Ok(num) = right_side.trim().parse::<i64>() {
535 let iterations = if num > 0 { num as f64 } else { 1.0 };
537 let iterations = iterations + 10.0;
539 return Some(iterations.max(1.0));
540 }
541 }
542 }
543 None
544 }
545
546 fn parse_collection_function(&self, expr: &str) -> Option<f64> {
548 let expr_lower = expr.to_lowercase();
549
550 if expr_lower.contains(".keys()") {
552 return Some(self.config.default_unwind_list_size);
554 }
555
556 if expr_lower.starts_with("values(") || expr_lower.contains(".values()") {
558 return Some(self.config.default_unwind_list_size);
559 }
560
561 if expr_lower.contains(".nodes()") {
563 return Some(self.config.default_unwind_list_size);
565 }
566
567 if expr_lower.starts_with("relationships(") || expr_lower.starts_with("rels(") {
569 return Some(self.config.default_unwind_list_size - 1.0); }
571
572 if expr_lower.starts_with("labels(") {
574 return Some(1.0); }
576
577 None
578 }
579
580 fn parse_split_function(&self, expr: &str) -> Option<f64> {
582 let expr_lower = expr.to_lowercase();
583
584 if expr_lower.contains("split(") || expr_lower.contains(".split(") {
586 return Some(self.config.default_unwind_list_size);
589 }
590
591 None
592 }
593
594 fn parse_collection_size(&self, expr: &str) -> u32 {
596 let inner = &expr[1..expr.len() - 1];
597 if inner.trim().is_empty() {
598 return 0;
599 }
600 let count = inner.split(',').count() as u32;
603 count.max(1)
604 }
605}
606
607#[cfg(test)]
608mod tests {
609 use super::*;
610
611 #[test]
612 fn test_parse_array_literal() {
613 let parser = ExpressionParser::new(CostModelConfig::default());
614 assert_eq!(parser.parse_list_size("[1, 2, 3]"), Some(3.0));
615 assert_eq!(parser.parse_list_size("[]"), Some(0.0));
616 assert_eq!(parser.parse_list_size("[a]"), Some(1.0));
617 }
618
619 #[test]
620 fn test_parse_range_function() {
621 let parser = ExpressionParser::new(CostModelConfig::default());
622 assert_eq!(parser.parse_list_size("range(1, 10)"), Some(9.0));
623 assert_eq!(parser.parse_list_size("range(1, 10, 2)"), Some(4.0));
624 }
625
626 #[test]
627 fn test_parse_range_expression() {
628 let parser = ExpressionParser::new(CostModelConfig::default());
629 assert_eq!(parser.parse_list_size("1..10"), Some(9.0));
630 assert_eq!(parser.parse_list_size("0..=5"), Some(6.0));
631 }
632
633 #[test]
634 fn test_parse_loop_iterations_number() {
635 let parser = ExpressionParser::new(CostModelConfig::default());
636 assert_eq!(parser.parse_loop_iterations("10"), Some(10)); assert_eq!(parser.parse_loop_iterations("0"), Some(1)); }
639
640 #[test]
641 fn test_parse_loop_iterations_comparison() {
642 let parser = ExpressionParser::new(CostModelConfig::default());
643 assert_eq!(parser.parse_loop_iterations("i < 10"), Some(20));
644 assert_eq!(parser.parse_loop_iterations("i <= 5"), Some(15));
645 }
646}