feat: generate module namespace binding on demand (#126)

* feat: generate module namespace binding on demand

* Fix comments

crates/rolldown/src/bundler/bundle/bundle.rs

@@ -66,9 +66,7 @@ impl<'a> Bundle<'a> {
       for module_id in chunk.modules.iter().copied() {
         match &self.graph.modules[module_id] {
           Module::Normal(module) => {
-            let linking_info = &self.graph.linking_infos[module_id];
-
-            for stmt_info in module.stmt_infos.iter().chain(linking_info.facade_stmt_infos.iter()) {
+            for stmt_info in module.stmt_infos.iter() {
               for declared in &stmt_info.declared_symbols {
                 // TODO: pass debug_assert!(self.graph.symbols.get(*declared).chunk_id.is_none());
                 // FIXME: I don't think this is correct, even though the assigned chunk_id is the same as the current chunk_id.
@@ -80,6 +78,10 @@ impl<'a> Bundle<'a> {
                 self.graph.symbols.get_mut(*declared).chunk_id = Some(chunk_id);
               }
 
+              if !stmt_info.is_included {
+                continue;
+              }
+
               for referenced in &stmt_info.referenced_symbols {
                 let canonical_ref = self.graph.symbols.get_canonical_ref(*referenced);
                 chunk_meta_imports.insert(canonical_ref);
@@ -105,6 +107,7 @@ impl<'a> Bundle<'a> {
         }
       }
     }
+
     for (chunk_id, chunk) in chunk_graph.chunks.iter_mut_enumerated() {
       let chunk_meta_imports = &chunk_meta_imports_vec[chunk_id];
       for import_ref in chunk_meta_imports.iter().copied() {

crates/rolldown/src/bundler/graph/linker.rs

@@ -1,14 +1,14 @@
 use index_vec::IndexVec;
 use oxc::span::Atom;
 use rolldown_common::{
-  ExportsKind, ImportKind, LocalOrReExport, ModuleId, StmtInfo, SymbolRef, WrapKind,
+  ExportsKind, ImportKind, LocalOrReExport, ModuleId, StmtInfo, StmtInfoId, SymbolRef, WrapKind,
 };
 use rustc_hash::FxHashMap;
 
 use super::{graph::Graph, symbols::NamespaceAlias};
 use crate::bundler::{
   graph::symbols::Symbols,
-  module::{normal_module::Resolution, Module},
+  module::{normal_module::Resolution, Module, NormalModule},
 };
 
 /// Store the linking info for module
@@ -20,7 +20,6 @@ pub struct LinkingInfo {
   pub facade_stmt_infos: Vec<StmtInfo>,
   pub resolved_exports: FxHashMap<Atom, SymbolRef>,
   pub resolved_star_exports: Vec<ModuleId>,
-  pub is_symbol_for_namespace_referenced: bool,
 }
 
 pub type LinkingInfoVec = IndexVec<ModuleId, LinkingInfo>;
@@ -34,6 +33,97 @@ impl<'graph> Linker<'graph> {
     Self { graph }
   }
 
+  fn include_statements(&mut self) {
+    use rayon::prelude::*;
+    struct Context<'a> {
+      graph: &'a Graph,
+      is_included_vec: &'a mut IndexVec<ModuleId, IndexVec<StmtInfoId, bool>>,
+    }
+
+    fn include_symbol(ctx: &mut Context, symbol_ref: SymbolRef) {
+      let mut canonical_ref = ctx.graph.symbols.par_get_canonical_ref(symbol_ref);
+      let canonical_ref_module = &ctx.graph.modules[canonical_ref.owner];
+      let canonical_ref_symbol = ctx.graph.symbols.get(canonical_ref);
+      if let Some(namespace_alias) = &canonical_ref_symbol.namespace_alias {
+        canonical_ref = namespace_alias.namespace_ref;
+      }
+      let Module::Normal(canonical_ref_module) = canonical_ref_module else {
+        return;
+      };
+      canonical_ref_module
+        .stmt_infos
+        .declared_stmts_by_symbol(&canonical_ref)
+        .iter()
+        .copied()
+        .for_each(|stmt_info_id| {
+          include_statement(ctx, canonical_ref_module, stmt_info_id);
+        });
+    }
+
+    fn include_statement(ctx: &mut Context, module: &NormalModule, stmt_info_id: StmtInfoId) {
+      let is_included = &mut ctx.is_included_vec[module.id][stmt_info_id];
+      if *is_included {
+        return;
+      }
+
+      // include the statement itself
+      *is_included = true;
+
+      let stmt_info = module.stmt_infos.get(stmt_info_id);
+
+      // include statements that are referenced by this statement
+      stmt_info.declared_symbols.iter().chain(stmt_info.referenced_symbols.iter()).for_each(
+        |symbol_ref| {
+          include_symbol(ctx, *symbol_ref);
+        },
+      );
+    }
+
+    let mut is_included_vec: IndexVec<ModuleId, IndexVec<StmtInfoId, bool>> = self
+      .graph
+      .modules
+      .iter()
+      .map(|m| match m {
+        Module::Normal(m) => {
+          m.stmt_infos.iter().map(|_| false).collect::<IndexVec<StmtInfoId, _>>()
+        }
+        Module::External(_) => IndexVec::default(),
+      })
+      .collect::<IndexVec<ModuleId, _>>();
+
+    let context = &mut Context { graph: self.graph, is_included_vec: &mut is_included_vec };
+
+    for module in &self.graph.modules {
+      match module {
+        Module::Normal(module) => {
+          let mut stmt_infos = module.stmt_infos.iter_enumerated();
+          // Skip the first one, because it's the namespace variable declaration.
+          // We want to include it on demand.
+          stmt_infos.next();
+          // Since we won't implement tree shaking, we just include all statements.
+          stmt_infos.for_each(|(stmt_info_id, _)| {
+            include_statement(context, module, stmt_info_id);
+          });
+          if module.is_entry {
+            let linking_info = &self.graph.linking_infos[module.id];
+            linking_info.resolved_exports.values().for_each(|symbol_ref| {
+              include_symbol(context, *symbol_ref);
+            });
+          }
+        }
+        Module::External(_) => {}
+      }
+    }
+    self.graph.modules.iter_mut().par_bridge().for_each(|module| {
+      let Module::Normal(module) = module else {
+        return;
+      };
+      is_included_vec[module.id].iter_enumerated().for_each(|(stmt_info_id, is_included)| {
+        module.stmt_infos.get_mut(stmt_info_id).is_included = *is_included;
+      });
+    });
+  }
+
   pub fn link(&mut self) {
     // Here take the symbols to avoid borrow graph and mut borrow graph at same time
     let mut symbols = std::mem::take(&mut self.graph.symbols);
@@ -68,7 +158,19 @@ impl<'graph> Linker<'graph> {
 
     // Set the symbols back and add linker modules to graph
     self.graph.symbols = symbols;
+
+    // FIXME: should move `linking_info.facade_stmt_infos` into a separate field
+    for (id, linking_info) in linking_infos.iter_mut_enumerated() {
+      std::mem::take(&mut linking_info.facade_stmt_infos).into_iter().for_each(|info| {
+        if let Module::Normal(module) = &mut self.graph.modules[id] {
+          module.stmt_infos.add_stmt_info(info);
+        }
+      });
+    }
+
     self.graph.linking_infos = linking_infos;
+
+    self.include_statements();
   }
 
   #[allow(clippy::too_many_lines)]
@@ -163,9 +265,6 @@ impl<'graph> Linker<'graph> {
     linking_infos: &mut LinkingInfoVec,
   ) {
     let linking_info = &mut linking_infos[target];
-    if let Module::Normal(module) = &self.graph.modules[target] {
-      module.initialize_namespace(linking_info);
-    }
     if linking_info.wrap_symbol.is_some() {
       return;
     }
@@ -195,32 +294,10 @@ impl<'graph> Linker<'graph> {
   }
 
   #[allow(clippy::needless_collect)]
-  fn mark_extra_symbols(&mut self, symbols: &mut Symbols, linking_infos: &mut LinkingInfoVec) {
-    // Determine if the namespace symbol need to be generated
+  fn mark_extra_symbols(&mut self, symbols: &mut Symbols, _linking_infos: &mut LinkingInfoVec) {
     for importer_id in &self.graph.sorted_modules {
       let importer = &self.graph.modules[*importer_id];
 
-      if let Module::Normal(importer) = importer {
-        let has_reexport_all_from_cjs = importer.get_star_exports_modules().any(|importee| matches!(&self.graph.modules[importee], Module::Normal(m) if m.exports_kind == ExportsKind::CommonJs));
-        if has_reexport_all_from_cjs {
-          self.graph.modules[*importer_id]
-            .mark_symbol_for_namespace_referenced(&mut linking_infos[*importer_id]);
-        }
-      }
-
-      importer
-        .import_records()
-        .iter()
-        .filter_map(|rec| {
-          ((rec.is_import_namespace || matches!(rec.kind, ImportKind::Require))
-            && rec.resolved_module.is_valid())
-          .then_some(rec.resolved_module)
-        })
-        .for_each(|importee| {
-          self.graph.modules[importee]
-            .mark_symbol_for_namespace_referenced(&mut linking_infos[importee]);
-        });
-
       // Create symbols for external module
       let mut extra_symbols = vec![];
       match importer {
@@ -306,9 +383,9 @@ impl<'graph> Linker<'graph> {
               local_binding.namespace_alias =
                 Some(NamespaceAlias { property_name: exported.clone(), namespace_ref: ns_ref });
               linking_info.facade_stmt_infos.push(StmtInfo {
-                stmt_idx: None,
                 declared_symbols: vec![local_binding_ref],
                 referenced_symbols: vec![ns_ref],
+                ..Default::default()
               });
               linking_info.resolved_exports.insert(exported.clone(), local_binding_ref);
             }
@@ -327,9 +404,9 @@ impl<'graph> Linker<'graph> {
             local_binding.namespace_alias =
               Some(NamespaceAlias { property_name: exported.clone(), namespace_ref: ns_ref });
             linking_info.facade_stmt_infos.push(StmtInfo {
-              stmt_idx: None,
               declared_symbols: vec![local_binding_ref],
               referenced_symbols: vec![ns_ref],
+              ..Default::default()
             });
           }
         });

crates/rolldown/src/bundler/graph/symbols.rs

@@ -42,7 +42,7 @@ impl SymbolMap {
     }
   }
 
-  pub fn create_symbol(&mut self, name: Atom) -> SymbolId {
+  pub fn _create_symbol(&mut self, name: Atom) -> SymbolId {
     if is_reserved_word(&name) {
       self.names.push(format!("_{name}").into())
     } else {

crates/rolldown/src/bundler/module/mod.rs

@@ -10,10 +10,7 @@ use string_wizard::MagicString;
 
 use self::external_module::ExternalModule;
 
-use super::{
-  chunk_graph::ChunkGraph,
-  graph::{graph::Graph, linker::LinkingInfo},
-};
+use super::{chunk_graph::ChunkGraph, graph::graph::Graph};
 
 pub type ModuleVec = IndexVec<ModuleId, Module>;
 
@@ -66,13 +63,6 @@ impl Module {
     }
   }
 
-  pub fn mark_symbol_for_namespace_referenced(&self, linking_info: &mut LinkingInfo) {
-    match self {
-      Self::Normal(m) => m.initialize_namespace(linking_info),
-      Self::External(_) => linking_info.is_symbol_for_namespace_referenced = true,
-    }
-  }
-
   pub fn render(&self, ctx: ModuleRenderContext) -> Option<MagicString<'_>> {
     match self {
       Self::Normal(m) => m.render(ctx),

crates/rolldown/src/bundler/module/normal_module.rs

@@ -7,7 +7,7 @@ use oxc::{
 };
 use rolldown_common::{
   ExportsKind, ImportRecord, ImportRecordId, LocalOrReExport, ModuleId, ModuleType, NamedImport,
-  ResourceId, StmtInfo, SymbolRef, WrapKind,
+  ResourceId, StmtInfo, StmtInfos, SymbolRef, WrapKind,
 };
 use rolldown_oxc::OxcProgram;
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -30,17 +30,18 @@ pub struct NormalModule {
   pub is_entry: bool,
   pub resource_id: ResourceId,
   pub module_type: ModuleType,
+  pub namespace_symbol: SymbolRef,
   pub ast: OxcProgram,
   pub named_imports: FxHashMap<SymbolId, NamedImport>,
   pub named_exports: FxHashMap<Atom, LocalOrReExport>,
-  pub stmt_infos: Vec<StmtInfo>,
+  /// `stmt_infos[0]` represents the namespace binding statement
+  pub stmt_infos: StmtInfos,
   pub import_records: IndexVec<ImportRecordId, ImportRecord>,
   pub imports: FxHashMap<Span, ImportRecordId>,
   // [[StarExportEntries]] in https://tc39.es/ecma262/#sec-source-text-module-records
   pub star_exports: Vec<ImportRecordId>,
   pub exports_kind: ExportsKind,
   pub scope: ScopeTree,
-  pub namespace_symbol: SymbolRef,
   pub default_export_symbol: Option<SymbolId>,
 }
 
@@ -56,6 +57,10 @@ pub enum Resolution {
 }
 
 impl NormalModule {
+  pub fn is_namespace_referenced(&self) -> bool {
+    self.stmt_infos.get(0.into()).is_included
+  }
+
   #[allow(clippy::needless_pass_by_value)]
   pub fn render(&self, ctx: ModuleRenderContext<'_>) -> Option<MagicString<'static>> {
     // FIXME: should not clone
@@ -87,15 +92,6 @@ impl NormalModule {
     }
   }
 
-  pub fn initialize_namespace(&self, self_linking_info: &mut LinkingInfo) {
-    if !self_linking_info.is_symbol_for_namespace_referenced {
-      self_linking_info.is_symbol_for_namespace_referenced = true;
-      self_linking_info
-        .facade_stmt_infos
-        .push(StmtInfo { declared_symbols: vec![self.namespace_symbol], ..Default::default() });
-    }
-  }
-
   // https://tc39.es/ecma262/#sec-getexportednames
   pub fn get_exported_names<'modules>(
     &'modules self,
@@ -325,7 +321,6 @@ impl NormalModule {
       self_linking_info
         .facade_stmt_infos
         .push(StmtInfo { declared_symbols: vec![(self.id, symbol).into()], ..Default::default() });
-      self.initialize_namespace(self_linking_info);
     }
   }
 

crates/rolldown/src/bundler/module/normal_module_builder.rs

@@ -5,13 +5,11 @@ use oxc::{
 };
 use rolldown_common::{
   ExportsKind, ImportRecord, ImportRecordId, LocalOrReExport, ModuleId, ModuleType, NamedImport,
-  ResourceId, StmtInfo, SymbolRef,
+  ResourceId, StmtInfos, SymbolRef,
 };
 use rolldown_oxc::OxcProgram;
 use rustc_hash::FxHashMap;
 
-use crate::bundler::graph::symbols::SymbolMap;
-
 use super::NormalModule;
 
 #[derive(Debug, Default)]
@@ -21,7 +19,7 @@ pub struct NormalModuleBuilder {
   pub ast: Option<OxcProgram>,
   pub named_imports: Option<FxHashMap<SymbolId, NamedImport>>,
   pub named_exports: Option<FxHashMap<Atom, LocalOrReExport>>,
-  pub stmt_infos: Option<Vec<StmtInfo>>,
+  pub stmt_infos: Option<StmtInfos>,
   pub import_records: Option<IndexVec<ImportRecordId, ImportRecord>>,
   pub imports: Option<FxHashMap<Span, ImportRecordId>>,
   pub star_exports: Option<Vec<ImportRecordId>>,
@@ -34,12 +32,6 @@ pub struct NormalModuleBuilder {
 }
 
 impl NormalModuleBuilder {
-  pub fn initialize_namespace_binding(&mut self, symbol_table: &mut SymbolMap) {
-    let name = format!("{}_ns", self.path.as_ref().unwrap().generate_unique_name());
-    let symbol_ref: SymbolRef = (self.id.unwrap(), symbol_table.create_symbol(name.into())).into();
-    self.namespace_symbol = Some(symbol_ref);
-  }
-
   pub fn build(self) -> NormalModule {
     NormalModule {
       exec_order: u32::MAX,