bundler: define __promiseAll when an unwrapped file awaits two async ESM wrappers

src/bundler/LinkerContext.rs

@@ -3538,29 +3538,6 @@ impl<'a> LinkerContext<'a> {
                 // This depends on the "__esm" symbol and declares the "init_foo" symbol
                 // for similar reasons to the CommonJS closure above.
 
-                // Count async dependencies to determine if we need __promiseAll
-                let mut async_import_count: usize = 0;
-                {
-                    let import_records =
-                        self.graph.ast.items_import_records()[source_index as usize].as_slice();
-                    let meta_flags = self.graph.meta.items_flags();
-
-                    for record in import_records {
-                        if !record.source_index.is_valid() {
-                            continue;
-                        }
-                        let other_flags = meta_flags[record.source_index.get() as usize];
-                        if other_flags.is_async_or_has_async_dependency {
-                            async_import_count += 1;
-                            if async_import_count >= 2 {
-                                break;
-                            }
-                        }
-                    }
-                }
-
-                let needs_promise_all = async_import_count >= 2;
-
                 let esm_parts: &[u32] = if wrapper_ref.is_valid()
                     && self.options.output_format != Format::InternalBakeDev
                 {
@@ -3569,30 +3546,14 @@ impl<'a> LinkerContext<'a> {
                     &[]
                 };
 
-                let promise_all_parts: &[u32] = if needs_promise_all
-                    && wrapper_ref.is_valid()
-                    && self.options.output_format != Format::InternalBakeDev
-                {
-                    self.top_level_symbols_to_parts_for_runtime(self.promise_all_runtime_ref)
-                } else {
-                    &[]
-                };
-
-                // generate a dummy part that depends on the "__esm" and optionally "__promiseAll" symbols
-                let mut dependencies =
-                    DependencyList::init_capacity(esm_parts.len() + promise_all_parts.len());
+                // generate a dummy part that depends on the "__esm" symbol
+                let mut dependencies = DependencyList::init_capacity(esm_parts.len());
                 for &part in esm_parts {
                     dependencies.append_assume_capacity(Dependency {
                         part_index: part,
                         source_index: bun_ast::Index::RUNTIME,
                     });
                 }
-                for &part in promise_all_parts {
-                    dependencies.append_assume_capacity(Dependency {
-                        part_index: part,
-                        source_index: bun_ast::Index::RUNTIME,
-                    });
-                }
 
                 let mut symbol_uses = PartSymbolUseMap::default();
                 symbol_uses
@@ -3626,19 +3587,6 @@ impl<'a> LinkerContext<'a> {
                             crate::Index::RUNTIME,
                         )
                         .expect("OOM");
-
-                    // Only mark __promiseAll as used if we have multiple async dependencies
-                    if needs_promise_all {
-                        self.graph
-                            .generate_symbol_import_and_use(
-                                source_index,
-                                part_index,
-                                self.promise_all_runtime_ref,
-                                1,
-                                crate::Index::RUNTIME,
-                            )
-                            .expect("OOM");
-                    }
                 }
             }
             WrapKind::None => {}

src/bundler/linker_context/scanImportsAndExports.rs

@@ -1068,11 +1068,14 @@ pub(crate) fn scan_imports_and_exports(
             );
 
             let parts_len = col_ref!(parts_list)[id].len() as usize;
+            // Per file, not per part: see `InsideWrapperPrefix::append_async_dependency`.
+            let mut async_esm_init_count: u32 = 0;
             for part_index in 0..parts_len {
                 let mut to_esm_uses: u32 = 0;
                 let mut to_common_js_uses: u32 = 0;
                 let mut runtime_require_uses: u32 = 0;
                 let mut preload_uses: u32 = 0;
+                let mut promise_all_uses: u32 = 0;
 
                 // Imports of wrapped files must depend on the wrapper
                 // Iterate by index so each iteration re-borrows
@@ -1205,6 +1208,16 @@ pub(crate) fn scan_imports_and_exports(
                                 1,
                                 Index::source(other_source_index),
                             )?;
+
+                            if kind == ImportKind::Stmt
+                                && other_flags.wrap == WrapKind::Esm
+                                && other_flags.is_async_or_has_async_dependency
+                            {
+                                async_esm_init_count += 1;
+                                if async_esm_init_count >= 2 {
+                                    promise_all_uses = 1;
+                                }
+                            }
                         }
 
                         // This is an ES6 import of a CommonJS module, so it needs the
@@ -1406,6 +1419,13 @@ pub(crate) fn scan_imports_and_exports(
                         b"__preload",
                         preload_uses,
                     )?;
+
+                    this.graph.generate_runtime_symbol_import_and_use(
+                        source_index,
+                        Index::part(part_index as u32),
+                        b"__promiseAll",
+                        promise_all_uses,
+                    )?;
                 }
             }
         }

test/bundler/bundler_promiseall_deadcode.test.ts

@@ -397,4 +397,136 @@ describe("bundler", () => {
       expect(bundled).not.toMatch(/await\s+__promiseAll\s*\(/);
     },
   });
