feat(rust): ensure span uniqueness of oxc AST (#1927)

hyf0 · web-flow · commit 9f164c69af43 · 2024-08-10T08:25:43.000Z
### Description Rolldown uses `Span` as the way to indentify AST ndoes. Now the `inject` feature would generate improt statements with empty Span. They might be duplicated, we need ensure it's uniqueness. Contexts: - Closes #1875 - #1926 - #1899 (comment)
diff --git a/crates/rolldown/src/utils/ecma_visitors/mod.rs b/crates/rolldown/src/utils/ecma_visitors/mod.rs
@@ -0,0 +1,48 @@
+use oxc::{
+  ast::{visit::walk_mut, VisitMut},
+  span::{GetSpanMut, Span},
+};
+use rustc_hash::FxHashSet;
+
+/// Make sure there aren't any duplicate spans in the AST.
+pub struct EnsureSpanUniqueness {
+  // visited_spans: FxHashMap</* start */ u32, /* ends */ FxHashSet<u32>>,
+  visited_spans: FxHashSet<Span>,
+  next_unique_span_start: u32,
+}
+
+impl<'a> VisitMut<'a> for EnsureSpanUniqueness {
+  fn visit_program(&mut self, it: &mut oxc::ast::ast::Program<'a>) {
+    self.next_unique_span_start = it.span.end + 1;
+    walk_mut::walk_program(self, it);
+  }
+
+  // TODO: it's better use `visit_span`, but it's not implemented yet by oxc. https://github.com/oxc-project/oxc/issues/4799
+  fn visit_module_declaration(&mut self, it: &mut oxc::ast::ast::ModuleDeclaration<'a>) {
+    self.ensure_uniqueness(it.span_mut());
+    walk_mut::walk_module_declaration(self, it);
+  }
+}
+
+impl EnsureSpanUniqueness {
+  pub fn new() -> Self {
+    Self { visited_spans: FxHashSet::default(), next_unique_span_start: 1 }
+  }
+
+  fn ensure_uniqueness(&mut self, span: &mut Span) {
+    if self.visited_spans.contains(span) {
+      *span = self.generate_unique_span();
+    }
+    self.visited_spans.insert(*span);
+  }
+
+  fn generate_unique_span(&mut self) -> Span {
+    let mut span_candidate = Span::new(self.next_unique_span_start, self.next_unique_span_start);
+    while self.visited_spans.contains(&span_candidate) {
+      self.next_unique_span_start += 1;
+      span_candidate = Span::new(self.next_unique_span_start, self.next_unique_span_start);
+    }
+    debug_assert!(span_candidate.is_empty());
+    span_candidate
+  }
+}
diff --git a/crates/rolldown/src/utils/mod.rs b/crates/rolldown/src/utils/mod.rs
@@ -9,6 +9,7 @@ pub mod apply_inner_plugins;
 pub mod augment_chunk_hash;
 pub mod call_expression_ext;
 pub mod chunk;
+pub mod ecma_visitors;
 pub mod extract_hash_pattern;
 pub mod extract_meaningful_input_name_from_path;
 pub mod hash_placeholder;
diff --git a/crates/rolldown/src/utils/pre_process_ecma_ast.rs b/crates/rolldown/src/utils/pre_process_ecma_ast.rs
@@ -1,5 +1,6 @@
 use std::path::Path;
 
+use oxc::ast::VisitMut;
 use oxc::minifier::{
   CompressOptions, Compressor, ReplaceGlobalDefines, ReplaceGlobalDefinesConfig,
 };
@@ -10,6 +11,7 @@ use rolldown_ecmascript::EcmaAst;
 
 use crate::types::oxc_parse_type::OxcParseType;
 
+use super::ecma_visitors::EnsureSpanUniqueness;
 use super::tweak_ast_for_scanning::tweak_ast_for_scanning;
 
 // #[allow(clippy::match_same_arms)]: `OxcParseType::Tsx` will have special logic to deal with ts compared to `OxcParseType::Jsx`
@@ -62,6 +64,10 @@ pub fn pre_process_ecma_ast(
 
   tweak_ast_for_scanning(&mut ast);
 
+  ast.program.with_mut(|fields| {
+    EnsureSpanUniqueness::new().visit_program(fields.program);
+  });
+
   // We have to re-create the symbol table and scope tree after the transformation so far to make sure they are up-to-date.
   let (symbols, scopes) = ast.make_symbol_table_and_scope_tree();
 
