diff --git a/asap-planner-rs/src/optimizer/aqe_extractor.rs b/asap-planner-rs/src/optimizer/aqe_extractor.rs index 7c2da85..1048ce0 100644 --- a/asap-planner-rs/src/optimizer/aqe_extractor.rs +++ b/asap-planner-rs/src/optimizer/aqe_extractor.rs @@ -11,6 +11,7 @@ use promql_utilities::query_logics::parsing::{ use tracing::warn; use crate::planner::patterns::build_patterns; +use crate::planner::promql::{parse_binary_arms, BinaryArm}; use super::solution::AQE; @@ -143,43 +144,18 @@ fn gcd(a: u64, b: u64) -> u64 { /// dropped — they contribute no AQE. Only arithmetic operators are split; /// comparison and set operators are left as-is (treated as opaque leaves). fn decompose_to_leaves(query: &str) -> Vec { - let ast = match promql_parser::parser::parse(query) { - Ok(a) => a, - Err(_) => return vec![query.to_string()], + let (lhs, rhs) = match parse_binary_arms(query) { + Some(arms) => arms, + None => return vec![query.to_string()], }; - if let promql_parser::parser::Expr::Binary(binary) = &ast { - if !binary.op.is_comparison_operator() && !binary.op.is_set_operator() { - let mut leaves = Vec::new(); - if let Some(lhs_str) = arm_to_query_string(binary.lhs.as_ref()) { - leaves.extend(decompose_to_leaves(&lhs_str)); - } - if let Some(rhs_str) = arm_to_query_string(binary.rhs.as_ref()) { - leaves.extend(decompose_to_leaves(&rhs_str)); - } - return leaves; + let mut leaves = Vec::new(); + for arm in [lhs, rhs] { + if let BinaryArm::Query(arm_query) = arm { + leaves.extend(decompose_to_leaves(&arm_query)); } } - - vec![query.to_string()] -} - -/// Convert one arm of a binary expression to a query string, returning `None` -/// for scalar literals (they don't map to AQEs). -fn arm_to_query_string(expr: &promql_parser::parser::Expr) -> Option { - let inner = strip_parens(expr); - match inner { - promql_parser::parser::Expr::NumberLiteral(_) => None, - other => Some(format!("{}", other)), - } -} - -fn strip_parens(expr: &promql_parser::parser::Expr) -> &promql_parser::parser::Expr { - if let promql_parser::parser::Expr::Paren(paren) = expr { - strip_parens(&paren.expr) - } else { - expr - } + leaves } /// Try to extract `QueryRequirements` from a single leaf PromQL query string. diff --git a/asap-planner-rs/src/planner/promql.rs b/asap-planner-rs/src/planner/promql.rs index 428622a..a0e952a 100644 --- a/asap-planner-rs/src/planner/promql.rs +++ b/asap-planner-rs/src/planner/promql.rs @@ -50,6 +50,25 @@ fn strip_parens(expr: &promql_parser::parser::Expr) -> &promql_parser::parser::E } } +/// Parse `query` and, if it is a top-level binary arithmetic expression, +/// return its two arms. Returns `None` if the query fails to parse, isn't a +/// `Binary` expression, or uses a comparison/set operator (e.g. `==`, `and`). +/// +/// This is the single shared entry point for binary-arm decomposition; both +/// `SingleQueryProcessor::get_binary_arm_queries` and the AQE extractor's +/// leaf decomposition build on it. +pub(crate) fn parse_binary_arms(query: &str) -> Option<(BinaryArm, BinaryArm)> { + let ast = promql_parser::parser::parse(query).ok()?; + if let promql_parser::parser::Expr::Binary(binary) = ast { + if !binary.op.is_comparison_operator() && !binary.op.is_set_operator() { + let lhs = expr_to_binary_arm(binary.lhs.as_ref()); + let rhs = expr_to_binary_arm(binary.rhs.as_ref()); + return Some((lhs, rhs)); + } + } + None +} + pub struct SingleQueryProcessor { query: String, t_repeat: u64, @@ -132,16 +151,7 @@ impl SingleQueryProcessor { /// (`BinaryArm::Scalar`). Returns `None` if the query is not a binary expression /// or cannot be parsed. pub fn get_binary_arm_queries(&self) -> Option<(BinaryArm, BinaryArm)> { - let ast = promql_parser::parser::parse(&self.query).ok()?; - if let promql_parser::parser::Expr::Binary(binary) = ast { - let lhs = expr_to_binary_arm(binary.lhs.as_ref()); - let rhs = expr_to_binary_arm(binary.rhs.as_ref()); - // Only handle arithmetic operators (not comparison or set operators) - if !binary.op.is_comparison_operator() && !binary.op.is_set_operator() { - return Some((lhs, rhs)); - } - } - None + parse_binary_arms(&self.query) } /// Create a new processor for an arm query, reusing all parameters from this processor.