diff --git a/crates/rspack_core/src/compilation/mod.rs b/crates/rspack_core/src/compilation/mod.rs
index 7f161d66c173..4bce2e1609de 100644
--- a/crates/rspack_core/src/compilation/mod.rs
+++ b/crates/rspack_core/src/compilation/mod.rs
@@ -4,7 +4,6 @@ use std::{
   collections::{VecDeque, hash_map},
   fmt::{self, Debug},
   hash::{BuildHasherDefault, Hash},
-  mem,
   sync::{
     Arc,
     atomic::{AtomicBool, AtomicU32, Ordering},
@@ -297,7 +296,7 @@ pub struct Compilation {
 
   pub modified_files: ArcPathSet,
   pub removed_files: ArcPathSet,
-  pub build_module_graph_artifact: BuildModuleGraphArtifact,
+  pub build_module_graph_artifact: DerefOption<BuildModuleGraphArtifact>,
   pub input_filesystem: Arc<dyn ReadableFileSystem>,
 
   pub intermediate_filesystem: Arc<dyn IntermediateFileSystem>,
@@ -427,7 +426,7 @@ impl Compilation {
       module_executor,
       in_finish_make: AtomicBool::new(false),
 
-      build_module_graph_artifact: BuildModuleGraphArtifact::default(),
+      build_module_graph_artifact: DerefOption::new(BuildModuleGraphArtifact::default()),
       modified_files,
       removed_files,
       input_filesystem,
@@ -458,14 +457,17 @@ impl Compilation {
     );
   }
   pub fn swap_build_module_graph_artifact(&mut self, make_artifact: &mut BuildModuleGraphArtifact) {
-    mem::swap(&mut self.build_module_graph_artifact, make_artifact);
+    self.build_module_graph_artifact.swap(make_artifact);
   }
   pub fn get_module_graph(&self) -> &ModuleGraph {
     self.build_module_graph_artifact.get_module_graph()
   }
 
-  // FIXME: find a better way to do this.
+  // it will return None during make phase since mg is incomplete
   pub fn module_by_identifier(&self, identifier: &ModuleIdentifier) -> Option<&BoxModule> {
+    if self.build_module_graph_artifact.is_none() {
+      return None;
+    }
     if let Some(module) = self.get_module_graph().module_by_identifier(identifier) {
       return Some(module);
     };
@@ -657,21 +659,23 @@ impl Compilation {
       self.add_entry(entry, options).await?;
     }
 
-    let make_artifact = mem::take(&mut self.build_module_graph_artifact);
-    self.build_module_graph_artifact = update_module_graph(
-      self,
-      make_artifact,
-      vec![UpdateParam::BuildEntry(
-        self
-          .entries
-          .values()
-          .flat_map(|item| item.all_dependencies())
-          .chain(self.global_entry.all_dependencies())
-          .copied()
-          .collect(),
-      )],
-    )
-    .await?;
+    let make_artifact = self.build_module_graph_artifact.take();
+    self.build_module_graph_artifact = DerefOption::new(
+      update_module_graph(
+        self,
+        make_artifact,
+        vec![UpdateParam::BuildEntry(
+          self
+            .entries
+            .values()
+            .flat_map(|item| item.all_dependencies())
+            .chain(self.global_entry.all_dependencies())
+            .copied()
+            .collect(),
+        )],
+      )
+      .await?,
+    );
     Ok(())
   }
 
@@ -706,21 +710,23 @@ impl Compilation {
 
     // Recheck entry and clean useless entry
     // This should before finish_modules hook is called, ensure providedExports effects on new added modules
-    let make_artifact = mem::take(&mut self.build_module_graph_artifact);
-    self.build_module_graph_artifact = update_module_graph(
-      self,
-      make_artifact,
-      vec![UpdateParam::BuildEntry(
-        self
-          .entries
-          .values()
-          .flat_map(|item| item.all_dependencies())
-          .chain(self.global_entry.all_dependencies())
-          .copied()
-          .collect(),
-      )],
-    )
-    .await?;
+    let make_artifact = self.build_module_graph_artifact.take();
+    self.build_module_graph_artifact = DerefOption::new(
+      update_module_graph(
+        self,
+        make_artifact,
+        vec![UpdateParam::BuildEntry(
+          self
+            .entries
+            .values()
+            .flat_map(|item| item.all_dependencies())
+            .chain(self.global_entry.all_dependencies())
+            .copied()
+            .collect(),
+        )],
+      )
+      .await?,
+    );
     Ok(())
   }
 
@@ -1016,8 +1022,8 @@ impl Compilation {
       self.module_executor = Some(module_executor);
     }
 
-    let artifact = std::mem::take(&mut self.build_module_graph_artifact);
-    self.build_module_graph_artifact = build_module_graph(self, artifact).await?;
+    let artifact = self.build_module_graph_artifact.take();
+    self.build_module_graph_artifact = DerefOption::new(build_module_graph(self, artifact).await?);
 
     self.in_finish_make.store(true, Ordering::Release);
 
@@ -1029,17 +1035,19 @@ impl Compilation {
     module_identifiers: IdentifierSet,
     f: impl Fn(Vec<&BoxModule>) -> T,
   ) -> Result<T> {
-    let artifact = std::mem::take(&mut self.build_module_graph_artifact);
+    let artifact = self.build_module_graph_artifact.take();
 
     // https://github.com/webpack/webpack/blob/19ca74127f7668aaf60d59f4af8fcaee7924541a/lib/Compilation.js#L2462C21-L2462C25
     self.module_graph_cache_artifact.unfreeze();
 
-    self.build_module_graph_artifact = update_module_graph(
-      self,
-      artifact,
-      vec![UpdateParam::ForceBuildModules(module_identifiers.clone())],
-    )
-    .await?;
+    self.build_module_graph_artifact = DerefOption::new(
+      update_module_graph(
+        self,
+        artifact,
+        vec![UpdateParam::ForceBuildModules(module_identifiers.clone())],
+      )
+      .await?,
+    );
 
     let module_graph = self.get_module_graph();
     Ok(f(module_identifiers
@@ -1434,8 +1442,9 @@ impl Compilation {
   pub async fn finish_build_module_graph(&mut self) -> Result<()> {
     self.in_finish_make.store(false, Ordering::Release);
     // clean up the entry deps
-    let make_artifact = std::mem::take(&mut self.build_module_graph_artifact);
-    self.build_module_graph_artifact = finish_build_module_graph(self, make_artifact).await?;
+    let make_artifact = self.build_module_graph_artifact.take();
+    self.build_module_graph_artifact =
+      DerefOption::new(finish_build_module_graph(self, make_artifact).await?);
     // sync assets to module graph from module_executor
     if let Some(module_executor) = &mut self.module_executor {
       let mut module_executor = std::mem::take(module_executor);
diff --git a/crates/rspack_core/src/utils/deref_option.rs b/crates/rspack_core/src/utils/deref_option.rs
index 6767ac9e1b3b..5016baf904d4 100644
--- a/crates/rspack_core/src/utils/deref_option.rs
+++ b/crates/rspack_core/src/utils/deref_option.rs
@@ -15,6 +15,9 @@ impl<T> DerefOption<T> {
   pub fn new(value: T) -> Self {
     Self(Some(value))
   }
+  pub fn is_none(&self) -> bool {
+    self.0.is_none()
+  }
   pub fn take(&mut self) -> T {
     self
       .0