risingwave_frontend/optimizer/plan_node/
plan_tree_node.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
15use smallvec::SmallVec;
16
17use super::PlanRef;
18use crate::utils::ColIndexMapping;
19
20/// The trait [`PlanNode`](super::PlanNode) really need about tree structure and used by optimizer
21/// framework. every plan node should impl it.
22///
23/// The trait [`PlanTreeNodeLeaf`], [`PlanTreeNodeUnary`] and [`PlanTreeNodeBinary`], is just
24/// special cases for [`PlanTreeNode`]. as long as you impl these trait for a plan node, we can
25/// easily impl the [`PlanTreeNode`] which is really need by framework with helper macros
26/// `impl_plan_tree_node_for_leaf`, `impl_plan_tree_node_for_unary` and
27/// `impl_plan_tree_node_for_binary`. We can't auto `impl PlanTreeNode for
28/// PlanTreeNodeLeaf/Unary/Binary`, because compiler doesn't know they are disjoint and thinks there
29/// are conflicting implementation.
30///
31/// And due to these three traits need not be used as dyn, it can return `Self` type, which is
32/// useful when implement rules and visitors. So we highly recommend not impl the [`PlanTreeNode`]
33/// trait directly, instead use these tree trait and impl [`PlanTreeNode`] use these helper
34/// macros.
35pub trait PlanTreeNode {
36    /// Get input nodes of the plan.
37    fn inputs(&self) -> SmallVec<[PlanRef; 2]>;
38
39    /// Clone the node with a list of new inputs.
40    fn clone_with_inputs(&self, inputs: &[PlanRef]) -> PlanRef;
41}
42
43/// See [`PlanTreeNode`](super)
44pub trait PlanTreeNodeLeaf: Clone {}
45
46/// See [`PlanTreeNode`](super)
47pub trait PlanTreeNodeUnary {
48    fn input(&self) -> PlanRef;
49    #[must_use]
50    fn clone_with_input(&self, input: PlanRef) -> Self;
51
52    /// Rewrites the plan node according to the schema change of its input node during rewriting.
53    ///
54    /// This function can be used to implement [`prune_col`](super::ColPrunable::prune_col) or
55    /// [`logical_rewrite_for_stream`](super::ToStream::logical_rewrite_for_stream)
56    #[must_use]
57    fn rewrite_with_input(
58        &self,
59        _input: PlanRef,
60        _input_col_change: ColIndexMapping,
61    ) -> (Self, ColIndexMapping)
62    where
63        Self: Sized,
64    {
65        unimplemented!()
66    }
67}
68
69/// See [`PlanTreeNode`](super)
70pub trait PlanTreeNodeBinary {
71    fn left(&self) -> PlanRef;
72    fn right(&self) -> PlanRef;
73
74    #[must_use]
75    fn clone_with_left_right(&self, left: PlanRef, right: PlanRef) -> Self;
76
77    /// Rewrites the plan node according to the schema change of its input nodes during rewriting.
78    ///
79    /// This function can be used to implement [`prune_col`](super::ColPrunable::prune_col) or
80    /// [`logical_rewrite_for_stream`](super::ToStream::logical_rewrite_for_stream)
81    #[must_use]
82    fn rewrite_with_left_right(
83        &self,
84        _left: PlanRef,
85        _left_col_change: ColIndexMapping,
86        _right: PlanRef,
87        _right_col_change: ColIndexMapping,
88    ) -> (Self, ColIndexMapping)
89    where
90        Self: Sized,
91    {
92        unimplemented!()
93    }
94}
95
96macro_rules! impl_plan_tree_node_for_leaf {
97    ($leaf_node_type:ident) => {
98        impl crate::optimizer::plan_node::PlanTreeNode for $leaf_node_type {
99            fn inputs(&self) -> smallvec::SmallVec<[crate::optimizer::PlanRef; 2]> {
100                smallvec::smallvec![]
101            }
102
103            /// Clone the node with a list of new inputs.
104            fn clone_with_inputs(
105                &self,
106                inputs: &[crate::optimizer::PlanRef],
107            ) -> crate::optimizer::PlanRef {
108                assert_eq!(inputs.len(), 0);
109                self.clone().into()
110            }
111        }
112    };
113}
114
115macro_rules! impl_plan_tree_node_for_unary {
116    ($unary_node_type:ident) => {
117        impl crate::optimizer::plan_node::PlanTreeNode for $unary_node_type {
118            fn inputs(&self) -> smallvec::SmallVec<[crate::optimizer::PlanRef; 2]> {
119                smallvec::smallvec![self.input()]
120            }
121
122            /// Clone the node with a list of new inputs.
123            fn clone_with_inputs(
124                &self,
125                inputs: &[crate::optimizer::PlanRef],
126            ) -> crate::optimizer::PlanRef {
127                assert_eq!(inputs.len(), 1);
128                self.clone_with_input(inputs[0].clone()).into()
129            }
130        }
131    };
132}
133
134macro_rules! impl_plan_tree_node_for_binary {
135    ($binary_node_type:ident) => {
136        impl crate::optimizer::plan_node::PlanTreeNode for $binary_node_type {
137            fn inputs(&self) -> smallvec::SmallVec<[crate::optimizer::PlanRef; 2]> {
138                smallvec::smallvec![self.left(), self.right()]
139            }
140
141            fn clone_with_inputs(
142                &self,
143                inputs: &[crate::optimizer::PlanRef],
144            ) -> crate::optimizer::PlanRef {
145                assert_eq!(inputs.len(), 2);
146                self.clone_with_left_right(inputs[0].clone(), inputs[1].clone())
147                    .into()
148            }
149        }
150    };
151}
risingwave_frontend/optimizer/plan_node/plan_tree_node.rs

risingwave_frontend/optimizer/plan_node/
plan_tree_node.rs
