From b0409698af1aa311afe0ce029a33cd64b717c56c Mon Sep 17 00:00:00 2001
From: pshu <pishu.spf@antfin.com>
Date: Wed, 7 Jan 2026 15:22:02 +0800
Subject: [PATCH] refactor: after code generation hook use RO compilation ref

---
 crates/rspack_core/src/compilation/mod.rs      |  7 +++++--
 crates/rspack_plugin_esm_library/src/link.rs   |  8 ++++++--
 crates/rspack_plugin_esm_library/src/plugin.rs | 10 +++++++---
 crates/rspack_plugin_rsdoctor/src/plugin.rs    |  6 +++++-
 4 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/crates/rspack_core/src/compilation/mod.rs b/crates/rspack_core/src/compilation/mod.rs
index 6313b6e995c9..ef032103b238 100644
--- a/crates/rspack_core/src/compilation/mod.rs
+++ b/crates/rspack_core/src/compilation/mod.rs
@@ -94,7 +94,7 @@ define_hook!(CompilationRuntimeRequirementInChunk: SeriesBail(compilation: &mut
 define_hook!(CompilationAdditionalTreeRuntimeRequirements: Series(compilation: &mut Compilation, chunk_ukey: &ChunkUkey, runtime_requirements: &mut RuntimeGlobals));
 define_hook!(CompilationRuntimeRequirementInTree: SeriesBail(compilation: &mut Compilation, chunk_ukey: &ChunkUkey, all_runtime_requirements: &RuntimeGlobals, runtime_requirements: &RuntimeGlobals, runtime_requirements_mut: &mut RuntimeGlobals));
 define_hook!(CompilationOptimizeCodeGeneration: Series(compilation: &mut Compilation));
-define_hook!(CompilationAfterCodeGeneration: Series(compilation: &mut Compilation));
+define_hook!(CompilationAfterCodeGeneration: Series(compilation: &Compilation, diagnostics: &mut Vec<Diagnostic>));
 define_hook!(CompilationChunkHash: Series(compilation: &Compilation, chunk_ukey: &ChunkUkey, hasher: &mut RspackHash),tracing=false);
 define_hook!(CompilationContentHash: Series(compilation: &Compilation, chunk_ukey: &ChunkUkey, hashes: &mut HashMap<SourceType, RspackHash>));
 define_hook!(CompilationDependentFullHash: SeriesBail(compilation: &Compilation, chunk_ukey: &ChunkUkey) -> bool);
@@ -1893,12 +1893,15 @@ impl Compilation {
     };
     self.code_generation(code_generation_modules).await?;
 
+    let mut diagnostics = vec![];
     plugin_driver
       .compilation_hooks
       .after_code_generation
-      .call(self)
+      .call(self, &mut diagnostics)
       .await
       .map_err(|e| e.wrap_err("caused by plugins in Compilation.hooks.afterCodeGeneration"))?;
+    self.extend_diagnostics(diagnostics);
+
     logger.time_end(start);
 
     let start = logger.time("runtime requirements");
diff --git a/crates/rspack_plugin_esm_library/src/link.rs b/crates/rspack_plugin_esm_library/src/link.rs
index 617a4135be37..ea7e2667076a 100644
--- a/crates/rspack_plugin_esm_library/src/link.rs
+++ b/crates/rspack_plugin_esm_library/src/link.rs
@@ -139,7 +139,11 @@ impl EsmLibraryPlugin {
     }
   }
 
-  pub(crate) async fn link(&self, compilation: &mut Compilation) -> Result<()> {
+  pub(crate) async fn link(
+    &self,
+    compilation: &Compilation,
+    diagnostics: &mut Vec<Diagnostic>,
+  ) -> Result<()> {
     let module_graph = compilation.get_module_graph();
 
     // codegen uses self.concatenated_modules_map_for_codegen which has hold another Arc, so
@@ -247,7 +251,7 @@ impl EsmLibraryPlugin {
 
     // link imported specifier with exported symbol
     let mut needed_namespace_objects_by_ukey = UkeyMap::default();
-    compilation.extend_diagnostics(self.link_imports_and_exports(
+    diagnostics.extend(self.link_imports_and_exports(
       compilation,
       &mut link,
       &mut concate_modules_map,
diff --git a/crates/rspack_plugin_esm_library/src/plugin.rs b/crates/rspack_plugin_esm_library/src/plugin.rs
index 1e51362f0679..f5b67db72a37 100644
--- a/crates/rspack_plugin_esm_library/src/plugin.rs
+++ b/crates/rspack_plugin_esm_library/src/plugin.rs
@@ -18,7 +18,7 @@ use rspack_core::{
   ParserOptions, Plugin, PrefetchExportsInfoMode, RuntimeGlobals, get_target, is_esm_dep_like,
   rspack_sources::{ReplaceSource, Source},
 };
-use rspack_error::Result;
+use rspack_error::{Diagnostic, Result};
 use rspack_hook::{plugin, plugin_hook};
 use rspack_plugin_javascript::{
   JavascriptModulesRenderChunkContent, JsPlugin, RenderSource,
@@ -303,7 +303,11 @@ async fn concatenation_scope(
 }
 
 #[plugin_hook(CompilationAfterCodeGeneration for EsmLibraryPlugin)]
-async fn after_code_generation(&self, compilation: &mut Compilation) -> Result<()> {
+async fn after_code_generation(
+  &self,
+  compilation: &Compilation,
+  diagnostics: &mut Vec<Diagnostic>,
+) -> Result<()> {
   let mut chunk_ids_to_ukey = FxHashMap::default();
 
   for chunk_ukey in compilation.chunk_by_ukey.keys() {
@@ -316,7 +320,7 @@ async fn after_code_generation(&self, compilation: &mut Compilation) -> Result<(
 
   *self.chunk_ids_to_ukey.borrow_mut() = chunk_ids_to_ukey;
 
-  self.link(compilation).await?;
+  self.link(compilation, diagnostics).await?;
   Ok(())
 }
 
diff --git a/crates/rspack_plugin_rsdoctor/src/plugin.rs b/crates/rspack_plugin_rsdoctor/src/plugin.rs
index 165982245414..30274f25f2a7 100644
--- a/crates/rspack_plugin_rsdoctor/src/plugin.rs
+++ b/crates/rspack_plugin_rsdoctor/src/plugin.rs
@@ -418,7 +418,11 @@ async fn module_ids(
 }
 
 #[plugin_hook(CompilationAfterCodeGeneration for RsdoctorPlugin, stage = 9999)]
-async fn after_code_generation(&self, compilation: &mut Compilation) -> Result<()> {
+async fn after_code_generation(
+  &self,
+  compilation: &Compilation,
+  _diagnostics: &mut Vec<Diagnostic>,
+) -> Result<()> {
   if !self.has_module_graph_feature(RsdoctorPluginModuleGraphFeature::ModuleSources) {
     return Ok(());
   }