Skip to content

Commit

Permalink
optimize server actions (#66523)
Browse files Browse the repository at this point in the history 
### What?

optimize allocations in server actions transform

In one edge case it reduces allocations from 30GB to 4.5MB and time from
760ms to 11ms.

### Why?

make it faster
  • Loading branch information
@sokra @lubieowoce
sokra authored and lubieowoce committed Sep 3, 2024
1 parent 562df1f commit c8bde3e
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 83 deletions.
141 changes: 71 additions & 70 deletions packages/next-swc/crates/next-custom-transforms/src/transforms/server_actions.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -437,7 +437,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
let old_in_default_export_decl = self.in_default_export_decl;
self.in_action_fn = is_action_fn;
self.in_module_level = false;
self.should_track_names = true;
self.should_track_names = is_action_fn || self.should_track_names;
self.in_export_decl = false;
self.in_default_export_decl = false;
f.visit_mut_children_with(self);
Expand All @@ -448,8 +448,14 @@ impl<C: Comments> VisitMut for ServerActions<C> {
self.in_default_export_decl = old_in_default_export_decl;
}

let mut child_names = self.names.clone();
self.names.extend(current_names);
let mut child_names = if self.should_track_names {
let names = take(&mut self.names);
self.names = current_names;
self.names.extend(names.iter().cloned());
names
} else {
take(&mut self.names)
};

if !is_action_fn {
return;
Expand Down Expand Up @@ -510,7 +516,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
fn visit_mut_fn_decl(&mut self, f: &mut FnDecl) {
let is_action_fn = self.get_action_info(f.function.body.as_mut(), true);

let current_declared_idents = self.declared_idents.clone();
let declared_idents_until = self.declared_idents.len();
let current_names = take(&mut self.names);

{
Expand All @@ -522,7 +528,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
let old_in_default_export_decl = self.in_default_export_decl;
self.in_action_fn = is_action_fn;
self.in_module_level = false;
self.should_track_names = true;
self.should_track_names = is_action_fn || self.should_track_names;
self.in_export_decl = false;
self.in_default_export_decl = false;
f.visit_mut_children_with(self);
Expand All @@ -533,8 +539,14 @@ impl<C: Comments> VisitMut for ServerActions<C> {
self.in_default_export_decl = old_in_default_export_decl;
}

let mut child_names = self.names.clone();
self.names.extend(current_names);
let mut child_names = if self.should_track_names {
let names = take(&mut self.names);
self.names = current_names;
self.names.extend(names.iter().cloned());
names
} else {
take(&mut self.names)
};

if !is_action_fn {
return;
Expand All @@ -551,7 +563,10 @@ impl<C: Comments> VisitMut for ServerActions<C> {
if !(self.in_action_file && self.in_export_decl) {
// Collect all the identifiers defined inside the closure and used
// in the action function. With deduplication.
retain_names_from_declared_idents(&mut child_names, &current_declared_idents);
retain_names_from_declared_idents(
&mut child_names,
&self.declared_idents[..declared_idents_until],
);

let maybe_new_expr =
self.maybe_hoist_and_create_proxy(child_names, Some(&mut f.function), None);
Expand Down Expand Up @@ -616,12 +631,12 @@ impl<C: Comments> VisitMut for ServerActions<C> {
let old_in_default_export_decl = self.in_default_export_decl;
self.in_action_fn = is_action_fn;
self.in_module_level = false;
self.should_track_names = true;
self.should_track_names = is_action_fn || self.should_track_names;
self.in_export_decl = false;
self.in_default_export_decl = false;
{
for n in &mut a.params {
collect_pat_idents(n, &mut self.declared_idents);
collect_idents_in_pat(n, &mut self.declared_idents);
}
}
a.visit_mut_children_with(self);
Expand All @@ -632,8 +647,14 @@ impl<C: Comments> VisitMut for ServerActions<C> {
self.in_default_export_decl = old_in_default_export_decl;
}

let mut child_names = self.names.clone();
self.names.extend(current_names);
let mut child_names = if self.should_track_names {
let names = take(&mut self.names);
self.names = current_names;
self.names.extend(names.iter().cloned());
names
} else {
take(&mut self.names)
};

if !is_action_fn {
return;
Expand Down Expand Up @@ -672,7 +693,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {

// If it's a closure (not in the module level), we need to collect
// identifiers defined in the closure.
self.declared_idents.extend(collect_decl_idents_in_stmt(n));
collect_decl_idents_in_stmt(n, &mut self.declared_idents);
}

fn visit_mut_param(&mut self, n: &mut Param) {
Expand All @@ -682,7 +703,7 @@ impl<C: Comments> VisitMut for ServerActions<C> {
return;
}

collect_pat_idents(&n.pat, &mut self.declared_idents);
collect_idents_in_pat(&n.pat, &mut self.declared_idents);
}

fn visit_mut_prop_or_spread(&mut self, n: &mut PropOrSpread) {
Expand Down Expand Up @@ -746,7 +767,8 @@ impl<C: Comments> VisitMut for ServerActions<C> {
}
Decl::Var(var) => {
// export const foo = 1
let ids: Vec<Id> = collect_idents_in_var_decls(&var.decls);
let mut ids: Vec<Id> = Vec::new();
collect_idents_in_var_decls(&var.decls, &mut ids);
self.exported_idents.extend(
ids.into_iter().map(|id| (id.clone(), id.0.to_string())),
);
Expand Down Expand Up @@ -1208,26 +1230,6 @@ fn attach_name_to_expr(ident: Ident, expr: Expr, extra_items: &mut Vec<ModuleIte
}
}

fn collect_pat_idents(pat: &Pat, closure_idents: &mut Vec<Id>) {
match &pat {
Pat::Ident(ident) => {
closure_idents.push(ident.id.to_id());
}
Pat::Array(array) => {
closure_idents.extend(collect_idents_in_array_pat(&array.elems));
}
Pat::Object(object) => {
closure_idents.extend(collect_idents_in_object_pat(&object.props));
}
Pat::Rest(rest) => {
if let Pat::Ident(ident) = &*rest.arg {
closure_idents.push(ident.id.to_id());
}
}
_ => {}
}
}

fn generate_action_id(file_name: &str, export_name: &str) -> String {
// Attach a checksum to the action using sha1:
// $$id = sha1('file_name' + ':' + 'export_name');
Expand Down Expand Up @@ -1489,35 +1491,32 @@ fn remove_server_directive_index_in_fn(
});
}

fn collect_idents_in_array_pat(elems: &[Option<Pat>]) -> Vec<Id> {
let mut ids = Vec::new();

fn collect_idents_in_array_pat(elems: &[Option<Pat>], ids: &mut Vec<Id>) {
for elem in elems.iter().flatten() {
match elem {
Pat::Ident(ident) => {
ids.push(ident.id.to_id());
}
Pat::Array(array) => {
ids.extend(collect_idents_in_array_pat(&array.elems));
collect_idents_in_array_pat(&array.elems, ids);
}
Pat::Object(object) => {
ids.extend(collect_idents_in_object_pat(&object.props));
collect_idents_in_object_pat(&object.props, ids);
}
Pat::Rest(rest) => {
if let Pat::Ident(ident) = &*rest.arg {
ids.push(ident.id.to_id());
}
}
_ => {}
Pat::Assign(AssignPat { left, .. }) => {
collect_idents_in_pat(left, ids);
}
Pat::Expr(..) | Pat::Invalid(..) => {}
}
}

ids
}

fn collect_idents_in_object_pat(props: &[ObjectPatProp]) -> Vec<Id> {
let mut ids = Vec::new();

fn collect_idents_in_object_pat(props: &[ObjectPatProp], ids: &mut Vec<Id>) {
for prop in props {
match prop {
ObjectPatProp::KeyValue(KeyValuePatProp { key, value }) => {
Expand All @@ -1530,10 +1529,10 @@ fn collect_idents_in_object_pat(props: &[ObjectPatProp]) -> Vec<Id> {
ids.push(ident.id.to_id());
}
Pat::Array(array) => {
ids.extend(collect_idents_in_array_pat(&array.elems));
collect_idents_in_array_pat(&array.elems, ids);
}
Pat::Object(object) => {
ids.extend(collect_idents_in_object_pat(&object.props));
collect_idents_in_object_pat(&object.props, ids);
}
_ => {}
}
Expand All @@ -1548,39 +1547,41 @@ fn collect_idents_in_object_pat(props: &[ObjectPatProp]) -> Vec<Id> {
}
}
}

ids
}

fn collect_idents_in_var_decls(decls: &[VarDeclarator]) -> Vec<Id> {
let mut ids = Vec::new();

fn collect_idents_in_var_decls(decls: &[VarDeclarator], ids: &mut Vec<Id>) {
for decl in decls {
match &decl.name {
Pat::Ident(ident) => {
collect_idents_in_pat(&decl.name, ids);
}
}

fn collect_idents_in_pat(pat: &Pat, ids: &mut Vec<Id>) {
match pat {
Pat::Ident(ident) => {
ids.push(ident.id.to_id());
}
Pat::Array(array) => {
collect_idents_in_array_pat(&array.elems, ids);
}
Pat::Object(object) => {
collect_idents_in_object_pat(&object.props, ids);
}
Pat::Assign(AssignPat { left, .. }) => {
collect_idents_in_pat(left, ids);
}
Pat::Rest(RestPat { arg, .. }) => {
if let Pat::Ident(ident) = &**arg {
ids.push(ident.id.to_id());
}
Pat::Array(array) => {
ids.extend(collect_idents_in_array_pat(&array.elems));
}
Pat::Object(object) => {
ids.extend(collect_idents_in_object_pat(&object.props));
}
_ => {}
}
Pat::Expr(..) | Pat::Invalid(..) => {}
}

ids
}

fn collect_decl_idents_in_stmt(stmt: &Stmt) -> Vec<Id> {
let mut ids = Vec::new();

fn collect_decl_idents_in_stmt(stmt: &Stmt, ids: &mut Vec<Id>) {
if let Stmt::Decl(Decl::Var(var)) = &stmt {
ids.extend(collect_idents_in_var_decls(&var.decls));
collect_idents_in_var_decls(&var.decls, ids);
}

ids
}

pub(crate) struct ClosureReplacer<'a> {
Expand Down
14 changes: 7 additions & 7 deletions ...s/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/28/output.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ import { encryptActionBoundArgs, decryptActionBoundArgs } from "private-next-rsc
let a, f;
function Comp(b, c, ...g) {
return registerServerReference("9878bfa39811ca7650992850a8751f9591b6a557", $$ACTION_2).bind(null, encryptActionBoundArgs("9878bfa39811ca7650992850a8751f9591b6a557", [
b,
c,
g,
b
g
]));
}
export async function $$ACTION_0($$ACTION_CLOSURE_BOUND, e) {
Expand All @@ -23,19 +23,19 @@ export async function $$ACTION_2($$ACTION_CLOSURE_BOUND, d) {
console.log(...window, {
window
});
console.log(a, $$ACTION_ARG_2, action2);
console.log(a, $$ACTION_ARG_0, action2);
var action2 = registerServerReference("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_0).bind(null, encryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", [
$$ACTION_ARG_0,
$$ACTION_ARG_1,
d,
f,
$$ACTION_ARG_1
$$ACTION_ARG_2
]));
return [
action2,
registerServerReference("188d5d945750dc32e2c842b93c75a65763d4a922", $$ACTION_1).bind(null, encryptActionBoundArgs("188d5d945750dc32e2c842b93c75a65763d4a922", [
action2,
$$ACTION_ARG_0,
$$ACTION_ARG_1,
d
]))
];
}
}
12 changes: 6 additions & 6 deletions ...s/next-swc/crates/next-custom-transforms/tests/fixture/server-actions/server/30/output.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ import { encryptActionBoundArgs, decryptActionBoundArgs } from "private-next-rsc
let a, f;
export async function action0(b, c, ...g) {
return registerServerReference("9878bfa39811ca7650992850a8751f9591b6a557", $$ACTION_2).bind(null, encryptActionBoundArgs("9878bfa39811ca7650992850a8751f9591b6a557", [
b,
c,
g,
b
g
]));
}
export async function $$ACTION_0($$ACTION_CLOSURE_BOUND, e) {
Expand All @@ -23,18 +23,18 @@ export async function $$ACTION_2($$ACTION_CLOSURE_BOUND, d) {
console.log(...window, {
window
});
console.log(a, $$ACTION_ARG_2, action2);
console.log(a, $$ACTION_ARG_0, action2);
var action2 = registerServerReference("6d53ce510b2e36499b8f56038817b9bad86cabb4", $$ACTION_0).bind(null, encryptActionBoundArgs("6d53ce510b2e36499b8f56038817b9bad86cabb4", [
$$ACTION_ARG_0,
$$ACTION_ARG_1,
d,
f,
$$ACTION_ARG_1
$$ACTION_ARG_2
]));
return [
action2,
registerServerReference("188d5d945750dc32e2c842b93c75a65763d4a922", $$ACTION_1).bind(null, encryptActionBoundArgs("188d5d945750dc32e2c842b93c75a65763d4a922", [
action2,
$$ACTION_ARG_0,
$$ACTION_ARG_1,
d
]))
];
Expand Down

0 comments on commit c8bde3e

Please sign in to comment.