Skip to main content

promql_parser/parser/
ast.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::label::{Labels, Matchers, METRIC_NAME};
16use crate::parser::token::{
17    self, token_display, T_BOTTOMK, T_COUNT_VALUES, T_END, T_QUANTILE, T_START, T_TOPK,
18};
19use crate::parser::token::{Token, TokenId, TokenType};
20use crate::parser::value::ValueType;
21use crate::parser::{indent, Function, FunctionArgs, Prettier, MAX_CHARACTERS_PER_LINE};
22use crate::util::display_duration;
23use chrono::{DateTime, Utc};
24use std::fmt::{self, Write};
25use std::ops::Neg;
26use std::sync::Arc;
27use std::time::{Duration, SystemTime};
28
29/// LabelModifier acts as
30///
31/// # Aggregation Modifier
32///
33/// - Exclude means `ignoring`
34/// - Include means `on`
35///
36/// # Vector Match Modifier
37///
38/// - Exclude means `without` removes the listed labels from the result vector,
39///   while all other labels are preserved in the output.
40///
41/// - Include means `by` does the opposite and drops labels that are not listed in the by clause,
42///   even if their label values are identical between all elements of the vector.
43///
44/// if empty listed labels, meaning no grouping
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum LabelModifier {
47    Include(Labels),
48    Exclude(Labels),
49}
50
51impl LabelModifier {
52    pub fn labels(&self) -> &Labels {
53        match self {
54            LabelModifier::Include(l) => l,
55            LabelModifier::Exclude(l) => l,
56        }
57    }
58
59    pub fn is_include(&self) -> bool {
60        matches!(*self, LabelModifier::Include(_))
61    }
62
63    pub fn include(ls: Vec<&str>) -> Self {
64        Self::Include(Labels::new(ls))
65    }
66
67    pub fn exclude(ls: Vec<&str>) -> Self {
68        Self::Exclude(Labels::new(ls))
69    }
70}
71
72/// The label list provided with the group_left or group_right modifier contains
73/// additional labels from the "one"-side to be included in the result metrics.
74#[derive(Debug, Clone, PartialEq, Eq)]
75#[cfg_attr(feature = "ser", derive(serde::Serialize))]
76#[cfg_attr(feature = "ser", serde(rename_all = "kebab-case"))]
77pub enum VectorMatchCardinality {
78    OneToOne,
79    ManyToOne(Labels),
80    OneToMany(Labels),
81    ManyToMany, // logical/set binary operators
82}
83
84impl VectorMatchCardinality {
85    pub fn labels(&self) -> Option<&Labels> {
86        match self {
87            VectorMatchCardinality::ManyToOne(l) => Some(l),
88            VectorMatchCardinality::OneToMany(l) => Some(l),
89            VectorMatchCardinality::ManyToMany => None,
90            VectorMatchCardinality::OneToOne => None,
91        }
92    }
93
94    pub fn many_to_one(ls: Vec<&str>) -> Self {
95        Self::ManyToOne(Labels::new(ls))
96    }
97
98    pub fn one_to_many(ls: Vec<&str>) -> Self {
99        Self::OneToMany(Labels::new(ls))
100    }
101}
102
103/// Binary Expr Modifier
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct BinModifier {
106    /// The matching behavior for the operation if both operands are Vectors.
107    /// If they are not this field is None.
108    pub card: VectorMatchCardinality,
109
110    /// on/ignoring on labels.
111    /// like a + b, no match modifier is needed.
112    pub matching: Option<LabelModifier>,
113    /// If a comparison operator, return 0/1 rather than filtering.
114    pub return_bool: bool,
115}
116
117impl fmt::Display for BinModifier {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        let mut s = String::from(self.bool_str());
120
121        if let Some(matching) = &self.matching {
122            match matching {
123                LabelModifier::Include(ls) => write!(s, "on ({ls}) ")?,
124                LabelModifier::Exclude(ls) if !ls.is_empty() => write!(s, "ignoring ({ls}) ")?,
125                _ => (),
126            }
127        }
128
129        match &self.card {
130            VectorMatchCardinality::ManyToOne(ls) => write!(s, "group_left ({ls}) ")?,
131            VectorMatchCardinality::OneToMany(ls) => write!(s, "group_right ({ls}) ")?,
132            _ => (),
133        }
134
135        if s.trim().is_empty() {
136            write!(f, "")
137        } else {
138            write!(f, " {}", s.trim_end()) // there is a leading space here
139        }
140    }
141}
142
143impl Default for BinModifier {
144    fn default() -> Self {
145        Self {
146            card: VectorMatchCardinality::OneToOne,
147            matching: None,
148            return_bool: false,
149        }
150    }
151}
152
153impl BinModifier {
154    pub fn with_card(mut self, card: VectorMatchCardinality) -> Self {
155        self.card = card;
156        self
157    }
158
159    pub fn with_matching(mut self, matching: Option<LabelModifier>) -> Self {
160        self.matching = matching;
161        self
162    }
163
164    pub fn with_return_bool(mut self, return_bool: bool) -> Self {
165        self.return_bool = return_bool;
166        self
167    }
168
169    pub fn is_labels_joint(&self) -> bool {
170        matches!(
171            (self.card.labels(), &self.matching),
172            (Some(labels), Some(matching)) if labels.is_joint(matching.labels())
173        )
174    }
175
176    pub fn intersect_labels(&self) -> Option<Vec<String>> {
177        if let Some(labels) = self.card.labels() {
178            if let Some(matching) = &self.matching {
179                return Some(labels.intersect(matching.labels()).labels);
180            }
181        };
182        None
183    }
184
185    pub fn is_matching_on(&self) -> bool {
186        matches!(&self.matching, Some(matching) if matching.is_include())
187    }
188
189    pub fn is_matching_labels_not_empty(&self) -> bool {
190        matches!(&self.matching, Some(matching) if !matching.labels().is_empty())
191    }
192
193    pub fn bool_str(&self) -> &str {
194        if self.return_bool {
195            "bool "
196        } else {
197            ""
198        }
199    }
200}
201
202#[cfg(feature = "ser")]
203pub(crate) fn serialize_grouping<S>(
204    this: &Option<LabelModifier>,
205    serializer: S,
206) -> Result<S::Ok, S::Error>
207where
208    S: serde::Serializer,
209{
210    use serde::ser::SerializeMap;
211    let mut map = serializer.serialize_map(Some(2))?;
212    match this {
213        Some(LabelModifier::Include(l)) => {
214            map.serialize_entry("grouping", l)?;
215            map.serialize_entry("without", &false)?;
216        }
217        Some(LabelModifier::Exclude(l)) => {
218            map.serialize_entry("grouping", l)?;
219            map.serialize_entry("without", &true)?;
220        }
221        None => {
222            map.serialize_entry("grouping", &(vec![] as Vec<String>))?;
223            map.serialize_entry("without", &false)?;
224        }
225    }
226
227    map.end()
228}
229
230#[cfg(feature = "ser")]
231pub(crate) fn serialize_bin_modifier<S>(
232    this: &Option<BinModifier>,
233    serializer: S,
234) -> Result<S::Ok, S::Error>
235where
236    S: serde::Serializer,
237{
238    use serde::ser::SerializeMap;
239    use serde_json::json;
240
241    let mut map = serializer.serialize_map(Some(2))?;
242
243    map.serialize_entry(
244        "bool",
245        &this.as_ref().map(|t| t.return_bool).unwrap_or(false),
246    )?;
247    if let Some(t) = this {
248        if let Some(labels) = &t.matching {
249            map.serialize_key("matching")?;
250
251            match labels {
252                LabelModifier::Include(labels) => {
253                    let value = json!({
254                        "card": t.card,
255                        "include": [],
256                        "labels": labels,
257                        "on": true,
258                    });
259                    map.serialize_value(&value)?;
260                }
261                LabelModifier::Exclude(labels) => {
262                    let value = json!({
263                        "card": t.card,
264                        "include": [],
265                        "labels": labels,
266                        "on": false,
267                    });
268                    map.serialize_value(&value)?;
269                }
270            }
271        } else {
272            let value = json!({
273                "card": t.card,
274                "include": [],
275                "labels": [],
276                "on": false,
277            });
278            map.serialize_entry("matching", &value)?;
279        }
280    } else {
281        map.serialize_entry("matching", &None::<bool>)?;
282    }
283
284    map.end()
285}
286
287#[cfg(feature = "ser")]
288pub(crate) fn serialize_at_modifier<S>(
289    this: &Option<AtModifier>,
290    serializer: S,
291) -> Result<S::Ok, S::Error>
292where
293    S: serde::Serializer,
294{
295    use serde::ser::SerializeMap;
296    let mut map = serializer.serialize_map(Some(2))?;
297    match this {
298        Some(AtModifier::Start) => {
299            map.serialize_entry("startOrEnd", &Some("start"))?;
300            map.serialize_entry("timestamp", &None::<u128>)?;
301        }
302        Some(AtModifier::End) => {
303            map.serialize_entry("startOrEnd", &Some("end"))?;
304            map.serialize_entry("timestamp", &None::<u128>)?;
305        }
306        Some(AtModifier::At(time)) => {
307            map.serialize_entry("startOrEnd", &None::<&str>)?;
308            map.serialize_entry(
309                "timestamp",
310                &time
311                    .duration_since(SystemTime::UNIX_EPOCH)
312                    .unwrap_or(Duration::ZERO)
313                    .as_millis(),
314            )?;
315        }
316        None => {
317            map.serialize_entry("startOrEnd", &None::<&str>)?;
318            map.serialize_entry("timestamp", &None::<u128>)?;
319        }
320    }
321
322    map.end()
323}
324
325#[derive(Debug, Clone, PartialEq, Eq)]
326pub enum Offset {
327    Pos(Duration),
328    Neg(Duration),
329}
330
331impl Offset {
332    #[cfg(feature = "ser")]
333    pub(crate) fn as_millis(&self) -> i128 {
334        match self {
335            Self::Pos(dur) => dur.as_millis() as i128,
336            Self::Neg(dur) => -(dur.as_millis() as i128),
337        }
338    }
339
340    #[cfg(feature = "ser")]
341    pub(crate) fn serialize_offset<S>(
342        offset: &Option<Self>,
343        serializer: S,
344    ) -> Result<S::Ok, S::Error>
345    where
346        S: serde::Serializer,
347    {
348        let value = offset.as_ref().map(|o| o.as_millis()).unwrap_or(0);
349        serializer.serialize_i128(value)
350    }
351}
352
353impl fmt::Display for Offset {
354    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
355        match self {
356            Offset::Pos(dur) => write!(f, "{}", display_duration(dur)),
357            Offset::Neg(dur) => write!(f, "-{}", display_duration(dur)),
358        }
359    }
360}
361
362#[derive(Debug, Clone, PartialEq, Eq)]
363pub enum AtModifier {
364    Start,
365    End,
366    /// at can be earlier than UNIX_EPOCH
367    At(SystemTime),
368}
369
370impl fmt::Display for AtModifier {
371    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
372        match self {
373            AtModifier::Start => write!(f, "@ {}()", token_display(T_START)),
374            AtModifier::End => write!(f, "@ {}()", token_display(T_END)),
375            AtModifier::At(time) => {
376                let d = time
377                    .duration_since(SystemTime::UNIX_EPOCH)
378                    .unwrap_or(Duration::ZERO); // This should not happen
379                write!(f, "@ {:.3}", d.as_secs_f64())
380            }
381        }
382    }
383}
384
385impl TryFrom<TokenId> for AtModifier {
386    type Error = String;
387
388    fn try_from(id: TokenId) -> Result<Self, Self::Error> {
389        match id {
390            T_START => Ok(AtModifier::Start),
391            T_END => Ok(AtModifier::End),
392            _ => Err(format!(
393                "invalid @ modifier preprocessor '{}', START or END is valid.",
394                token::token_display(id)
395            )),
396        }
397    }
398}
399
400impl TryFrom<Token> for AtModifier {
401    type Error = String;
402
403    fn try_from(token: Token) -> Result<Self, Self::Error> {
404        AtModifier::try_from(token.id())
405    }
406}
407
408impl TryFrom<NumberLiteral> for AtModifier {
409    type Error = String;
410
411    fn try_from(num: NumberLiteral) -> Result<Self, Self::Error> {
412        AtModifier::try_from(num.val)
413    }
414}
415
416impl TryFrom<Expr> for AtModifier {
417    type Error = String;
418
419    fn try_from(ex: Expr) -> Result<Self, Self::Error> {
420        match ex {
421            Expr::NumberLiteral(nl) => AtModifier::try_from(nl),
422            _ => Err("invalid float value after @ modifier".into()),
423        }
424    }
425}
426
427impl TryFrom<f64> for AtModifier {
428    type Error = String;
429
430    fn try_from(secs: f64) -> Result<Self, Self::Error> {
431        let err_info = format!("timestamp out of bounds for @ modifier: {secs}");
432
433        if secs.is_nan() || secs.is_infinite() || secs >= f64::MAX || secs <= f64::MIN {
434            return Err(err_info);
435        }
436        let milli = (secs * 1000f64).round().abs() as u64;
437
438        let duration = Duration::from_millis(milli);
439        let mut st = Some(SystemTime::UNIX_EPOCH);
440        if secs.is_sign_positive() {
441            st = SystemTime::UNIX_EPOCH.checked_add(duration);
442        }
443        if secs.is_sign_negative() {
444            st = SystemTime::UNIX_EPOCH.checked_sub(duration);
445        }
446
447        st.map(Self::At).ok_or(err_info)
448    }
449}
450
451/// EvalStmt holds an expression and information on the range it should
452/// be evaluated on.
453#[allow(rustdoc::broken_intra_doc_links)]
454#[derive(Debug, Clone)]
455pub struct EvalStmt {
456    /// Expression to be evaluated.
457    pub expr: Expr,
458
459    /// The time boundaries for the evaluation. If start equals end an instant
460    /// is evaluated.
461    pub start: SystemTime,
462    pub end: SystemTime,
463    /// Time between two evaluated instants for the range [start:end].
464    pub interval: Duration,
465    /// Lookback delta to use for this evaluation.
466    pub lookback_delta: Duration,
467}
468
469impl fmt::Display for EvalStmt {
470    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
471        write!(
472            f,
473            "[{}] @ [{}, {}, {}, {}]",
474            self.expr,
475            DateTime::<Utc>::from(self.start).to_rfc3339(),
476            DateTime::<Utc>::from(self.end).to_rfc3339(),
477            display_duration(&self.interval),
478            display_duration(&self.lookback_delta)
479        )
480    }
481}
482
483/// Grammar:
484/// ``` norust
485/// <aggr-op> [without|by (<label list>)] ([parameter,] <vector expression>)
486/// <aggr-op>([parameter,] <vector expression>) [without|by (<label list>)]
487/// ```
488///
489/// parameter is only required for `count_values`, `quantile`, `topk` and `bottomk`.
490#[derive(Debug, Clone, PartialEq, Eq)]
491#[cfg_attr(feature = "ser", derive(serde::Serialize))]
492pub struct AggregateExpr {
493    /// The used aggregation operation.
494    pub op: TokenType,
495    /// The Vector expression over which is aggregated.
496    pub expr: Box<Expr>,
497    /// Parameter used by some aggregators.
498    pub param: Option<Box<Expr>>,
499    /// modifier is optional for some aggregation operators, like sum.
500    #[cfg_attr(feature = "ser", serde(flatten))]
501    #[cfg_attr(feature = "ser", serde(serialize_with = "serialize_grouping"))]
502    pub modifier: Option<LabelModifier>,
503}
504
505impl AggregateExpr {
506    fn get_op_string(&self) -> String {
507        let mut s = self.op.to_string();
508
509        if let Some(modifier) = &self.modifier {
510            match modifier {
511                LabelModifier::Exclude(ls) => write!(s, " without ({ls}) ").unwrap(),
512                LabelModifier::Include(ls) if !ls.is_empty() => write!(s, " by ({ls}) ").unwrap(),
513                _ => (),
514            }
515        }
516        s
517    }
518}
519
520impl fmt::Display for AggregateExpr {
521    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
522        write!(f, "{}", self.get_op_string())?;
523
524        write!(f, "(")?;
525        if let Some(param) = &self.param {
526            write!(f, "{param}, ")?;
527        }
528        write!(f, "{})", self.expr)?;
529
530        Ok(())
531    }
532}
533
534impl Prettier for AggregateExpr {
535    fn format(&self, level: usize, max: usize) -> String {
536        let mut s = format!("{}{}(\n", indent(level), self.get_op_string());
537        if let Some(param) = &self.param {
538            writeln!(s, "{},", param.pretty(level + 1, max)).unwrap();
539        }
540        writeln!(s, "{}", self.expr.pretty(level + 1, max)).unwrap();
541        write!(s, "{})", indent(level)).unwrap();
542        s
543    }
544}
545
546/// UnaryExpr will negate the expr
547#[derive(Debug, Clone, PartialEq, Eq)]
548pub struct UnaryExpr {
549    pub expr: Box<Expr>,
550}
551
552impl fmt::Display for UnaryExpr {
553    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
554        write!(f, "-{}", self.expr)
555    }
556}
557
558impl Prettier for UnaryExpr {
559    fn pretty(&self, level: usize, max: usize) -> String {
560        format!(
561            "{}-{}",
562            indent(level),
563            self.expr.pretty(level, max).trim_start()
564        )
565    }
566}
567
568#[cfg(feature = "ser")]
569impl serde::Serialize for UnaryExpr {
570    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
571    where
572        S: serde::Serializer,
573    {
574        use serde::ser::SerializeMap;
575        let mut map = serializer.serialize_map(Some(2))?;
576        map.serialize_entry("op", "-")?;
577        map.serialize_entry("expr", &self.expr)?;
578
579        map.end()
580    }
581}
582
583/// Grammar:
584/// ``` norust
585/// <vector expr> <bin-op> ignoring(<label list>) group_left(<label list>) <vector expr>
586/// <vector expr> <bin-op> ignoring(<label list>) group_right(<label list>) <vector expr>
587/// <vector expr> <bin-op> on(<label list>) group_left(<label list>) <vector expr>
588/// <vector expr> <bin-op> on(<label list>) group_right(<label list>) <vector expr>
589/// ```
590#[derive(Debug, Clone, PartialEq, Eq)]
591#[cfg_attr(feature = "ser", derive(serde::Serialize))]
592pub struct BinaryExpr {
593    /// The operation of the expression.
594    pub op: TokenType,
595    /// The operands on the left sides of the operator.
596    pub lhs: Box<Expr>,
597    /// The operands on the right sides of the operator.
598    pub rhs: Box<Expr>,
599    #[cfg_attr(feature = "ser", serde(flatten))]
600    #[cfg_attr(feature = "ser", serde(serialize_with = "serialize_bin_modifier"))]
601    pub modifier: Option<BinModifier>,
602}
603
604impl BinaryExpr {
605    pub fn is_matching_on(&self) -> bool {
606        matches!(&self.modifier, Some(modifier) if modifier.is_matching_on())
607    }
608
609    pub fn is_matching_labels_not_empty(&self) -> bool {
610        matches!(&self.modifier, Some(modifier) if modifier.is_matching_labels_not_empty())
611    }
612
613    pub fn return_bool(&self) -> bool {
614        matches!(&self.modifier, Some(modifier) if modifier.return_bool)
615    }
616
617    /// check if labels of card and matching are joint
618    pub fn is_labels_joint(&self) -> bool {
619        matches!(&self.modifier, Some(modifier) if modifier.is_labels_joint())
620    }
621
622    /// intersect labels of card and matching
623    pub fn intersect_labels(&self) -> Option<Vec<String>> {
624        self.modifier
625            .as_ref()
626            .and_then(|modifier| modifier.intersect_labels())
627    }
628
629    fn get_op_matching_string(&self) -> String {
630        match &self.modifier {
631            Some(modifier) => format!("{}{modifier}", self.op),
632            None => self.op.to_string(),
633        }
634    }
635}
636
637impl fmt::Display for BinaryExpr {
638    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
639        write!(
640            f,
641            "{} {} {}",
642            self.lhs,
643            self.get_op_matching_string(),
644            self.rhs
645        )
646    }
647}
648
649impl Prettier for BinaryExpr {
650    fn format(&self, level: usize, max: usize) -> String {
651        format!(
652            "{}\n{}{}\n{}",
653            self.lhs.pretty(level + 1, max),
654            indent(level),
655            self.get_op_matching_string(),
656            self.rhs.pretty(level + 1, max)
657        )
658    }
659}
660
661#[derive(Debug, Clone, PartialEq, Eq)]
662#[cfg_attr(feature = "ser", derive(serde::Serialize))]
663pub struct ParenExpr {
664    pub expr: Box<Expr>,
665}
666
667impl fmt::Display for ParenExpr {
668    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
669        write!(f, "({})", self.expr)
670    }
671}
672
673impl Prettier for ParenExpr {
674    fn format(&self, level: usize, max: usize) -> String {
675        format!(
676            "{}(\n{}\n{})",
677            indent(level),
678            self.expr.pretty(level + 1, max),
679            indent(level)
680        )
681    }
682}
683
684/// Grammar:
685/// ```norust
686/// <instant_query> '[' <range> ':' [<resolution>] ']' [ @ <float_literal> ] [ offset <duration> ]
687/// ```
688#[derive(Debug, Clone, PartialEq, Eq)]
689#[cfg_attr(feature = "ser", derive(serde::Serialize))]
690pub struct SubqueryExpr {
691    pub expr: Box<Expr>,
692    #[cfg_attr(feature = "ser", serde(serialize_with = "Offset::serialize_offset"))]
693    pub offset: Option<Offset>,
694    #[cfg_attr(feature = "ser", serde(flatten))]
695    #[cfg_attr(feature = "ser", serde(serialize_with = "serialize_at_modifier"))]
696    pub at: Option<AtModifier>,
697    #[cfg_attr(
698        feature = "ser",
699        serde(serialize_with = "crate::util::duration::serialize_duration")
700    )]
701    pub range: Duration,
702    /// Default is the global evaluation interval.
703    #[cfg_attr(
704        feature = "ser",
705        serde(serialize_with = "crate::util::duration::serialize_duration_opt")
706    )]
707    pub step: Option<Duration>,
708}
709
710impl SubqueryExpr {
711    fn get_time_suffix_string(&self) -> String {
712        let step = match &self.step {
713            Some(step) => display_duration(step),
714            None => String::from(""),
715        };
716        let range = display_duration(&self.range);
717
718        let mut s = format!("[{range}:{step}]");
719
720        if let Some(at) = &self.at {
721            write!(s, " {at}").unwrap();
722        }
723
724        if let Some(offset) = &self.offset {
725            write!(s, " offset {offset}").unwrap();
726        }
727        s
728    }
729}
730
731impl fmt::Display for SubqueryExpr {
732    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
733        write!(f, "{}{}", self.expr, self.get_time_suffix_string())
734    }
735}
736
737impl Prettier for SubqueryExpr {
738    fn pretty(&self, level: usize, max: usize) -> String {
739        format!(
740            "{}{}",
741            self.expr.pretty(level, max),
742            self.get_time_suffix_string()
743        )
744    }
745}
746
747#[derive(Debug, Clone)]
748pub struct NumberLiteral {
749    pub val: f64,
750}
751
752impl NumberLiteral {
753    pub fn new(val: f64) -> Self {
754        Self { val }
755    }
756}
757
758impl PartialEq for NumberLiteral {
759    fn eq(&self, other: &Self) -> bool {
760        self.val == other.val || self.val.is_nan() && other.val.is_nan()
761    }
762}
763
764impl Eq for NumberLiteral {}
765
766impl Neg for NumberLiteral {
767    type Output = Self;
768
769    fn neg(self) -> Self::Output {
770        NumberLiteral { val: -self.val }
771    }
772}
773
774impl fmt::Display for NumberLiteral {
775    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
776        if self.val == f64::INFINITY {
777            write!(f, "Inf")
778        } else if self.val == f64::NEG_INFINITY {
779            write!(f, "-Inf")
780        } else if f64::is_nan(self.val) {
781            write!(f, "NaN")
782        } else {
783            write!(f, "{}", self.val)
784        }
785    }
786}
787
788#[cfg(feature = "ser")]
789impl serde::Serialize for NumberLiteral {
790    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
791    where
792        S: serde::Serializer,
793    {
794        use serde::ser::SerializeMap;
795        let mut map = serializer.serialize_map(Some(1))?;
796        map.serialize_entry("val", &self.to_string())?;
797
798        map.end()
799    }
800}
801
802impl Prettier for NumberLiteral {
803    fn needs_split(&self, _max: usize) -> bool {
804        false
805    }
806}
807
808#[derive(Debug, Clone, PartialEq, Eq)]
809#[cfg_attr(feature = "ser", derive(serde::Serialize))]
810pub struct StringLiteral {
811    pub val: String,
812}
813
814impl fmt::Display for StringLiteral {
815    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
816        write!(f, "\"{}\"", self.val)
817    }
818}
819
820impl Prettier for StringLiteral {
821    fn needs_split(&self, _max: usize) -> bool {
822        false
823    }
824}
825
826#[derive(Debug, Clone, PartialEq, Eq)]
827#[cfg_attr(feature = "ser", derive(serde::Serialize))]
828pub struct VectorSelector {
829    pub name: Option<String>,
830    #[cfg_attr(feature = "ser", serde(flatten))]
831    pub matchers: Matchers,
832    #[cfg_attr(feature = "ser", serde(serialize_with = "Offset::serialize_offset"))]
833    pub offset: Option<Offset>,
834    #[cfg_attr(feature = "ser", serde(flatten))]
835    #[cfg_attr(feature = "ser", serde(serialize_with = "serialize_at_modifier"))]
836    pub at: Option<AtModifier>,
837}
838
839impl VectorSelector {
840    pub fn new(name: Option<String>, matchers: Matchers) -> Self {
841        VectorSelector {
842            name,
843            matchers,
844            offset: None,
845            at: None,
846        }
847    }
848}
849
850impl Default for VectorSelector {
851    fn default() -> Self {
852        Self {
853            name: None,
854            matchers: Matchers::empty(),
855            offset: None,
856            at: None,
857        }
858    }
859}
860
861impl From<String> for VectorSelector {
862    fn from(name: String) -> Self {
863        VectorSelector {
864            name: Some(name),
865            offset: None,
866            at: None,
867            matchers: Matchers::empty(),
868        }
869    }
870}
871
872/// directly create an instant vector with only METRIC_NAME matcher.
873///
874/// # Examples
875///
876/// Basic usage:
877///
878/// ``` rust
879/// use promql_parser::label::Matchers;
880/// use promql_parser::parser::VectorSelector;
881///
882/// let vs = VectorSelector {
883///     name: Some(String::from("foo")),
884///     offset: None,
885///     at: None,
886///     matchers: Matchers::empty(),
887/// };
888///
889/// assert_eq!(VectorSelector::from("foo"), vs);
890/// ```
891impl From<&str> for VectorSelector {
892    fn from(name: &str) -> Self {
893        VectorSelector::from(name.to_string())
894    }
895}
896
897impl Neg for VectorSelector {
898    type Output = UnaryExpr;
899
900    fn neg(self) -> Self::Output {
901        let ex = Expr::VectorSelector(self);
902        UnaryExpr { expr: Box::new(ex) }
903    }
904}
905
906impl fmt::Display for VectorSelector {
907    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
908        if let Some(name) = &self.name {
909            write!(f, "{name}")?;
910        }
911        let matchers = &self.matchers.to_string();
912        if !matchers.is_empty() {
913            write!(f, "{{{matchers}}}")?;
914        }
915        if let Some(at) = &self.at {
916            write!(f, " {at}")?;
917        }
918        if let Some(offset) = &self.offset {
919            write!(f, " offset {offset}")?;
920        }
921        Ok(())
922    }
923}
924
925impl Prettier for VectorSelector {
926    fn needs_split(&self, _max: usize) -> bool {
927        false
928    }
929}
930
931#[derive(Debug, Clone, PartialEq, Eq)]
932#[cfg_attr(feature = "ser", derive(serde::Serialize))]
933pub struct MatrixSelector {
934    #[cfg_attr(feature = "ser", serde(flatten))]
935    pub vs: VectorSelector,
936    #[cfg_attr(
937        feature = "ser",
938        serde(serialize_with = "crate::util::duration::serialize_duration")
939    )]
940    pub range: Duration,
941}
942
943impl fmt::Display for MatrixSelector {
944    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
945        if let Some(name) = &self.vs.name {
946            write!(f, "{name}")?;
947        }
948
949        let matchers = &self.vs.matchers.to_string();
950        if !matchers.is_empty() {
951            write!(f, "{{{matchers}}}")?;
952        }
953
954        write!(f, "[{}]", display_duration(&self.range))?;
955
956        if let Some(at) = &self.vs.at {
957            write!(f, " {at}")?;
958        }
959
960        if let Some(offset) = &self.vs.offset {
961            write!(f, " offset {offset}")?;
962        }
963
964        Ok(())
965    }
966}
967
968impl Prettier for MatrixSelector {
969    fn needs_split(&self, _max: usize) -> bool {
970        false
971    }
972}
973
974/// Call represents Prometheus Function.
975/// Some functions have special cases:
976///
977/// ## exp
978///
979/// exp(v instant-vector) calculates the exponential function for all elements in v.
980/// Special cases are:
981///
982/// ```promql
983/// Exp(+Inf) = +Inf
984/// Exp(NaN) = NaN
985/// ```
986///
987/// ## ln
988///
989/// ln(v instant-vector) calculates the natural logarithm for all elements in v.
990/// Special cases are:
991///
992/// ```promql
993/// ln(+Inf) = +Inf
994/// ln(0) = -Inf
995/// ln(x < 0) = NaN
996/// ln(NaN) = NaN
997/// ```
998///
999/// TODO: support more special cases of function call
1000///
1001///  - acos()
1002///  - acosh()
1003///  - asin()
1004///  - asinh()
1005///  - atan()
1006///  - atanh()
1007///  - cos()
1008///  - cosh()
1009///  - sin()
1010///  - sinh()
1011///  - tan()
1012///  - tanh()
1013#[derive(Debug, Clone, PartialEq, Eq)]
1014#[cfg_attr(feature = "ser", derive(serde::Serialize))]
1015pub struct Call {
1016    pub func: Function,
1017    #[cfg_attr(feature = "ser", serde(flatten))]
1018    pub args: FunctionArgs,
1019}
1020
1021impl fmt::Display for Call {
1022    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1023        write!(f, "{}({})", self.func.name, self.args)
1024    }
1025}
1026
1027impl Prettier for Call {
1028    fn format(&self, level: usize, max: usize) -> String {
1029        format!(
1030            "{}{}(\n{}\n{})",
1031            indent(level),
1032            self.func.name,
1033            self.args.pretty(level + 1, max),
1034            indent(level)
1035        )
1036    }
1037}
1038
1039/// Node for extending the AST. [Extension] won't be generate by this parser itself.
1040#[derive(Debug, Clone)]
1041pub struct Extension {
1042    pub expr: Arc<dyn ExtensionExpr>,
1043}
1044
1045/// The interface for extending the AST with custom expression node.
1046pub trait ExtensionExpr: std::fmt::Debug + Send + Sync {
1047    fn as_any(&self) -> &dyn std::any::Any;
1048
1049    fn name(&self) -> &str;
1050
1051    fn value_type(&self) -> ValueType;
1052
1053    fn children(&self) -> &[Expr];
1054}
1055
1056impl PartialEq for Extension {
1057    fn eq(&self, other: &Self) -> bool {
1058        format!("{self:?}") == format!("{other:?}")
1059    }
1060}
1061
1062impl Eq for Extension {}
1063
1064#[derive(Debug, Clone, PartialEq, Eq)]
1065#[cfg_attr(feature = "ser", derive(serde::Serialize))]
1066#[cfg_attr(feature = "ser", serde(tag = "type", rename_all = "camelCase"))]
1067pub enum Expr {
1068    /// Aggregate represents an aggregation operation on a Vector.
1069    #[cfg_attr(feature = "ser", serde(rename = "aggregation"))]
1070    Aggregate(AggregateExpr),
1071
1072    /// Unary represents a unary operation on another expression.
1073    /// Currently unary operations are only supported for Scalars.
1074    #[cfg_attr(feature = "ser", serde(rename = "unaryExpr"))]
1075    Unary(UnaryExpr),
1076
1077    /// Binary represents a binary expression between two child expressions.
1078    #[cfg_attr(feature = "ser", serde(rename = "binaryExpr"))]
1079    Binary(BinaryExpr),
1080
1081    /// Paren wraps an expression so it cannot be disassembled as a consequence
1082    /// of operator precedence.
1083    #[cfg_attr(feature = "ser", serde(rename = "parenExpr"))]
1084    Paren(ParenExpr),
1085
1086    /// SubqueryExpr represents a subquery.
1087    Subquery(SubqueryExpr),
1088
1089    /// NumberLiteral represents a number.
1090    NumberLiteral(NumberLiteral),
1091
1092    /// StringLiteral represents a string.
1093    StringLiteral(StringLiteral),
1094
1095    /// VectorSelector represents a Vector selection.
1096    VectorSelector(VectorSelector),
1097
1098    /// MatrixSelector represents a Matrix selection.
1099    MatrixSelector(MatrixSelector),
1100
1101    /// Call represents a function call.
1102    Call(Call),
1103
1104    /// Extension represents an extension expression. It is for user to attach additional
1105    /// information to the AST. This parser won't generate Extension node.
1106    #[cfg_attr(feature = "ser", serde(skip))]
1107    Extension(Extension),
1108}
1109
1110impl Expr {
1111    pub(crate) fn new_vector_selector(
1112        name: Option<String>,
1113        matchers: Matchers,
1114    ) -> Result<Self, String> {
1115        let vs = VectorSelector::new(name, matchers);
1116        Ok(Self::VectorSelector(vs))
1117    }
1118
1119    pub(crate) fn new_unary_expr(expr: Expr) -> Result<Self, String> {
1120        match expr {
1121            Expr::StringLiteral(_) => Err("unary expression only allowed on expressions of type scalar or vector, got: string".into()),
1122            Expr::MatrixSelector(_) => Err("unary expression only allowed on expressions of type scalar or vector, got: matrix".into()),
1123            _ => Ok(-expr),
1124        }
1125    }
1126
1127    pub(crate) fn new_subquery_expr(
1128        expr: Expr,
1129        range: Duration,
1130        step: Option<Duration>,
1131    ) -> Result<Self, String> {
1132        let se = Expr::Subquery(SubqueryExpr {
1133            expr: Box::new(expr),
1134            offset: None,
1135            at: None,
1136            range,
1137            step,
1138        });
1139        Ok(se)
1140    }
1141
1142    pub(crate) fn new_paren_expr(expr: Expr) -> Result<Self, String> {
1143        let ex = Expr::Paren(ParenExpr {
1144            expr: Box::new(expr),
1145        });
1146        Ok(ex)
1147    }
1148
1149    /// NOTE: @ and offset is not set here.
1150    pub(crate) fn new_matrix_selector(expr: Expr, range: Duration) -> Result<Self, String> {
1151        match expr {
1152            Expr::VectorSelector(VectorSelector {
1153                offset: Some(_), ..
1154            }) => Err("no offset modifiers allowed before range".into()),
1155            Expr::VectorSelector(VectorSelector { at: Some(_), .. }) => {
1156                Err("no @ modifiers allowed before range".into())
1157            }
1158            Expr::VectorSelector(vs) => {
1159                let ms = Expr::MatrixSelector(MatrixSelector { vs, range });
1160                Ok(ms)
1161            }
1162            _ => Err("ranges only allowed for vector selectors".into()),
1163        }
1164    }
1165
1166    pub(crate) fn at_expr(self, at: AtModifier) -> Result<Self, String> {
1167        let already_set_err = Err("@ <timestamp> may not be set multiple times".into());
1168        match self {
1169            Expr::VectorSelector(mut vs) => match vs.at {
1170                None => {
1171                    vs.at = Some(at);
1172                    Ok(Expr::VectorSelector(vs))
1173                }
1174                Some(_) => already_set_err,
1175            },
1176            Expr::MatrixSelector(mut ms) => match ms.vs.at {
1177                None => {
1178                    ms.vs.at = Some(at);
1179                    Ok(Expr::MatrixSelector(ms))
1180                }
1181                Some(_) => already_set_err,
1182            },
1183            Expr::Subquery(mut s) => match s.at {
1184                None => {
1185                    s.at = Some(at);
1186                    Ok(Expr::Subquery(s))
1187                }
1188                Some(_) => already_set_err,
1189            },
1190            _ => {
1191                Err("@ modifier must be preceded by an vector selector or matrix selector or a subquery".into())
1192            }
1193        }
1194    }
1195
1196    /// set offset field for specified Expr, but CAN ONLY be set once.
1197    pub(crate) fn offset_expr(self, offset: Offset) -> Result<Self, String> {
1198        let already_set_err = Err("offset may not be set multiple times".into());
1199        match self {
1200            Expr::VectorSelector(mut vs) => match vs.offset {
1201                None => {
1202                    vs.offset = Some(offset);
1203                    Ok(Expr::VectorSelector(vs))
1204                }
1205                Some(_) => already_set_err,
1206            },
1207            Expr::MatrixSelector(mut ms) => match ms.vs.offset {
1208                None => {
1209                    ms.vs.offset = Some(offset);
1210                    Ok(Expr::MatrixSelector(ms))
1211                }
1212                Some(_) => already_set_err,
1213            },
1214            Expr::Subquery(mut s) => match s.offset {
1215                None => {
1216                    s.offset = Some(offset);
1217                    Ok(Expr::Subquery(s))
1218                }
1219                Some(_) => already_set_err,
1220            },
1221            _ => {
1222                Err("offset modifier must be preceded by an vector selector or matrix selector or a subquery".into())
1223            }
1224        }
1225    }
1226
1227    pub(crate) fn new_call(func: Function, args: FunctionArgs) -> Result<Expr, String> {
1228        Ok(Expr::Call(Call { func, args }))
1229    }
1230
1231    pub(crate) fn new_binary_expr(
1232        lhs: Expr,
1233        op: TokenId,
1234        modifier: Option<BinModifier>,
1235        rhs: Expr,
1236    ) -> Result<Expr, String> {
1237        let ex = BinaryExpr {
1238            op: TokenType::new(op),
1239            lhs: Box::new(lhs),
1240            rhs: Box::new(rhs),
1241            modifier,
1242        };
1243        Ok(Expr::Binary(ex))
1244    }
1245
1246    pub(crate) fn new_aggregate_expr(
1247        op: TokenId,
1248        modifier: Option<LabelModifier>,
1249        args: FunctionArgs,
1250    ) -> Result<Expr, String> {
1251        let op = TokenType::new(op);
1252        if args.is_empty() {
1253            return Err(format!(
1254                "no arguments for aggregate expression '{op}' provided"
1255            ));
1256        }
1257        let mut desired_args_count = 1;
1258        let mut param = None;
1259        if op.is_aggregator_with_param() {
1260            desired_args_count = 2;
1261            param = args.first();
1262        }
1263        if args.len() != desired_args_count {
1264            return Err(format!(
1265                "wrong number of arguments for aggregate expression provided, expected {}, got {}",
1266                desired_args_count,
1267                args.len()
1268            ));
1269        }
1270
1271        match args.last() {
1272            Some(expr) => Ok(Expr::Aggregate(AggregateExpr {
1273                op,
1274                expr,
1275                param,
1276                modifier,
1277            })),
1278            None => Err(
1279                "aggregate operation needs a single instant vector parameter, but found none"
1280                    .into(),
1281            ),
1282        }
1283    }
1284
1285    pub fn value_type(&self) -> ValueType {
1286        match self {
1287            Expr::Aggregate(_) => ValueType::Vector,
1288            Expr::Unary(ex) => ex.expr.value_type(),
1289            Expr::Binary(ex) => {
1290                if ex.lhs.value_type() == ValueType::Scalar
1291                    && ex.rhs.value_type() == ValueType::Scalar
1292                {
1293                    ValueType::Scalar
1294                } else {
1295                    ValueType::Vector
1296                }
1297            }
1298            Expr::Paren(ex) => ex.expr.value_type(),
1299            Expr::Subquery(_) => ValueType::Matrix,
1300            Expr::NumberLiteral(_) => ValueType::Scalar,
1301            Expr::StringLiteral(_) => ValueType::String,
1302            Expr::VectorSelector(_) => ValueType::Vector,
1303            Expr::MatrixSelector(_) => ValueType::Matrix,
1304            Expr::Call(ex) => ex.func.return_type,
1305            Expr::Extension(ex) => ex.expr.value_type(),
1306        }
1307    }
1308
1309    /// only Some if expr is [Expr::NumberLiteral]
1310    pub(crate) fn scalar_value(&self) -> Option<f64> {
1311        match self {
1312            Expr::NumberLiteral(nl) => Some(nl.val),
1313            _ => None,
1314        }
1315    }
1316
1317    pub fn prettify(&self) -> String {
1318        self.pretty(0, MAX_CHARACTERS_PER_LINE)
1319    }
1320}
1321
1322impl From<String> for Expr {
1323    fn from(val: String) -> Self {
1324        Expr::StringLiteral(StringLiteral { val })
1325    }
1326}
1327
1328impl From<&str> for Expr {
1329    fn from(s: &str) -> Self {
1330        Expr::StringLiteral(StringLiteral { val: s.into() })
1331    }
1332}
1333
1334impl From<f64> for Expr {
1335    fn from(val: f64) -> Self {
1336        Expr::NumberLiteral(NumberLiteral { val })
1337    }
1338}
1339
1340/// directly create an Expr::VectorSelector from instant vector
1341///
1342/// # Examples
1343///
1344/// Basic usage:
1345///
1346/// ``` rust
1347/// use promql_parser::label::Matchers;
1348/// use promql_parser::parser::{Expr, VectorSelector};
1349///
1350/// let name = String::from("foo");
1351/// let vs = VectorSelector::new(Some(name), Matchers::empty());
1352///
1353/// assert_eq!(Expr::VectorSelector(vs), Expr::from(VectorSelector::from("foo")));
1354/// ```
1355impl From<VectorSelector> for Expr {
1356    fn from(vs: VectorSelector) -> Self {
1357        Expr::VectorSelector(vs)
1358    }
1359}
1360
1361impl Neg for Expr {
1362    type Output = Self;
1363
1364    fn neg(self) -> Self::Output {
1365        match self {
1366            Expr::NumberLiteral(nl) => Expr::NumberLiteral(-nl),
1367            _ => Expr::Unary(UnaryExpr {
1368                expr: Box::new(self),
1369            }),
1370        }
1371    }
1372}
1373
1374impl fmt::Display for Expr {
1375    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1376        match self {
1377            Expr::Aggregate(ex) => write!(f, "{ex}"),
1378            Expr::Unary(ex) => write!(f, "{ex}"),
1379            Expr::Binary(ex) => write!(f, "{ex}"),
1380            Expr::Paren(ex) => write!(f, "{ex}"),
1381            Expr::Subquery(ex) => write!(f, "{ex}"),
1382            Expr::NumberLiteral(ex) => write!(f, "{ex}"),
1383            Expr::StringLiteral(ex) => write!(f, "{ex}"),
1384            Expr::VectorSelector(ex) => write!(f, "{ex}"),
1385            Expr::MatrixSelector(ex) => write!(f, "{ex}"),
1386            Expr::Call(ex) => write!(f, "{ex}"),
1387            Expr::Extension(ext) => write!(f, "{ext:?}"),
1388        }
1389    }
1390}
1391
1392impl Prettier for Expr {
1393    fn pretty(&self, level: usize, max: usize) -> String {
1394        match self {
1395            Expr::Aggregate(ex) => ex.pretty(level, max),
1396            Expr::Unary(ex) => ex.pretty(level, max),
1397            Expr::Binary(ex) => ex.pretty(level, max),
1398            Expr::Paren(ex) => ex.pretty(level, max),
1399            Expr::Subquery(ex) => ex.pretty(level, max),
1400            Expr::NumberLiteral(ex) => ex.pretty(level, max),
1401            Expr::StringLiteral(ex) => ex.pretty(level, max),
1402            Expr::VectorSelector(ex) => ex.pretty(level, max),
1403            Expr::MatrixSelector(ex) => ex.pretty(level, max),
1404            Expr::Call(ex) => ex.pretty(level, max),
1405            Expr::Extension(ext) => format!("{ext:?}"),
1406        }
1407    }
1408}
1409
1410/// check_ast checks the validity of the provided AST. This includes type checking.
1411/// Recursively check correct typing for child nodes and raise errors in case of bad typing.
1412pub(crate) fn check_ast(expr: Expr) -> Result<Expr, String> {
1413    match expr {
1414        Expr::Binary(ex) => check_ast_for_binary_expr(ex),
1415        Expr::Aggregate(ex) => check_ast_for_aggregate_expr(ex),
1416        Expr::Call(ex) => check_ast_for_call(ex),
1417        Expr::Unary(ex) => check_ast_for_unary(ex),
1418        Expr::Subquery(ex) => check_ast_for_subquery(ex),
1419        Expr::VectorSelector(ex) => check_ast_for_vector_selector(ex),
1420        Expr::Paren(_) => Ok(expr),
1421        Expr::NumberLiteral(_) => Ok(expr),
1422        Expr::StringLiteral(_) => Ok(expr),
1423        Expr::MatrixSelector(_) => Ok(expr),
1424        Expr::Extension(_) => Ok(expr),
1425    }
1426}
1427
1428fn expect_type(
1429    expected: ValueType,
1430    actual: Option<ValueType>,
1431    context: &str,
1432) -> Result<bool, String> {
1433    match actual {
1434        Some(actual) => {
1435            if actual == expected {
1436                Ok(true)
1437            } else {
1438                Err(format!(
1439                    "expected type {expected} in {context}, got {actual}"
1440                ))
1441            }
1442        }
1443        None => Err(format!("expected type {expected} in {context}, got None")),
1444    }
1445}
1446
1447/// the original logic is redundant in prometheus, and the following coding blocks
1448/// have been optimized for readability, but all logic SHOULD be covered.
1449fn check_ast_for_binary_expr(mut ex: BinaryExpr) -> Result<Expr, String> {
1450    if !ex.op.is_operator() {
1451        return Err(format!(
1452            "binary expression does not support operator '{}'",
1453            ex.op
1454        ));
1455    }
1456
1457    if ex.return_bool() && !ex.op.is_comparison_operator() {
1458        return Err("bool modifier can only be used on comparison operators".into());
1459    }
1460
1461    if ex.op.is_comparison_operator()
1462        && ex.lhs.value_type() == ValueType::Scalar
1463        && ex.rhs.value_type() == ValueType::Scalar
1464        && !ex.return_bool()
1465    {
1466        return Err("comparisons between scalars must use BOOL modifier".into());
1467    }
1468
1469    // For `on` matching, a label can only appear in one of the lists.
1470    // Every time series of the result vector must be uniquely identifiable.
1471    if ex.is_matching_on() && ex.is_labels_joint() {
1472        if let Some(labels) = ex.intersect_labels() {
1473            if let Some(label) = labels.first() {
1474                return Err(format!(
1475                    "label '{label}' must not occur in ON and GROUP clause at once"
1476                ));
1477            }
1478        };
1479    }
1480
1481    if ex.op.is_set_operator() {
1482        if ex.lhs.value_type() == ValueType::Scalar || ex.rhs.value_type() == ValueType::Scalar {
1483            return Err(format!(
1484                "set operator '{}' not allowed in binary scalar expression",
1485                ex.op
1486            ));
1487        }
1488
1489        if ex.lhs.value_type() == ValueType::Vector && ex.rhs.value_type() == ValueType::Vector {
1490            if let Some(ref modifier) = ex.modifier {
1491                if matches!(modifier.card, VectorMatchCardinality::OneToMany(_))
1492                    || matches!(modifier.card, VectorMatchCardinality::ManyToOne(_))
1493                {
1494                    return Err(format!("no grouping allowed for '{}' operation", ex.op));
1495                }
1496            };
1497        }
1498
1499        match &mut ex.modifier {
1500            Some(modifier) => {
1501                if modifier.card == VectorMatchCardinality::OneToOne {
1502                    modifier.card = VectorMatchCardinality::ManyToMany;
1503                }
1504            }
1505            None => {
1506                ex.modifier =
1507                    Some(BinModifier::default().with_card(VectorMatchCardinality::ManyToMany));
1508            }
1509        }
1510    }
1511
1512    if ex.lhs.value_type() != ValueType::Scalar && ex.lhs.value_type() != ValueType::Vector {
1513        return Err("binary expression must contain only scalar and instant vector types".into());
1514    }
1515    if ex.rhs.value_type() != ValueType::Scalar && ex.rhs.value_type() != ValueType::Vector {
1516        return Err("binary expression must contain only scalar and instant vector types".into());
1517    }
1518
1519    if (ex.lhs.value_type() != ValueType::Vector || ex.rhs.value_type() != ValueType::Vector)
1520        && ex.is_matching_labels_not_empty()
1521    {
1522        return Err("vector matching only allowed between vectors".into());
1523    }
1524
1525    Ok(Expr::Binary(ex))
1526}
1527
1528fn check_ast_for_aggregate_expr(ex: AggregateExpr) -> Result<Expr, String> {
1529    if !ex.op.is_aggregator() {
1530        return Err(format!(
1531            "aggregation operator expected in aggregation expression but got '{}'",
1532            ex.op
1533        ));
1534    }
1535
1536    expect_type(
1537        ValueType::Vector,
1538        Some(ex.expr.value_type()),
1539        "aggregation expression",
1540    )?;
1541
1542    if matches!(ex.op.id(), T_TOPK | T_BOTTOMK | T_QUANTILE) {
1543        expect_type(
1544            ValueType::Scalar,
1545            ex.param.as_ref().map(|ex| ex.value_type()),
1546            "aggregation expression",
1547        )?;
1548    }
1549
1550    if ex.op.id() == T_COUNT_VALUES {
1551        expect_type(
1552            ValueType::String,
1553            ex.param.as_ref().map(|ex| ex.value_type()),
1554            "aggregation expression",
1555        )?;
1556    }
1557
1558    Ok(Expr::Aggregate(ex))
1559}
1560
1561fn check_ast_for_call(ex: Call) -> Result<Expr, String> {
1562    let expected_args_len = ex.func.arg_types.len();
1563    let name = ex.func.name;
1564    let actual_args_len = ex.args.len();
1565
1566    if ex.func.variadic == 0 {
1567        if expected_args_len != actual_args_len {
1568            return Err(format!(
1569                "expected {expected_args_len} argument(s) in call to '{name}', got {actual_args_len}"
1570            ));
1571        }
1572    } else {
1573        let expected_args_len_without_default = expected_args_len.saturating_sub(1);
1574        if expected_args_len_without_default > actual_args_len {
1575            return Err(format!(
1576                "expected at least {expected_args_len_without_default} argument(s) in call to '{name}', got {actual_args_len}"
1577            ));
1578        }
1579
1580        if ex.func.variadic > 0 {
1581            let expected_max_args_len =
1582                expected_args_len_without_default + ex.func.variadic as usize;
1583            if expected_max_args_len < actual_args_len {
1584                return Err(format!(
1585                    "expected at most {expected_max_args_len} argument(s) in call to '{name}', got {actual_args_len}"
1586                ));
1587            }
1588        }
1589    }
1590
1591    // special cases from https://prometheus.io/docs/prometheus/latest/querying/functions
1592    if name.eq("exp") {
1593        if let Some(val) = ex.args.first().and_then(|ex| ex.scalar_value()) {
1594            if val.is_nan() || val.is_infinite() {
1595                return Ok(Expr::Call(ex));
1596            }
1597        }
1598    } else if name.eq("ln") || name.eq("log2") || name.eq("log10") {
1599        if let Some(val) = ex.args.first().and_then(|ex| ex.scalar_value()) {
1600            if val.is_nan() || val.is_infinite() || val <= 0.0 {
1601                return Ok(Expr::Call(ex));
1602            }
1603        }
1604    }
1605
1606    for (mut idx, actual_arg) in ex.args.args.iter().enumerate() {
1607        // this only happens when function args are variadic
1608        if idx >= ex.func.arg_types.len() {
1609            idx = ex.func.arg_types.len() - 1;
1610        }
1611
1612        expect_type(
1613            ex.func.arg_types[idx],
1614            Some(actual_arg.value_type()),
1615            &format!("call to function '{name}'"),
1616        )?;
1617    }
1618
1619    Ok(Expr::Call(ex))
1620}
1621
1622fn check_ast_for_unary(ex: UnaryExpr) -> Result<Expr, String> {
1623    let value_type = ex.expr.value_type();
1624    if value_type != ValueType::Scalar && value_type != ValueType::Vector {
1625        return Err(format!(
1626            "unary expression only allowed on expressions of type scalar or vector, got {value_type}"
1627        ));
1628    }
1629
1630    Ok(Expr::Unary(ex))
1631}
1632
1633fn check_ast_for_subquery(ex: SubqueryExpr) -> Result<Expr, String> {
1634    let value_type = ex.expr.value_type();
1635    if value_type != ValueType::Vector {
1636        return Err(format!(
1637            "subquery is only allowed on vector, got {value_type} instead"
1638        ));
1639    }
1640
1641    Ok(Expr::Subquery(ex))
1642}
1643
1644fn check_ast_for_vector_selector(ex: VectorSelector) -> Result<Expr, String> {
1645    match ex.name {
1646        Some(ref name) => match ex.matchers.find_matcher_value(METRIC_NAME) {
1647            Some(val) => Err(format!(
1648                "metric name must not be set twice: '{name}' or '{val}'"
1649            )),
1650            None => Ok(Expr::VectorSelector(ex)),
1651        },
1652        None if ex.matchers.is_empty_matchers() => {
1653            // When name is None, a vector selector must contain at least one non-empty matcher
1654            // to prevent implicit selection of all metrics (e.g. by a typo).
1655            Err("vector selector must contain at least one non-empty matcher".into())
1656        }
1657        _ => Ok(Expr::VectorSelector(ex)),
1658    }
1659}
1660
1661#[cfg(test)]
1662mod tests {
1663    use super::*;
1664    use crate::label::{MatchOp, Matcher};
1665
1666    #[test]
1667    fn test_valid_at_modifier() {
1668        let cases = vec![
1669            // tuple: (seconds, elapsed milliseconds before or after UNIX_EPOCH)
1670            (0.0, 0),
1671            (1000.3, 1000300),    // after UNIX_EPOCH
1672            (1000.9, 1000900),    // after UNIX_EPOCH
1673            (1000.9991, 1000999), // after UNIX_EPOCH
1674            (1000.9999, 1001000), // after UNIX_EPOCH
1675            (-1000.3, 1000300),   // before UNIX_EPOCH
1676            (-1000.9, 1000900),   // before UNIX_EPOCH
1677        ];
1678
1679        for (secs, elapsed) in cases {
1680            match AtModifier::try_from(secs).unwrap() {
1681                AtModifier::At(st) => {
1682                    if secs.is_sign_positive() || secs == 0.0 {
1683                        assert_eq!(
1684                            elapsed,
1685                            st.duration_since(SystemTime::UNIX_EPOCH)
1686                                .unwrap()
1687                                .as_millis()
1688                        )
1689                    } else if secs.is_sign_negative() {
1690                        assert_eq!(
1691                            elapsed,
1692                            SystemTime::UNIX_EPOCH
1693                                .duration_since(st)
1694                                .unwrap()
1695                                .as_millis()
1696                        )
1697                    }
1698                }
1699                _ => panic!(),
1700            }
1701        }
1702
1703        assert_eq!(
1704            AtModifier::try_from(Expr::from(1.0)),
1705            AtModifier::try_from(1.0),
1706        );
1707    }
1708
1709    #[test]
1710    fn test_invalid_at_modifier() {
1711        let cases = vec![
1712            f64::MAX,
1713            f64::MIN,
1714            f64::NAN,
1715            f64::INFINITY,
1716            f64::NEG_INFINITY,
1717        ];
1718
1719        for secs in cases {
1720            assert!(AtModifier::try_from(secs).is_err())
1721        }
1722
1723        assert_eq!(
1724            AtModifier::try_from(token::T_ADD),
1725            Err("invalid @ modifier preprocessor '+', START or END is valid.".into())
1726        );
1727
1728        assert_eq!(
1729            AtModifier::try_from(Expr::from("string literal")),
1730            Err("invalid float value after @ modifier".into())
1731        );
1732    }
1733
1734    #[test]
1735    fn test_binary_labels() {
1736        assert_eq!(
1737            &Labels::new(vec!["foo", "bar"]),
1738            LabelModifier::Include(Labels::new(vec!["foo", "bar"])).labels()
1739        );
1740
1741        assert_eq!(
1742            &Labels::new(vec!["foo", "bar"]),
1743            LabelModifier::Exclude(Labels::new(vec!["foo", "bar"])).labels()
1744        );
1745
1746        assert_eq!(
1747            &Labels::new(vec!["foo", "bar"]),
1748            VectorMatchCardinality::OneToMany(Labels::new(vec!["foo", "bar"]))
1749                .labels()
1750                .unwrap()
1751        );
1752
1753        assert_eq!(
1754            &Labels::new(vec!["foo", "bar"]),
1755            VectorMatchCardinality::ManyToOne(Labels::new(vec!["foo", "bar"]))
1756                .labels()
1757                .unwrap()
1758        );
1759
1760        assert_eq!(VectorMatchCardinality::OneToOne.labels(), None);
1761        assert_eq!(VectorMatchCardinality::ManyToMany.labels(), None);
1762    }
1763
1764    #[test]
1765    fn test_neg() {
1766        assert_eq!(
1767            -VectorSelector::from("foo"),
1768            UnaryExpr {
1769                expr: Box::new(Expr::from(VectorSelector::from("foo")))
1770            }
1771        )
1772    }
1773
1774    #[test]
1775    fn test_scalar_value() {
1776        assert_eq!(Some(1.0), Expr::from(1.0).scalar_value());
1777        assert_eq!(None, Expr::from("1.0").scalar_value());
1778    }
1779
1780    #[test]
1781    fn test_at_expr() {
1782        assert_eq!(
1783            "@ <timestamp> may not be set multiple times",
1784            Expr::from(VectorSelector::from("foo"))
1785                .at_expr(AtModifier::try_from(1.0).unwrap())
1786                .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1787                .unwrap_err()
1788        );
1789
1790        assert_eq!(
1791            "@ <timestamp> may not be set multiple times",
1792            Expr::new_matrix_selector(
1793                Expr::from(VectorSelector::from("foo")),
1794                Duration::from_secs(1),
1795            )
1796            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1797            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1798            .unwrap_err()
1799        );
1800
1801        assert_eq!(
1802            "@ <timestamp> may not be set multiple times",
1803            Expr::new_subquery_expr(
1804                Expr::from(VectorSelector::from("foo")),
1805                Duration::from_secs(1),
1806                None,
1807            )
1808            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1809            .and_then(|ex| ex.at_expr(AtModifier::try_from(1.0).unwrap()))
1810            .unwrap_err()
1811        )
1812    }
1813
1814    #[test]
1815    fn test_offset_expr() {
1816        assert_eq!(
1817            "offset may not be set multiple times",
1818            Expr::from(VectorSelector::from("foo"))
1819                .offset_expr(Offset::Pos(Duration::from_secs(1000)))
1820                .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1821                .unwrap_err()
1822        );
1823
1824        assert_eq!(
1825            "offset may not be set multiple times",
1826            Expr::new_matrix_selector(
1827                Expr::from(VectorSelector::from("foo")),
1828                Duration::from_secs(1),
1829            )
1830            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1831            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1832            .unwrap_err()
1833        );
1834
1835        assert_eq!(
1836            "offset may not be set multiple times",
1837            Expr::new_subquery_expr(
1838                Expr::from(VectorSelector::from("foo")),
1839                Duration::from_secs(1),
1840                None,
1841            )
1842            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1843            .and_then(|ex| ex.offset_expr(Offset::Pos(Duration::from_secs(1000))))
1844            .unwrap_err()
1845        );
1846    }
1847
1848    #[test]
1849    fn test_expr_to_string() {
1850        let mut cases = vec![
1851            ("1", "1"),
1852            ("- 1", "-1"),
1853            ("+ 1", "1"),
1854            ("Inf", "Inf"),
1855            ("inf", "Inf"),
1856            ("+Inf", "Inf"),
1857            ("- Inf", "-Inf"),
1858            (".5", "0.5"),
1859            ("5.", "5"),
1860            ("123.4567", "123.4567"),
1861            ("5e-3", "0.005"),
1862            ("5e3", "5000"),
1863            ("0xc", "12"),
1864            ("0755", "493"),
1865            ("08", "8"),
1866            ("+5.5e-3", "0.0055"),
1867            ("-0755", "-493"),
1868            ("NaN", "NaN"),
1869            ("NAN", "NaN"),
1870            ("- 1^2", "-1 ^ 2"),
1871            ("+1 + -2 * 1", "1 + -2 * 1"),
1872            ("1 + 2/(3*1)", "1 + 2 / (3 * 1)"),
1873            ("foo*sum", "foo * sum"),
1874            ("foo * on(test,blub) bar", "foo * on (test, blub) bar"),
1875            (
1876                r#"up{job="hi", instance="in"} offset 5m @ 100"#,
1877                r#"up{instance="in",job="hi"} @ 100.000 offset 5m"#,
1878            ),
1879            (
1880                r#"up{job="hi", instance="in"}"#,
1881                r#"up{instance="in",job="hi"}"#,
1882            ),
1883            ("sum (up) by (job,instance)", "sum by (job, instance) (up)"),
1884            (
1885                "foo / on(test,blub) group_left(bar) bar",
1886                "foo / on (test, blub) group_left (bar) bar",
1887            ),
1888            (
1889                "foo / on(test,blub) group_right(bar) bar",
1890                "foo / on (test, blub) group_right (bar) bar",
1891            ),
1892            (
1893                r#"foo{a="b",foo!="bar",test=~"test",bar!~"baz"}"#,
1894                r#"foo{a="b",bar!~"baz",foo!="bar",test=~"test"}"#,
1895            ),
1896            (
1897                r#"{__name__=~"foo.+",__name__=~".*bar"}"#,
1898                r#"{__name__=~".*bar",__name__=~"foo.+"}"#,
1899            ),
1900            (
1901                r#"test{a="b"}[5y] OFFSET 3d"#,
1902                r#"test{a="b"}[5y] offset 3d"#,
1903            ),
1904            (
1905                r#"{a="b"}[5y] OFFSET 3d"#,
1906                r#"{a="b"}[5y] offset 3d"#,
1907            ),
1908            (
1909                "sum(some_metric) without(and, by, avg, count, alert, annotations)",
1910                "sum without (and, by, avg, count, alert, annotations) (some_metric)",
1911            ),
1912            (
1913                r#"floor(some_metric{foo!="bar"})"#,
1914                r#"floor(some_metric{foo!="bar"})"#,
1915            ),
1916            (
1917                "sum(rate(http_request_duration_seconds[10m])) / count(rate(http_request_duration_seconds[10m]))",
1918                "sum(rate(http_request_duration_seconds[10m])) / count(rate(http_request_duration_seconds[10m]))",
1919            ),
1920            ("rate(some_metric[5m])", "rate(some_metric[5m])"),
1921            ("round(some_metric,5)", "round(some_metric, 5)"),
1922            (
1923                r#"absent(sum(nonexistent{job="myjob"}))"#,
1924                r#"absent(sum(nonexistent{job="myjob"}))"#,
1925            ),
1926            (
1927                "histogram_quantile(0.9,rate(http_request_duration_seconds_bucket[10m]))",
1928                "histogram_quantile(0.9, rate(http_request_duration_seconds_bucket[10m]))",
1929            ),
1930            (
1931                "histogram_quantile(0.9,sum(rate(http_request_duration_seconds_bucket[10m])) by(job,le))",
1932                "histogram_quantile(0.9, sum by (job, le) (rate(http_request_duration_seconds_bucket[10m])))",
1933            ),
1934            (
1935                r#"label_join(up{job="api-server",src1="a",src2="b",src3="c"}, "foo", ",", "src1", "src2", "src3")"#,
1936                r#"label_join(up{job="api-server",src1="a",src2="b",src3="c"}, "foo", ",", "src1", "src2", "src3")"#,
1937            ),
1938            (
1939                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s] @ 100"#,
1940                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s] @ 100.000"#,
1941            ),
1942            (
1943                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s]"#,
1944                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:])[4m:3s]"#,
1945            ),
1946            (
1947                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:] offset 4m)[4m:3s]"#,
1948                r#"min_over_time(rate(foo{bar="baz"}[2s])[5m:] offset 4m)[4m:3s]"#,
1949            ),
1950            ("some_metric OFFSET 1m [10m:5s]", "some_metric offset 1m[10m:5s]"),
1951            ("some_metric @123 [10m:5s]", "some_metric @ 123.000[10m:5s]"),
1952            ("some_metric <= 1ms", "some_metric <= 0.001"),
1953        ];
1954
1955        // the following cases are from https://github.com/prometheus/prometheus/blob/main/promql/parser/printer_test.go
1956        let mut cases1 = vec![
1957            (
1958                r#"sum by() (task:errors:rate10s{job="s"})"#,
1959                r#"sum(task:errors:rate10s{job="s"})"#,
1960            ),
1961            (
1962                r#"sum by(code) (task:errors:rate10s{job="s"})"#,
1963                r#"sum by (code) (task:errors:rate10s{job="s"})"#,
1964            ),
1965            (
1966                r#"sum without() (task:errors:rate10s{job="s"})"#,
1967                r#"sum without () (task:errors:rate10s{job="s"})"#,
1968            ),
1969            (
1970                r#"sum without(instance) (task:errors:rate10s{job="s"})"#,
1971                r#"sum without (instance) (task:errors:rate10s{job="s"})"#,
1972            ),
1973            (
1974                r#"topk(5, task:errors:rate10s{job="s"})"#,
1975                r#"topk(5, task:errors:rate10s{job="s"})"#,
1976            ),
1977            (
1978                r#"count_values("value", task:errors:rate10s{job="s"})"#,
1979                r#"count_values("value", task:errors:rate10s{job="s"})"#,
1980            ),
1981            ("a - on() c", "a - on () c"),
1982            ("a - on(b) c", "a - on (b) c"),
1983            ("a - on(b) group_left(x) c", "a - on (b) group_left (x) c"),
1984            (
1985                "a - on(b) group_left(x, y) c",
1986                "a - on (b) group_left (x, y) c",
1987            ),
1988            ("a - on(b) group_left c", "a - on (b) group_left () c"),
1989            ("a - ignoring(b) c", "a - ignoring (b) c"),
1990            ("a - ignoring() c", "a - c"),
1991            ("up > bool 0", "up > bool 0"),
1992            ("a offset 1m", "a offset 1m"),
1993            ("a offset -7m", "a offset -7m"),
1994            (r#"a{c="d"}[5m] offset 1m"#, r#"a{c="d"}[5m] offset 1m"#),
1995            ("a[5m] offset 1m", "a[5m] offset 1m"),
1996            ("a[12m] offset -3m", "a[12m] offset -3m"),
1997            ("a[1h:5m] offset 1m", "a[1h:5m] offset 1m"),
1998            (r#"{__name__="a"}"#, r#"{__name__="a"}"#),
1999            (r#"a{b!="c"}[1m]"#, r#"a{b!="c"}[1m]"#),
2000            (r#"a{b=~"c"}[1m]"#, r#"a{b=~"c"}[1m]"#),
2001            (r#"a{b!~"c"}[1m]"#, r#"a{b!~"c"}[1m]"#),
2002            ("a @ 10", "a @ 10.000"),
2003            ("a[1m] @ 10", "a[1m] @ 10.000"),
2004            ("a @ start()", "a @ start()"),
2005            ("a @ end()", "a @ end()"),
2006            ("a[1m] @ start()", "a[1m] @ start()"),
2007            ("a[1m] @ end()", "a[1m] @ end()"),
2008        ];
2009
2010        // the following cases copy the tests from the following: https://github.com/prometheus/prometheus/pull/9138
2011        let mut cases2 = vec![
2012            (
2013                r#"test{a="b"}[5y] OFFSET 3d"#,
2014                r#"test{a="b"}[5y] offset 3d"#,
2015            ),
2016            (
2017                r#"test{a="b"}[5m] OFFSET 3600"#,
2018                r#"test{a="b"}[5m] offset 1h"#,
2019            ),
2020            ("foo[3ms] @ 2.345", "foo[3ms] @ 2.345"),
2021            ("foo[4s180ms] @ 2.345", "foo[4s180ms] @ 2.345"),
2022            ("foo[4.18] @ 2.345", "foo[4s180ms] @ 2.345"),
2023            ("foo[4s18ms] @ 2.345", "foo[4s18ms] @ 2.345"),
2024            ("foo[4.018] @ 2.345", "foo[4s18ms] @ 2.345"),
2025            ("test[5]", "test[5s]"),
2026            ("some_metric[5m] @ 1m", "some_metric[5m] @ 60.000"),
2027            ("metric @ 100s", "metric @ 100.000"),
2028            ("metric @ 1m40s", "metric @ 100.000"),
2029            ("metric @ 100 offset 50", "metric @ 100.000 offset 50s"),
2030            ("metric offset 50 @ 100", "metric @ 100.000 offset 50s"),
2031            ("metric @ 0 offset -50", "metric @ 0.000 offset -50s"),
2032            ("metric offset -50 @ 0", "metric @ 0.000 offset -50s"),
2033            (
2034                r#"sum_over_time(metric{job="1"}[100] @ 100 offset 50)"#,
2035                r#"sum_over_time(metric{job="1"}[1m40s] @ 100.000 offset 50s)"#,
2036            ),
2037            (
2038                r#"sum_over_time(metric{job="1"}[100] offset 50s @ 100)"#,
2039                r#"sum_over_time(metric{job="1"}[1m40s] @ 100.000 offset 50s)"#,
2040            ),
2041            (
2042                r#"sum_over_time(metric{job="1"}[100] @ 100) + label_replace(sum_over_time(metric{job="2"}[100] @ 100), "job", "1", "", "")"#,
2043                r#"sum_over_time(metric{job="1"}[1m40s] @ 100.000) + label_replace(sum_over_time(metric{job="2"}[1m40s] @ 100.000), "job", "1", "", "")"#,
2044            ),
2045            (
2046                r#"sum_over_time(metric{job="1"}[100:1] offset 20 @ 100)"#,
2047                r#"sum_over_time(metric{job="1"}[1m40s:1s] @ 100.000 offset 20s)"#,
2048            ),
2049        ];
2050
2051        cases.append(&mut cases1);
2052        cases.append(&mut cases2);
2053        for (input, expected) in cases {
2054            let expr = crate::parser::parse(input).unwrap();
2055            assert_eq!(expected, expr.to_string())
2056        }
2057    }
2058
2059    #[test]
2060    fn test_vector_selector_to_string() {
2061        let cases = vec![
2062            (VectorSelector::default(), ""),
2063            (VectorSelector::from("foobar"), "foobar"),
2064            (
2065                {
2066                    let name = Some(String::from("foobar"));
2067                    let matchers = Matchers::one(Matcher::new(MatchOp::Equal, "a", "x"));
2068                    VectorSelector::new(name, matchers)
2069                },
2070                r#"foobar{a="x"}"#,
2071            ),
2072            (
2073                {
2074                    let matchers = Matchers::new(vec![
2075                        Matcher::new(MatchOp::Equal, "a", "x"),
2076                        Matcher::new(MatchOp::Equal, "b", "y"),
2077                    ]);
2078                    VectorSelector::new(None, matchers)
2079                },
2080                r#"{a="x",b="y"}"#,
2081            ),
2082            (
2083                {
2084                    let matchers =
2085                        Matchers::one(Matcher::new(MatchOp::Equal, METRIC_NAME, "foobar"));
2086                    VectorSelector::new(None, matchers)
2087                },
2088                r#"{__name__="foobar"}"#,
2089            ),
2090        ];
2091
2092        for (vs, expect) in cases {
2093            assert_eq!(expect, vs.to_string())
2094        }
2095    }
2096
2097    #[test]
2098    fn test_aggregate_expr_pretty() {
2099        let cases = vec![
2100            ("sum(foo)", "sum(foo)"),
2101            (
2102                r#"sum by() (task:errors:rate10s{job="s"})"#,
2103                r#"sum(
2104  task:errors:rate10s{job="s"}
2105)"#,
2106            ),
2107            (
2108                r#"sum without(job,foo) (task:errors:rate10s{job="s"})"#,
2109                r#"sum without (job, foo) (
2110  task:errors:rate10s{job="s"}
2111)"#,
2112            ),
2113            (
2114                r#"sum(task:errors:rate10s{job="s"}) without(job,foo)"#,
2115                r#"sum without (job, foo) (
2116  task:errors:rate10s{job="s"}
2117)"#,
2118            ),
2119            (
2120                r#"sum by(job,foo) (task:errors:rate10s{job="s"})"#,
2121                r#"sum by (job, foo) (
2122  task:errors:rate10s{job="s"}
2123)"#,
2124            ),
2125            (
2126                r#"sum (task:errors:rate10s{job="s"}) by(job,foo)"#,
2127                r#"sum by (job, foo) (
2128  task:errors:rate10s{job="s"}
2129)"#,
2130            ),
2131            (
2132                r#"topk(10, ask:errors:rate10s{job="s"})"#,
2133                r#"topk(
2134  10,
2135  ask:errors:rate10s{job="s"}
2136)"#,
2137            ),
2138            (
2139                r#"sum by(job,foo) (sum by(job,foo) (task:errors:rate10s{job="s"}))"#,
2140                r#"sum by (job, foo) (
2141  sum by (job, foo) (
2142    task:errors:rate10s{job="s"}
2143  )
2144)"#,
2145            ),
2146            (
2147                r#"sum by(job,foo) (sum by(job,foo) (sum by(job,foo) (task:errors:rate10s{job="s"})))"#,
2148                r#"sum by (job, foo) (
2149  sum by (job, foo) (
2150    sum by (job, foo) (
2151      task:errors:rate10s{job="s"}
2152    )
2153  )
2154)"#,
2155            ),
2156            (
2157                r#"sum by(job,foo)
2158(sum by(job,foo) (task:errors:rate10s{job="s"}))"#,
2159                r#"sum by (job, foo) (
2160  sum by (job, foo) (
2161    task:errors:rate10s{job="s"}
2162  )
2163)"#,
2164            ),
2165            (
2166                r#"sum by(job,foo)
2167(sum(task:errors:rate10s{job="s"}) without(job,foo))"#,
2168                r#"sum by (job, foo) (
2169  sum without (job, foo) (
2170    task:errors:rate10s{job="s"}
2171  )
2172)"#,
2173            ),
2174            (
2175                r#"sum by(job,foo) # Comment 1.
2176(sum by(job,foo) ( # Comment 2.
2177task:errors:rate10s{job="s"}))"#,
2178                r#"sum by (job, foo) (
2179  sum by (job, foo) (
2180    task:errors:rate10s{job="s"}
2181  )
2182)"#,
2183            ),
2184        ];
2185
2186        for (input, expect) in cases {
2187            let expr = crate::parser::parse(input);
2188            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2189        }
2190    }
2191
2192    #[test]
2193    fn test_binary_expr_pretty() {
2194        let cases = vec![
2195            ("a+b", "a + b"),
2196            (
2197                "a == bool 1",
2198                "  a
2199== bool
2200  1",
2201            ),
2202            (
2203                "a == 1024000",
2204                "  a
2205==
2206  1024000",
2207            ),
2208            (
2209                "a + ignoring(job) b",
2210                "  a
2211+ ignoring (job)
2212  b",
2213            ),
2214            (
2215                "foo_1 + foo_2",
2216                "  foo_1
2217+
2218  foo_2",
2219            ),
2220            (
2221                "foo_1 + foo_2 + foo_3",
2222                "    foo_1
2223  +
2224    foo_2
2225+
2226  foo_3",
2227            ),
2228            (
2229                "foo + baar + foo_3",
2230                "  foo + baar
2231+
2232  foo_3",
2233            ),
2234            (
2235                "foo_1 + foo_2 + foo_3 + foo_4",
2236                "      foo_1
2237    +
2238      foo_2
2239  +
2240    foo_3
2241+
2242  foo_4",
2243            ),
2244            (
2245                "foo_1 + ignoring(foo) foo_2 + ignoring(job) group_left foo_3 + on(instance) group_right foo_4",
2246
2247                 "      foo_1
2248    + ignoring (foo)
2249      foo_2
2250  + ignoring (job) group_left ()
2251    foo_3
2252+ on (instance) group_right ()
2253  foo_4",
2254            ),
2255        ];
2256
2257        for (input, expect) in cases {
2258            let expr = crate::parser::parse(input);
2259            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2260        }
2261    }
2262
2263    #[test]
2264    fn test_call_expr_pretty() {
2265        let cases = vec![
2266            (
2267                "rate(foo[1m])",
2268                "rate(
2269  foo[1m]
2270)",
2271            ),
2272            (
2273                "sum_over_time(foo[1m])",
2274                "sum_over_time(
2275  foo[1m]
2276)",
2277            ),
2278            (
2279                "rate(long_vector_selector[10m:1m] @ start() offset 1m)",
2280                "rate(
2281  long_vector_selector[10m:1m] @ start() offset 1m
2282)",
2283            ),
2284            (
2285                "histogram_quantile(0.9, rate(foo[1m]))",
2286                "histogram_quantile(
2287  0.9,
2288  rate(
2289    foo[1m]
2290  )
2291)",
2292            ),
2293            (
2294                "histogram_quantile(0.9, rate(foo[1m] @ start()))",
2295                "histogram_quantile(
2296  0.9,
2297  rate(
2298    foo[1m] @ start()
2299  )
2300)",
2301            ),
2302            (
2303                "max_over_time(rate(demo_api_request_duration_seconds_count[1m])[1m:] @ start() offset 1m)",
2304                "max_over_time(
2305  rate(
2306    demo_api_request_duration_seconds_count[1m]
2307  )[1m:] @ start() offset 1m
2308)",
2309            ),
2310            (
2311                r#"label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*")"#,
2312                r#"label_replace(
2313  up{job="api-server",service="a:c"},
2314  "foo",
2315  "$1",
2316  "service",
2317  "(.*):.*"
2318)"#,
2319            ),
2320            (
2321                r#"label_replace(label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*"), "foo", "$1", "service", "(.*):.*")"#,
2322                r#"label_replace(
2323  label_replace(
2324    up{job="api-server",service="a:c"},
2325    "foo",
2326    "$1",
2327    "service",
2328    "(.*):.*"
2329  ),
2330  "foo",
2331  "$1",
2332  "service",
2333  "(.*):.*"
2334)"#,
2335            ),
2336        ];
2337
2338        for (input, expect) in cases {
2339            let expr = crate::parser::parse(input);
2340            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2341        }
2342    }
2343
2344    #[test]
2345    fn test_paren_expr_pretty() {
2346        let cases = vec![
2347            ("(foo)", "(foo)"),
2348            (
2349                "(_foo_long_)",
2350                "(
2351  _foo_long_
2352)",
2353            ),
2354            (
2355                "((foo_long))",
2356                "(
2357  (foo_long)
2358)",
2359            ),
2360            (
2361                "((_foo_long_))",
2362                "(
2363  (
2364    _foo_long_
2365  )
2366)",
2367            ),
2368            (
2369                "(((foo_long)))",
2370                "(
2371  (
2372    (foo_long)
2373  )
2374)",
2375            ),
2376            ("(1 + 2)", "(1 + 2)"),
2377            (
2378                "(foo + bar)",
2379                "(
2380  foo + bar
2381)",
2382            ),
2383            (
2384                "(foo_long + bar_long)",
2385                "(
2386    foo_long
2387  +
2388    bar_long
2389)",
2390            ),
2391            (
2392                "(foo_long + bar_long + bar_2_long)",
2393                "(
2394      foo_long
2395    +
2396      bar_long
2397  +
2398    bar_2_long
2399)",
2400            ),
2401            (
2402                "((foo_long + bar_long) + bar_2_long)",
2403                "(
2404    (
2405        foo_long
2406      +
2407        bar_long
2408    )
2409  +
2410    bar_2_long
2411)",
2412            ),
2413            (
2414                "(1111 + 2222)",
2415                "(
2416    1111
2417  +
2418    2222
2419)",
2420            ),
2421            (
2422                "(sum_over_time(foo[1m]))",
2423                "(
2424  sum_over_time(
2425    foo[1m]
2426  )
2427)",
2428            ),
2429            (
2430                r#"(label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*"))"#,
2431                r#"(
2432  label_replace(
2433    up{job="api-server",service="a:c"},
2434    "foo",
2435    "$1",
2436    "service",
2437    "(.*):.*"
2438  )
2439)"#,
2440            ),
2441            (
2442                r#"(label_replace(label_replace(up{job="api-server",service="a:c"}, "foo", "$1", "service", "(.*):.*"), "foo", "$1", "service", "(.*):.*"))"#,
2443                r#"(
2444  label_replace(
2445    label_replace(
2446      up{job="api-server",service="a:c"},
2447      "foo",
2448      "$1",
2449      "service",
2450      "(.*):.*"
2451    ),
2452    "foo",
2453    "$1",
2454    "service",
2455    "(.*):.*"
2456  )
2457)"#,
2458            ),
2459            (
2460                r#"(label_replace(label_replace((up{job="api-server",service="a:c"}), "foo", "$1", "service", "(.*):.*"), "foo", "$1", "service", "(.*):.*"))"#,
2461                r#"(
2462  label_replace(
2463    label_replace(
2464      (
2465        up{job="api-server",service="a:c"}
2466      ),
2467      "foo",
2468      "$1",
2469      "service",
2470      "(.*):.*"
2471    ),
2472    "foo",
2473    "$1",
2474    "service",
2475    "(.*):.*"
2476  )
2477)"#,
2478            ),
2479        ];
2480
2481        for (input, expect) in cases {
2482            let expr = crate::parser::parse(input);
2483            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2484        }
2485    }
2486
2487    #[test]
2488    fn test_unary_expr_pretty() {
2489        let cases = vec![
2490            ("-1", "-1"),
2491            ("-vector_selector", "-vector_selector"),
2492            (
2493                "(-vector_selector)",
2494                "(
2495  -vector_selector
2496)",
2497            ),
2498            (
2499                "-histogram_quantile(0.9,rate(foo[1m]))",
2500                "-histogram_quantile(
2501  0.9,
2502  rate(
2503    foo[1m]
2504  )
2505)",
2506            ),
2507            (
2508                "-histogram_quantile(0.99, sum by (le) (rate(foo[1m])))",
2509                "-histogram_quantile(
2510  0.99,
2511  sum by (le) (
2512    rate(
2513      foo[1m]
2514    )
2515  )
2516)",
2517            ),
2518            (
2519                "-histogram_quantile(0.9, -rate(foo[1m] @ start()))",
2520                "-histogram_quantile(
2521  0.9,
2522  -rate(
2523    foo[1m] @ start()
2524  )
2525)",
2526            ),
2527            (
2528                "(-histogram_quantile(0.9, -rate(foo[1m] @ start())))",
2529                "(
2530  -histogram_quantile(
2531    0.9,
2532    -rate(
2533      foo[1m] @ start()
2534    )
2535  )
2536)",
2537            ),
2538        ];
2539
2540        for (input, expect) in cases {
2541            let expr = crate::parser::parse(input);
2542            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2543        }
2544    }
2545
2546    #[test]
2547    fn test_expr_pretty() {
2548        // Following queries have been taken from https://monitoring.mixins.dev/
2549        let cases = vec![
2550            (
2551                r#"(node_filesystem_avail_bytes{job="node",fstype!=""} / node_filesystem_size_bytes{job="node",fstype!=""} * 100 < 40 and predict_linear(node_filesystem_avail_bytes{job="node",fstype!=""}[6h], 24*60*60) < 0 and node_filesystem_readonly{job="node",fstype!=""} == 0)"#,
2552                r#"(
2553            node_filesystem_avail_bytes{fstype!="",job="node"}
2554          /
2555            node_filesystem_size_bytes{fstype!="",job="node"}
2556        *
2557          100
2558      <
2559        40
2560    and
2561        predict_linear(
2562          node_filesystem_avail_bytes{fstype!="",job="node"}[6h],
2563            24 * 60
2564          *
2565            60
2566        )
2567      <
2568        0
2569  and
2570      node_filesystem_readonly{fstype!="",job="node"}
2571    ==
2572      0
2573)"#,
2574            ),
2575            (
2576                r#"(node_filesystem_avail_bytes{job="node",fstype!=""} / node_filesystem_size_bytes{job="node",fstype!=""} * 100 < 20 and predict_linear(node_filesystem_avail_bytes{job="node",fstype!=""}[6h], 4*60*60) < 0 and node_filesystem_readonly{job="node",fstype!=""} == 0)"#,
2577                r#"(
2578            node_filesystem_avail_bytes{fstype!="",job="node"}
2579          /
2580            node_filesystem_size_bytes{fstype!="",job="node"}
2581        *
2582          100
2583      <
2584        20
2585    and
2586        predict_linear(
2587          node_filesystem_avail_bytes{fstype!="",job="node"}[6h],
2588            4 * 60
2589          *
2590            60
2591        )
2592      <
2593        0
2594  and
2595      node_filesystem_readonly{fstype!="",job="node"}
2596    ==
2597      0
2598)"#,
2599            ),
2600            (
2601                r#"(node_timex_offset_seconds > 0.05 and deriv(node_timex_offset_seconds[5m]) >= 0) or (node_timex_offset_seconds < -0.05 and deriv(node_timex_offset_seconds[5m]) <= 0)"#,
2602                r#"  (
2603        node_timex_offset_seconds
2604      >
2605        0.05
2606    and
2607        deriv(
2608          node_timex_offset_seconds[5m]
2609        )
2610      >=
2611        0
2612  )
2613or
2614  (
2615        node_timex_offset_seconds
2616      <
2617        -0.05
2618    and
2619        deriv(
2620          node_timex_offset_seconds[5m]
2621        )
2622      <=
2623        0
2624  )"#,
2625            ),
2626            (
2627                r#"1 - ((node_memory_MemAvailable_bytes{job="node"} or (node_memory_Buffers_bytes{job="node"} + node_memory_Cached_bytes{job="node"} + node_memory_MemFree_bytes{job="node"} + node_memory_Slab_bytes{job="node"}) ) / node_memory_MemTotal_bytes{job="node"})"#,
2628                r#"  1
2629-
2630  (
2631      (
2632          node_memory_MemAvailable_bytes{job="node"}
2633        or
2634          (
2635                  node_memory_Buffers_bytes{job="node"}
2636                +
2637                  node_memory_Cached_bytes{job="node"}
2638              +
2639                node_memory_MemFree_bytes{job="node"}
2640            +
2641              node_memory_Slab_bytes{job="node"}
2642          )
2643      )
2644    /
2645      node_memory_MemTotal_bytes{job="node"}
2646  )"#,
2647            ),
2648            (
2649                r#"min by (job, integration) (rate(alertmanager_notifications_failed_total{job="alertmanager", integration=~".*"}[5m]) / rate(alertmanager_notifications_total{job="alertmanager", integration="~.*"}[5m])) > 0.01"#,
2650                r#"  min by (job, integration) (
2651      rate(
2652        alertmanager_notifications_failed_total{integration=~".*",job="alertmanager"}[5m]
2653      )
2654    /
2655      rate(
2656        alertmanager_notifications_total{integration="~.*",job="alertmanager"}[5m]
2657      )
2658  )
2659>
2660  0.01"#,
2661            ),
2662            (
2663                r#"(count by (job) (changes(process_start_time_seconds{job="alertmanager"}[10m]) > 4) / count by (job) (up{job="alertmanager"})) >= 0.5"#,
2664                r#"  (
2665      count by (job) (
2666          changes(
2667            process_start_time_seconds{job="alertmanager"}[10m]
2668          )
2669        >
2670          4
2671      )
2672    /
2673      count by (job) (
2674        up{job="alertmanager"}
2675      )
2676  )
2677>=
2678  0.5"#,
2679            ),
2680        ];
2681
2682        for (input, expect) in cases {
2683            let expr = crate::parser::parse(input);
2684            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2685        }
2686    }
2687
2688    #[test]
2689    fn test_step_invariant_pretty() {
2690        let cases = vec![
2691            ("a @ 1", "a @ 1.000"),
2692            ("a @ start()", "a @ start()"),
2693            ("vector_selector @ start()", "vector_selector @ start()"),
2694        ];
2695
2696        for (input, expect) in cases {
2697            let expr = crate::parser::parse(input);
2698            assert_eq!(expect, expr.unwrap().pretty(0, 10));
2699        }
2700    }
2701
2702    #[test]
2703    fn test_prettify() {
2704        let cases = vec![
2705            ("vector_selector", "vector_selector"),
2706            (
2707                r#"vector_selector{fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",process_name="alertmanager"}"#,
2708                r#"vector_selector{barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",process_name="alertmanager"}"#,
2709            ),
2710            (
2711                r#"matrix_selector{fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",process_name="alertmanager"}[1y2w3d]"#,
2712                r#"matrix_selector{barrrrrrrrrrrrrrrrrrr="fooooooooooooooooo",fooooooooooooooooo="barrrrrrrrrrrrrrrrrrr",process_name="alertmanager"}[382d]"#,
2713            ),
2714        ];
2715
2716        for (input, expect) in cases {
2717            assert_eq!(expect, crate::parser::parse(input).unwrap().prettify());
2718        }
2719    }
2720
2721    #[test]
2722    fn test_eval_stmt_to_string() {
2723        let query = r#"http_requests_total{job="apiserver", handler="/api/comments"}[5m]"#;
2724        let start = "2024-10-08T07:15:00.022978+00:00";
2725        let end = "2024-10-08T07:15:30.012978+00:00";
2726        let expect = r#"[http_requests_total{handler="/api/comments",job="apiserver"}[5m]] @ [2024-10-08T07:15:00.022978+00:00, 2024-10-08T07:15:30.012978+00:00, 1m, 5m]"#;
2727
2728        let stmt = EvalStmt {
2729            expr: crate::parser::parse(query).unwrap(),
2730            start: DateTime::parse_from_rfc3339(start)
2731                .unwrap()
2732                .with_timezone(&Utc)
2733                .into(),
2734            end: DateTime::parse_from_rfc3339(end)
2735                .unwrap()
2736                .with_timezone(&Utc)
2737                .into(),
2738            interval: Duration::from_secs(60),
2739            lookback_delta: Duration::from_secs(300),
2740        };
2741
2742        assert_eq!(expect, stmt.to_string());
2743    }
2744
2745    #[test]
2746    fn test_prettify_with_utf8_labels() {
2747        // Test that labels with special characters are properly quoted in display
2748        let cases = vec![
2749            // (input, expected_display)
2750            (r#"{"some.metric"}"#, r#"{__name__="some.metric"}"#),
2751            (
2752                r#"foo{"label.with.dots"="value"}"#,
2753                r#"foo{"label.with.dots"="value"}"#,
2754            ),
2755            (
2756                r#"bar{"label-with-dashes"="test"}"#,
2757                r#"bar{"label-with-dashes"="test"}"#,
2758            ),
2759            (
2760                r#"baz{"label:with:colons"="data"}"#,
2761                r#"baz{"label:with:colons"="data"}"#,
2762            ),
2763            (
2764                r#"sum by ("service.version", foo) ({"some.metric"})"#,
2765                r#"sum by ("service.version", foo) ({__name__="some.metric"})"#,
2766            ),
2767            (
2768                r#"sum by (`service.version`, foo) ({"some.metric"})"#,
2769                r#"sum by ("service.version", foo) ({__name__="some.metric"})"#,
2770            ),
2771            // Regular labels should not be quoted
2772            (r#"foo{job="web"}"#, r#"foo{job="web"}"#),
2773            (
2774                r#"bar{instance_id="server1"}"#,
2775                r#"bar{instance_id="server1"}"#,
2776            ),
2777        ];
2778
2779        for (input, expected) in cases {
2780            let parsed = crate::parser::parse(input).unwrap();
2781            let prettified = parsed.prettify();
2782            assert_eq!(prettified, expected);
2783        }
2784    }
2785}