rusty_promql_parser/
ast.rs

1//! Abstract Syntax Tree (AST) types for PromQL expressions.
2//!
3//! This module defines all the types that make up the parsed representation
4//! of a PromQL query. The main type is [`Expr`], which represents any valid
5//! PromQL expression.
6//!
7//! # Overview
8//!
9//! PromQL expressions can be:
10//! - **Literals**: Numbers ([`Expr::Number`]) and strings ([`Expr::String`])
11//! - **Selectors**: Instant vectors ([`Expr::VectorSelector`]) and range vectors ([`Expr::MatrixSelector`])
12//! - **Operations**: Binary ([`Expr::Binary`]) and unary ([`Expr::Unary`]) expressions
13//! - **Functions**: Function calls ([`Expr::Call`]) and aggregations ([`Expr::Aggregation`])
14//! - **Subqueries**: Range queries over instant vectors ([`Expr::Subquery`])
15//!
16//! # Display
17//!
18//! All AST types implement [`std::fmt::Display`] to convert back to PromQL syntax:
19//!
20//! ```rust
21//! use rusty_promql_parser::expr;
22//!
23//! let (_, ast) = expr("sum(rate(http_requests[5m]))").unwrap();
24//! println!("{}", ast); // Prints: sum(rate(http_requests[5m]))
25//! ```
26
27use std::fmt;
28
29use crate::lexer::duration::Duration;
30use crate::parser::aggregation::Grouping;
31use crate::parser::selector::{AtModifier, MatrixSelector, VectorSelector};
32
33/// Root expression type for PromQL AST
34#[derive(Debug, Clone, PartialEq)]
35pub enum Expr {
36    /// Numeric literal: `42`, `3.14`, `0x1F`, `1e-10`, `Inf`, `NaN`
37    Number(f64),
38
39    /// String literal: `"hello"`, `'world'`, `` `raw` ``
40    String(String),
41
42    /// Instant vector selector: `http_requests{job="api"}`
43    VectorSelector(VectorSelector),
44
45    /// Range vector selector: `http_requests{job="api"}[5m]`
46    MatrixSelector(MatrixSelector),
47
48    /// Function call: `rate(http_requests[5m])`
49    Call(Call),
50
51    /// Aggregation: `sum by (job) (http_requests)`
52    Aggregation(Box<Aggregation>),
53
54    /// Binary operation: `foo + bar`, `foo / on(job) bar`
55    Binary(Box<BinaryExpr>),
56
57    /// Unary operation: `-foo`, `+bar`
58    Unary(Box<UnaryExpr>),
59
60    /// Parenthesized: `(foo + bar)`
61    Paren(Box<Expr>),
62
63    /// Subquery: `rate(http_requests[5m])[30m:1m]`
64    Subquery(Box<SubqueryExpr>),
65}
66
67impl Expr {
68    /// Check if this is a scalar expression (number literal)
69    pub fn is_scalar(&self) -> bool {
70        matches!(self, Expr::Number(_))
71    }
72
73    /// Check if this is a string literal
74    pub fn is_string(&self) -> bool {
75        matches!(self, Expr::String(_))
76    }
77
78    /// Check if this expression produces an instant vector
79    pub fn is_instant_vector(&self) -> bool {
80        matches!(
81            self,
82            Expr::VectorSelector(_)
83                | Expr::Call(_)
84                | Expr::Aggregation(_)
85                | Expr::Binary(_)
86                | Expr::Unary(_)
87        ) || matches!(self, Expr::Paren(e) if e.is_instant_vector())
88    }
89
90    /// Check if this expression produces a range vector
91    pub fn is_range_vector(&self) -> bool {
92        matches!(self, Expr::MatrixSelector(_) | Expr::Subquery(_))
93    }
94
95    /// Unwrap parentheses to get the inner expression
96    pub fn unwrap_parens(&self) -> &Expr {
97        match self {
98            Expr::Paren(inner) => inner.unwrap_parens(),
99            other => other,
100        }
101    }
102}
103
104impl fmt::Display for Expr {
105    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
106        match self {
107            Expr::Number(n) => {
108                if n.is_nan() {
109                    write!(f, "NaN")
110                } else if n.is_infinite() {
111                    if *n > 0.0 {
112                        write!(f, "Inf")
113                    } else {
114                        write!(f, "-Inf")
115                    }
116                } else {
117                    write!(f, "{}", n)
118                }
119            }
120            Expr::String(s) => write!(f, "\"{}\"", s.escape_default()),
121            Expr::VectorSelector(v) => write!(f, "{}", v),
122            Expr::MatrixSelector(m) => write!(f, "{}", m),
123            Expr::Call(c) => write!(f, "{}", c),
124            Expr::Aggregation(a) => write!(f, "{}", a),
125            Expr::Binary(b) => write!(f, "{}", b),
126            Expr::Unary(u) => write!(f, "{}", u),
127            Expr::Paren(e) => write!(f, "({})", e),
128            Expr::Subquery(s) => write!(f, "{}", s),
129        }
130    }
131}
132
133/// Function call expression
134#[derive(Debug, Clone, PartialEq)]
135pub struct Call {
136    /// Function name
137    pub name: String,
138    /// Function arguments
139    pub args: Vec<Expr>,
140}
141
142impl Call {
143    /// Create a new function call
144    pub fn new(name: impl Into<String>, args: Vec<Expr>) -> Self {
145        Self {
146            name: name.into(),
147            args,
148        }
149    }
150}
151
152impl fmt::Display for Call {
153    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154        write!(f, "{}(", self.name)?;
155        for (i, arg) in self.args.iter().enumerate() {
156            if i > 0 {
157                write!(f, ", ")?;
158            }
159            write!(f, "{}", arg)?;
160        }
161        write!(f, ")")
162    }
163}
164
165/// Aggregation expression
166#[derive(Debug, Clone, PartialEq)]
167pub struct Aggregation {
168    /// The aggregation operator name
169    pub op: String,
170    /// The expression to aggregate
171    pub expr: Expr,
172    /// Parameter for parametric aggregations (topk, quantile, etc.)
173    pub param: Option<Expr>,
174    /// Optional grouping clause (by/without)
175    pub grouping: Option<Grouping>,
176}
177
178impl Aggregation {
179    /// Create a new aggregation
180    pub fn new(op: impl Into<String>, expr: Expr) -> Self {
181        Self {
182            op: op.into(),
183            expr,
184            param: None,
185            grouping: None,
186        }
187    }
188
189    /// Create a new aggregation with a parameter
190    pub fn with_param(op: impl Into<String>, param: Expr, expr: Expr) -> Self {
191        Self {
192            op: op.into(),
193            expr,
194            param: Some(param),
195            grouping: None,
196        }
197    }
198
199    /// Set the grouping clause
200    pub fn with_grouping(mut self, grouping: Grouping) -> Self {
201        self.grouping = Some(grouping);
202        self
203    }
204}
205
206impl fmt::Display for Aggregation {
207    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
208        write!(f, "{}", self.op)?;
209        if let Some(ref grouping) = self.grouping {
210            write!(f, " {} ", grouping)?;
211        }
212        write!(f, "(")?;
213        if let Some(ref param) = self.param {
214            write!(f, "{}, ", param)?;
215        }
216        write!(f, "{})", self.expr)
217    }
218}
219
220/// Binary operators
221#[derive(Debug, Clone, Copy, PartialEq, Eq)]
222pub enum BinaryOp {
223    // Arithmetic
224    Add,   // +
225    Sub,   // -
226    Mul,   // *
227    Div,   // /
228    Mod,   // %
229    Pow,   // ^
230    Atan2, // atan2
231
232    // Comparison
233    Eq, // ==
234    Ne, // !=
235    Lt, // <
236    Le, // <=
237    Gt, // >
238    Ge, // >=
239
240    // Set operations
241    And,    // and
242    Or,     // or
243    Unless, // unless
244}
245
246impl BinaryOp {
247    /// Get the operator as a string
248    pub fn as_str(&self) -> &'static str {
249        match self {
250            BinaryOp::Add => "+",
251            BinaryOp::Sub => "-",
252            BinaryOp::Mul => "*",
253            BinaryOp::Div => "/",
254            BinaryOp::Mod => "%",
255            BinaryOp::Pow => "^",
256            BinaryOp::Atan2 => "atan2",
257            BinaryOp::Eq => "==",
258            BinaryOp::Ne => "!=",
259            BinaryOp::Lt => "<",
260            BinaryOp::Le => "<=",
261            BinaryOp::Gt => ">",
262            BinaryOp::Ge => ">=",
263            BinaryOp::And => "and",
264            BinaryOp::Or => "or",
265            BinaryOp::Unless => "unless",
266        }
267    }
268
269    /// Get the precedence of this operator (higher = binds tighter)
270    pub fn precedence(&self) -> u8 {
271        match self {
272            BinaryOp::Or => 1,                     // Lowest
273            BinaryOp::And | BinaryOp::Unless => 2, // Set intersection/difference
274            BinaryOp::Eq
275            | BinaryOp::Ne
276            | BinaryOp::Lt
277            | BinaryOp::Le
278            | BinaryOp::Gt
279            | BinaryOp::Ge => 3, // Comparison
280            BinaryOp::Add | BinaryOp::Sub => 4,    // Addition/subtraction
281            BinaryOp::Mul | BinaryOp::Div | BinaryOp::Mod | BinaryOp::Atan2 => 5, // Multiplication/division
282            BinaryOp::Pow => 6,                                                   // Highest
283        }
284    }
285
286    /// Check if this operator is right-associative
287    pub fn is_right_associative(&self) -> bool {
288        matches!(self, BinaryOp::Pow)
289    }
290
291    /// Check if this is a comparison operator
292    pub fn is_comparison(&self) -> bool {
293        matches!(
294            self,
295            BinaryOp::Eq | BinaryOp::Ne | BinaryOp::Lt | BinaryOp::Le | BinaryOp::Gt | BinaryOp::Ge
296        )
297    }
298
299    /// Check if this is a set operator
300    pub fn is_set_operator(&self) -> bool {
301        matches!(self, BinaryOp::And | BinaryOp::Or | BinaryOp::Unless)
302    }
303
304    /// Check if this is an arithmetic operator
305    pub fn is_arithmetic(&self) -> bool {
306        matches!(
307            self,
308            BinaryOp::Add
309                | BinaryOp::Sub
310                | BinaryOp::Mul
311                | BinaryOp::Div
312                | BinaryOp::Mod
313                | BinaryOp::Pow
314                | BinaryOp::Atan2
315        )
316    }
317}
318
319impl fmt::Display for BinaryOp {
320    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
321        write!(f, "{}", self.as_str())
322    }
323}
324
325/// Vector matching for binary operations
326#[derive(Debug, Clone, Copy, PartialEq, Eq)]
327pub enum VectorMatchingOp {
328    On,       // on (label1, label2)
329    Ignoring, // ignoring (label1, label2)
330}
331
332impl fmt::Display for VectorMatchingOp {
333    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
334        match self {
335            VectorMatchingOp::On => write!(f, "on"),
336            VectorMatchingOp::Ignoring => write!(f, "ignoring"),
337        }
338    }
339}
340
341/// Group modifier side
342#[derive(Debug, Clone, Copy, PartialEq, Eq)]
343pub enum GroupSide {
344    Left,  // group_left
345    Right, // group_right
346}
347
348impl fmt::Display for GroupSide {
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        match self {
351            GroupSide::Left => write!(f, "group_left"),
352            GroupSide::Right => write!(f, "group_right"),
353        }
354    }
355}
356
357/// Group modifier for many-to-one/one-to-many matching
358#[derive(Debug, Clone, PartialEq, Eq)]
359pub struct GroupModifier {
360    /// Which side to group (left or right)
361    pub side: GroupSide,
362    /// Additional labels to include from the "one" side
363    pub labels: Vec<String>,
364}
365
366impl fmt::Display for GroupModifier {
367    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
368        write!(f, "{}", self.side)?;
369        if !self.labels.is_empty() {
370            write!(f, " (")?;
371            for (i, label) in self.labels.iter().enumerate() {
372                if i > 0 {
373                    write!(f, ", ")?;
374                }
375                write!(f, "{}", label)?;
376            }
377            write!(f, ")")?;
378        }
379        Ok(())
380    }
381}
382
383/// Vector matching specification for binary operations
384#[derive(Debug, Clone, PartialEq, Eq)]
385pub struct VectorMatching {
386    /// The matching operation (on or ignoring)
387    pub op: VectorMatchingOp,
388    /// Labels to match on or ignore
389    pub labels: Vec<String>,
390    /// Optional group modifier
391    pub group: Option<GroupModifier>,
392}
393
394impl fmt::Display for VectorMatching {
395    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
396        write!(f, "{} (", self.op)?;
397        for (i, label) in self.labels.iter().enumerate() {
398            if i > 0 {
399                write!(f, ", ")?;
400            }
401            write!(f, "{}", label)?;
402        }
403        write!(f, ")")?;
404        if let Some(ref group) = self.group {
405            write!(f, " {}", group)?;
406        }
407        Ok(())
408    }
409}
410
411/// Modifier for binary operations
412#[derive(Debug, Clone, PartialEq, Eq, Default)]
413pub struct BinaryModifier {
414    /// Whether to return bool (0/1) instead of filtering for comparisons
415    pub return_bool: bool,
416    /// Vector matching specification
417    pub matching: Option<VectorMatching>,
418}
419
420impl BinaryModifier {
421    /// Create a modifier with just the bool flag
422    pub fn with_bool() -> Self {
423        Self {
424            return_bool: true,
425            matching: None,
426        }
427    }
428
429    /// Create a modifier with vector matching
430    pub fn with_matching(matching: VectorMatching) -> Self {
431        Self {
432            return_bool: false,
433            matching: Some(matching),
434        }
435    }
436
437    /// Check if this modifier has any settings
438    pub fn is_empty(&self) -> bool {
439        !self.return_bool && self.matching.is_none()
440    }
441}
442
443impl fmt::Display for BinaryModifier {
444    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
445        if self.return_bool {
446            write!(f, "bool")?;
447            if self.matching.is_some() {
448                write!(f, " ")?;
449            }
450        }
451        if let Some(ref matching) = self.matching {
452            write!(f, "{}", matching)?;
453        }
454        Ok(())
455    }
456}
457
458/// Binary expression
459#[derive(Debug, Clone, PartialEq)]
460pub struct BinaryExpr {
461    /// The binary operator
462    pub op: BinaryOp,
463    /// Left-hand side expression
464    pub lhs: Expr,
465    /// Right-hand side expression
466    pub rhs: Expr,
467    /// Optional modifier (bool, on, ignoring, group_left, group_right)
468    pub modifier: Option<BinaryModifier>,
469}
470
471impl BinaryExpr {
472    /// Create a new binary expression
473    pub fn new(op: BinaryOp, lhs: Expr, rhs: Expr) -> Self {
474        Self {
475            op,
476            lhs,
477            rhs,
478            modifier: None,
479        }
480    }
481
482    /// Create a new binary expression with a modifier
483    pub fn with_modifier(op: BinaryOp, lhs: Expr, rhs: Expr, modifier: BinaryModifier) -> Self {
484        Self {
485            op,
486            lhs,
487            rhs,
488            modifier: Some(modifier),
489        }
490    }
491}
492
493impl fmt::Display for BinaryExpr {
494    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
495        write!(f, "{} {}", self.lhs, self.op)?;
496        if let Some(ref modifier) = self.modifier
497            && !modifier.is_empty()
498        {
499            write!(f, " {}", modifier)?;
500        }
501        write!(f, " {}", self.rhs)
502    }
503}
504
505/// Unary operators
506#[derive(Debug, Clone, Copy, PartialEq, Eq)]
507pub enum UnaryOp {
508    /// Unary plus (no-op)
509    Plus,
510    /// Unary minus (negation)
511    Minus,
512}
513
514impl UnaryOp {
515    /// Get the operator as a string
516    pub fn as_str(&self) -> &'static str {
517        match self {
518            UnaryOp::Plus => "+",
519            UnaryOp::Minus => "-",
520        }
521    }
522}
523
524impl fmt::Display for UnaryOp {
525    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
526        write!(f, "{}", self.as_str())
527    }
528}
529
530/// Unary expression
531#[derive(Debug, Clone, PartialEq)]
532pub struct UnaryExpr {
533    /// The unary operator
534    pub op: UnaryOp,
535    /// The operand expression
536    pub expr: Expr,
537}
538
539impl UnaryExpr {
540    /// Create a new unary expression
541    pub fn new(op: UnaryOp, expr: Expr) -> Self {
542        Self { op, expr }
543    }
544}
545
546impl fmt::Display for UnaryExpr {
547    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
548        write!(f, "{}{}", self.op, self.expr)
549    }
550}
551
552/// Subquery expression
553#[derive(Debug, Clone, PartialEq)]
554pub struct SubqueryExpr {
555    /// The inner expression to evaluate as a subquery
556    pub expr: Expr,
557    /// The time range of the subquery
558    pub range: Duration,
559    /// Optional step/resolution (if None, uses default evaluation interval)
560    pub step: Option<Duration>,
561    /// Offset modifier
562    pub offset: Option<Duration>,
563    /// @ modifier for timestamp pinning
564    pub at: Option<AtModifier>,
565}
566
567impl SubqueryExpr {
568    /// Create a new subquery expression
569    pub fn new(expr: Expr, range: Duration) -> Self {
570        Self {
571            expr,
572            range,
573            step: None,
574            offset: None,
575            at: None,
576        }
577    }
578
579    /// Create a new subquery expression with step
580    pub fn with_step(expr: Expr, range: Duration, step: Duration) -> Self {
581        Self {
582            expr,
583            range,
584            step: Some(step),
585            offset: None,
586            at: None,
587        }
588    }
589}
590
591impl fmt::Display for SubqueryExpr {
592    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
593        write!(f, "{}[{}:", self.expr, self.range)?;
594        if let Some(ref step) = self.step {
595            write!(f, "{}", step)?;
596        }
597        write!(f, "]")?;
598        if let Some(ref at) = self.at {
599            write!(f, " {}", at)?;
600        }
601        if let Some(ref offset) = self.offset {
602            write!(f, " offset {}", offset)?;
603        }
604        Ok(())
605    }
606}
607
608#[cfg(test)]
609mod tests {
610    use super::*;
611
612    #[test]
613    fn test_binary_op_precedence() {
614        // or < and/unless < comparison < +/- < */% < ^
615        assert!(BinaryOp::Or.precedence() < BinaryOp::And.precedence());
616        assert!(BinaryOp::And.precedence() < BinaryOp::Eq.precedence());
617        assert!(BinaryOp::Eq.precedence() < BinaryOp::Add.precedence());
618        assert!(BinaryOp::Add.precedence() < BinaryOp::Mul.precedence());
619        assert!(BinaryOp::Mul.precedence() < BinaryOp::Pow.precedence());
620    }
621
622    #[test]
623    fn test_binary_op_associativity() {
624        assert!(!BinaryOp::Add.is_right_associative());
625        assert!(!BinaryOp::Mul.is_right_associative());
626        assert!(BinaryOp::Pow.is_right_associative());
627    }
628
629    #[test]
630    fn test_binary_op_categories() {
631        assert!(BinaryOp::Add.is_arithmetic());
632        assert!(BinaryOp::Eq.is_comparison());
633        assert!(BinaryOp::And.is_set_operator());
634    }
635
636    #[test]
637    fn test_expr_display_number() {
638        assert_eq!(Expr::Number(42.0).to_string(), "42");
639        assert_eq!(Expr::Number(3.5).to_string(), "3.5");
640        assert_eq!(Expr::Number(f64::INFINITY).to_string(), "Inf");
641        assert_eq!(Expr::Number(f64::NEG_INFINITY).to_string(), "-Inf");
642        assert_eq!(Expr::Number(f64::NAN).to_string(), "NaN");
643    }
644
645    #[test]
646    fn test_expr_display_string() {
647        assert_eq!(Expr::String("hello".to_string()).to_string(), "\"hello\"");
648    }
649
650    #[test]
651    fn test_unary_expr_display() {
652        let expr = UnaryExpr::new(UnaryOp::Minus, Expr::Number(42.0));
653        assert_eq!(expr.to_string(), "-42");
654
655        let expr = UnaryExpr::new(UnaryOp::Plus, Expr::Number(42.0));
656        assert_eq!(expr.to_string(), "+42");
657    }
658
659    #[test]
660    fn test_binary_expr_display() {
661        let expr = BinaryExpr::new(BinaryOp::Add, Expr::Number(1.0), Expr::Number(2.0));
662        assert_eq!(expr.to_string(), "1 + 2");
663    }
664
665    #[test]
666    fn test_call_display() {
667        let call = Call::new(
668            "rate",
669            vec![Expr::VectorSelector(VectorSelector::new("http_requests"))],
670        );
671        assert_eq!(call.to_string(), "rate(http_requests)");
672    }
673
674    #[test]
675    fn test_aggregation_display() {
676        let agg = Aggregation::new("sum", Expr::VectorSelector(VectorSelector::new("metric")));
677        assert_eq!(agg.to_string(), "sum(metric)");
678    }
679
680    #[test]
681    fn test_expr_is_scalar() {
682        assert!(Expr::Number(42.0).is_scalar());
683        assert!(!Expr::String("test".to_string()).is_scalar());
684    }
685
686    #[test]
687    fn test_expr_unwrap_parens() {
688        let inner = Expr::Number(42.0);
689        let paren = Expr::Paren(Box::new(inner.clone()));
690        let double_paren = Expr::Paren(Box::new(paren.clone()));
691
692        assert_eq!(*double_paren.unwrap_parens(), inner);
693    }
694}
rusty_promql_parser/ast.rs

rusty_promql_parser/
ast.rs
