Repository: vercel/next.js
Pull request: perf(turbopack): Optimize `next_page_static_info`

### What?

Avoid using `FxHashMap` for the `next_page_static_info` pass and use a custom trait instead. The goal of using a custom trait is to keep it general and make the responsibility of each code. 

### Why?

We do not need to allocate at all. 

Also, visitor calls are expensive, and we don't have to pay for them in this case.

### How?

 - Closes SWC-655



Closes PACK-4604

Changed files:
filename: crates/next-core/src/next_shared/transforms/next_page_static_info.rs
patch: @@ -1,10 +1,13 @@
 use anyhow::Result;
 use async_trait::async_trait;
 use next_custom_transforms::transforms::page_static_info::{
-    Const, collect_exports, extract_exported_const_values,
+    Const, collect_exported_const_visitor::GetMut, collect_exports, extract_exported_const_values,
 };
 use serde_json::Value;
-use swc_core::{atoms::atom, ecma::ast::Program};
+use swc_core::{
+    atoms::{Atom, atom},
+    ecma::ast::Program,
+};
 use turbo_tasks::{ResolvedVc, Vc};
 use turbo_tasks_fs::FileSystemPath;
 use turbopack::module_options::{ModuleRule, ModuleRuleEffect};
@@ -45,15 +48,27 @@ struct NextPageStaticInfo {
     client_context: Option<ClientContextType>,
 }
 
+#[derive(Default)]
+struct PropertiesToExtract {
+    config: Option<Const>,
+}
+impl GetMut<Atom, Option<Const>> for PropertiesToExtract {
+    fn get_mut(&mut self, key: &Atom) -> Option<&mut Option<Const>> {
+        if key == &atom!("config") {
+            Some(&mut self.config)
+        } else {
+            None
+        }
+    }
+}
 #[async_trait]
 impl CustomTransformer for NextPageStaticInfo {
     #[tracing::instrument(level = tracing::Level::TRACE, name = "next_page_static_info", skip_all)]
     async fn transform(&self, program: &mut Program, ctx: &TransformContext<'_>) -> Result<()> {
         if let Some(collected_exports) = collect_exports(program)? {
-            let mut properties_to_extract = collected_exports.extra_properties.clone();
-            properties_to_extract.insert(atom!("config"));
+            let mut properties_to_extract = PropertiesToExtract::default();
 
-            let extracted = extract_exported_const_values(program, properties_to_extract);
+            extract_exported_const_values(program, &mut properties_to_extract);
 
             let is_server_layer_page = matches!(
                 self.server_context,
@@ -82,8 +97,7 @@ impl CustomTransformer for NextPageStaticInfo {
             }
 
             if is_app_page {
-                if let Some(Some(Const::Value(Value::Object(config_obj)))) =
-                    extracted.get(&atom!("config"))
+                if let Some(Const::Value(Value::Object(config_obj))) = properties_to_extract.config
                 {
                     let mut messages = vec![format!(
                         "Page config in {} is deprecated. Replace `export const config=…` with \

filename: crates/next-custom-transforms/src/transforms/page_static_info/collect_exported_const_visitor.rs
patch: @@ -1,15 +1,13 @@
-use rustc_hash::{FxHashMap, FxHashSet};
 use serde_json::{Map, Number, Value};
 use swc_core::{
     atoms::Atom,
     common::{Mark, SyntaxContext},
     ecma::{
         ast::{
-            BindingIdent, Decl, ExportDecl, Expr, Lit, ModuleDecl, ModuleItem, Pat, Prop, PropName,
-            PropOrSpread, VarDecl, VarDeclKind, VarDeclarator,
+            BindingIdent, Decl, ExportDecl, Expr, Lit, Module, ModuleDecl, ModuleItem, Pat,
+            Program, Prop, PropName, PropOrSpread, VarDecl, VarDeclKind, VarDeclarator,
         },
         utils::{ExprCtx, ExprExt},
-        visit::{Visit, VisitWith},
     },
 };
 
@@ -21,18 +19,21 @@ pub enum Const {
     Unsupported(String),
 }
 
-pub(crate) struct CollectExportedConstVisitor {
-    pub properties: FxHashMap<Atom, Option<Const>>,
+pub(crate) struct CollectExportedConstVisitor<'a, M>
+where
+    M: GetMut<Atom, Option<Const>>,
+{
+    pub properties: &'a mut M,
     expr_ctx: ExprCtx,
 }
 
-impl CollectExportedConstVisitor {
-    pub fn new(properties_to_extract: FxHashSet<Atom>) -> Self {
+impl<'a, M> CollectExportedConstVisitor<'a, M>
+where
+    M: GetMut<Atom, Option<Const>>,
+{
+    pub fn new(properties: &'a mut M) -> Self {
         Self {
-            properties: properties_to_extract
-                .into_iter()
-                .map(|p| (p, None))
-                .collect(),
+            properties,
             expr_ctx: ExprCtx {
                 unresolved_ctxt: SyntaxContext::empty().apply_mark(Mark::new()),
                 is_unresolved_ref_safe: false,
@@ -41,11 +42,13 @@ impl CollectExportedConstVisitor {
             },
         }
     }
-}
 
-impl Visit for CollectExportedConstVisitor {
-    fn visit_module_items(&mut self, module_items: &[ModuleItem]) {
-        for module_item in module_items {
+    pub fn check_program(&mut self, program: &Program) {
+        let Program::Module(Module { body, .. }) = program else {
+            return;
+        };
+
+        for module_item in body {
             if let ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(ExportDecl {
                 decl: Decl::Var(decl),
                 ..
@@ -69,11 +72,13 @@ impl Visit for CollectExportedConstVisitor {
                 }
             }
         }
-
-        module_items.visit_children_with(self);
     }
 }
 
+pub trait GetMut<K, V> {
+    fn get_mut(&mut self, key: &K) -> Option<&mut V>;
+}
+
 /// Coerece the actual value of the given ast node.
 fn extract_value(ctx: ExprCtx, init: &Expr, id: String) -> Option<Const> {
     match init {

filename: crates/next-custom-transforms/src/transforms/page_static_info/mod.rs
patch: @@ -3,7 +3,7 @@ pub use collect_exported_const_visitor::Const;
 use collect_exports_visitor::CollectExportsVisitor;
 use once_cell::sync::Lazy;
 use regex::Regex;
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashSet;
 use serde::{Deserialize, Serialize};
 use swc_core::{
     atoms::Atom,
@@ -12,6 +12,8 @@ use swc_core::{
     ecma::{ast::Program, visit::VisitWith},
 };
 
+use crate::transforms::page_static_info::collect_exported_const_visitor::GetMut;
+
 pub mod collect_exported_const_visitor;
 pub mod collect_exports_visitor;
 
@@ -209,15 +211,13 @@ pub fn collect_rsc_module_info(
 /// error.
 pub fn extract_exported_const_values(
     source_ast: &Program,
-    properties_to_extract: FxHashSet<Atom>,
-) -> FxHashMap<Atom, Option<Const>> {
+    properties_to_extract: &mut impl GetMut<Atom, Option<Const>>,
+) {
     GLOBALS.set(&Default::default(), || {
         let mut visitor =
             collect_exported_const_visitor::CollectExportedConstVisitor::new(properties_to_extract);
 
-        source_ast.visit_with(&mut visitor);
-
-        visitor.properties
+        visitor.check_program(source_ast);
     })
 }