diff --git a/crates/oxc_ast/src/lib.rs b/crates/oxc_ast/src/lib.rs index 60e2ebfc86210..293aa4024dedf 100644 --- a/crates/oxc_ast/src/lib.rs +++ b/crates/oxc_ast/src/lib.rs @@ -42,7 +42,7 @@ pub use num_bigint::BigUint; pub use crate::{ ast_builder::AstBuilder, ast_kind::{AstKind, AstType}, - trivia::{Comment, CommentKind, Trivias, TriviasMap}, + trivia::{Comment, CommentKind, SortedComments, Trivias}, visit::{Visit, VisitMut}, }; diff --git a/crates/oxc_ast/src/trivia.rs b/crates/oxc_ast/src/trivia.rs index b9353b9c51e19..57333c6029a68 100644 --- a/crates/oxc_ast/src/trivia.rs +++ b/crates/oxc_ast/src/trivia.rs @@ -1,8 +1,8 @@ //! Trivias such as comments and irregular whitespaces use std::{ - collections::btree_map::{BTreeMap, Range}, - ops::{Deref, RangeBounds}, + iter::FusedIterator, + ops::{Bound, Deref, RangeBounds}, sync::Arc, }; @@ -37,15 +37,16 @@ impl CommentKind { } } -pub type TriviasMap = BTreeMap; +/// Sorted set of unique trivia comments, in ascending order by starting position. +pub type SortedComments = Box<[(u32, Comment)]>; #[derive(Debug, Clone, Default)] pub struct Trivias(Arc); #[derive(Debug, Default)] pub struct TriviasImpl { - /// Keyed by span.start - comments: TriviasMap, + /// Unique comments, ordered by increasing span-start. + comments: SortedComments, irregular_whitespaces: Vec, } @@ -59,7 +60,7 @@ impl Deref for Trivias { } impl Trivias { - pub fn new(comments: TriviasMap, irregular_whitespaces: Vec) -> Trivias { + pub fn new(comments: SortedComments, irregular_whitespaces: Vec) -> Trivias { Self(Arc::new(TriviasImpl { comments, irregular_whitespaces })) } @@ -67,18 +68,116 @@ impl Trivias { self.comments.iter().map(|(start, comment)| (comment.kind, Span::new(*start, comment.end))) } - pub fn comments_range(&self, range: R) -> Range<'_, u32, Comment> + pub fn comments_range(&self, range: R) -> CommentsRange<'_> where R: RangeBounds, { - self.comments.range(range) + CommentsRange::new(&self.comments, range.start_bound().cloned(), range.end_bound().cloned()) } pub fn has_comments_between(&self, span: Span) -> bool { - self.comments.range(span.start..span.end).count() > 0 + self.comments_range(span.start..span.end).count() > 0 } pub fn irregular_whitespaces(&self) -> &Vec { &self.irregular_whitespaces } } + +/// Double-ended iterator over a range of comments, by starting position. +pub struct CommentsRange<'a> { + comments: &'a [(u32, Comment)], + range: (Bound, Bound), + current_start: usize, + current_end: usize, +} + +impl<'a> CommentsRange<'a> { + fn new(comments: &'a [(u32, Comment)], start: Bound, end: Bound) -> Self { + // Directly skip all comments that are already known to start + // outside the requested range. + let partition_start = { + let range_start = match start { + Bound::Unbounded => 0, + Bound::Included(x) => x, + Bound::Excluded(x) => x.saturating_add(1), + }; + comments.partition_point(|(start, _)| *start < range_start) + }; + let partition_end = { + let range_end = match end { + Bound::Unbounded => u32::MAX, + Bound::Included(x) => x, + Bound::Excluded(x) => x.saturating_sub(1), + }; + comments.partition_point(|(start, _)| *start <= range_end) + }; + Self { + comments, + range: (start, end), + current_start: partition_start, + current_end: partition_end, + } + } +} + +impl<'c> Iterator for CommentsRange<'c> { + type Item = (&'c u32, &'c Comment); + + fn next(&mut self) -> Option { + if self.current_start < self.current_end { + for (start, comment) in &self.comments[self.current_start..self.current_end] { + self.current_start = self.current_start.saturating_add(1); + if self.range.contains(start) { + return Some((start, comment)); + } + } + } + None + } + + fn size_hint(&self) -> (usize, Option) { + let max_remaining = self.current_end.saturating_sub(self.current_start); + (0, Some(max_remaining)) + } +} + +impl<'c> DoubleEndedIterator for CommentsRange<'c> { + fn next_back(&mut self) -> Option { + if self.current_start < self.current_end { + for (start, comment) in self.comments[self.current_start..self.current_end].iter().rev() + { + self.current_end = self.current_end.saturating_sub(1); + if self.range.contains(start) { + return Some((start, comment)); + } + } + } + None + } +} + +impl FusedIterator for CommentsRange<'_> {} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn test_comments_range() { + let comments: SortedComments = vec![ + (0, Comment { end: 4, kind: CommentKind::SingleLine }), + (5, Comment { end: 9, kind: CommentKind::SingleLine }), + (10, Comment { end: 13, kind: CommentKind::SingleLine }), + (14, Comment { end: 17, kind: CommentKind::SingleLine }), + (18, Comment { end: 23, kind: CommentKind::SingleLine }), + ] + .into_boxed_slice(); + let full_len = comments.len(); + let trivias = Trivias::new(comments, vec![]); + assert_eq!(trivias.comments_range(..).count(), full_len); + assert_eq!(trivias.comments_range(1..).count(), full_len.saturating_sub(1)); + assert_eq!(trivias.comments_range(..18).count(), full_len.saturating_sub(1)); + assert_eq!(trivias.comments_range(..=18).count(), full_len); + } +} diff --git a/crates/oxc_parser/src/lexer/trivia_builder.rs b/crates/oxc_parser/src/lexer/trivia_builder.rs index 91fde3627cb58..696cb23e24daa 100644 --- a/crates/oxc_parser/src/lexer/trivia_builder.rs +++ b/crates/oxc_parser/src/lexer/trivia_builder.rs @@ -1,16 +1,18 @@ -use oxc_ast::{Comment, CommentKind, Trivias, TriviasMap}; +use oxc_ast::{Comment, CommentKind, SortedComments, Trivias}; use oxc_span::Span; #[derive(Debug, Default)] pub struct TriviaBuilder { - // Duplicated comments can be added from rewind, use `BTreeMap` to ensure uniqueness + // NOTE(lucab): This is a set of unique comments. Duplicated + // comments could be generated in case of rewind; they are + // filtered out at insertion time. comments: Vec<(u32, Comment)>, irregular_whitespaces: Vec, } impl TriviaBuilder { pub fn build(self) -> Trivias { - let comments = TriviasMap::from_iter(self.comments); + let comments = SortedComments::from_iter(self.comments); Trivias::new(comments, self.irregular_whitespaces) } @@ -26,9 +28,9 @@ impl TriviaBuilder { fn add_comment(&mut self, start: u32, comment: Comment) { // The comments array is an ordered vec, only add the comment if its not added before, - // to avoid situations where the parser needs to rewind and reinsert the comment. - if let Some(comment) = self.comments.last_mut() { - if start <= comment.0 { + // to avoid situations where the parser needs to rewind and tries to reinsert the comment. + if let Some((last_start, _)) = self.comments.last() { + if start <= *last_start { return; } }