Skip to content

Commit 2f06d1e

Browse files
DunqingDunqing
committed
refactor(formatter): simplify AstNodeIterator and expand impl_ast_node_vec macro
1 parent 9c3931a commit 2f06d1e

File tree

1 file changed

+98
-180
lines changed

1 file changed

+98
-180
lines changed

‎crates/oxc_formatter/src/ast_nodes/iterator.rs‎

Lines changed: 98 additions & 180 deletions
Original file line numberDiff line numberDiff line change
@@ -12,42 +12,111 @@ use oxc_span::GetSpan;
1212

1313
use super::{AstNode, AstNodes};
1414

15+
/// Iterator for `AstNode<Vec<T>>`.
1516
pub struct AstNodeIterator<'a, T> {
1617
inner: std::iter::Peekable<std::slice::Iter<'a, T>>,
1718
parent: AstNodes<'a>,
1819
allocator: &'a Allocator,
20+
/// The `following_span_start` for the last element when there's no next element in this iterator.
21+
///
22+
/// This is essential for [`Comments::get_trailing_comments`] to correctly distinguish trailing
23+
/// comments from leading comments of the following sibling. When `following_span_start` is 0,
24+
/// comments after the last element are treated as its trailing comments. But when set to
25+
/// the next sibling's span start, `get_trailing_comments` can properly determine which
26+
/// comments belong to the current node vs the following sibling outside this iterator.
27+
///
28+
/// Example: For directives, without this field, comments between the last directive and
29+
/// first statement would be incorrectly treated as trailing comments of the directive,
30+
/// when they should be leading comments of the statement.
31+
///
32+
/// See [`Comments::get_trailing_comments`] in `crates/oxc_formatter/src/formatter/comments.rs`
33+
/// for the detailed handling logic.
34+
///
35+
/// [`Comments::get_trailing_comments`]: crate::formatter::Comments::get_trailing_comments
36+
following_span_start: u32,
37+
/// Function to compute `following_span_start` from the next element.
38+
get_following_span_start: fn(&T) -> u32,
39+
}
40+
41+
/// Custom span getter for Statement that handles decorated exports.
42+
/// <https://github.com/oxc-project/oxc/issues/10409>
43+
fn get_statement_span(stmt: &Statement<'_>) -> u32 {
44+
match stmt {
45+
Statement::ExportDefaultDeclaration(export) => {
46+
if let ExportDefaultDeclarationKind::ClassDeclaration(class) = &export.declaration
47+
&& let Some(decorator) = class.decorators.first()
48+
{
49+
min(decorator.span.start, export.span.start)
50+
} else {
51+
export.span.start
52+
}
53+
}
54+
Statement::ExportNamedDeclaration(export) => {
55+
if let Some(Declaration::ClassDeclaration(class)) = &export.declaration
56+
&& let Some(decorator) = class.decorators.first()
57+
{
58+
min(decorator.span.start, export.span.start)
59+
} else {
60+
export.span.start
61+
}
62+
}
63+
_ => stmt.span().start,
64+
}
1965
}
2066

