rspack_plugin_javascript/plugin/
infer_async_modules_plugin.rs

1use linked_hash_set::LinkedHashSet;
2use rayon::prelude::*;
3use rspack_collections::{IdentifierMap, IdentifierSet};
4use rspack_core::{
5  Compilation, CompilationFinishModules, DependencyType, Logger, ModuleGraph, ModuleIdentifier,
6  Plugin,
7  incremental::{IncrementalPasses, Mutation, Mutations},
8};
9use rspack_error::Result;
10use rspack_hook::{plugin, plugin_hook};
11
12#[plugin]
13#[derive(Debug, Default)]
14pub struct InferAsyncModulesPlugin;
15
16#[plugin_hook(CompilationFinishModules for InferAsyncModulesPlugin)]
17async fn finish_modules(&self, compilation: &mut Compilation) -> Result<()> {
18  if let Some(mutations) = compilation
19    .incremental
20    .mutations_read(IncrementalPasses::INFER_ASYNC_MODULES)
21  {
22    mutations
23      .iter()
24      .filter_map(|mutation| {
25        if let Mutation::ModuleRemove { module } = mutation {
26          Some(module)
27        } else {
28          None
29        }
30      })
31      .for_each(|module| {
32        compilation.async_modules_artifact.remove(module);
33      });
34  }
35
36  let module_graph = compilation.get_module_graph();
37  let modules = module_graph.modules();
38  let mut sync_modules = LinkedHashSet::default();
39  let mut async_modules = LinkedHashSet::default();
40  for (module_identifier, module) in modules {
41    let build_meta = module.build_meta();
42    if build_meta.has_top_level_await {
43      async_modules.insert(module_identifier);
44    } else {
45      sync_modules.insert(module_identifier);
46    }
47  }
48
49  let mut mutations = compilation
50    .incremental
51    .mutations_writeable()
52    .then(Mutations::default);
53
54  set_sync_modules(compilation, sync_modules, &mut mutations);
55  set_async_modules(compilation, async_modules, &mut mutations);
56
57  if compilation
58    .incremental
59    .mutations_readable(IncrementalPasses::INFER_ASYNC_MODULES)
60    && let Some(mutations) = &mutations
61  {
62    let logger = compilation.get_logger("rspack.incremental.inferAsyncModules");
63    logger.log(format!(
64      "{} modules are updated by set_async",
65      mutations.len()
66    ));
67  }
68
69  if let Some(compilation_mutations) = compilation.incremental.mutations_write()
70    && let Some(mutations) = mutations
71  {
72    compilation_mutations.extend(mutations);
73  }
74
75  Ok(())
76}
77
78fn set_sync_modules(
79  compilation: &mut Compilation,
80  modules: LinkedHashSet<ModuleIdentifier>,
81  mutations: &mut Option<Mutations>,
82) {
83  let module_graph = compilation.get_module_graph();
84  let outgoing_connections = modules
85    .iter()
86    .par_bridge()
87    .map(|mid| {
88      (
89        *mid,
90        module_graph
91          .get_outgoing_connections(mid)
92          .filter_map(|con| module_graph.module_identifier_by_dependency_id(&con.dependency_id))
93          .filter(|&out| out != mid)
94          .copied()
95          .collect::<Vec<_>>(),
96      )
97    })
98    .collect::<IdentifierMap<_>>();
99
100  let mut queue = modules;
101  while let Some(module) = queue.pop_front() {
102    if outgoing_connections
103      .get(&module)
104      .cloned()
105      .unwrap_or_else(|| {
106        let module_graph = compilation.get_module_graph();
107        module_graph
108          .get_outgoing_connections(&module)
109          .filter_map(|con| module_graph.module_identifier_by_dependency_id(&con.dependency_id))
110          .filter(|&out| &module != out)
111          .copied()
112          .collect::<Vec<_>>()
113      })
114      .iter()
115      .any(|out| ModuleGraph::is_async(compilation, out))
116    {
117      // We can't safely reset is_async to false if there are any outgoing module is async
118      continue;
119    }
120    // The module is_async = false will also decide its parent module is_async, so if the module is_async = false
121    // is not changed, this means its parent module will be not affected, so we stop the infer at here.
122    if ModuleGraph::set_async(compilation, module, false) {
123      if let Some(mutations) = mutations {
124        mutations.add(Mutation::ModuleSetAsync { module });
125      }
126      let module_graph = compilation.get_module_graph();
127      module_graph
128        .get_incoming_connections(&module)
129        .filter(|con| {
130          module_graph
131            .dependency_by_id(&con.dependency_id)
132            .map(|dep| {
133              matches!(
134                dep.dependency_type(),
135                DependencyType::EsmImport | DependencyType::EsmExport
136              )
137            })
138            .unwrap_or_default()
139        })
140        .for_each(|con| {
141          if let Some(id) = con.original_module_identifier {
142            queue.insert(id);
143          }
144        });
145    }
146  }
147}
148
149fn set_async_modules(
150  compilation: &mut Compilation,
151  modules: LinkedHashSet<ModuleIdentifier>,
152  mutations: &mut Option<Mutations>,
153) {
154  let mut queue = modules;
155  let mut visited = IdentifierSet::from_iter(queue.iter().copied());
156
157  while let Some(module) = queue.pop_front() {
158    if ModuleGraph::set_async(compilation, module, true)
159      && let Some(mutations) = mutations
160    {
161      mutations.add(Mutation::ModuleSetAsync { module });
162    }
163    let module_graph = compilation.get_module_graph();
164    module_graph
165      .get_incoming_connections(&module)
166      .filter(|con| {
167        module_graph
168          .dependency_by_id(&con.dependency_id)
169          .map(|dep| {
170            matches!(
171              dep.dependency_type(),
172              DependencyType::EsmImport | DependencyType::EsmExport
173            )
174          })
175          .unwrap_or_default()
176      })
177      .for_each(|con| {
178        if let Some(id) = con.original_module_identifier
179          && visited.insert(id)
180        {
181          queue.insert(id);
182        }
183      });
184  }
185}
186
187impl Plugin for InferAsyncModulesPlugin {
188  fn name(&self) -> &'static str {
189    "InferAsyncModulesPlugin"
190  }
191
192  fn apply(&self, ctx: &mut rspack_core::ApplyContext<'_>) -> Result<()> {
193    ctx
194      .compilation_hooks
195      .finish_modules
196      .tap(finish_modules::new(self));
197    Ok(())
198  }
199}
rspack_plugin_javascript/plugin/infer_async_modules_plugin.rs

rspack_plugin_javascript/plugin/
infer_async_modules_plugin.rs
