Optimise isAssetReferenced in packager (#858)

Author: marcins

.changeset/lovely-pens-love.md

@@ -0,0 +1,8 @@
+---
+'@atlaspack/types-internal': minor
+'@atlaspack/feature-flags': minor
+'@atlaspack/packager-js': minor
+'@atlaspack/core': minor
+---
+
+Introduce a new `getReferencedAssets(bundle)` method to the BundleGraph to pre-compute referenced assets, this is used by the scope hoisting packager behind a new `precomputeReferencedAssets` feature flag.

packages/core/core/src/BundleGraph.ts

@@ -1435,6 +1435,158 @@ export default class BundleGraph {
     });
   }
 
+  // New method: Fast checks only (no caching of results)
+  isAssetReferencedFastCheck(bundle: Bundle, asset: Asset): boolean | null {
+    // Fast Check #1: If asset is in multiple bundles in same target, it's referenced
+    let bundlesWithAsset = this.getBundlesWithAsset(asset).filter(
+      (b) =>
+        b.target.name === bundle.target.name &&
+        b.target.distDir === bundle.target.distDir,
+    );
+
+    if (bundlesWithAsset.length > 1) {
+      return true;
+    }
+
+    // Fast Check #2: If asset is referenced by any async/conditional dependency, it's referenced
+    let assetNodeId = nullthrows(this._graph.getNodeIdByContentKey(asset.id));
+
+    if (
+      this._graph
+        .getNodeIdsConnectedTo(assetNodeId, bundleGraphEdgeTypes.references)
+        .map((id) => this._graph.getNode(id))
+        .some(
+          (node) =>
+            node?.type === 'dependency' &&
+            (node.value.priority === Priority.lazy ||
+              node.value.priority === Priority.conditional) &&
+            node.value.specifierType !== SpecifierType.url,
+        )
+    ) {
+      return true;
+    }
+
+    // Fast checks failed - return null to indicate expensive computation needed
+    return null;
+  }
+
+  getReferencedAssets(bundle: Bundle): Set<Asset> {
+    let referencedAssets = new Set<Asset>();
+
+    // Build a map of all assets in this bundle with their dependencies
+    // This allows us to check all assets in a single traversal
+    let assetDependenciesMap = new Map<Asset, Array<Dependency>>();
+
+    this.traverseAssets(bundle, (asset) => {
+      // Always do fast checks (no caching)
+      let fastCheckResult = this.isAssetReferencedFastCheck(bundle, asset);
+
+      if (fastCheckResult === true) {
+        referencedAssets.add(asset);
+        return;
+      }
+
+      // Fast checks failed (fastCheckResult === null), need expensive computation
+      // Check if it's actually referenced via traversal
+
+      // Store dependencies for later batch checking
+      let dependencies = this._graph
+        .getNodeIdsConnectedTo(
+          nullthrows(this._graph.getNodeIdByContentKey(asset.id)),
+        )
+        .map((id) => nullthrows(this._graph.getNode(id)))
+        .filter((node) => node.type === 'dependency')
+        .map((node) => {
+          invariant(node.type === 'dependency');
+          return node.value;
+        });
+
+      if (dependencies.length > 0) {
+        assetDependenciesMap.set(asset, dependencies);
+      }
+    });
+
+    // If no assets need the expensive check, return early
+    if (assetDependenciesMap.size === 0) {
+      return referencedAssets;
+    }
+
+    // Get the assets we need to check once
+    let assetsToCheck = Array.from(assetDependenciesMap.keys());
+
+    // Helper function to check if all assets from assetDependenciesMap are in referencedAssets
+    const allAssetsReferenced = (): boolean =>
+      assetsToCheck.length <= referencedAssets.size &&
+      assetsToCheck.every((asset) => referencedAssets.has(asset));
+
+    // Do ONE traversal to check all remaining assets
+    // We can share visitedBundles across all assets because we check every asset
+    // against every visited bundle, which matches the original per-asset behavior
+    let siblingBundles = new Set(
+      this.getBundleGroupsContainingBundle(bundle).flatMap((bundleGroup) =>
+        this.getBundlesInBundleGroup(bundleGroup, {includeInline: true}),
+      ),
+    );
+
+    let visitedBundles: Set<Bundle> = new Set();
+
+    // Single traversal from all referencers
+    for (let referencer of siblingBundles) {
+      this.traverseBundles((descendant, _, actions) => {
+        if (descendant.id === bundle.id) {
+          return;
+        }
+
+        if (visitedBundles.has(descendant)) {
+          actions.skipChildren();
+          return;
+        }
+
+        visitedBundles.add(descendant);
+
+        if (
+          descendant.type !== bundle.type ||
+          fromEnvironmentId(descendant.env).context !==
+            fromEnvironmentId(bundle.env).context
+        ) {
+          // Don't skip children - they might be the right type!
+          return;
+        }
+
+        // Check ALL assets at once in this descendant bundle
+        for (let [asset, dependencies] of assetDependenciesMap) {
+          // Skip if already marked as referenced
+          if (referencedAssets.has(asset)) {
+            continue;
+          }
+
+          // Check if this descendant bundle references the asset
+          if (
+            !this.bundleHasAsset(descendant, asset) &&
+            dependencies.some((dependency) =>
+              this.bundleHasDependency(descendant, dependency),
+            )
+          ) {
+            referencedAssets.add(asset);
+          }
+        }
+
+        // If all assets from assetDependenciesMap are now marked as referenced, we can stop early
+        if (allAssetsReferenced()) {
+          actions.stop();
+          return;
+        }
+      }, referencer);
+
+      // If all assets from assetDependenciesMap are referenced, no need to check more sibling bundles
+      if (allAssetsReferenced()) {
+        break;
+      }
+    }
+
+    return referencedAssets;
+  }
+
   hasParentBundleOfType(bundle: Bundle, type: string): boolean {
     let parents = this.getParentBundles(bundle);
     return (

packages/core/core/src/public/BundleGraph.ts

@@ -198,6 +198,27 @@ export default class BundleGraph<TBundle extends IBundle>
     );
   }
 
