Fix Server Actions closure idents tracking (#66601)

shuding · web-flow · commit 04ce445105ee · 2024-06-06T18:08:10.000+02:00
This PR fixes the same case mention in #66464. Instead of collecting all values eagerly, here we merge fields (on any level of depth) of the same value and skip methods. For example: ```ts foo.bar foo.bar.baz qux.fn() ``` Previously we're (wrongly) collecting `[foo.bar, foo.bar.baz, qux.fn]`, and now it will be just `[foo.bar, qux]`. Merging of fields is critical for collecting methods correctly because in theory we can't tell if an object member is a method or not: ```ts data.push.call(data, 1) // or inside a function that does the same: doPush(data.push, data) ``` If we don't merge fields we'll collect `[data.push, data]` which still fails.
diff --git a/packages/next-swc/crates/next-custom-transforms/src/transforms/server_actions.rs b/packages/next-swc/crates/next-custom-transforms/src/transforms/server_actions.rs
@@ -48,6 +48,7 @@ pub fn server_actions<C: Comments>(
         in_action_file: false,
         in_export_decl: false,
         in_default_export_decl: false,
+        in_callee: false,
         has_action: false,
 
         reference_index: 0,
@@ -89,6 +90,7 @@ struct ServerActions<C: Comments> {
     in_action_file: bool,
     in_export_decl: bool,
     in_default_export_decl: bool,
+    in_callee: bool,
     has_action: bool,
 
     reference_index: u32,
@@ -709,9 +711,24 @@ impl<C: Comments> VisitMut for ServerActions<C> {
         n.visit_mut_children_with(self);
     }
 
+    fn visit_mut_callee(&mut self, n: &mut Callee) {
+        let old_in_callee = self.in_callee;
+        self.in_callee = true;
+        n.visit_mut_children_with(self);
+        self.in_callee = old_in_callee;
+    }
+
     fn visit_mut_expr(&mut self, n: &mut Expr) {
         if !self.in_module_level && self.should_track_names {
-            if let Ok(name) = Name::try_from(&*n) {
+            if let Ok(mut name) = Name::try_from(&*n) {
+                if self.in_callee {
+                    // This is a callee i.e. `foo.bar()`,
+                    // we need to track the actual value instead of the method name.
+                    if !name.1.is_empty() {
+                        name.1.pop();
+                    }
+                }
+
                 self.names.push(name);
                 self.should_track_names = false;
                 n.visit_mut_children_with(self);
@@ -1168,19 +1185,46 @@ impl<C: Comments> VisitMut for ServerActions<C> {
 }
 
 fn retain_names_from_declared_idents(child_names: &mut Vec<Name>, current_declared_idents: &[Id]) {
-    // Collect all the identifiers defined inside the closure and used
-    // in the action function. With deduplication.
-    let mut added_names = Vec::new();
-    child_names.retain(|name| {
-        if added_names.contains(name) {
-            false
-        } else if current_declared_idents.contains(&name.0) {
-            added_names.push(name.clone());
-            true
-        } else {
-            false
+    // Collect the names to retain in a separate vector
+    let mut retained_names = Vec::new();
+
+    for name in child_names.iter() {
+        let mut should_retain = true;
+
+        // Merge child_names. For example if both `foo.bar` and `foo.bar.baz` are used,
+        // we only need to keep `foo.bar` as it covers the other.
+
+        // Currently this is O(n^2) and we can potentially improve this to O(n log n)
+        // by sorting or using a hashset.
+        for another_name in child_names.iter() {
+            if name != another_name
+                && name.0 == another_name.0
+                && name.1.len() >= another_name.1.len()
+            {
+                let mut is_prefix = true;
+                for i in 0..another_name.1.len() {
+                    if name.1[i] != another_name.1[i] {
+                        is_prefix = false;
+                        break;
+                    }
+                }
+                if is_prefix {
+                    should_retain = false;
+                    break;
+                }
+            }
         }
-    });
+
+        if should_retain
+            && current_declared_idents.contains(&name.0)
+            && !retained_names.contains(name)
+        {
+            retained_names.push(name.clone());
+        }
+    }
+
+    // Replace the original child_names with the retained names
+    *child_names = retained_names;
 }
 
 fn gen_ident(cnt: &mut u32) -> JsWord {
diff --git a/packages/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/32/input.js b/packages/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/32/input.js
@@ -0,0 +1,14 @@
+export function Component() {
+  const data = [1, 2, 3]
+  const foo = [2, 3, 4]
+  const baz = { value: { current: 1 } }
+
+  async function action() {
+    'use server'
+
+    console.log(data.at(1), baz.value, baz.value.current)
+    console.log(foo.push.call(foo, 5))
+  }
+
+  return <form action={action} />
+}
diff --git a/packages/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/32/output.js b/packages/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/32/output.js
@@ -0,0 +1,30 @@
+/* __next_internal_action_entry_do_not_use__ {"6d53ce510b2e36499b8f56038817b9bad86cabb4":"$$ACTION_0"} */ import { registerServerReference } from "private-next-rsc-server-reference";
+import { encryptActionBoundArgs, decryptActionBoundArgs } from "private-next-rsc-action-encryption";
+export function Component() {
+    const data = [
+        1,
+        2,
+        3
+    ];
+    const foo = [
+        2,
+        3,
+        4
+    ];
+    const baz = {
+        value: {
+            current: 1
+        }
+    };
+    var action = registerServerReference("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_0).bind(null, encryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", [
+        data,
+        baz.value,
+        foo
+    ]));
+    return <form action={action}/>;
+}
+export async function $$ACTION_0($$ACTION_CLOSURE_BOUND) {
+    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2] = await decryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_CLOSURE_BOUND);
+    console.log($$ACTION_ARG_0.at(1), $$ACTION_ARG_1, $$ACTION_ARG_1.current);
+    console.log($$ACTION_ARG_2.push.call($$ACTION_ARG_2, 5));
+}
diff --git a/packages/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/7/output.js b/packages/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/7/output.js
@@ -3,61 +3,49 @@ import { encryptActionBoundArgs, decryptActionBoundArgs } from "private-next-rsc
 import deleteFromDb from 'db';
 export function Item1(product, foo, bar) {
     const a = registerServerReference("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_0).bind(null, encryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", [
-        product.id,
-        product?.foo,
-        product.bar.baz,
         product,
         foo,
         bar
     ]));
     return <Button action={a}>Delete</Button>;
 }
 export async function $$ACTION_0($$ACTION_CLOSURE_BOUND) {
-    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2, $$ACTION_ARG_3, $$ACTION_ARG_4, $$ACTION_ARG_5] = await decryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_CLOSURE_BOUND);
-    await deleteFromDb($$ACTION_ARG_3.id, $$ACTION_ARG_3?.foo, $$ACTION_ARG_3.bar.baz, $$ACTION_ARG_3[$$ACTION_ARG_4, $$ACTION_ARG_5]);
+    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2] = await decryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_CLOSURE_BOUND);
+    await deleteFromDb($$ACTION_ARG_0.id, $$ACTION_ARG_0?.foo, $$ACTION_ARG_0.bar.baz, $$ACTION_ARG_0[$$ACTION_ARG_1, $$ACTION_ARG_2]);
 }
 export function Item2(product, foo, bar) {
     var deleteItem2 = registerServerReference("188d5d945750dc32e2c842b93c75a65763d4a922", $$ACTION_1).bind(null, encryptActionBoundArgs("188d5d945750dc32e2c842b93c75a65763d4a922", [
-        product.id,
-        product?.foo,
-        product.bar.baz,
         product,
         foo,
         bar
     ]));
     return <Button action={deleteItem2}>Delete</Button>;
 }
 export async function $$ACTION_1($$ACTION_CLOSURE_BOUND) {
-    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2, $$ACTION_ARG_3, $$ACTION_ARG_4, $$ACTION_ARG_5] = await decryptActionBoundArgs("188d5d945750dc32e2c842b93c75a65763d4a922", $$ACTION_CLOSURE_BOUND);
-    await deleteFromDb($$ACTION_ARG_3.id, $$ACTION_ARG_3?.foo, $$ACTION_ARG_3.bar.baz, $$ACTION_ARG_3[$$ACTION_ARG_4, $$ACTION_ARG_5]);
+    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2] = await decryptActionBoundArgs("188d5d945750dc32e2c842b93c75a65763d4a922", $$ACTION_CLOSURE_BOUND);
+    await deleteFromDb($$ACTION_ARG_0.id, $$ACTION_ARG_0?.foo, $$ACTION_ARG_0.bar.baz, $$ACTION_ARG_0[$$ACTION_ARG_1, $$ACTION_ARG_2]);
 }
 export function Item3(product, foo, bar) {
     const deleteItem3 = registerServerReference("56a859f462d35a297c46a1bbd1e6a9058c104ab8", $$ACTION_3).bind(null, encryptActionBoundArgs("56a859f462d35a297c46a1bbd1e6a9058c104ab8", [
-        product.id,
-        product?.foo,
-        product.bar.baz,
         product,
         foo,
         bar
     ]));
     return <Button action={deleteItem3}>Delete</Button>;
 }
 export async function $$ACTION_3($$ACTION_CLOSURE_BOUND) {
-    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2, $$ACTION_ARG_3, $$ACTION_ARG_4, $$ACTION_ARG_5] = await decryptActionBoundArgs("56a859f462d35a297c46a1bbd1e6a9058c104ab8", $$ACTION_CLOSURE_BOUND);
-    await deleteFromDb($$ACTION_ARG_3.id, $$ACTION_ARG_3?.foo, $$ACTION_ARG_3.bar.baz, $$ACTION_ARG_3[$$ACTION_ARG_4, $$ACTION_ARG_5]);
+    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2] = await decryptActionBoundArgs("56a859f462d35a297c46a1bbd1e6a9058c104ab8", $$ACTION_CLOSURE_BOUND);
+    await deleteFromDb($$ACTION_ARG_0.id, $$ACTION_ARG_0?.foo, $$ACTION_ARG_0.bar.baz, $$ACTION_ARG_0[$$ACTION_ARG_1, $$ACTION_ARG_2]);
 }
 export function Item4(product, foo, bar) {
     const deleteItem4 = registerServerReference("9c0dd1f7c2b3f41d32e10f5c437de3d67ad32c6c", $$ACTION_4).bind(null, encryptActionBoundArgs("9c0dd1f7c2b3f41d32e10f5c437de3d67ad32c6c", [
-        product.id,
-        product?.foo,
-        product.bar.baz,
         product,
         foo,
         bar
     ]));
     return <Button action={deleteItem4}>Delete</Button>;
 }
 export async function $$ACTION_4($$ACTION_CLOSURE_BOUND) {
-    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2, $$ACTION_ARG_3, $$ACTION_ARG_4, $$ACTION_ARG_5] = await decryptActionBoundArgs("9c0dd1f7c2b3f41d32e10f5c437de3d67ad32c6c", $$ACTION_CLOSURE_BOUND);
-    await deleteFromDb($$ACTION_ARG_3.id, $$ACTION_ARG_3?.foo, $$ACTION_ARG_3.bar.baz, $$ACTION_ARG_3[$$ACTION_ARG_4, $$ACTION_ARG_5]);
+    var [$$ACTION_ARG_0, $$ACTION_ARG_1, $$ACTION_ARG_2] = await decryptActionBoundArgs("9c0dd1f7c2b3f41d32e10f5c437de3d67ad32c6c", $$ACTION_CLOSURE_BOUND);
+    await deleteFromDb($$ACTION_ARG_0.id, $$ACTION_ARG_0?.foo, $$ACTION_ARG_0.bar.baz, $$ACTION_ARG_0[$$ACTION_ARG_1, $$ACTION_ARG_2]);
 }