2167
macro_rules! impl_ast_node_vec {
2268
($type:ty) => {
69+
impl_ast_node_vec!($type, false, |n: &$type| n.span().start);
70+
};
71+
($type:ty, has_following_span_in_the_last_item) => {
72+
impl_ast_node_vec!($type, true, |n: &$type| n.span().start);
73+
};
74+
($type:ty, $has_following_span_in_the_last_item:tt, $get_span:expr) => {
2375
impl<'a> AstNode<'a, Vec<'a, $type>> {
2476
pub fn iter(&self) -> AstNodeIterator<'a, $type> {
2577
AstNodeIterator {
2678
inner: self.inner.iter().peekable(),
2779
parent: self.parent,
2880
allocator: self.allocator,
81+
following_span_start: if $has_following_span_in_the_last_item {
82+
self.following_span_start
83+
} else {
84+
0
85+
},
86+
get_following_span_start: $get_span,
2987
}
3088
}
3189

3290
pub fn first(&self) -> Option<&'a AstNode<'a, $type>> {
91+
let following = if $has_following_span_in_the_last_item {
92+
self.following_span_start
93+
} else {
94+
0
95+
};
96+
let get_span: fn(&$type) -> u32 = $get_span;
3397
let mut inner_iter = self.inner.iter();
3498
self.allocator
3599
.alloc(inner_iter.next().map(|inner| AstNode {
36100
inner,
37101
parent: self.parent,
38102
allocator: self.allocator,
39-
following_span_start: inner_iter.next().map_or(0, |n| n.span().start),
103+
following_span_start: inner_iter.next().map_or(following, get_span),
40104
}))
41105
.as_ref()
42106
}
43107

44108
pub fn last(&self) -> Option<&'a AstNode<'a, $type>> {
109+
let following = if $has_following_span_in_the_last_item {
110+
self.following_span_start
111+
} else {
112+
0
113+
};
45114
self.allocator
46115
.alloc(self.inner.last().map(|inner| AstNode {
47116
inner,
48117
parent: self.parent,
49118
allocator: self.allocator,
50-
following_span_start: 0,
119+
following_span_start: following,
51120
}))
52121
.as_ref()
53122
}
@@ -57,12 +126,14 @@ macro_rules! impl_ast_node_vec {
57126
type Item = &'a AstNode<'a, $type>;
58127
fn next(&mut self) -> Option<Self::Item> {
59128
let allocator = self.allocator;
129+
let following = self.following_span_start;
130+
let get_span = self.get_following_span_start;
60131
allocator
61132
.alloc(self.inner.next().map(|inner| AstNode {
62133
parent: self.parent,
63134
inner,
64135
allocator,
65-
following_span_start: self.inner.peek().map_or(0, |n| n.span().start),
136+
following_span_start: self.inner.peek().map_or(following, |n| get_span(*n)),
66137
}))
67138
.as_ref()
68139
}
@@ -72,10 +143,16 @@ macro_rules! impl_ast_node_vec {
72143
type Item = &'a AstNode<'a, $type>;
73144
type IntoIter = AstNodeIterator<'a, $type>;
74145
fn into_iter(self) -> Self::IntoIter {
75-
AstNodeIterator::<$type> {
146+
AstNodeIterator {
76147
inner: self.inner.iter().peekable(),
77148
parent: self.parent,
78149
allocator: self.allocator,
150+
following_span_start: if $has_following_span_in_the_last_item {
151+
self.following_span_start
152+
} else {
153+
0
154+
},
155+
get_following_span_start: $get_span,
79156
}
80157
}
81158
}
@@ -90,6 +167,8 @@ macro_rules! impl_ast_node_vec_for_option {
90167
inner: self.inner.iter().peekable(),
91168
parent: self.parent,
92169
allocator: self.allocator,
170+
following_span_start: 0,
171+
get_following_span_start: |opt| opt.as_ref().map_or(0, |n| n.span().start),
93172
}
94173
}
95174

@@ -126,18 +205,14 @@ macro_rules! impl_ast_node_vec_for_option {
126205
type Item = &'a AstNode<'a, $type>;
127206
fn next(&mut self) -> Option<Self::Item> {
128207
let allocator = self.allocator;
208+
let following = self.following_span_start;
209+
let get_span = self.get_following_span_start;
129210
allocator
130-
.alloc(self.inner.next().map(|inner| {
131-
AstNode {
132-
parent: self.parent,
133-
inner,
134-
allocator,
135-
following_span_start: self
136-
.inner
137-
.peek()
138-
.and_then(|opt| opt.as_ref().map(|n| n.span().start))
139-
.unwrap_or(0),
140-
}
211+
.alloc(self.inner.next().map(|inner| AstNode {
212+
parent: self.parent,
213+
inner,
214+
allocator,
215+
following_span_start: self.inner.peek().map_or(following, |n| get_span(*n)),
141216
}))
142217
.as_ref()
143218
}
@@ -147,17 +222,18 @@ macro_rules! impl_ast_node_vec_for_option {
147222
type Item = &'a AstNode<'a, $type>;
148223
type IntoIter = AstNodeIterator<'a, $type>;
149224
fn into_iter(self) -> Self::IntoIter {
150-
AstNodeIterator::<$type> {
225+
AstNodeIterator {
151226
inner: self.inner.iter().peekable(),
152227
parent: self.parent,
153228
allocator: self.allocator,
229+
following_span_start: 0,
230+
get_following_span_start: |opt| opt.as_ref().map_or(0, |n| n.span().start),
154231
}
155232
}
156233
}
157234
};
158235
}
159236

160-
// Concrete implementations for all Vec types used in the AST
161237
impl_ast_node_vec!(Expression<'a>);
162238
impl_ast_node_vec!(ArrayExpressionElement<'a>);
163239
impl_ast_node_vec!(ObjectPropertyKind<'a>);
@@ -186,167 +262,9 @@ impl_ast_node_vec!(Decorator<'a>);
186262
impl_ast_node_vec_for_option!(Option<AssignmentTargetMaybeDefault<'a>>);
187263
impl_ast_node_vec_for_option!(Option<BindingPattern<'a>>);
188264

189-
// Manual implementation for Vec<Statement> because we have to handle span
190-
// for ExportDefaultDeclaration and ExportNamedDeclaration that may include
191-
// decorators.
265+
// Directive needs `following_span_start` to distinguish trailing comments from leading comments
266+
// of the first statement. See the struct field comment for `following_span_start` for details.
267+
impl_ast_node_vec!(Directive<'a>, has_following_span_in_the_last_item);
268+
// Custom get_span for Statement to handle decorated exports.
192269
// <https://github.com/oxc-project/oxc/issues/10409>
193-
impl<'a> AstNode<'a, Vec<'a, Statement<'a>>> {
194-
pub fn iter(&self) -> AstNodeIterator<'a, Statement<'a>> {
195-
AstNodeIterator {
196-
inner: self.inner.iter().peekable(),
197-
parent: self.parent,
198-
allocator: self.allocator,
199-
}
200-
}
201-
pub fn first(&self) -> Option<&'a AstNode<'a, Statement<'a>>> {
202-
let mut inner_iter = self.inner.iter();
203-
self.allocator
204-
.alloc(inner_iter.next().map(|inner| AstNode {
205-
inner,
206-
parent: self.parent,
207-
allocator: self.allocator,
208-
following_span_start: inner_iter.next().map_or(0, |n| n.span().start),
209-
}))
210-
.as_ref()
211-
}
212-
pub fn last(&self) -> Option<&'a AstNode<'a, Statement<'a>>> {
213-
self.allocator
214-
.alloc(self.inner.last().map(|inner| AstNode {
215-
inner,
216-
parent: self.parent,
217-
allocator: self.allocator,
218-
following_span_start: 0,
219-
}))
220-
.as_ref()
221-
}
222-
}
223-
impl<'a> Iterator for AstNodeIterator<'a, Statement<'a>> {
224-
type Item = &'a AstNode<'a, Statement<'a>>;
225-
fn next(&mut self) -> Option<Self::Item> {
226-
let allocator = self.allocator;
227-
allocator
228-
.alloc(self.inner.next().map(|inner| AstNode {
229-
parent: self.parent,
230-
inner,
231-
allocator,
232-
following_span_start: {
233-
match self.inner.peek() {
234-
// `@decorator export default class A {}`
235-
// Get the span start of the decorator.
236-
Some(Statement::ExportDefaultDeclaration(export)) => {
237-
if let ExportDefaultDeclarationKind::ClassDeclaration(class) =
238-
&export.declaration
239-
&& let Some(decorator) = class.decorators.first()
240-
{
241-
min(decorator.span.start, export.span.start)
242-
} else {
243-
export.span.start
244-
}
245-
}
246-
// `@decorator export class A {}`
247-
// Get the span start of the decorator.
248-
Some(Statement::ExportNamedDeclaration(export)) => {
249-
if let Some(Declaration::ClassDeclaration(class)) = &export.declaration
250-
&& let Some(decorator) = class.decorators.first()
251-
{
252-
min(decorator.span.start, export.span.start)
253-
} else {
254-
export.span.start
255-
}
256-
}
257-
Some(next) => next.span().start,
258-
None => 0,
259-
}
260-
},
261-
}))
262-
.as_ref()
263-
}
264-
}
265-
impl<'a> IntoIterator for &AstNode<'a, Vec<'a, Statement<'a>>> {
266-
type Item = &'a AstNode<'a, Statement<'a>>;
267-
type IntoIter = AstNodeIterator<'a, Statement<'a>>;
268-
fn into_iter(self) -> Self::IntoIter {
269-
AstNodeIterator::<Statement<'a>> {
270-
inner: self.inner.iter().peekable(),
271-
parent: self.parent,
272-
allocator: self.allocator,
273-
}
274-
}
275-
}
276-
277-
fn get_following_span_start_for_directive_parent(parent: &AstNodes<'_>) -> u32 {
278-
match parent {
279-
AstNodes::Program(program) => program.body().first().map_or(0, |n| n.span().start),
280-
AstNodes::FunctionBody(function_body) => {
281-
function_body.statements().first().map_or(0, |n| n.span().start)
282-
}
283-
AstNodes::TSModuleBlock(ts_module_block) => {
284-
ts_module_block.body().first().map_or(0, |n| n.span().start)
285-
}
286-
_ => 0,
287-
}
288-
}
289-
290-
// Manual implementation for Vec<Directive> because we need to handle
291-
// following_span for the last directive in Program.body.
292-
impl<'a> AstNode<'a, Vec<'a, Directive<'a>>> {
293-
pub fn iter(&self) -> AstNodeIterator<'a, Directive<'a>> {
294-
AstNodeIterator {
295-
inner: self.inner.iter().peekable(),
296-
parent: self.parent,
297-
allocator: self.allocator,
298-
}
299-
}
300-
pub fn first(&self) -> Option<&'a AstNode<'a, Directive<'a>>> {
301-
let mut inner_iter = self.inner.iter();
302-
self.allocator
303-
.alloc(inner_iter.next().map(|inner| AstNode {
304-
inner,
305-
parent: self.parent,
306-
allocator: self.allocator,
307-
following_span_start: inner_iter.next().map_or_else(
308-
|| get_following_span_start_for_directive_parent(&self.parent),
309-
|n| n.span().start,
310-
),
311-
}))
312-
.as_ref()
313-
}
314-
pub fn last(&self) -> Option<&'a AstNode<'a, Directive<'a>>> {
315-
self.allocator
316-
.alloc(self.inner.last().map(|inner| AstNode {
317-
inner,
318-
parent: self.parent,
319-
allocator: self.allocator,
320-
following_span_start: get_following_span_start_for_directive_parent(&self.parent),
321-
}))
322-
.as_ref()
323-
}
324-
}
325-
impl<'a> Iterator for AstNodeIterator<'a, Directive<'a>> {
326-
type Item = &'a AstNode<'a, Directive<'a>>;
327-
fn next(&mut self) -> Option<Self::Item> {
328-
let allocator = self.allocator;
329-
allocator
330-
.alloc(self.inner.next().map(|inner| AstNode {
331-
parent: self.parent,
332-
inner,
333-
allocator,
334-
following_span_start: self.inner.peek().map_or_else(
335-
|| get_following_span_start_for_directive_parent(&self.parent),
336-
|n| n.span().start,
337-
),
338-
}))
339-
.as_ref()
340-
}
341-
}
342-
impl<'a> IntoIterator for &AstNode<'a, Vec<'a, Directive<'a>>> {
343-
type Item = &'a AstNode<'a, Directive<'a>>;
344-
type IntoIter = AstNodeIterator<'a, Directive<'a>>;
345-
fn into_iter(self) -> Self::IntoIter {
346-
AstNodeIterator::<Directive<'a>> {
347-
inner: self.inner.iter().peekable(),
348-
parent: self.parent,
349-
allocator: self.allocator,
350-
}
351-
}
352-
}
270+
impl_ast_node_vec!(Statement<'a>, false, get_statement_span);

0 commit comments

Comments
 (0)