refactor(planner): dedup PromQL binary-arm decomposition logic

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<String> {
-    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<String> {
-    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.

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.