+  isAssetReferencedFastCheck(bundle: IBundle, asset: IAsset): boolean | null {
+    return this.#graph.isAssetReferencedFastCheck(
+      bundleToInternalBundle(bundle),
+      assetToAssetValue(asset),
+    );
+  }
+
+  getReferencedAssets(bundle: IBundle): Set<IAsset> {
+    let internalReferencedAssets = this.#graph.getReferencedAssets(
+      bundleToInternalBundle(bundle),
+    );
+
+    // Convert internal assets to public assets
+    let publicReferencedAssets = new Set<IAsset>();
+    for (let internalAsset of internalReferencedAssets) {
+      publicReferencedAssets.add(assetFromValue(internalAsset, this.#options));
+    }
+
+    return publicReferencedAssets;
+  }
+
   hasParentBundleOfType(bundle: IBundle, type: string): boolean {
     return this.#graph.hasParentBundleOfType(
       bundleToInternalBundle(bundle),

packages/core/feature-flags/src/index.ts

@@ -326,6 +326,16 @@ export const DEFAULT_FEATURE_FLAGS = {
    * @since 2025-11-05
    */
   nestedPromiseImportFix: process.env.ATLASPACK_BUILD_ENV === 'test',
+
+  /**
+   * Precompute referenced assets in bundles to avoid repeated traversals
+   * during scope hoisting packaging. This optimization caches which assets
+   * are referenced in a bundle, reducing O(N*M) calls to O(1).
+   *
+   * @author Marcin Szczepanski <[email protected]>
+   * @since 2025-10-24
+   */
+  precomputeReferencedAssets: process.env.ATLASPACK_BUILD_ENV === 'test',
 };
 
 export type FeatureFlags = typeof DEFAULT_FEATURE_FLAGS;

packages/core/types-internal/src/index.ts

@@ -1690,6 +1690,13 @@ export interface BundleGraph<TBundle extends Bundle> {
   isAssetReachableFromBundle(asset: Asset, bundle: Bundle): boolean;
   /** Returns whether an asset is referenced outside the given bundle. */
   isAssetReferenced(bundle: Bundle, asset: Asset): boolean;
+  /**
+   * Fast checks only for asset reference status. Returns true if fast checks succeed,
+   * null if expensive traversal computation is needed.
+   */
+  isAssetReferencedFastCheck(bundle: Bundle, asset: Asset): boolean | null;
+  /** Returns a set of all assets that are referenced outside the given bundle. */
+  getReferencedAssets(bundle: Bundle): Set<Asset>;
   /**
    * Resolves the export `symbol` of `asset` to the source,
    * stopping at the first asset after leaving `bundle`.

packages/packagers/js/src/ScopeHoistingPackager.ts

@@ -131,6 +131,7 @@ export class ScopeHoistingPackager {
   useBothScopeHoistingImprovements: boolean =
     getFeatureFlag('applyScopeHoistingImprovementV2') ||
     getFeatureFlag('applyScopeHoistingImprovement');
+  referencedAssetsCache: Map<string, Set<Asset>> = new Map();
 
   constructor(
     options: PluginOptions,
@@ -412,6 +413,44 @@ export class ScopeHoistingPackager {
     return `$parcel$global.rwr(${params.join(', ')});`;
   }
 
+  // Helper to check if an asset is referenced, with cache-first + fast-check hybrid approach
+  isAssetReferencedInBundle(bundle: NamedBundle, asset: Asset): boolean {
+    // STEP 1: Check expensive computation cache first (fastest when it hits)
+    if (getFeatureFlag('precomputeReferencedAssets')) {
+      let bundleId = bundle.id;
+      let referencedAssets = this.referencedAssetsCache.get(bundleId);
+
+      if (referencedAssets) {
+        // Cache hit - fastest path (~0.001ms)
+        return referencedAssets.has(asset);
+      }
+    }
+
+    // STEP 2: Cache miss - try fast checks (~0.01ms)
+    let fastCheckResult = this.bundleGraph.isAssetReferencedFastCheck(
+      bundle,
+      asset,
+    );
+
+    if (fastCheckResult === true) {
+      // Fast check succeeded - asset is referenced
+      return true;
+    }
+
+    // STEP 3: Need expensive computation (~20-2000ms)
+    if (getFeatureFlag('precomputeReferencedAssets')) {
+      // Compute and cache expensive results for this bundle
+      let bundleId = bundle.id;
+      let referencedAssets = this.bundleGraph.getReferencedAssets(bundle);
+      this.referencedAssetsCache.set(bundleId, referencedAssets);
+
+      return referencedAssets.has(asset);
+    } else {
+      // No caching - fall back to original expensive method
+      return this.bundleGraph.isAssetReferenced(bundle, asset);
+    }
+  }
+
   async loadAssets() {
     type QueueItem = [Asset, {code: string; map: Buffer | undefined | null}];
     let queue = new PromiseQueue<QueueItem>({
@@ -431,7 +470,7 @@ export class ScopeHoistingPackager {
       if (
         asset.meta.shouldWrap ||
         this.bundle.env.sourceType === 'script' ||
-        this.bundleGraph.isAssetReferenced(this.bundle, asset) ||
+        this.isAssetReferencedInBundle(this.bundle, asset) ||
         this.bundleGraph
           .getIncomingDependencies(asset)
           .some((dep) => dep.meta.shouldWrap && dep.specifierType !== 'url')
@@ -993,7 +1032,7 @@ ${code}
           referencedBundle &&
           referencedBundle.getMainEntry() === resolved &&
           referencedBundle.type === 'js' &&
-          !this.bundleGraph.isAssetReferenced(referencedBundle, resolved)
+          !this.isAssetReferencedInBundle(referencedBundle, resolved)
         ) {
           this.addExternal(dep, replacements, referencedBundle);
           this.externalAssets.add(resolved);
@@ -1792,7 +1831,7 @@ ${code}
     return (
       asset.sideEffects === false &&
       nullthrows(this.bundleGraph.getUsedSymbols(asset)).size == 0 &&
-      !this.bundleGraph.isAssetReferenced(this.bundle, asset)
+      !this.isAssetReferencedInBundle(this.bundle, asset)
     );
   }