turbopack: drop ASTs of tree-shaken modules

@@ -1012,9 +1012,12 @@ impl AppProject {                 let (full, binding_usage_info) = if remove_unused_imports {                     let full_with_unused_references =                         ModuleGraph::from_graphs(graphs.clone(), None);+                    // This is a per-route graph: a module that is unused here might still be code+                    // generated via another graph, so dropping its data would only cause re-parses.                     let binding_usage_info = compute_binding_usage_info(                         full_with_unused_references,                         should_read_binding_usage,+                        /* drop_unused_module_data */ false,                     );                     (                         ModuleGraph::from_graphs(graphs, Some(binding_usage_info)),

@@ -919,7 +919,12 @@ impl PageEndpoint {              let graph = if remove_unused_imports {                 let graph = ModuleGraph::from_graphs(graphs.clone(), None);-                let binding_usage_info = compute_binding_usage_info(graph, true);+                // This is a per-route graph: a module that is unused here might still be code+                // generated via another graph, so dropping its data would only cause re-parses.+                let binding_usage_info = compute_binding_usage_info(+                    graph, /* remove_unused_imports */ true,+                    /* drop_unused_module_data */ false,+                );                 ModuleGraph::from_graphs(graphs, Some(binding_usage_info))             } else {                 ModuleGraph::from_graphs(graphs, None)

@@ -2942,7 +2942,10 @@ async fn whole_app_module_graph_operation(         let base = if turbopack_remove_unused_imports {             // TODO suboptimal that we do compute_binding_usage_info twice (once for the base             // graph and later for the full graph)-            let binding_usage_info = compute_binding_usage_info(base, true);+            let binding_usage_info = compute_binding_usage_info(+                base, /* remove_unused_imports */ true,+                /* drop_unused_module_data */ false,+            );             ModuleGraph::from_graphs(vec![base_single_module_graph], Some(binding_usage_info))         } else {             base@@ -2989,7 +2992,11 @@ async fn whole_app_module_graph_operation(          let (full, binding_usage_info) = if turbopack_remove_unused_imports {             let full_with_unused_references = ModuleGraph::from_graphs(graphs.clone(), None);-            let binding_usage_info = compute_binding_usage_info(full_with_unused_references, true);+            let binding_usage_info = compute_binding_usage_info(+                full_with_unused_references,+                /* remove_unused_imports */ true,+                /* drop_unused_module_data */ true,+            );             (                 ModuleGraph::from_graphs(graphs, Some(binding_usage_info)),                 Some(binding_usage_info),

@@ -1189,6 +1189,19 @@ impl TurboTasksBackend {         Ok(ReadOutcome::Scheduled(listener))     } +    fn drop_task_cell_content_hint(+        &self,+        task_id: TaskId,+        cell: CellId,+        turbo_tasks: &TurboTasks<TurboTasksBackend>,+    ) {+        let mut ctx = self.execute_context(turbo_tasks);+        let mut task = ctx.task(task_id, TaskDataCategory::Data);+        // Only drop the content, if it's there. A later read recomputes the task, exactly like+        // after a `final_read_hint` read.+        let _ = task.remove_cell_data(&cell, &get_value_type(cell.type_id()).persistence);+    }+     fn listen_to_cell(         &self,         task: &mut TaskGuard<'_>,@@ -3957,6 +3970,15 @@ impl Backend for TurboTasksBackend {         self.try_read_own_task_cell(task_id, cell, turbo_tasks)     } +    fn drop_task_cell_content_hint(+        &self,+        task_id: TaskId,+        cell: CellId,+        turbo_tasks: &TurboTasks<Self>,+    ) {+        self.drop_task_cell_content_hint(task_id, cell, turbo_tasks)+    }+     fn read_task_collectibles(         &self,         task_id: TaskId,

@@ -0,0 +1,92 @@+#![feature(arbitrary_self_types)]+#![feature(arbitrary_self_types_pointers)]+#![allow(clippy::needless_return)] // tokio macro-generated code doesn't respect this++use std::{+    sync::atomic::{AtomicU32, Ordering},+    time::Duration,+};++use anyhow::Result;+use turbo_tasks::Vc;+use turbo_tasks_testing::{Registration, register, run_once_without_cache_check};++static REGISTRATION: Registration = register!();++static COMPUTE_COUNT: AtomicU32 = AtomicU32::new(0);++#[tokio::test(flavor = "multi_thread", worker_threads = 2)]+async fn drop_content_hint() {+    // The test counts executions, so it must run only once, on a fresh instance.+    run_once_without_cache_check(&REGISTRATION, async {+        // Initial read computes the value.+        assert_eq!(*read_operation(1).read_strongly_consistent().await?, 42);+        assert_eq!(COMPUTE_COUNT.load(Ordering::SeqCst), 1);++        // Dropping present content doesn't execute the task.+        drop_operation(2).read_strongly_consistent().await?;+        // Dropping content that is not present anymore doesn't recompute it either.+        drop_operation(3).read_strongly_consistent().await?;+        // Give a (wrongly) scheduled recomputation a chance to run.+        tokio::time::sleep(Duration::from_millis(100)).await;+        assert_eq!(COMPUTE_COUNT.load(Ordering::SeqCst), 1);++        // A later read recomputes the dropped content.+        assert_eq!(*read_operation(4).read_strongly_consistent().await?, 42);+        assert_eq!(COMPUTE_COUNT.load(Ordering::SeqCst), 2);++        // And the content is kept again after that read.+        assert_eq!(*read_operation(5).read_strongly_consistent().await?, 42);+        assert_eq!(COMPUTE_COUNT.load(Ordering::SeqCst), 2);++        // Dropping through a chain of task outputs drops the cell the chain ends in.+        drop_through_output_operation(6)+            .read_strongly_consistent()+            .await?;+        assert_eq!(*read_operation(7).read_strongly_consistent().await?, 42);+        assert_eq!(COMPUTE_COUNT.load(Ordering::SeqCst), 3);++        anyhow::Ok(())+    })+    .await+    .unwrap()+}++/// Not serializable, so the content can't be restored from the persistent cache and must be+/// recomputed after it has been dropped.+#[turbo_tasks::value(serialization = "skip")]+struct Content(u32);++#[turbo_tasks::function]+fn compute() -> Vc<Content> {+    COMPUTE_COUNT.fetch_add(1, Ordering::SeqCst);+    Content(42).cell()+}++/// Returns the output of another task, so its own output points to a task output.+#[turbo_tasks::function]+fn forward() -> Vc<Content> {+    compute()+}++/// `nonce` makes every call a new task, so that it actually reads the cell again. (It has to be+/// used, as unused arguments are not part of the task key.)+#[turbo_tasks::function(operation, root)]+async fn read_operation(nonce: u32) -> Result<Vc<u32>> {+    let _ = nonce;+    Ok(Vc::cell(compute().await?.0))+}++#[turbo_tasks::function(operation, root)]+async fn drop_operation(nonce: u32) -> Result<Vc<()>> {+    let _ = nonce;+    compute().drop_content_hint().await?;+    Ok(Vc::cell(()))+}++#[turbo_tasks::function(operation, root)]+async fn drop_through_output_operation(nonce: u32) -> Result<Vc<()>> {+    let _ = nonce;+    forward().drop_content_hint().await?;+    Ok(Vc::cell(()))+}

@@ -2,7 +2,7 @@ use proc_macro::TokenStream; use proc_macro2::{Ident, TokenStream as TokenStream2}; use quote::{ToTokens, quote}; use syn::{-    Error, Expr, ExprLit, Generics, ImplItem, ImplItemFn, ItemImpl, Lit, LitStr, Meta,+    Attribute, Error, Expr, ExprLit, Generics, ImplItem, ImplItemFn, ItemImpl, Lit, LitStr, Meta,     MetaNameValue, Path, Token, Type,     parse::{Parse, ParseStream},     parse_macro_input,@@ -167,6 +167,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {     }      fn trait_value_impl(+        attrs: &[Attribute],         ty: &Type,         generics: &Generics,         ty_ident: &Ident,@@ -315,6 +316,10 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {             #[automatically_derived]             unsafe impl #impl_generics turbo_tasks::UpcastStrict<::std::boxed::Box<dyn #trait_path>> for #ty #where_clause {} +            // Forward the impl's attributes, so that e.g. `#[async_trait]` placed below+            // `#[turbo_tasks::value_impl]` applies to the plain (non-turbo-tasks) methods. The+            // turbo-tasks functions are already lowered to sync signatures at this point.+            #(#attrs)*             impl #impl_generics #trait_path for #ty #where_clause {                 #(#trait_items)*                 #(#trait_functions)*@@ -340,6 +345,7 @@ pub fn value_impl(args: TokenStream, input: TokenStream) -> TokenStream {     match &item.trait_ {         None => inherent_value_impl(&item.self_ty, &ty_ident, &item.items).into(),         Some((_, trait_path, _)) => trait_value_impl(+            &item.attrs,             &item.self_ty,             &item.generics,             &ty_ident,

@@ -676,6 +676,22 @@ pub trait Backend: Sized + Sync + Send {         turbo_tasks: &TurboTasks<Self>,     ) -> Result<TypedCellContent>; +    /// Hint that the content of the cell `cell` of task `task_id` is no longer needed, so the+    /// backend may drop it to free memory. A later read of a dropped cell recomputes the owning+    /// task, like after a `final_read_hint` read.+    ///+    /// The task must exist. The cell content doesn't need to be present (e.g. it was already+    /// dropped); this never reads or recomputes it.+    ///+    /// This doesn't track any dependency. The default implementation does nothing.+    fn drop_task_cell_content_hint(+        &self,+        _task_id: TaskId,+        _cell: CellId,+        _turbo_tasks: &TurboTasks<Self>,+    ) {+    }+     /// INVALIDATION: Be careful with this, when reader is None, it will not track dependencies, so     /// using it could break cache invalidation.     fn read_task_collectibles(

@@ -197,6 +197,9 @@ pub trait TurboTasksApi: TurboTasksCallApi + Sync + Send {         index: CellId,     ) -> Result<TypedCellContent>; +    /// See [`Backend::drop_task_cell_content_hint`].+    fn drop_task_cell_content_hint(&self, task: TaskId, index: CellId);+     fn read_own_task_cell(&self, task: TaskId, index: CellId) -> Result<TypedCellContent>;     fn update_own_task_cell(         &self,@@ -1951,6 +1954,10 @@ impl<B: Backend + 'static> TurboTasksApi for TurboTasks<B> {             .try_read_own_task_cell(current_task, index, self)     } +    fn drop_task_cell_content_hint(&self, task: TaskId, index: CellId) {+        self.backend.drop_task_cell_content_hint(task, index, self)+    }+     #[track_caller]     fn try_read_local_output(         &self,

@@ -439,6 +439,24 @@ where         }     } +    /// Hint that the cell content this `Vc` points to is no longer needed, so it may be dropped to+    /// free memory. Future reads might need to recompute the value.+    ///+    /// This resolves `self` like [`Vc::to_resolved`] does (tracking the dependencies on task+    /// outputs, and waiting for them if needed), but unlike [`Vc::final_read_hint`] it never reads+    /// the cell content. If the content isn't present (e.g. it was already dropped, or it was never+    /// computed because its dependents were restored from a persistent cache), this does nothing,+    /// and in particular doesn't recompute it.+    ///+    /// This does not track any dependency on the cell.+    pub async fn drop_content_hint(self) -> Result<()> {+        let RawVcUnpacked::TaskCell(task, index) = self.node.resolve().await?.unpack() else {+            unreachable!("resolve always returns a TaskCell");+        };+        crate::turbo_tasks().drop_task_cell_content_hint(task, index);+        Ok(())+    }+     /// Returns `true` if the reference is resolved, meaning the underlying [`RawVc`] uses the     /// [`RawVcUnpacked::TaskCell`] representation.     ///

@@ -320,7 +320,11 @@ async fn build_internal(         true,     );     let mut module_graph = ModuleGraph::from_graphs(vec![single_graph], None);-    let binding_usage = compute_binding_usage_info(module_graph, true);+    let binding_usage = compute_binding_usage_info(+        module_graph,+        /* remove_unused_imports */ true,+        /* drop_unused_module_data */ true,+    );     let unused_references = binding_usage         .connect()         .unused_references()

@@ -1,3 +1,5 @@+use anyhow::Result;+use async_trait::async_trait; use turbo_rcstr::RcStr; use turbo_tasks::{ResolvedVc, ValueToString, Vc}; @@ -37,6 +39,7 @@ pub enum ModuleSideEffects { /// /// For documentation about where this is used and how it fits into the rest of Turbopack, see /// [`crate::_layers`].+#[async_trait] #[turbo_tasks::value_trait] pub trait Module {     /// The identifier of the [`Module`]. It's expected to be unique and capture all properties of@@ -77,6 +80,17 @@ pub trait Module {     /// [packagejson]: https://webpack.js.org/guides/tree-shaking/#mark-the-file-as-side-effect-free     #[turbo_tasks::function]     fn side_effects(self: Vc<Self>) -> Vc<ModuleSideEffects>;++    /// Hints that the module won't be code generated (e.g. because tree shaking removed it), so+    /// data that it only keeps in memory for code generation (e.g. its parsed AST) can be dropped+    /// to free memory.+    ///+    /// Implementations must not compute that data if it isn't in memory, and must stay correct if+    /// the module is code generated nevertheless (e.g. by recomputing the data). This is not a+    /// turbo-tasks function, so calling it doesn't create a task. The default does nothing.+    async fn drop_code_gen_data_hint(&self) -> Result<()> {+        Ok(())+    } }  #[turbo_tasks::value_trait]

@@ -5,7 +5,7 @@ use auto_hash_map::AutoSet; use rustc_hash::{FxHashMap, FxHashSet}; use tracing::Instrument; use turbo_rcstr::RcStr;-use turbo_tasks::{OperationVc, ResolvedVc, Vc};+use turbo_tasks::{OperationVc, ResolvedVc, TryJoinIterExt, Vc};  use crate::{     chunk::chunking_context::UnusedReferences,@@ -113,10 +113,17 @@ impl BindingUsageInfo {     } } +/// Computes which exports and references of the modules in `graph` are used.+///+/// With `drop_unused_module_data`, modules that are not reachable via used references (i.e. that+/// remove_unused_imports removes) are hinted to drop data they only retain for code generation, see+/// [`Module::drop_code_gen_data_hint`]. Only enable this when this graph decides which modules+/// are code generated, otherwise the data might be needed again and has to be recomputed. #[turbo_tasks::function(operation)] pub async fn compute_binding_usage_info(     graph: OperationVc<ModuleGraph>,     remove_unused_imports: bool,+    drop_unused_module_data: bool, ) -> Result<Vc<BindingUsageInfo>> {     let span_outer = tracing::info_span!(         "compute binding usage info",@@ -374,6 +381,23 @@ pub async fn compute_binding_usage_info(         span.record("visit_count", visit_count);         span.record("unused_reference_count", unused_references.len()); +        if drop_unused_module_data {+            // The traversal visits every module that is reachable via used references (and sets+            // its usage), so the remaining ones are unused.+            graph_ref+                .iter_nodes()+                .filter(|module| !used_exports.contains_key(module))+                .map(async |module| {+                    module+                        .into_trait_ref()+                        .await?+                        .drop_code_gen_data_hint()+                        .await+                })+                .try_join()+                .await?;+        }+         #[cfg(debug_assertions)]         {             use std::sync::LazyLock;@@ -713,7 +737,11 @@ mod graph_tests {             true,         );         let graph = ModuleGraph::from_graphs(vec![single_graph], None);-        let binding_usage = compute_binding_usage_info(graph, false).connect().await?;+        let binding_usage = compute_binding_usage_info(+            graph, /* remove_unused_imports */ false, /* drop_unused_module_data */ false,+        )+        .connect()+        .await?;          let escapes = names             .iter()

@@ -885,7 +885,16 @@ impl EcmascriptModuleAsset { }  #[turbo_tasks::value_impl]+#[async_trait::async_trait] impl Module for EcmascriptModuleAsset {+    /// Drops the parsed AST. Analysis keeps it in memory for code generation, which drops it with+    /// `final_read_hint` after its final read.+    async fn drop_code_gen_data_hint(&self) -> Result<()> {+        // Analysis and code generation read the AST through `failsafe_parse`, which returns this+        // cell, unless parsing failed and the last successful parse is kept.+        self.parse().await?.drop_content_hint().await+    }+     #[turbo_tasks::function]     async fn ident(&self) -> Result<Vc<AssetIdent>> {         let mut ident = self.source.ident().owned().await?;

@@ -598,6 +598,7 @@ async fn run_test_operation(prepared_test: ResolvedVc<PreparedTest>) -> Result<V         Some(compute_binding_usage_info(             module_graph,             options.remove_unused_imports,+            /* drop_unused_module_data */ false,         ))     } else {         None

@@ -509,6 +509,7 @@ async fn run_test_operation(resource: RcStr) -> Result<Vc<FileSystemPath>> {         Some(compute_binding_usage_info(             module_graph,             options.remove_unused_imports,+            /* drop_unused_module_data */ false,         ))     } else {         None