From 67a3a59e9030846424cfe77656d7a9a3da3a29a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lichu=20Acu=C3=B1a?= <lichu.acuna@vercel.com>
Date: Mon, 19 Aug 2024 15:04:36 -0700
Subject: [PATCH] Fixed async loader module id strategy bug

---
 .../src/global_module_id_strategy.rs          | 131 ++++++++++++++++--
 1 file changed, 116 insertions(+), 15 deletions(-)

diff --git a/turbopack/crates/turbopack-ecmascript/src/global_module_id_strategy.rs b/turbopack/crates/turbopack-ecmascript/src/global_module_id_strategy.rs
index 20596165ca9b1..e6c1f82c827a3 100644
--- a/turbopack/crates/turbopack-ecmascript/src/global_module_id_strategy.rs
+++ b/turbopack/crates/turbopack-ecmascript/src/global_module_id_strategy.rs
@@ -1,17 +1,19 @@
 use std::collections::{HashMap, HashSet};
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use turbo_tasks::{
     graph::{AdjacencyMap, GraphTraversal},
-    RcStr, ValueToString, Vc,
+    RcStr, TryJoinIterExt, ValueToString, Vc,
 };
 use turbo_tasks_hash::hash_xxh3_hash64;
 use turbopack_core::{
     chunk::ModuleId,
     module::{Module, Modules},
-    reference::primary_referenced_modules,
+    reference::ModuleReference,
 };
 
+use crate::references::esm::EsmAsyncAssetReference;
+
 #[turbo_tasks::value]
 pub struct PreprocessedChildrenIdents {
     // ident.to_string() -> full hash
@@ -20,12 +22,93 @@ pub struct PreprocessedChildrenIdents {
     modules_idents: HashMap<RcStr, u64>,
 }
 
+#[derive(Clone, Hash)]
+#[turbo_tasks::value(shared)]
+pub enum ReferencedModule {
+    Module(Vc<Box<dyn Module>>),
+    AsyncLoaderModule(Vc<Box<dyn Module>>),
+}
+
+impl ReferencedModule {
+    fn module(&self) -> Vc<Box<dyn Module>> {
+        match *self {
+            ReferencedModule::Module(module) => module,
+            ReferencedModule::AsyncLoaderModule(module) => module,
+        }
+    }
+}
+
+// TODO(LichuAcu): Reduce type complexity
+#[allow(clippy::type_complexity)]
+type ModulesAndAsyncLoaders = Vec<(Vec<Vc<Box<dyn Module>>>, Option<Vc<Box<dyn Module>>>)>;
+
+#[turbo_tasks::value(transparent)]
+pub struct ReferencedModules(Vec<Vc<ReferencedModule>>);
+
+#[turbo_tasks::function]
+async fn referenced_modules(module: Vc<Box<dyn Module>>) -> Result<Vc<ReferencedModules>> {
+    let references = module.references().await?;
+
+    let mut set = HashSet::new();
+    let modules_and_async_loaders: ModulesAndAsyncLoaders = references
+        .iter()
+        .map(|reference| async move {
+            let mut async_loader = None;
+            if Vc::try_resolve_downcast_type::<EsmAsyncAssetReference>(*reference)
+                .await?
+                .is_some()
+            {
+                async_loader = Some(
+                    reference
+                        .resolve_reference()
+                        .resolve()
+                        .await?
+                        .first_module()
+                        .await?
+                        .context("No module found")?,
+                );
+            }
+
+            let modules = reference
+                .resolve_reference()
+                .resolve()
+                .await?
+                .primary_modules()
+                .await?
+                .clone_value();
+
+            Ok((modules, async_loader))
+        })
+        .try_join()
+        .await?;
+
+    let mut modules = Vec::new();
+
+    for (module_list, async_loader) in modules_and_async_loaders {
+        for module in module_list {
+            if set.insert(module) {
+                modules.push(ReferencedModule::Module(module).cell());
+            }
+        }
+        if let Some(async_loader_module) = async_loader {
+            if set.insert(async_loader_module) {
+                modules.push(ReferencedModule::AsyncLoaderModule(async_loader_module).cell());
+            }
+        }
+    }
+
+    Ok(Vc::cell(modules))
+}
+
 pub async fn get_children_modules(
-    parent: Vc<Box<dyn Module>>,
-) -> Result<impl Iterator<Item = Vc<Box<dyn Module>>> + Send> {
-    let mut primary_modules = primary_referenced_modules(parent).await?.clone_value();
-    primary_modules.extend(parent.additional_layers_modules().await?.clone_value());
-    Ok(primary_modules.into_iter())
+    parent: Vc<ReferencedModule>,
+) -> Result<impl Iterator<Item = Vc<ReferencedModule>> + Send> {
+    let parent_module = parent.await?.module();
+    let mut modules = referenced_modules(parent_module).await?.clone_value();
+    for module in parent_module.additional_layers_modules().await? {
+        modules.push(ReferencedModule::Module(*module).cell());
+    }
+    Ok(modules.into_iter())
 }
 
 // NOTE(LichuAcu) Called on endpoint.root_modules(). It would probably be better if this was called
@@ -38,7 +121,14 @@ pub async fn children_modules_idents(
 ) -> Result<Vc<PreprocessedChildrenIdents>> {
     let children_modules_iter = AdjacencyMap::new()
         .skip_duplicates()
-        .visit(root_modules.await?.iter().copied(), get_children_modules)
+        .visit(
+            root_modules
+                .await?
+                .iter()
+                .map(|module| ReferencedModule::Module(*module).cell())
+                .collect::<Vec<_>>(),
+            get_children_modules,
+        )
         .await
         .completed()?
         .into_inner()
@@ -46,12 +136,23 @@ pub async fn children_modules_idents(
 
     // module_id -> full hash
     let mut modules_idents = HashMap::new();
-
-    for module in children_modules_iter {
-        let module_ident = module.ident();
-        let ident_str = module_ident.to_string().await?.clone_value();
-        let hash = hash_xxh3_hash64(&ident_str);
-        modules_idents.insert(ident_str, hash);
+    for child_module in children_modules_iter {
+        match *child_module.await? {
+            ReferencedModule::Module(module) => {
+                let module_ident = module.ident();
+                let ident_str = module_ident.to_string().await?.clone_value();
+                let hash = hash_xxh3_hash64(&ident_str);
+                modules_idents.insert(ident_str, hash);
+            }
+            ReferencedModule::AsyncLoaderModule(async_loader_module) => {
+                let async_loader_ident = async_loader_module
+                    .ident()
+                    .with_modifier(Vc::cell("async loader".into()));
+                let ident_str = async_loader_ident.to_string().await?.clone_value();
+                let hash = hash_xxh3_hash64(&ident_str);
+                modules_idents.insert(ident_str, hash);
+            }
+        }
     }
 
     Ok(PreprocessedChildrenIdents { modules_idents }.cell())