+
+  // The importer below is not wrapped (nothing `import()`s or `require()`s it),
+  // so its `await __promiseAll([...])` is printed at the top level of the chunk.
+  // The `import()` calls are what make p.js and q.js lazily wrapped (`init_p`,
+  // `init_q`), and their top-level await makes those wrappers async.
+  itBundled("bundler/__promiseAll is defined when an unwrapped entry awaits two async ESM wrappers", {
+    files: {
+      "/entry.js": `
+        import "./p.js";
+        import "./q.js";
+        import("./p.js");
+        import("./q.js");
+        console.log("entry");
+      `,
+      "/p.js": `
+        await Promise.resolve();
+        console.log("p");
+        export const p = 1;
+      `,
+      "/q.js": `
+        await Promise.resolve();
+        console.log("q");
+        export const q = 1;
+      `,
+    },
+    run: {
+      stdout: "p\nq\nentry",
+    },
+    onAfterBundle(api) {
+      const bundled = api.readFile("out.js");
+      expect(bundled).toMatch(/await\s+__promiseAll\s*\(\s*\[\s*init_p\(\),\s*init_q\(\)\s*\]\s*\)/);
+      expect(bundled).toContain("var __promiseAll = ");
+    },
+  });
+
+  itBundled("bundler/__promiseAll is defined when an unwrapped entry awaits three async ESM wrappers", {
+    files: {
+      "/entry.js": `
+        import { a } from "./a.js";
+        import { b } from "./b.js";
+        import { c } from "./c.js";
+        console.log(a, b, c);
+        const mods = await Promise.all([import("./a.js"), import("./b.js"), import("./c.js")]);
+        console.log(mods.map(m => Object.keys(m)).join());
+      `,
+      "/a.js": `export const a = await Promise.resolve("A");`,
+      "/b.js": `export const b = await Promise.resolve("B");`,
+      "/c.js": `export const c = await Promise.resolve("C");`,
+    },
+    run: {
+      stdout: "A B C\na,b,c",
+    },
+    onAfterBundle(api) {
+      const bundled = api.readFile("out.js");
+      expect(bundled).toMatch(/await\s+__promiseAll\s*\(\s*\[\s*init_a\(\),\s*init_b\(\),\s*init_c\(\)\s*\]\s*\)/);
+      expect(bundled).toContain("var __promiseAll = ");
+    },
+  });
+
+  itBundled("bundler/__promiseAll is defined when an unwrapped non-entry file awaits two async ESM wrappers", {
+    files: {
+      "/entry.js": `
+        import { joined } from "./mid.js";
+        console.log(joined);
+        const [{ a }, { b }] = await Promise.all([import("./a.js"), import("./b.js")]);
+        console.log(a + b === joined);
+      `,
+      "/mid.js": `
+        import { a } from "./a.js";
+        import { b } from "./b.js";
+        export const joined = a + b;
+      `,
+      "/a.js": `export const a = await Promise.resolve("A");`,
+      "/b.js": `export const b = await Promise.resolve("B");`,
+    },
+    run: {
+      stdout: "AB\ntrue",
+    },
+    onAfterBundle(api) {
+      const bundled = api.readFile("out.js");
+      expect(bundled).toMatch(/\/\/ mid\.js\s+await\s+__promiseAll\s*\(\s*\[\s*init_a\(\),\s*init_b\(\)\s*\]\s*\)/);
+      expect(bundled).toContain("var __promiseAll = ");
+    },
+  });
+
+  itBundled("bundler/__promiseAll is tree-shaken when an unwrapped entry awaits one async ESM wrapper", {
+    files: {
+      "/entry.js": `
+        import { a } from "./a.js";
+        console.log(a);
+        console.log(Object.keys(await import("./a.js")).join());
+      `,
+      "/a.js": `export const a = await Promise.resolve("A");`,
+    },
+    run: {
+      stdout: "A\na",
+    },
+    onAfterBundle(api) {
+      const bundled = api.readFile("out.js");
+      expect(bundled).toMatch(/await\s+init_a\(\);/);
+      expect(bundled).not.toContain("__promiseAll");
+    },
+  });
+
+  // Only `import` statements join the `await __promiseAll([...])`. An `import()`
+  // of a second async module calls `init_c()` where the expression is and does
+  // not need the helper.
+  itBundled("bundler/__promiseAll is tree-shaken when a wrapper's second async dependency is an import()", {
+    files: {
+      "/entry.js": `
+        const { run } = await import("./lazy.js");
+        await run();
+      `,
+      "/lazy.js": `
+        import { a } from "./a.js";
+        export async function run() {
+          const { c } = await import("./c.js");
+          console.log(a, c);
+        }
+      `,
+      "/a.js": `export const a = await Promise.resolve("A");`,
+      "/c.js": `export const c = await Promise.resolve("C");`,
+    },
+    run: {
+      stdout: "A C",
+    },
+    onAfterBundle(api) {
+      const bundled = api.readFile("out.js");
+      expect(bundled).toMatch(/var init_lazy = __esm\(async \(\) => \{\s*await\s+init_a\(\);/);
+      expect(bundled).not.toContain("__promiseAll");
+    },
+  });
 });