risingwave_sqlsmith/sqlreduce/
path.rs

1// Copyright 2025 RisingWave Labs
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
15//! Path-based AST navigation for SQL reduction.
16//!
17//! This module provides utilities for navigating and modifying SQL ASTs using
18//! path-based addressing for precise AST manipulation.
19
20use std::fmt;
21
22use risingwave_sqlparser::ast::*;
23use strum::EnumDiscriminants;
24
25/// Represents all possible AST field names that can be navigated.
26/// This provides compile-time safety for field access.
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum AstField {
29    // Statement fields
30    Query,
31    Name,
32    Columns,
33
34    // Query fields
35    Body,
36    With,
37    OrderBy,
38    Limit,
39    Offset,
40
41    // Select fields
42    Projection,
43    Selection,
44    From,
45    GroupBy,
46    Having,
47    Distinct,
48
49    // Expression fields
50    Left,
51    Right,
52    Operand,
53    ElseResult,
54    Subquery,
55    Inner,
56    Expr,
57    Low,
58    High,
59
60    // TableWithJoins fields
61    Relation,
62    Joins,
63
64    // Join fields
65    JoinOperator,
66
67    // SelectItem fields
68    Alias,
69
70    // OrderByExpr fields
71    Asc,
72    NullsFirst,
73
74    // With clause fields
75    CteTable,
76    Recursive,
77
78    // CTE fields
79    CteInner,
80
81    // TableFactor fields
82    Lateral,
83}
84
85impl fmt::Display for AstField {
86    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
87        let s = match self {
88            AstField::Query => "query",
89            AstField::Name => "name",
90            AstField::Columns => "columns",
91            AstField::Body => "body",
92            AstField::With => "with",
93            AstField::OrderBy => "order_by",
94            AstField::Limit => "limit",
95            AstField::Offset => "offset",
96            AstField::Projection => "projection",
97            AstField::Selection => "selection",
98            AstField::From => "from",
99            AstField::GroupBy => "group_by",
100            AstField::Having => "having",
101            AstField::Distinct => "distinct",
102            AstField::Left => "left",
103            AstField::Right => "right",
104            AstField::Operand => "operand",
105            AstField::ElseResult => "else_result",
106            AstField::Subquery => "subquery",
107            AstField::Inner => "inner",
108            AstField::Expr => "expr",
109            AstField::Low => "low",
110            AstField::High => "high",
111            AstField::Relation => "relation",
112            AstField::Joins => "joins",
113            AstField::JoinOperator => "join_operator",
114            AstField::Alias => "alias",
115            AstField::Asc => "asc",
116            AstField::NullsFirst => "nulls_first",
117            AstField::CteTable => "cte_tables",
118            AstField::Recursive => "recursive",
119            AstField::CteInner => "cte_inner",
120            AstField::Lateral => "lateral",
121        };
122        write!(f, "{}", s)
123    }
124}
125
126/// Represents a path component in an AST navigation path.
127/// Components that make up a path through the AST.
128/// Enables precise navigation to any AST node.
129#[derive(Debug, Clone, PartialEq, Eq)]
130pub enum PathComponent {
131    /// Field access by AST field enum (type-safe)
132    Field(AstField),
133    /// Array/Vec index access
134    Index(usize),
135}
136
137impl PathComponent {
138    /// Create a Field `PathComponent` from `AstField` enum (preferred)
139    pub fn field(field: AstField) -> Self {
140        PathComponent::Field(field)
141    }
142
143    /// Get the field name as owned String for compatibility
144    pub fn field_name(&self) -> Option<String> {
145        match self {
146            PathComponent::Field(field) => Some(field.to_string()),
147            PathComponent::Index(_) => None,
148        }
149    }
150
151    /// Get the `AstField` enum if this is a field
152    pub fn as_ast_field(&self) -> Option<&AstField> {
153        match self {
154            PathComponent::Field(field) => Some(field),
155            PathComponent::Index(_) => None,
156        }
157    }
158
159    // Convenience constructors for common fields
160    pub fn query() -> Self {
161        Self::field(AstField::Query)
162    }
163
164    pub fn body() -> Self {
165        Self::field(AstField::Body)
166    }
167
168    pub fn selection() -> Self {
169        Self::field(AstField::Selection)
170    }
171
172    pub fn projection() -> Self {
173        Self::field(AstField::Projection)
174    }
175
176    pub fn from_clause() -> Self {
177        Self::field(AstField::From)
178    }
179
180    pub fn group_by() -> Self {
181        Self::field(AstField::GroupBy)
182    }
183
184    pub fn having() -> Self {
185        Self::field(AstField::Having)
186    }
187
188    pub fn with_clause() -> Self {
189        Self::field(AstField::With)
190    }
191
192    pub fn order_by() -> Self {
193        Self::field(AstField::OrderBy)
194    }
195
196    pub fn left() -> Self {
197        Self::field(AstField::Left)
198    }
199
200    pub fn right() -> Self {
201        Self::field(AstField::Right)
202    }
203
204    pub fn operand() -> Self {
205        Self::field(AstField::Operand)
206    }
207
208    pub fn else_result() -> Self {
209        Self::field(AstField::ElseResult)
210    }
211
212    pub fn relation() -> Self {
213        Self::field(AstField::Relation)
214    }
215
216    pub fn joins() -> Self {
217        Self::field(AstField::Joins)
218    }
219
220    pub fn subquery() -> Self {
221        Self::field(AstField::Subquery)
222    }
223
224    pub fn cte_tables() -> Self {
225        Self::field(AstField::CteTable)
226    }
227
228    pub fn cte_inner() -> Self {
229        Self::field(AstField::CteInner)
230    }
231}
232
233impl fmt::Display for PathComponent {
234    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235        match self {
236            PathComponent::Field(field) => write!(f, ".{}", field),
237            PathComponent::Index(idx) => write!(f, "[{}]", idx),
238        }
239    }
240}
241
242/// A path through the AST for precise node identification.
243/// This allows us to precisely identify and modify any node in the tree.
244pub type AstPath = Vec<PathComponent>;
245
246// Note: Display implementation moved to helper function to avoid orphan rule
247pub fn display_ast_path(path: &AstPath) -> String {
248    path.iter().map(|c| c.to_string()).collect::<String>()
249}
250
251/// Represents a node in the AST that can be navigated and modified.
252/// This is a simplified representation focusing on the most commonly
253/// reduced SQL constructs.
254#[derive(Debug, Clone, EnumDiscriminants)]
255#[strum_discriminants(derive(strum::Display))]
256#[strum_discriminants(name(AstNodeType))]
257pub enum AstNode {
258    Statement(Statement),
259    Query(Box<Query>),
260    Select(Box<Select>),
261    Expr(Expr),
262    SelectItem(SelectItem),
263    TableWithJoins(TableWithJoins),
264    Join(Join),
265    TableFactor(TableFactor),
266    OrderByExpr(OrderByExpr),
267    With(With),
268    Cte(Cte),
269    ExprList(Vec<Expr>),
270    SelectItemList(Vec<SelectItem>),
271    TableList(Vec<TableWithJoins>),
272    JoinList(Vec<Join>),
273    OrderByList(Vec<OrderByExpr>),
274    CteList(Vec<Cte>),
275    Option(Option<Box<AstNode>>),
276}
277
278impl AstNode {
279    /// Navigate to a child node using a path component.
280    pub fn get_child(&self, component: &PathComponent) -> Option<AstNode> {
281        match (self, component) {
282            // Statement navigation
283            (AstNode::Statement(Statement::Query(query)), PathComponent::Field(field))
284                if *field == AstField::Query =>
285            {
286                Some(AstNode::Query(query.clone()))
287            }
288
289            (
290                AstNode::Statement(Statement::CreateView { query, .. }),
291                PathComponent::Field(field),
292            ) if *field == AstField::Query => Some(AstNode::Query(query.clone())),
293
294            // More CreateView field access
295            (AstNode::Statement(Statement::CreateView { .. }), PathComponent::Field(field)) => {
296                match field {
297                    AstField::Name => None,    // ObjectName is complex, skip for now
298                    AstField::Columns => None, // Column list is complex, skip for now
299                    _ => None,
300                }
301            }
302
303            // Query navigation - enhanced
304            (AstNode::Query(query), PathComponent::Field(field)) => match field {
305                AstField::Body => match &query.body {
306                    SetExpr::Select(select) => Some(AstNode::Select(select.clone())),
307                    SetExpr::Query(subquery) => Some(AstNode::Query(subquery.clone())),
308                    SetExpr::SetOperation { left, .. } => {
309                        // For SetOperation, recursively handle the left side
310                        match left.as_ref() {
311                            SetExpr::Select(select) => Some(AstNode::Select(select.clone())),
312                            SetExpr::Query(subquery) => Some(AstNode::Query(subquery.clone())),
313                            _ => None,
314                        }
315                    }
316                    _ => None,
317                },
318                AstField::With => query.with.as_ref().map(|w| AstNode::With(w.clone())),
319                AstField::OrderBy => {
320                    if query.order_by.is_empty() {
321                        None
322                    } else {
323                        Some(AstNode::OrderByList(query.order_by.clone()))
324                    }
325                }
326                AstField::Limit => query.limit.as_ref().map(|e| AstNode::Expr(e.clone())),
327                AstField::Offset => None, // offset is Option<String>, not navigable as expression
328                _ => None,
329            },
330
331            // Select navigation - enhanced
332            (AstNode::Select(select), PathComponent::Field(field)) => match field {
333                AstField::Projection => {
334                    if select.projection.is_empty() {
335                        None
336                    } else {
337                        Some(AstNode::SelectItemList(select.projection.clone()))
338                    }
339                }
340                AstField::Selection => select.selection.as_ref().map(|e| AstNode::Expr(e.clone())),
341                AstField::From => {
342                    if select.from.is_empty() {
343                        None
344                    } else {
345                        Some(AstNode::TableList(select.from.clone()))
346                    }
347                }
348                AstField::GroupBy => {
349                    if select.group_by.is_empty() {
350                        None
351                    } else {
352                        Some(AstNode::ExprList(select.group_by.clone()))
353                    }
354                }
355                AstField::Having => select.having.as_ref().map(|e| AstNode::Expr(e.clone())),
356                AstField::Distinct => None, // Distinct is an enum, handle separately if needed
357                _ => None,
358            },
359
360            // List navigation
361            (AstNode::SelectItemList(items), PathComponent::Index(idx)) => items
362                .get(*idx)
363                .map(|item| AstNode::SelectItem(item.clone())),
364            (AstNode::ExprList(exprs), PathComponent::Index(idx)) => {
365                exprs.get(*idx).map(|expr| AstNode::Expr(expr.clone()))
366            }
367            (AstNode::TableList(tables), PathComponent::Index(idx)) => tables
368                .get(*idx)
369                .map(|table| AstNode::TableWithJoins(table.clone())),
370            (AstNode::OrderByList(orders), PathComponent::Index(idx)) => orders
371                .get(*idx)
372                .map(|order| AstNode::OrderByExpr(order.clone())),
373
374            // Expression navigation (for pullup operations) - enhanced
375            (AstNode::Expr(expr), PathComponent::Field(field)) => match (expr, field) {
376                (Expr::BinaryOp { left, .. }, AstField::Left) => Some(AstNode::Expr(*left.clone())),
377                (Expr::BinaryOp { right, .. }, AstField::Right) => {
378                    Some(AstNode::Expr(*right.clone()))
379                }
380                (Expr::Case { operand, .. }, AstField::Operand) => {
381                    operand.as_ref().map(|e| AstNode::Expr(*e.clone()))
382                }
383                (Expr::Case { else_result, .. }, AstField::ElseResult) => {
384                    else_result.as_ref().map(|e| AstNode::Expr(*e.clone()))
385                }
386                (Expr::Exists(subquery), AstField::Subquery) => {
387                    Some(AstNode::Query(subquery.clone()))
388                }
389                (Expr::Subquery(subquery), AstField::Subquery) => {
390                    Some(AstNode::Query(subquery.clone()))
391                }
392                (Expr::Function(_func), AstField::Name) => None, // ObjectName is complex
393                (Expr::Nested(inner), AstField::Inner) => Some(AstNode::Expr(*inner.clone())),
394                (Expr::UnaryOp { expr, .. }, AstField::Expr) => Some(AstNode::Expr(*expr.clone())),
395                (Expr::Cast { expr, .. }, AstField::Expr) => Some(AstNode::Expr(*expr.clone())),
396                (Expr::IsNull(expr), AstField::Expr) => Some(AstNode::Expr(*expr.clone())),
397                (Expr::IsNotNull(expr), AstField::Expr) => Some(AstNode::Expr(*expr.clone())),
398                (Expr::Between { expr, .. }, AstField::Expr) => Some(AstNode::Expr(*expr.clone())),
399                (Expr::Between { low, .. }, AstField::Low) => Some(AstNode::Expr(*low.clone())),
400                (Expr::Between { high, .. }, AstField::High) => Some(AstNode::Expr(*high.clone())),
401                _ => None,
402            },
403
404            // TableWithJoins navigation
405            (AstNode::TableWithJoins(table_with_joins), PathComponent::Field(field)) => match field
406            {
407                AstField::Relation => Some(AstNode::TableFactor(table_with_joins.relation.clone())),
408                AstField::Joins => {
409                    if !table_with_joins.joins.is_empty() {
410                        Some(AstNode::JoinList(table_with_joins.joins.clone()))
411                    } else {
412                        None
413                    }
414                }
415                _ => None,
416            },
417
418            // Join navigation
419            (AstNode::Join(join), PathComponent::Field(field)) => match field {
420                AstField::Relation => Some(AstNode::TableFactor(join.relation.clone())),
421                AstField::JoinOperator => None, // JoinOperator is simple enum, skip navigation
422                _ => None,
423            },
424
425            // TableFactor navigation
426            (AstNode::TableFactor(table_factor), PathComponent::Field(field)) => {
427                match (table_factor, field) {
428                    (TableFactor::Table { .. }, _) => None, // Table references are terminal
429                    (TableFactor::Derived { subquery, .. }, AstField::Subquery) => {
430                        Some(AstNode::Query(subquery.clone()))
431                    }
432                    (TableFactor::TableFunction { .. }, _) => None, // Function calls are complex
433                    _ => None,
434                }
435            }
436
437            // JoinList navigation (for Vec<Join>)
438            (AstNode::JoinList(joins), PathComponent::Index(idx)) => {
439                joins.get(*idx).map(|join| AstNode::Join(join.clone()))
440            }
441
442            // SelectItem navigation
443            (AstNode::SelectItem(select_item), PathComponent::Field(field)) => {
444                match (select_item, field) {
445                    (SelectItem::UnnamedExpr(expr), AstField::Expr) => {
446                        Some(AstNode::Expr(expr.clone()))
447                    }
448                    (SelectItem::ExprWithAlias { expr, .. }, AstField::Expr) => {
449                        Some(AstNode::Expr(expr.clone()))
450                    }
451                    (SelectItem::QualifiedWildcard(..), _) => None,
452                    (SelectItem::Wildcard(..), _) => None,
453                    _ => None,
454                }
455            }
456
457            // OrderByExpr navigation
458            (AstNode::OrderByExpr(order_by), PathComponent::Field(field)) => match field {
459                AstField::Expr => Some(AstNode::Expr(order_by.expr.clone())),
460                AstField::Asc => None,        // Boolean, not navigable
461                AstField::NullsFirst => None, // Option<bool>, not navigable
462                _ => None,
463            },
464
465            // With clause navigation
466            (AstNode::With(with_clause), PathComponent::Field(field)) => match field {
467                AstField::CteTable => {
468                    if with_clause.cte_tables.is_empty() {
469                        None
470                    } else {
471                        Some(AstNode::CteList(with_clause.cte_tables.clone()))
472                    }
473                }
474                AstField::Recursive => None, // Boolean, not navigable
475                _ => None,
476            },
477
478            // CTE list navigation
479            (AstNode::CteList(ctes), PathComponent::Index(idx)) => {
480                if *idx < ctes.len() {
481                    Some(AstNode::Cte(ctes[*idx].clone()))
482                } else {
483                    None
484                }
485            }
486
487            // CTE navigation
488            (AstNode::Cte(cte), PathComponent::Field(field)) => match field {
489                AstField::Alias => None, // TableAlias is complex, skip for now
490                AstField::CteInner => match &cte.cte_inner {
491                    CteInner::Query(query) => Some(AstNode::Query(query.clone())),
492                    CteInner::ChangeLog(_) => None, // ObjectName is complex, skip for now
493                },
494                _ => None,
495            },
496
497            _ => {
498                // Add debug logging for unmatched cases
499                tracing::debug!(
500                    "get_child: No match for {:?} with component {:?}",
501                    std::mem::discriminant(self),
502                    component
503                );
504                None
505            }
506        }
507    }
508
509    /// Set a child node using a path component.
510    /// Returns a new `AstNode` with the modification applied.
511    pub fn set_child(
512        &self,
513        component: &PathComponent,
514        new_child: Option<AstNode>,
515    ) -> Option<AstNode> {
516        match (self, component) {
517            // Handle Statement::Query (root query statements)
518            (AstNode::Statement(Statement::Query(_)), PathComponent::Field(field))
519                if *field == AstField::Query =>
520            {
521                if let Some(AstNode::Query(new_query)) = new_child {
522                    Some(AstNode::Statement(Statement::Query(new_query)))
523                } else {
524                    None
525                }
526            }
527
528            (
529                AstNode::Statement(Statement::CreateView {
530                    name,
531                    columns,
532                    query: _,
533                    or_replace,
534                    materialized,
535                    if_not_exists,
536                    emit_mode,
537                    with_options,
538                }),
539                PathComponent::Field(field),
540            ) => match field {
541                AstField::Query => {
542                    if let Some(AstNode::Query(new_query)) = new_child {
543                        let new_stmt = Statement::CreateView {
544                            or_replace: *or_replace,
545                            materialized: *materialized,
546                            if_not_exists: *if_not_exists,
547                            name: name.clone(),
548                            columns: columns.clone(),
549                            query: new_query,
550                            emit_mode: emit_mode.clone(),
551                            with_options: with_options.clone(),
552                        };
553                        Some(AstNode::Statement(new_stmt))
554                    } else {
555                        None
556                    }
557                }
558                _ => None,
559            },
560
561            // Query field modifications
562            (AstNode::Query(query), PathComponent::Field(field)) => {
563                let mut new_query = (**query).clone();
564                match field {
565                    AstField::Body => {
566                        if let Some(AstNode::Select(select)) = new_child {
567                            new_query.body = SetExpr::Select(select);
568                            Some(AstNode::Query(Box::new(new_query)))
569                        } else {
570                            None
571                        }
572                    }
573                    AstField::OrderBy => {
574                        if let Some(AstNode::OrderByList(orders)) = new_child {
575                            new_query.order_by = orders;
576                        } else {
577                            new_query.order_by = vec![];
578                        }
579                        Some(AstNode::Query(Box::new(new_query)))
580                    }
581                    AstField::Limit => {
582                        new_query.limit = new_child.and_then(|n| match n {
583                            AstNode::Expr(e) => Some(e),
584                            _ => None,
585                        });
586                        Some(AstNode::Query(Box::new(new_query)))
587                    }
588                    AstField::With => {
589                        new_query.with = new_child.and_then(|n| match n {
590                            AstNode::With(w) => Some(w),
591                            _ => None,
592                        });
593                        Some(AstNode::Query(Box::new(new_query)))
594                    }
595                    _ => None,
596                }
597            }
598
599            // Select field modifications
600            (AstNode::Select(select), PathComponent::Field(field)) => {
601                let mut new_select = (**select).clone();
602                match field {
603                    AstField::Selection => {
604                        new_select.selection = new_child.and_then(|n| match n {
605                            AstNode::Expr(e) => Some(e),
606                            _ => None,
607                        });
608                    }
609                    AstField::Having => {
610                        new_select.having = new_child.and_then(|n| match n {
611                            AstNode::Expr(e) => Some(e),
612                            _ => None,
613                        });
614                    }
615                    AstField::Projection => {
616                        if let Some(AstNode::SelectItemList(items)) = new_child {
617                            new_select.projection = items;
618                        } else {
619                            // Remove projection by setting to empty vec (SELECT without columns is invalid, but we try it)
620                            new_select.projection = vec![];
621                        }
622                    }
623                    AstField::From => {
624                        if let Some(AstNode::TableList(tables)) = new_child {
625                            new_select.from = tables;
626                        } else {
627                            // Remove FROM clause by setting to empty vec
628                            new_select.from = vec![];
629                        }
630                    }
631                    AstField::GroupBy => {
632                        if let Some(AstNode::ExprList(exprs)) = new_child {
633                            new_select.group_by = exprs;
634                        } else {
635                            // Remove GROUP BY by setting to empty vec
636                            new_select.group_by = vec![];
637                        }
638                    }
639                    _ => return None,
640                }
641                Some(AstNode::Select(Box::new(new_select)))
642            }
643
644            // List modifications
645            (AstNode::SelectItemList(items), PathComponent::Index(idx)) => {
646                let mut new_items = items.clone();
647                if *idx < new_items.len() {
648                    if let Some(AstNode::SelectItem(item)) = new_child {
649                        new_items[*idx] = item;
650                    } else {
651                        new_items.remove(*idx);
652                    }
653                    Some(AstNode::SelectItemList(new_items))
654                } else {
655                    None
656                }
657            }
658
659            (AstNode::ExprList(exprs), PathComponent::Index(idx)) => {
660                let mut new_exprs = exprs.clone();
661                if *idx < new_exprs.len() {
662                    if let Some(AstNode::Expr(expr)) = new_child {
663                        new_exprs[*idx] = expr;
664                    } else {
665                        new_exprs.remove(*idx);
666                    }
667                    Some(AstNode::ExprList(new_exprs))
668                } else {
669                    None
670                }
671            }
672
673            (AstNode::TableList(tables), PathComponent::Index(idx)) => {
674                let mut new_tables = tables.clone();
675                if *idx < new_tables.len() {
676                    if let Some(AstNode::TableWithJoins(table)) = new_child {
677                        new_tables[*idx] = table;
678                    } else {
679                        new_tables.remove(*idx);
680                    }
681                    Some(AstNode::TableList(new_tables))
682                } else {
683                    None
684                }
685            }
686
687            // TableWithJoins modifications
688            (AstNode::TableWithJoins(table_with_joins), PathComponent::Field(field)) => {
689                match field {
690                    AstField::Relation => {
691                        if let Some(AstNode::TableFactor(new_relation)) = new_child {
692                            Some(AstNode::TableWithJoins(TableWithJoins {
693                                relation: new_relation,
694                                joins: table_with_joins.joins.clone(),
695                            }))
696                        } else {
697                            None
698                        }
699                    }
700                    AstField::Joins => {
701                        if let Some(AstNode::JoinList(new_joins)) = new_child {
702                            Some(AstNode::TableWithJoins(TableWithJoins {
703                                relation: table_with_joins.relation.clone(),
704                                joins: new_joins,
705                            }))
706                        } else {
707                            // Allow removing joins by setting to empty list
708                            Some(AstNode::TableWithJoins(TableWithJoins {
709                                relation: table_with_joins.relation.clone(),
710                                joins: vec![],
711                            }))
712                        }
713                    }
714                    _ => None,
715                }
716            }
717
718            // Join modifications
719            (AstNode::Join(join), PathComponent::Field(field)) => match field {
720                AstField::Relation => {
721                    if let Some(AstNode::TableFactor(new_relation)) = new_child {
722                        Some(AstNode::Join(Join {
723                            relation: new_relation,
724                            join_operator: join.join_operator.clone(),
725                        }))
726                    } else {
727                        None
728                    }
729                }
730                _ => None,
731            },
732
733            // TableFactor modifications
734            (AstNode::TableFactor(table_factor), PathComponent::Field(field)) => {
735                match (table_factor, field) {
736                    (TableFactor::Derived { lateral, alias, .. }, AstField::Subquery) => {
737                        if let Some(AstNode::Query(new_subquery)) = new_child {
738                            Some(AstNode::TableFactor(TableFactor::Derived {
739                                lateral: *lateral,
740                                subquery: new_subquery,
741                                alias: alias.clone(),
742                            }))
743                        } else {
744                            None
745                        }
746                    }
747                    _ => None,
748                }
749            }
750
751            // JoinList modifications
752            (AstNode::JoinList(joins), PathComponent::Index(idx)) => {
753                let mut new_joins = joins.clone();
754                if *idx < new_joins.len() {
755                    if let Some(AstNode::Join(new_join)) = new_child {
756                        new_joins[*idx] = new_join;
757                    } else {
758                        new_joins.remove(*idx);
759                    }
760                    Some(AstNode::JoinList(new_joins))
761                } else {
762                    None
763                }
764            }
765
766            (AstNode::OrderByList(orders), PathComponent::Index(idx)) => {
767                let mut new_orders = orders.clone();
768                if *idx < new_orders.len() {
769                    if let Some(AstNode::OrderByExpr(order)) = new_child {
770                        new_orders[*idx] = order;
771                    } else {
772                        new_orders.remove(*idx);
773                    }
774                    Some(AstNode::OrderByList(new_orders))
775                } else {
776                    None
777                }
778            }
779
780            // Expression field modifications
781            (AstNode::Expr(expr), PathComponent::Field(field)) => match (expr, field) {
782                (Expr::BinaryOp { left: _, op, right }, AstField::Left) => {
783                    if let Some(AstNode::Expr(new_left)) = new_child {
784                        Some(AstNode::Expr(Expr::BinaryOp {
785                            left: Box::new(new_left),
786                            op: op.clone(),
787                            right: right.clone(),
788                        }))
789                    } else {
790                        None
791                    }
792                }
793                (Expr::BinaryOp { left, op, right: _ }, AstField::Right) => {
794                    if let Some(AstNode::Expr(new_right)) = new_child {
795                        Some(AstNode::Expr(Expr::BinaryOp {
796                            left: left.clone(),
797                            op: op.clone(),
798                            right: Box::new(new_right),
799                        }))
800                    } else {
801                        None
802                    }
803                }
804                (Expr::Nested(_), AstField::Inner) => {
805                    if let Some(AstNode::Expr(new_inner)) = new_child {
806                        Some(AstNode::Expr(Expr::Nested(Box::new(new_inner))))
807                    } else {
808                        None
809                    }
810                }
811                _ => None,
812            },
813
814            // SelectItem field modifications
815            (AstNode::SelectItem(select_item), PathComponent::Field(field)) => {
816                match (select_item, field) {
817                    (SelectItem::UnnamedExpr(_), AstField::Expr) => {
818                        if let Some(AstNode::Expr(new_expr)) = new_child {
819                            Some(AstNode::SelectItem(SelectItem::UnnamedExpr(new_expr)))
820                        } else {
821                            None
822                        }
823                    }
824                    (SelectItem::ExprWithAlias { alias, .. }, AstField::Expr) => {
825                        if let Some(AstNode::Expr(new_expr)) = new_child {
826                            Some(AstNode::SelectItem(SelectItem::ExprWithAlias {
827                                expr: new_expr,
828                                alias: alias.clone(),
829                            }))
830                        } else {
831                            None
832                        }
833                    }
834                    _ => None,
835                }
836            }
837
838            // OrderByExpr field modifications
839            (AstNode::OrderByExpr(order_by), PathComponent::Field(field)) => match field {
840                AstField::Expr => {
841                    if let Some(AstNode::Expr(new_expr)) = new_child {
842                        Some(AstNode::OrderByExpr(OrderByExpr {
843                            expr: new_expr,
844                            asc: order_by.asc,
845                            nulls_first: order_by.nulls_first,
846                        }))
847                    } else {
848                        None
849                    }
850                }
851                _ => None,
852            },
853
854            // With clause modifications
855            (AstNode::With(with_clause), PathComponent::Field(field)) => match field {
856                AstField::CteTable => {
857                    if let Some(AstNode::CteList(new_ctes)) = new_child {
858                        let mut new_with = with_clause.clone();
859                        new_with.cte_tables = new_ctes;
860                        Some(AstNode::With(new_with))
861                    } else {
862                        // Remove all CTEs by setting to empty vec
863                        let mut new_with = with_clause.clone();
864                        new_with.cte_tables = vec![];
865                        Some(AstNode::With(new_with))
866                    }
867                }
868                _ => None,
869            },
870
871            // CTE list modifications
872            (AstNode::CteList(ctes), PathComponent::Index(idx)) => {
873                if let Some(AstNode::Cte(new_cte)) = new_child {
874                    if *idx < ctes.len() {
875                        let mut new_ctes = ctes.clone();
876                        new_ctes[*idx] = new_cte;
877                        Some(AstNode::CteList(new_ctes))
878                    } else {
879                        None
880                    }
881                } else {
882                    None
883                }
884            }
885
886            // CTE modifications
887            (AstNode::Cte(cte), PathComponent::Field(field)) => match field {
888                AstField::CteInner => {
889                    if let Some(AstNode::Query(new_query)) = new_child {
890                        let mut new_cte = cte.clone();
891                        new_cte.cte_inner = CteInner::Query(new_query);
892                        Some(AstNode::Cte(new_cte))
893                    } else {
894                        None
895                    }
896                }
897                _ => None,
898            },
899
900            _ => {
901                // Add debug logging for unmatched cases
902                tracing::debug!(
903                    "set_child: No match for {} ({:?}) with component {:?}",
904                    AstNodeType::from(self),
905                    std::mem::discriminant(self),
906                    component
907                );
908                None
909            }
910        }
911    }
912
913    /// Convert back to a Statement if this is the root node.
914    pub fn to_statement(&self) -> Option<Statement> {
915        match self {
916            AstNode::Statement(stmt) => Some(stmt.clone()),
917            _ => None,
918        }
919    }
920}
921
922/// Navigate to a node in the AST using the given path.
923/// Enables precise AST node retrieval.
924pub fn get_node_at_path(root: &AstNode, path: &AstPath) -> Option<AstNode> {
925    let mut current = root.clone();
926    for component in path {
927        current = current.get_child(component)?;
928    }
929    Some(current)
930}
931
932/// Set a node in the AST at the given path.
933/// Enables precise AST node modification.
934pub fn set_node_at_path(
935    root: &AstNode,
936    path: &AstPath,
937    new_node: Option<AstNode>,
938) -> Option<AstNode> {
939    if path.is_empty() {
940        return new_node;
941    }
942
943    let result = root.clone();
944    let mut current_path = Vec::new();
945
946    // Navigate to the parent of the target node
947    for component in &path[..path.len() - 1] {
948        current_path.push(component.clone());
949    }
950
951    // Get the parent node
952    let parent = get_node_at_path(&result, &current_path)?;
953
954    // Apply the modification to the parent
955    let modified_parent = parent.set_child(&path[path.len() - 1], new_node)?;
956
957    // Now we need to set this modified parent back in the tree
958    if current_path.is_empty() {
959        Some(modified_parent)
960    } else {
961        set_node_at_path(&result, &current_path, Some(modified_parent))
962    }
963}
964
965/// Helper function to get a child node and recurse if it exists.
966fn explore_child_field(
967    node: &AstNode,
968    field: AstField,
969    current_path: &AstPath,
970    paths: &mut Vec<AstPath>,
971) {
972    let field_component = PathComponent::field(field);
973    let child_path = [current_path.clone(), vec![field_component.clone()]].concat();
974    let relative_path = vec![field_component];
975
976    if let Some(child_node) = get_node_at_path(node, &relative_path) {
977        // Collect child paths but don't add them yet (for outer-first ordering)
978        let child_paths = enumerate_reduction_paths(&child_node, child_path);
979        paths.extend(child_paths);
980    }
981}
982
983/// Calculate the depth of a path (number of components).
984/// Used for outer-first ordering: shallower paths (outer queries) come first.
985fn path_depth(path: &AstPath) -> usize {
986    path.len()
987}
988
989/// Enumerate all interesting paths in the AST for reduction.
990/// Systematically discovers all reducible AST locations.
991pub fn enumerate_reduction_paths(node: &AstNode, current_path: AstPath) -> Vec<AstPath> {
992    let mut paths = vec![current_path.clone()];
993
994    tracing::debug!(
995        "Enumerating paths for node {:?} at path {}",
996        get_node_type_name(node),
997        display_ast_path(&current_path)
998    );
999
1000    match node {
1001        AstNode::Statement(Statement::Query(_)) => {
1002            explore_child_field(node, AstField::Query, &current_path, &mut paths);
1003        }
1004
1005        AstNode::Statement(Statement::CreateView { .. }) => {
1006            explore_child_field(node, AstField::Query, &current_path, &mut paths);
1007        }
1008
1009        AstNode::Query(_query) => {
1010            explore_child_field(node, AstField::Body, &current_path, &mut paths);
1011            explore_child_field(node, AstField::With, &current_path, &mut paths);
1012            explore_child_field(node, AstField::OrderBy, &current_path, &mut paths);
1013        }
1014
1015        AstNode::Select(_) => {
1016            explore_child_field(node, AstField::Projection, &current_path, &mut paths);
1017            explore_child_field(node, AstField::From, &current_path, &mut paths);
1018            explore_child_field(node, AstField::GroupBy, &current_path, &mut paths);
1019
1020            // For optional fields, just add the path if they exist
1021            let selection_path = [current_path.clone(), vec![PathComponent::selection()]].concat();
1022            if get_node_at_path(node, &vec![PathComponent::selection()]).is_some() {
1023                paths.push(selection_path);
1024            }
1025
1026            let having_path = [current_path.clone(), vec![PathComponent::having()]].concat();
1027            if get_node_at_path(node, &vec![PathComponent::having()]).is_some() {
1028                paths.push(having_path);
1029            }
1030        }
1031
1032        // For lists, enumerate individual elements
1033        AstNode::SelectItemList(items) => {
1034            for i in 0..items.len() {
1035                let item_path = [current_path.clone(), vec![PathComponent::Index(i)]].concat();
1036                paths.push(item_path);
1037            }
1038        }
1039
1040        AstNode::ExprList(exprs) => {
1041            for i in 0..exprs.len() {
1042                let expr_path = [current_path.clone(), vec![PathComponent::Index(i)]].concat();
1043                paths.push(expr_path.clone());
1044                // Also descend into expressions for pullup opportunities
1045                let relative_path = vec![PathComponent::Index(i)];
1046                if let Some(expr_node) = get_node_at_path(node, &relative_path) {
1047                    paths.extend(enumerate_reduction_paths(&expr_node, expr_path));
1048                }
1049            }
1050        }
1051
1052        AstNode::TableList(tables) => {
1053            for i in 0..tables.len() {
1054                let table_path = [current_path.clone(), vec![PathComponent::Index(i)]].concat();
1055                paths.push(table_path.clone());
1056                // Recursively explore TableWithJoins
1057                if let Some(table_node) = node.get_child(&PathComponent::Index(i)) {
1058                    paths.extend(enumerate_reduction_paths(&table_node, table_path));
1059                }
1060            }
1061        }
1062
1063        // TableWithJoins path enumeration - key for JOIN reduction
1064        AstNode::TableWithJoins(_) => {
1065            explore_child_field(node, AstField::Relation, &current_path, &mut paths);
1066            explore_child_field(node, AstField::Joins, &current_path, &mut paths);
1067        }
1068
1069        // Join path enumeration
1070        AstNode::Join(_) => {
1071            explore_child_field(node, AstField::Relation, &current_path, &mut paths);
1072        }
1073
1074        // TableFactor path enumeration
1075        AstNode::TableFactor(_) => {
1076            // For Derived tables (subqueries), explore the subquery
1077            explore_child_field(node, AstField::Subquery, &current_path, &mut paths);
1078        }
1079
1080        // JoinList path enumeration
1081        AstNode::JoinList(joins) => {
1082            for i in 0..joins.len() {
1083                let join_path = [current_path.clone(), vec![PathComponent::Index(i)]].concat();
1084                paths.push(join_path.clone());
1085                // Recursively explore Join
1086                if let Some(join_node) = node.get_child(&PathComponent::Index(i)) {
1087                    paths.extend(enumerate_reduction_paths(&join_node, join_path));
1088                }
1089            }
1090        }
1091
1092        AstNode::OrderByList(orders) => {
1093            for i in 0..orders.len() {
1094                let order_path = [current_path.clone(), vec![PathComponent::Index(i)]].concat();
1095                paths.push(order_path);
1096            }
1097        }
1098
1099        // For expressions, look for pullup opportunities
1100        AstNode::Expr(expr) => match expr {
1101            Expr::BinaryOp { .. } => {
1102                let left_path = [current_path.clone(), vec![PathComponent::left()]].concat();
1103                let right_path = [current_path.clone(), vec![PathComponent::right()]].concat();
1104                paths.push(left_path);
1105                paths.push(right_path);
1106            }
1107            Expr::Case {
1108                operand,
1109                else_result,
1110                ..
1111            } => {
1112                if operand.is_some() {
1113                    let operand_path =
1114                        [current_path.clone(), vec![PathComponent::operand()]].concat();
1115                    paths.push(operand_path);
1116                }
1117                if else_result.is_some() {
1118                    let else_path =
1119                        [current_path.clone(), vec![PathComponent::else_result()]].concat();
1120                    paths.push(else_path);
1121                }
1122            }
1123            Expr::Exists(_) => {
1124                explore_child_field(node, AstField::Subquery, &current_path, &mut paths);
1125            }
1126            Expr::Subquery(_) => {
1127                explore_child_field(node, AstField::Subquery, &current_path, &mut paths);
1128            }
1129            _ => {}
1130        },
1131
1132        // WITH clause enumeration
1133        AstNode::With(_) => {
1134            explore_child_field(node, AstField::CteTable, &current_path, &mut paths);
1135        }
1136
1137        // CTE list enumeration
1138        AstNode::CteList(ctes) => {
1139            for i in 0..ctes.len() {
1140                let cte_path = [current_path.clone(), vec![PathComponent::Index(i)]].concat();
1141                paths.push(cte_path.clone());
1142
1143                // Recursively enumerate paths within each CTE
1144                if let Some(cte_node) = get_node_at_path(node, &vec![PathComponent::Index(i)]) {
1145                    paths.extend(enumerate_reduction_paths(&cte_node, cte_path));
1146                }
1147            }
1148        }
1149
1150        // CTE enumeration
1151        AstNode::Cte(_) => {
1152            explore_child_field(node, AstField::CteInner, &current_path, &mut paths);
1153        }
1154
1155        _ => {}
1156    }
1157
1158    // Sort paths by depth (outer-first): shallower paths come first
1159    // This ensures outer queries are reduced before inner subqueries
1160    // For example: SELECT (SELECT ...) will reduce the outer SELECT first
1161    paths.sort_by_key(path_depth);
1162
1163    paths
1164}
1165
1166/// Convert a Statement to an `AstNode` for path-based operations.
1167pub fn statement_to_ast_node(stmt: &Statement) -> AstNode {
1168    AstNode::Statement(stmt.clone())
1169}
1170
1171/// Extract a Statement from an `AstNode`.
1172pub fn ast_node_to_statement(node: &AstNode) -> Option<Statement> {
1173    match node {
1174        AstNode::Statement(stmt) => Some(stmt.clone()),
1175        AstNode::Query(query) => Some(Statement::Query(Box::new(query.as_ref().clone()))),
1176        _ => None,
1177    }
1178}
1179
1180/// Get a human-readable name for an AST node type.
1181pub fn get_node_type_name(node: &AstNode) -> String {
1182    AstNodeType::from(node).to_string()
1183}
1184
1185#[cfg(test)]
1186mod tests {
1187    use risingwave_sqlparser::parser::Parser;
1188
1189    use super::*;
1190
1191    #[test]
1192    fn test_path_enumeration() {
1193        let sql = "SELECT a FROM b;";
1194        let parsed = Parser::parse_sql(sql).expect("Failed to parse SQL");
1195        let stmt = &parsed[0];
1196        let ast_node = statement_to_ast_node(stmt);
1197
1198        let paths = enumerate_reduction_paths(&ast_node, vec![]);
1199
1200        // Should have at least a few paths for this simple query
1201        assert!(paths.len() >= 3);
1202        println!("Found {} paths for simple query", paths.len());
1203    }
1204
1205    #[test]
1206    fn test_create_materialized_view() {
1207        let sql = "CREATE MATERIALIZED VIEW stream_query AS SELECT min((tumble_0.c5 + tumble_0.c5)) AS col_0, false AS col_1, TIME '07:30:48' AS col_2, tumble_0.c14 AS col_3 FROM tumble(alltypes1, alltypes1.c11, INTERVAL '10') AS tumble_0 WHERE tumble_0.c1 GROUP BY tumble_0.c10, tumble_0.c9, tumble_0.c13, tumble_0.c1, tumble_0.c16, tumble_0.c14, tumble_0.c5, tumble_0.c8;";
1208        let parsed = Parser::parse_sql(sql).expect("Failed to parse SQL");
1209        let stmt = &parsed[0];
1210        let ast_node = statement_to_ast_node(stmt);
1211
1212        let paths = enumerate_reduction_paths(&ast_node, vec![]);
1213
1214        println!("Found {} paths for CREATE MATERIALIZED VIEW", paths.len());
1215        for (i, path) in paths.iter().enumerate() {
1216            println!("Path {}: {}", i, display_ast_path(path));
1217        }
1218
1219        // Should have paths for the SELECT statement inside the MV
1220        assert!(!paths.is_empty());
1221
1222        // Should be able to get some node
1223        assert!(get_node_at_path(&ast_node, &paths[0]).is_some());
1224    }
1225}
risingwave_sqlsmith/sqlreduce/path.rs

risingwave_sqlsmith/sqlreduce/
path.rs
