Repository: vercel/next.js
Pull request: Replace tuples with better suited data structures for server references

This PR introduces the `ServerReferenceExport` struct to replace `(Ident, Atom, Atom)` tuples for tracking exported server references. It also replaces the `Vec<(Atom, Atom)>` `export_actions` with an `FxIndexMap` for `reference_ids_by_export_name`, enabling more efficient lookups in upstack PRs. There are no functional changes.

Changed files:
filename: crates/next-custom-transforms/src/transforms/server_actions.rs
patch: @@ -70,6 +70,12 @@ enum ThisStatus {
     Forbidden { directive: Directive },
 }
 
+struct ServerReferenceExport {
+    ident: Ident,
+    export_name: Atom,
+    reference_id: Atom,
+}
+
 #[derive(Clone, Debug)]
 enum ServerActionsErrorKind {
     ExportedSyncFunction {
@@ -130,10 +136,6 @@ enum ServerActionsErrorKind {
     },
 }
 
-/// A mapping of hashed action id to the action's exported function name.
-// Using BTreeMap to ensure the order of the actions is deterministic.
-pub type ActionsMap = BTreeMap<Atom, Atom>;
-
 #[tracing::instrument(level = tracing::Level::TRACE, skip_all)]
 pub fn server_actions<C: Comments>(
     file_name: &FileName,
@@ -169,8 +171,6 @@ pub fn server_actions<C: Comments>(
         names: Default::default(),
         declared_idents: Default::default(),
 
-        exported_idents: Default::default(),
-
         // This flag allows us to rewrite `function foo() {}` to `const foo = createProxy(...)`.
         rewrite_fn_decl_to_proxy_decl: None,
         rewrite_default_fn_expr_to_proxy_expr: None,
@@ -179,7 +179,8 @@ pub fn server_actions<C: Comments>(
         annotations: Default::default(),
         extra_items: Default::default(),
         hoisted_extra_items: Default::default(),
-        export_actions: Default::default(),
+        reference_ids_by_export_name: Default::default(),
+        server_reference_exports: Default::default(),
 
         private_ctxt: SyntaxContext::empty().apply_mark(Mark::new()),
 
@@ -190,18 +191,22 @@ pub fn server_actions<C: Comments>(
     })
 }
 
-/// Serializes the Server Actions into a magic comment prefixed by
+/// Serializes the Server References into a magic comment prefixed by
 /// `__next_internal_action_entry_do_not_use__`.
-fn generate_server_actions_comment(
-    actions: &ActionsMap,
+fn generate_server_references_comment(
+    export_names_ordered_by_reference_id: &BTreeMap<&Atom, &Atom>,
     entry_path_query: Option<(&str, &str)>,
 ) -> String {
     format!(
         " __next_internal_action_entry_do_not_use__ {} ",
         if let Some(entry_path_query) = entry_path_query {
-            serde_json::to_string(&(actions, entry_path_query.0, entry_path_query.1))
+            serde_json::to_string(&(
+                export_names_ordered_by_reference_id,
+                entry_path_query.0,
+                entry_path_query.1,
+            ))
         } else {
-            serde_json::to_string(&actions)
+            serde_json::to_string(&export_names_ordered_by_reference_id)
         }
         .unwrap()
     )
@@ -239,16 +244,15 @@ struct ServerActions<C: Comments> {
     rewrite_default_fn_expr_to_proxy_expr: Option<Box<Expr>>,
     rewrite_expr_to_proxy_expr: Option<Box<Expr>>,
 
-    exported_idents: Vec<(
-        /* ident */ Ident,
-        /* name */ Atom,
-        /* id */ Atom,
-    )>,
-
     annotations: Vec<Stmt>,
     extra_items: Vec<ModuleItem>,
     hoisted_extra_items: Vec<ModuleItem>,
-    export_actions: Vec<(/* name */ Atom, /* id */ Atom)>,
+
+    /// A map of all server references (inline + exported): export_name -> reference_id
+    reference_ids_by_export_name: FxIndexMap<Atom, Atom>,
+
+    /// A list of server references for originally exported server functions only.
+    server_reference_exports: Vec<ServerReferenceExport>,
 
     private_ctxt: SyntaxContext,
 
@@ -460,8 +464,8 @@ impl<C: Comments> ServerActions<C> {
         let action_id = self.generate_server_reference_id(&action_name, false, Some(&new_params));
 
         self.has_action = true;
-        self.export_actions
-            .push((action_name.clone(), action_id.clone()));
+        self.reference_ids_by_export_name
+            .insert(action_name.clone(), action_id.clone());
 
         if let BlockStmtOrExpr::BlockStmt(block) = &mut *arrow.body {
             block.visit_mut_with(&mut ClosureReplacer {
@@ -612,8 +616,8 @@ impl<C: Comments> ServerActions<C> {
         let action_id = self.generate_server_reference_id(&action_name, false, Some(&new_params));
 
         self.has_action = true;
-        self.export_actions
-            .push((action_name.clone(), action_id.clone()));
+        self.reference_ids_by_export_name
+            .insert(action_name.clone(), action_id.clone());
 
         function.body.visit_mut_with(&mut ClosureReplacer {
             used_ids: &ids_from_closure,
@@ -742,8 +746,8 @@ impl<C: Comments> ServerActions<C> {
         let reference_id = self.generate_server_reference_id(&export_name, true, Some(&new_params));
 
         self.has_cache = true;
-        self.export_actions
-            .push((export_name.clone(), reference_id.clone()));
+        self.reference_ids_by_export_name
+            .insert(export_name.clone(), reference_id.clone());
 
         if let BlockStmtOrExpr::BlockStmt(block) = &mut *arrow.body {
             block.visit_mut_with(&mut ClosureReplacer {
@@ -826,8 +830,8 @@ impl<C: Comments> ServerActions<C> {
         let reference_id = self.generate_server_reference_id(&cache_name, true, Some(&new_params));
 
         self.has_cache = true;
-        self.export_actions
-            .push((cache_name.clone(), reference_id.clone()));
+        self.reference_ids_by_export_name
+            .insert(cache_name.clone(), reference_id.clone());
 
         function.body.visit_mut_with(&mut ClosureReplacer {
             used_ids: &ids_from_closure,
@@ -1419,7 +1423,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                 let (is_action_fn, is_cache_fn) =
                                     has_body_directive(&f.function.body);
 
-                                let ref_id = if is_action_fn {
+                                let is_cache = if is_action_fn {
                                     false
                                 } else if is_cache_fn {
                                     true
@@ -1433,15 +1437,15 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                 // TODO(shu): This is a workaround. We should have a better way
                                 // to skip self-annotated exports here.
                                 if !(is_cache_fn && self.config.is_react_server_layer) {
-                                    self.exported_idents.push((
-                                        f.ident.clone(),
-                                        f.ident.sym.clone(),
-                                        self.generate_server_reference_id(
+                                    self.server_reference_exports.push(ServerReferenceExport {
+                                        ident: f.ident.clone(),
+                                        export_name: f.ident.sym.clone(),
+                                        reference_id: self.generate_server_reference_id(
                                             f.ident.sym.as_ref(),
-                                            ref_id,
+                                            is_cache,
                                             Some(&f.function.params),
                                         ),
-                                    ));
+                                    });
                                 }
                             }
                             Decl::Var(var) => {
@@ -1450,15 +1454,15 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                 collect_idents_in_var_decls(&var.decls, &mut idents);
 
                                 for ident in &idents {
-                                    self.exported_idents.push((
-                                        ident.clone(),
-                                        ident.sym.clone(),
-                                        self.generate_server_reference_id(
+                                    self.server_reference_exports.push(ServerReferenceExport {
+                                        ident: ident.clone(),
+                                        export_name: ident.sym.clone(),
+                                        reference_id: self.generate_server_reference_id(
                                             ident.sym.as_ref(),
                                             in_cache_file,
                                             None,
                                         ),
-                                    ));
+                                    });
                                 }
 
                                 for decl in &mut var.decls {
@@ -1505,40 +1509,49 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                                 }) = export_name
                                                 {
                                                     // export { foo as bar }
-                                                    self.exported_idents.push((
-                                                        ident.clone(),
-                                                        sym.clone(),
-                                                        self.generate_server_reference_id(
-                                                            sym.as_ref(),
-                                                            in_cache_file,
-                                                            None,
-                                                        ),
-                                                    ));
+                                                    self.server_reference_exports.push(
+                                                        ServerReferenceExport {
+                                                            ident: ident.clone(),
+                                                            export_name: sym.clone(),
+                                                            reference_id: self
+                                                                .generate_server_reference_id(
+                                                                    sym.as_ref(),
+                                                                    in_cache_file,
+                                                                    None,
+                                                                ),
+                                                        },
+                                                    );
                                                 } else if let ModuleExportName::Str(str) =
                                                     export_name
                                                 {
                                                     // export { foo as "bar" }
-                                                    self.exported_idents.push((
-                                                        ident.clone(),
-                                                        str.value.clone(),
-                                                        self.generate_server_reference_id(
-                                                            str.value.as_ref(),
-                                                            in_cache_file,
-                                                            None,
-                                                        ),
-                                                    ));
+                                                    self.server_reference_exports.push(
+                                                        ServerReferenceExport {
+                                                            ident: ident.clone(),
+                                                            export_name: str.value.clone(),
+                                                            reference_id: self
+                                                                .generate_server_reference_id(
+                                                                    str.value.as_ref(),
+                                                                    in_cache_file,
+                                                                    None,
+                                                                ),
+                                                        },
+                                                    );
                                                 }
                                             } else {
                                                 // export { foo }
-                                                self.exported_idents.push((
-                                                    ident.clone(),
-                                                    ident.sym.clone(),
-                                                    self.generate_server_reference_id(
-                                                        ident.sym.as_ref(),
-                                                        in_cache_file,
-                                                        None,
-                                                    ),
-                                                ));
+                                                self.server_reference_exports.push(
+                                                    ServerReferenceExport {
+                                                        ident: ident.clone(),
+                                                        export_name: ident.sym.clone(),
+                                                        reference_id: self
+                                                            .generate_server_reference_id(
+                                                                ident.sym.as_ref(),
+                                                                in_cache_file,
+                                                                None,
+                                                            ),
+                                                    },
+                                                );
                                             }
                                         }
                                     } else {
@@ -1577,11 +1590,11 @@ impl<C: Comments> VisitMut for ServerActions<C> {
 
                                 if let Some(ident) = &f.ident {
                                     // export default function foo() {}
-                                    self.exported_idents.push((
-                                        ident.clone(),
-                                        atom!("default"),
-                                        ref_id,
-                                    ));
+                                    self.server_reference_exports.push(ServerReferenceExport {
+                                        ident: ident.clone(),
+                                        export_name: atom!("default"),
+                                        reference_id: ref_id,
+                                    });
                                 } else {
                                     // export default function() {}
                                     // Use the span from the function expression
@@ -1595,11 +1608,11 @@ impl<C: Comments> VisitMut for ServerActions<C> {
 
                                     f.ident = Some(new_ident.clone());
 
-                                    self.exported_idents.push((
-                                        new_ident.clone(),
-                                        atom!("default"),
-                                        ref_id,
-                                    ));
+                                    self.server_reference_exports.push(ServerReferenceExport {
+                                        ident: new_ident.clone(),
+                                        export_name: atom!("default"),
+                                        reference_id: ref_id,
+                                    });
 
                                     assign_name_to_ident(
                                         &new_ident,
@@ -1651,10 +1664,10 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                         self.private_ctxt,
                                     );
 
-                                    self.exported_idents.push((
-                                        new_ident.clone(),
-                                        atom!("default"),
-                                        self.generate_server_reference_id(
+                                    self.server_reference_exports.push(ServerReferenceExport {
+                                        ident: new_ident.clone(),
+                                        export_name: atom!("default"),
+                                        reference_id: self.generate_server_reference_id(
                                             "default",
                                             is_cache,
                                             Some(
@@ -1665,7 +1678,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                                     .collect(),
                                             ),
                                         ),
-                                    ));
+                                    });
 
                                     create_var_declarator(&new_ident, &mut self.extra_items);
                                     assign_name_to_ident(
@@ -1680,15 +1693,15 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                             }
                             Expr::Ident(ident) => {
                                 // export default foo
-                                self.exported_idents.push((
-                                    ident.clone(),
-                                    atom!("default"),
-                                    self.generate_server_reference_id(
+                                self.server_reference_exports.push(ServerReferenceExport {
+                                    ident: ident.clone(),
+                                    export_name: atom!("default"),
+                                    reference_id: self.generate_server_reference_id(
                                         "default",
                                         in_cache_file,
                                         None,
                                     ),
-                                ));
+                                });
                             }
                             Expr::Call(call) => {
                                 // export default fn()
@@ -1698,15 +1711,15 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                 let new_ident =
                                     Ident::new(self.gen_action_ident(), span, self.private_ctxt);
 
-                                self.exported_idents.push((
-                                    new_ident.clone(),
-                                    atom!("default"),
-                                    self.generate_server_reference_id(
+                                self.server_reference_exports.push(ServerReferenceExport {
+                                    ident: new_ident.clone(),
+                                    export_name: atom!("default"),
+                                    reference_id: self.generate_server_reference_id(
                                         "default",
                                         in_cache_file,
                                         None,
                                     ),
-                                ));
+                                });
 
                                 create_var_declarator(&new_ident, &mut self.extra_items);
                                 assign_name_to_ident(&new_ident, "default", &mut self.extra_items);
@@ -1763,27 +1776,19 @@ impl<C: Comments> VisitMut for ServerActions<C> {
             }
         }
 
-        let mut actions = self.export_actions.take();
-
         if in_action_file || in_cache_file && !self.config.is_react_server_layer {
-            actions.extend(
-                self.exported_idents
+            self.reference_ids_by_export_name.extend(
+                self.server_reference_exports
                     .iter()
-                    .map(|e| (e.1.clone(), e.2.clone())),
+                    .map(|e| (e.export_name.clone(), e.reference_id.clone())),
             );
 
-            if !actions.is_empty() {
+            if !self.reference_ids_by_export_name.is_empty() {
                 self.has_action |= in_action_file;
                 self.has_cache |= in_cache_file;
             }
         };
 
-        // Make it a hashmap of id -> name.
-        let actions = actions
-            .into_iter()
-            .map(|a| (a.1, a.0))
-            .collect::<ActionsMap>();
-
         // If it's compiled in the client layer, each export field needs to be
         // wrapped by a reference creation call.
         let create_ref_ident = private_ident!("createServerReference");
@@ -1836,7 +1841,12 @@ impl<C: Comments> VisitMut for ServerActions<C> {
 
         // If it's a "use server" or a "use cache" file, all exports need to be annotated.
         if should_track_exports {
-            for (ident, export_name, ref_id) in self.exported_idents.iter() {
+            for ServerReferenceExport {
+                ident,
+                export_name,
+                reference_id: ref_id,
+            } in &self.server_reference_exports
+            {
                 if !self.config.is_react_server_layer {
                     if export_name == "default" {
                         let export_expr = ModuleItem::ModuleDecl(ModuleDecl::ExportDefaultExpr(
@@ -1943,7 +1953,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                 new.append(&mut self.extra_items);
 
                 // For "use cache" files, there's no need to do extra annotations.
-                if !in_cache_file && !self.exported_idents.is_empty() {
+                if !in_cache_file && !self.server_reference_exports.is_empty() {
                     let ensure_ident = private_ident!("ensureServerEntryExports");
                     new.push(ModuleItem::ModuleDecl(ModuleDecl::Import(ImportDecl {
                         span: DUMMY_SP,
@@ -1972,9 +1982,9 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                 expr: Box::new(Expr::Array(ArrayLit {
                                     span: DUMMY_SP,
                                     elems: self
-                                        .exported_idents
+                                        .server_reference_exports
                                         .iter()
-                                        .map(|(ident, _, _)| {
+                                        .map(|ServerReferenceExport { ident, .. }| {
                                             Some(ExprOrSpread {
                                                 spread: None,
                                                 expr: Box::new(Expr::Ident(ident.clone())),
@@ -2096,15 +2106,23 @@ impl<C: Comments> VisitMut for ServerActions<C> {
         }
 
         if self.has_action || self.has_cache {
+            // Flip the map and convert it to a BTreeMap for deterministic
+            // ordering in the server references comment.
+            let export_names_ordered_by_reference_id = self
+                .reference_ids_by_export_name
+                .iter()
+                .map(|(export_name, reference_id)| (reference_id, export_name))
+                .collect::<BTreeMap<_, _>>();
+
             if self.config.is_react_server_layer {
                 // Prepend a special comment to the top of the file.
                 self.comments.add_leading(
                     self.start_pos,
                     Comment {
                         span: DUMMY_SP,
                         kind: CommentKind::Block,
-                        text: generate_server_actions_comment(
-                            &actions,
+                        text: generate_server_references_comment(
+                            &export_names_ordered_by_reference_id,
                             match self.mode {
                                 ServerActionsMode::Webpack => None,
                                 ServerActionsMode::Turbopack => Some(("", "")),
@@ -2121,7 +2139,11 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                             Comment {
                                 span: DUMMY_SP,
                                 kind: CommentKind::Block,
-                                text: generate_server_actions_comment(&actions, None).into(),
+                                text: generate_server_references_comment(
+                                    &export_names_ordered_by_reference_id,
+                                    None,
+                                )
+                                .into(),
                             },
                         );
                         new.push(client_layer_import.unwrap());
@@ -2165,8 +2187,8 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                                             Comment {
                                                 span: DUMMY_SP,
                                                 kind: CommentKind::Block,
-                                                text: generate_server_actions_comment(
-                                                    &std::iter::once((ref_id, export)).collect(),
+                                                text: generate_server_references_comment(
+                                                    &std::iter::once((&ref_id, &export)).collect(),
                                                     Some((
                                                         &self.file_name,
                                                         self.file_query.as_ref().map_or("", |v| v),