From 9f374dec2042017d06bb6a0b02d696edf1f45c63 Mon Sep 17 00:00:00 2001 From: Hendrik Liebau Date: Fri, 18 Oct 2024 23:56:32 +0200 Subject: [PATCH] Fix handling of local functions in `"use cache"` modules --- .../src/transforms/server_actions.rs | 65 ++++++++++++++++--- .../fixture/server-actions/server/34/input.js | 21 +++++- .../server-actions/server/34/output.js | 23 +++++-- .../fixture/server-actions/server/45/input.js | 12 ++++ .../server-actions/server/45/output.js | 13 ++++ 5 files changed, 117 insertions(+), 17 deletions(-) create mode 100644 crates/next-custom-transforms/tests/fixture/server-actions/server/45/input.js create mode 100644 crates/next-custom-transforms/tests/fixture/server-actions/server/45/output.js diff --git a/crates/next-custom-transforms/src/transforms/server_actions.rs b/crates/next-custom-transforms/src/transforms/server_actions.rs index f6f8de78e3d9e..e674646d61c86 100644 --- a/crates/next-custom-transforms/src/transforms/server_actions.rs +++ b/crates/next-custom-transforms/src/transforms/server_actions.rs @@ -1,5 +1,5 @@ use std::{ - collections::BTreeMap, + collections::{BTreeMap, HashSet}, convert::{TryFrom, TryInto}, mem::take, }; @@ -75,6 +75,7 @@ pub fn server_actions( private_ctxt: SyntaxContext::empty().apply_mark(Mark::new()), arrow_or_fn_expr_ident: None, + exported_local_ids: HashSet::new(), }) } @@ -125,10 +126,11 @@ struct ServerActions { private_ctxt: SyntaxContext, arrow_or_fn_expr_ident: Option, + exported_local_ids: HashSet, } impl ServerActions { - // Check if the function or arrow function is an action function + // Check if the function or arrow function is an action or cache function fn get_body_info(&mut self, maybe_body: Option<&mut BlockStmt>) -> (bool, Option) { let mut is_action_fn = false; let mut cache_type = None; @@ -174,13 +176,12 @@ impl ServerActions { } } - if self.in_export_decl && self.in_action_file { - // All export functions in a server file are actions - is_action_fn = true; - } - - if self.in_module_level { - if let Some(cache_file_type) = &self.in_cache_file { + if self.in_export_decl { + if self.in_action_file { + // All export functions in a server file are actions + is_action_fn = true; + } else if let Some(cache_file_type) = &self.in_cache_file { + // All export functions in a cache file are cache functions cache_type = Some(cache_file_type.clone()); } } @@ -1012,6 +1013,12 @@ impl VisitMut for ServerActions { } fn visit_mut_fn_decl(&mut self, f: &mut FnDecl) { + let old_in_export_decl = self.in_export_decl; + + if self.in_module_level && self.exported_local_ids.contains(&f.ident.to_id()) { + self.in_export_decl = true + } + let (is_action_fn, cache_type) = self.get_body_info(f.function.body.as_mut()); let declared_idents_until = self.declared_idents.len(); @@ -1036,6 +1043,8 @@ impl VisitMut for ServerActions { } if !is_action_fn && cache_type.is_none() || !self.config.is_react_server_layer { + self.in_export_decl = old_in_export_decl; + return; } @@ -1059,6 +1068,8 @@ impl VisitMut for ServerActions { .emit(); }); + self.in_export_decl = old_in_export_decl; + return; } @@ -1126,6 +1137,8 @@ impl VisitMut for ServerActions { }); } } + + self.in_export_decl = old_in_export_decl; } fn visit_mut_method_prop(&mut self, m: &mut MethodProp) { @@ -1332,6 +1345,32 @@ impl VisitMut for ServerActions { self.config.enabled, ); + // If we're in a "use cache" file, collect all original IDs from export + // specifiers in a pre-pass so that we know which functions are + // exported, e.g. for this case: + // ``` + // "use cache" + // function Foo() {} + // export { Foo } + // ``` + if self.in_cache_file.is_some() { + for stmt in stmts.iter() { + if let ModuleItem::ModuleDecl(ModuleDecl::ExportNamed(named)) = stmt { + if named.src.is_none() { + for spec in &named.specifiers { + if let ExportSpecifier::Named(ExportNamedSpecifier { + orig: ModuleExportName::Ident(ident), + .. + }) = spec + { + self.exported_local_ids.insert(ident.to_id()); + } + } + } + } + } + } + let old_annotations = self.annotations.take(); let mut new = Vec::with_capacity(stmts.len()); @@ -1594,7 +1633,7 @@ impl VisitMut for ServerActions { new.rotate_right(1); } - // If it's a "use server" file, all exports need to be annotated as actions. + // If it's a "use server" or a "use cache" file, all exports need to be annotated. if self.in_action_file || self.in_cache_file.is_some() { for (ident, export_name) in self.exported_idents.iter() { if !self.config.is_react_server_layer { @@ -1881,16 +1920,22 @@ impl VisitMut for ServerActions { } fn visit_mut_var_declarator(&mut self, var_declarator: &mut VarDeclarator) { + let old_in_export_decl = self.in_export_decl; let old_arrow_expr_ident = self.arrow_or_fn_expr_ident.take(); if let (Pat::Ident(ident), Some(box Expr::Arrow(_) | box Expr::Fn(_))) = (&var_declarator.name, &var_declarator.init) { + if self.in_module_level && self.exported_local_ids.contains(&ident.to_id()) { + self.in_export_decl = true + } + self.arrow_or_fn_expr_ident = Some(ident.id.clone()); } var_declarator.visit_mut_children_with(self); + self.in_export_decl = old_in_export_decl; self.arrow_or_fn_expr_ident = old_arrow_expr_ident; } diff --git a/crates/next-custom-transforms/tests/fixture/server-actions/server/34/input.js b/crates/next-custom-transforms/tests/fixture/server-actions/server/34/input.js index 50b9988d8867b..3abfe00850efa 100644 --- a/crates/next-custom-transforms/tests/fixture/server-actions/server/34/input.js +++ b/crates/next-custom-transforms/tests/fixture/server-actions/server/34/input.js @@ -1,5 +1,22 @@ 'use cache' -export async function foo() {} -const bar = async () => {} +const foo = async () => { + return 'foo' +} + export { bar } + +async function bar() { + return 'bar' +} + +// Should not be wrapped in $$cache__. +const qux = async function qux() { + return 'qux' +} + +const baz = async function () { + return qux() + 'baz' +} + +export { foo, baz } diff --git a/crates/next-custom-transforms/tests/fixture/server-actions/server/34/output.js b/crates/next-custom-transforms/tests/fixture/server-actions/server/34/output.js index 5ac2d6476edd8..004ff4e5b57fd 100644 --- a/crates/next-custom-transforms/tests/fixture/server-actions/server/34/output.js +++ b/crates/next-custom-transforms/tests/fixture/server-actions/server/34/output.js @@ -1,8 +1,21 @@ -/* __next_internal_action_entry_do_not_use__ {"3128060c414d59f8552e4788b846c0d2b7f74743":"$$RSC_SERVER_CACHE_0","951c375b4a6a6e89d67b743ec5808127cfde405d":"$$RSC_SERVER_CACHE_1"} */ import { registerServerReference } from "private-next-rsc-server-reference"; +/* __next_internal_action_entry_do_not_use__ {"12a8d21b6362b4cc8f5b15560525095bc48dba80":"$$RSC_SERVER_CACHE_3","3128060c414d59f8552e4788b846c0d2b7f74743":"$$RSC_SERVER_CACHE_0","951c375b4a6a6e89d67b743ec5808127cfde405d":"$$RSC_SERVER_CACHE_1"} */ import { registerServerReference } from "private-next-rsc-server-reference"; import { encryptActionBoundArgs, decryptActionBoundArgs } from "private-next-rsc-action-encryption"; import { cache as $$cache__ } from "private-next-rsc-cache-wrapper"; -export var $$RSC_SERVER_CACHE_0 = $$cache__("default", "3128060c414d59f8552e4788b846c0d2b7f74743", async function foo() {}); -export var foo = registerServerReference($$RSC_SERVER_CACHE_0, "3128060c414d59f8552e4788b846c0d2b7f74743", null); -export var $$RSC_SERVER_CACHE_1 = $$cache__("default", "951c375b4a6a6e89d67b743ec5808127cfde405d", async function() {}); -const bar = registerServerReference($$RSC_SERVER_CACHE_1, "951c375b4a6a6e89d67b743ec5808127cfde405d", null); +export var $$RSC_SERVER_CACHE_0 = $$cache__("default", "3128060c414d59f8552e4788b846c0d2b7f74743", async function() { + return 'foo'; +}); +const foo = registerServerReference($$RSC_SERVER_CACHE_0, "3128060c414d59f8552e4788b846c0d2b7f74743", null); export { bar }; +export var $$RSC_SERVER_CACHE_1 = $$cache__("default", "951c375b4a6a6e89d67b743ec5808127cfde405d", async function bar() { + return 'bar'; +}); +var bar = registerServerReference($$RSC_SERVER_CACHE_1, "951c375b4a6a6e89d67b743ec5808127cfde405d", null); +// Should not be wrapped in $$cache__. +const qux = async function qux() { + return 'qux'; +}; +export var $$RSC_SERVER_CACHE_3 = $$cache__("default", "12a8d21b6362b4cc8f5b15560525095bc48dba80", async function $$RSC_SERVER_CACHE_2() { + return qux() + 'baz'; +}); +const baz = registerServerReference($$RSC_SERVER_CACHE_3, "12a8d21b6362b4cc8f5b15560525095bc48dba80", null); +export { foo, baz }; diff --git a/crates/next-custom-transforms/tests/fixture/server-actions/server/45/input.js b/crates/next-custom-transforms/tests/fixture/server-actions/server/45/input.js new file mode 100644 index 0000000000000..1ca81397c8e70 --- /dev/null +++ b/crates/next-custom-transforms/tests/fixture/server-actions/server/45/input.js @@ -0,0 +1,12 @@ +'use cache' + +// Expect no error here, this is allowed to be sync because it's not exported. +function Foo() { + const v = Math.random() + console.log(v) + return v +} + +export async function bar() { + return +} diff --git a/crates/next-custom-transforms/tests/fixture/server-actions/server/45/output.js b/crates/next-custom-transforms/tests/fixture/server-actions/server/45/output.js new file mode 100644 index 0000000000000..f093af51af6f0 --- /dev/null +++ b/crates/next-custom-transforms/tests/fixture/server-actions/server/45/output.js @@ -0,0 +1,13 @@ +/* __next_internal_action_entry_do_not_use__ {"3128060c414d59f8552e4788b846c0d2b7f74743":"$$RSC_SERVER_CACHE_0"} */ import { registerServerReference } from "private-next-rsc-server-reference"; +import { encryptActionBoundArgs, decryptActionBoundArgs } from "private-next-rsc-action-encryption"; +import { cache as $$cache__ } from "private-next-rsc-cache-wrapper"; +// Expect no error here, this is allowed to be sync because it's not exported. +function Foo() { + const v = Math.random(); + console.log(v); + return v; +} +export var $$RSC_SERVER_CACHE_0 = $$cache__("default", "3128060c414d59f8552e4788b846c0d2b7f74743", async function bar() { + return ; +}); +export var bar = registerServerReference($$RSC_SERVER_CACHE_0, "3128060c414d59f8552e4788b846c0d2b7f74743", null);