Support mutation and rewriting during AST traversal

[scintillavoy]

First of all, thank you for maintaining this project! From my experience, this crate is the only reliable PromQL parser available in Rust right now.

Problem:

Currently, walk_expr provides a great way to traverse the PromQL AST immutably. In the official Go implementation of the PromQL parser, traversal functions like Walk and Inspect naturally allow for AST mutation because Go passes pointers without strict mutability constraints.

However, walk_expr relies entirely on shared references (&Expr). This means if we need to manipulate the AST, such as injecting labels into selectors or optimizing queries, we currently lack built-in traversal patterns for mutation and restructuring.

---

Here are some thoughts on how to implement this feature:

1. Mutation

To support modifying internal fields in-place without allocating a new tree, we can introduce an ExprMutVisitor trait and a walk_expr_mut function (which takes &mut Expr).

However, duplicating the traversal logic across walk_expr and walk_expr_mut creates a maintenance burden, as they easily fall out of sync when new variants are added to Expr. Instead, we can write a single declarative macro that defines the AST destructuring logic once and generates both the & and &mut walker functions.

2. Rewriting

For structural changes where we need to replace or delete entire nodes, we can introduce an ExprRewriter trait and a walk_and_rewrite function (which takes ownership of Expr and returns a new one).

Unlike the reference walkers, the restructuring logic here is fundamentally different, making it impractical to generate all three walking functions with a single macro. Writing a separate, complex macro just for one single function definition feels like overkill, so it would likely be better to just implement walk_and_rewrite manually.

---

Questions for maintainers:

• Are you open to adding an ExprMutVisitor and an ExprRewriter trait?
• How do you feel about using a macro to deduplicate implementations?

[evenyag]

Makes sense to me.

The rust sqlparser also uses two traits. One difference is that sqlparser defines procedure macros for Visit/VisitMut, not declarative macros. But I'm open to both approaches if declarative macros work.

• https://docs.rs/sqlparser/latest/sqlparser/ast/trait.VisitMut.html
• https://docs.rs/sqlparser/latest/sqlparser/ast/trait.Visit.html

Would you be interested in submitting PRs for these features?

Should we choose ExprMutVisitor or ExprVisitorMut?

[scintillavoy]

@evenyag Thanks for the links! Looking at sqlparser, it actually uses four traits (Visit, VisitMut, Visitor, and VisitorMut) to separate the traversal logic from the callback logic across every AST node using procedural macros (#[derive(Visit, VisitMut)]).

My intention for using a macro was actually much more lightweight. If we were to support mutation, there would be walk_expr and walk_expr_mut. To deduplicate the match block inside them, I think a simple declarative macro would work here.

It would look like this. It is basically abstracting the & and &mut reference types and the trait bounds so we can generate both versions from a single implementation:

pub trait ExprVisitor {
    type Error;

    fn pre_visit(&mut self, plan: &Expr) -> Result<bool, Self::Error>;
    fn post_visit(&mut self, _plan: &Expr) -> Result<bool, Self::Error> {
        Ok(true)
    }
}

pub trait ExprVisitorMut {
    type Error;

    fn pre_visit(&mut self, plan: &mut Expr) -> Result<bool, Self::Error>;
    fn post_visit(&mut self, _plan: &mut Expr) -> Result<bool, Self::Error> {
        Ok(true)
    }
}

macro_rules! generate_walker {
    (
        $(#[$meta:meta])*
        $func_name:ident, 
        $trait_name:ident, 
        $ref_type:tt
    ) => {
        $(#[$meta])*
        pub fn $func_name<V: $trait_name>(visitor: &mut V, expr: $ref_type Expr) -> Result<bool, V::Error> {
            if !visitor.pre_visit(expr)? {
                return Ok(false);
            }

            let recurse = match expr {
                Expr::Aggregate(AggregateExpr { expr, .. }) => $func_name(visitor, expr)?,
                // ... other match arms
            };

            if !recurse {
                return Ok(false);
            }

            if !visitor.post_visit(expr)? {
                return Ok(false);
            }

            Ok(true)
        }
    }
}

generate_walker!(
    /// A util function that traverses an AST [Expr] immutably in depth-first order.
    /// ...
    walk_expr, ExprVisitor, &
);

generate_walker!(
    /// A util function that traverses an AST [Expr] mutably in depth-first order.
    /// ...
    walk_expr_mut, ExprVisitorMut, &mut
);

Would you be interested in submitting PRs for these features?

Yes, I would be glad to submit PRs for these features.

Should we choose ExprMutVisitor or ExprVisitorMut?

ExprVisitorMut seems better because it aligns well with standard Rust conventions.

[evenyag]

Thanks! Feel free to have a try. I wonder if tt can match both & and &mut.

As an initial version, I think duplicating is acceptable as well. cc @sunng87 @waynexia

[sunng87]

Agree. I would prefer readability to tricky macros.

[scintillavoy]

I wonder if tt can match both & and &mut.

Thanks for pointing out this. &mut consists of two separate tokens so we should use a quantifier if we were to use macros.

As suggested, I prioritized readability and went with the manual duplication for now. I submitted a PR with the initial implementation here: #139, so please have a look!

[scintillavoy]

The mutable visitor has been implemented in #139. Thank you for the quick feedback throughout the review process.

Regarding the remaining points:

• Rewriter: I don't have an immediate need for this feature.
• Macro deduplication: This turned out to be more complex than expected, and I'm not sure it is worthwhile.

I wanted to mention them for completeness, but I don't plan to work on them myself for now. Feel free to close this issue or keep it open depending on how you prefer to track them.