From 186f3a95034d576079ecebe2f684a50470e6f0fe Mon Sep 17 00:00:00 2001
From: Tobias Koppers <tobias.koppers@googlemail.com>
Date: Mon, 6 May 2024 18:27:39 +0200
Subject: [PATCH] improve versioned content map

---
 .../next-api/src/versioned_content_map.rs     | 85 ++++++++++---------
 1 file changed, 47 insertions(+), 38 deletions(-)

diff --git a/packages/next-swc/crates/next-api/src/versioned_content_map.rs b/packages/next-swc/crates/next-api/src/versioned_content_map.rs
index f2deaaf8e13d9..357f2969fa1fc 100644
--- a/packages/next-swc/crates/next-api/src/versioned_content_map.rs
+++ b/packages/next-swc/crates/next-api/src/versioned_content_map.rs
@@ -23,31 +23,30 @@ use turbopack_binding::{
 #[turbo_tasks::value(transparent)]
 pub struct OutputAssetsOperation(Vc<OutputAssets>);
 
-#[derive(
-    Clone, Copy, TraceRawVcs, PartialEq, Eq, ValueDebugFormat, Serialize, Deserialize, Debug,
-)]
+#[derive(Clone, TraceRawVcs, PartialEq, Eq, ValueDebugFormat, Serialize, Deserialize, Debug)]
 struct MapEntry {
     assets_operation: Vc<OutputAssets>,
     side_effects: Vc<Completion>,
+    path_to_asset: HashMap<Vc<FileSystemPath>, Vc<Box<dyn OutputAsset>>>,
 }
 
 #[turbo_tasks::value(transparent)]
 struct OptionMapEntry(Option<MapEntry>);
 
 type PathToOutputOperation = HashMap<Vc<FileSystemPath>, Vc<OutputAssets>>;
-type OutputOperationToSideEffects = HashMap<Vc<OutputAssets>, Vc<Completion>>;
+type OutputOperationToComputeEntry = HashMap<Vc<OutputAssets>, Vc<OptionMapEntry>>;
 
 #[turbo_tasks::value]
 pub struct VersionedContentMap {
     map_path_to_op: State<PathToOutputOperation>,
-    map_op_to_side_effects: State<OutputOperationToSideEffects>,
+    map_op_to_compute_entry: State<OutputOperationToComputeEntry>,
 }
 
 impl ValueDefault for VersionedContentMap {
     fn value_default() -> Vc<Self> {
         VersionedContentMap {
             map_path_to_op: State::new(HashMap::new()),
-            map_op_to_side_effects: State::new(HashMap::new()),
+            map_op_to_compute_entry: State::new(HashMap::new()),
         }
         .cell()
     }
@@ -71,31 +70,35 @@ impl VersionedContentMap {
         client_output_path: Vc<FileSystemPath>,
     ) -> Result<Vc<Completion>> {
         let this = self.await?;
-        let side_effects =
-            self.output_side_effects(assets_operation, client_relative_path, client_output_path);
+        let compute_entry =
+            self.compute_entry(assets_operation, client_relative_path, client_output_path);
         let assets = *assets_operation.await?;
-        this.map_op_to_side_effects
-            .update_conditionally(|map| map.insert(assets, side_effects) != Some(side_effects));
-        Ok(side_effects)
+        this.map_op_to_compute_entry
+            .update_conditionally(|map| map.insert(assets, compute_entry) != Some(compute_entry));
+        let Some(entry) = &*compute_entry.await? else {
+            unreachable!("compute_entry always returns Some(MapEntry)")
+        };
+        Ok(entry.side_effects)
     }
 
     #[turbo_tasks::function]
-    async fn output_side_effects(
+    async fn compute_entry(
         self: Vc<Self>,
         assets_operation: Vc<OutputAssetsOperation>,
         client_relative_path: Vc<FileSystemPath>,
         client_output_path: Vc<FileSystemPath>,
-    ) -> Result<Vc<Completion>> {
+    ) -> Result<Vc<OptionMapEntry>> {
         let assets = *assets_operation.await?;
         let entries: Vec<_> = assets
             .await?
             .iter()
-            .map(|&asset| async move { Ok((asset.ident().path().resolve().await?, assets)) })
+            .map(|&asset| async move { Ok((asset.ident().path().resolve().await?, asset, assets)) })
             .try_join()
             .await?;
-        self.await?.map_path_to_op.update_conditionally(move |map| {
+
+        self.await?.map_path_to_op.update_conditionally(|map| {
             let mut changed = false;
-            for (k, v) in entries {
+            for &(k, _, v) in entries.iter() {
                 if map.insert(k, v) != Some(v) {
                     changed = true;
                 }
@@ -103,11 +106,13 @@ impl VersionedContentMap {
             changed
         });
         // Make sure all written client assets are up-to-date
-        Ok(emit_client_assets(
-            assets,
-            client_relative_path,
-            client_output_path,
-        ))
+        let side_effects = emit_client_assets(assets, client_relative_path, client_output_path);
+        let map_entry = Vc::cell(Some(MapEntry {
+            assets_operation: assets,
+            side_effects,
+            path_to_asset: entries.into_iter().map(|(k, v, _)| (k, v)).collect(),
+        }));
+        Ok(map_entry)
     }
 
     #[turbo_tasks::function]
@@ -142,20 +147,21 @@ impl VersionedContentMap {
     ) -> Result<Vc<Box<dyn OutputAsset>>> {
         let result = self.raw_get(path).await?;
         if let Some(MapEntry {
-            assets_operation,
+            assets_operation: _,
             side_effects,
-        }) = *result
+            path_to_asset,
+        }) = &*result
         {
-            // NOTE(alexkirsz) This is necessary to mark the task as active again.
-            Vc::connect(assets_operation);
-            Vc::connect(side_effects);
-
             side_effects.await?;
 
-            for &asset in assets_operation.await?.iter() {
-                if asset.ident().path().resolve().await? == path {
-                    return Ok(asset);
-                }
+            if let Some(asset) = path_to_asset.get(&path) {
+                return Ok(*asset);
+            } else {
+                let path = path.to_string().await?;
+                bail!(
+                    "could not find asset for path {} (asset has been removed)",
+                    path
+                );
             }
         }
         let path = path.to_string().await?;
@@ -186,16 +192,19 @@ impl VersionedContentMap {
         let Some(assets) = assets else {
             return Ok(Vc::cell(None));
         };
-        let side_effects = {
-            let map = self.map_op_to_side_effects.get();
+        // Need to reconnect the operation to the map
+        Vc::connect(assets);
+
+        let compute_entry = {
+            let map = self.map_op_to_compute_entry.get();
             map.get(&assets).copied()
         };
-        let Some(side_effects) = side_effects else {
+        let Some(compute_entry) = compute_entry else {
             return Ok(Vc::cell(None));
         };
-        Ok(Vc::cell(Some(MapEntry {
-            assets_operation: assets,
-            side_effects,
-        })))
+        // Need to reconnect the operation to the map
+        Vc::connect(compute_entry);
+
+        Ok(compute_entry)
     }
 }