Fix handling of local functions in "use cache" modules

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<C: Comments>(
         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<C: Comments> {
     private_ctxt: SyntaxContext,
 
     arrow_or_fn_expr_ident: Option<Ident>,
+    exported_local_ids: HashSet<Id>,
 }
 
 impl<C: Comments> ServerActions<C> {
-    // 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<String>) {
         let mut is_action_fn = false;
         let mut cache_type = None;
@@ -174,13 +176,12 @@ impl<C: Comments> ServerActions<C> {
             }
         }
 
-        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<C: Comments> VisitMut for ServerActions<C> {
     }
 
     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<C: Comments> VisitMut for ServerActions<C> {
         }
 
         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<C: Comments> VisitMut for ServerActions<C> {
                         .emit();
                 });
 
+                self.in_export_decl = old_in_export_decl;
+
                 return;
             }
 
@@ -1126,6 +1137,8 @@ impl<C: Comments> VisitMut for ServerActions<C> {
                 });
             }
         }
+
+        self.in_export_decl = old_in_export_decl;
     }
 
     fn visit_mut_method_prop(&mut self, m: &mut MethodProp) {
@@ -1332,6 +1345,32 @@ impl<C: Comments> VisitMut for ServerActions<C> {
             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<C: Comments> VisitMut for ServerActions<C> {
             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<C: Comments> VisitMut for ServerActions<C> {
     }
 
     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;
     }
 

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 }

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 };

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 <Foo />
+}

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 <Foo/>;
+});
+export var bar = registerServerReference($$RSC_SERVER_CACHE_0, "3128060c414d59f8552e4788b846c0d2b7f74743", null);