| 1 | // Gating rewrite logic for compiled functions. |
| 2 | // |
| 3 | // When gating is enabled, the compiled function is wrapped in a conditional: |
| 4 | // `gating() ? optimized_fn : original_fn` |
| 5 | // |
| 6 | // For function declarations referenced before their declaration, a special |
| 7 | // hoisting pattern is used (see `insert_additional_function_declaration`). |
| 8 | // |
| 9 | // Ported from `Entrypoint/Gating.ts`. |
| 10 | |
| 11 | use react_compiler_ast::common::BaseNode; |
| 12 | use react_compiler_ast::expressions::*; |
| 13 | use react_compiler_ast::patterns::PatternLike; |
| 14 | use react_compiler_ast::statements::*; |
| 15 | use react_compiler_diagnostics::CompilerDiagnostic; |
| 16 | use react_compiler_diagnostics::ErrorCategory; |
| 17 | |
| 18 | use super::imports::ProgramContext; |
| 19 | use super::plugin_options::GatingConfig; |
| 20 | |
| 21 | /// A compiled function node, can be any function type. |
| 22 | #[derive(Debug, Clone)] |
| 23 | pub enum CompiledFunctionNode { |
| 24 | FunctionDeclaration(FunctionDeclaration), |
| 25 | FunctionExpression(FunctionExpression), |
| 26 | ArrowFunctionExpression(ArrowFunctionExpression), |
| 27 | } |
| 28 | |
| 29 | /// Represents a compiled function that needs gating. |
| 30 | /// In the Rust version, we work with indices into the program body |
| 31 | /// rather than Babel paths. |
| 32 | pub struct GatingRewrite { |
| 33 | /// Index in program.body where the original function is |
| 34 | pub original_index: usize, |
| 35 | /// The compiled function AST node |
| 36 | pub compiled_fn: CompiledFunctionNode, |
| 37 | /// The gating config |
| 38 | pub gating: GatingConfig, |
| 39 | /// Whether the function is referenced before its declaration at top level |
| 40 | pub referenced_before_declared: bool, |
| 41 | /// Whether the parent statement is an ExportDefaultDeclaration |
| 42 | pub is_export_default: bool, |
| 43 | } |
| 44 | |
| 45 | /// Apply gating rewrites to the program. |
| 46 | /// This modifies program.body by replacing/inserting statements. |
| 47 | /// |
| 48 | /// Corresponds to `insertGatedFunctionDeclaration` in the TS version, |
| 49 | /// but batched: all rewrites are collected first, then applied in reverse |
| 50 | /// index order to maintain validity of earlier indices. |
| 51 | pub fn apply_gating_rewrites( |
| 52 | program: &mut react_compiler_ast::Program, |
| 53 | mut rewrites: Vec<GatingRewrite>, |
| 54 | context: &mut ProgramContext, |
| 55 | ) -> Result<(), CompilerDiagnostic> { |
| 56 | // Sort rewrites in reverse order by original_index so that insertions |
| 57 | // at higher indices don't invalidate lower indices. |
| 58 | rewrites.sort_by(|a, b| b.original_index.cmp(&a.original_index)); |
| 59 | |
| 60 | for rewrite in rewrites { |
| 61 | let gating_imported_name = context |
| 62 | .add_import_specifier( |
| 63 | &rewrite.gating.source, |
| 64 | &rewrite.gating.import_specifier_name, |
| 65 | None, |
| 66 | ) |
| 67 | .name |
| 68 | .clone(); |
| 69 | |
| 70 | if rewrite.referenced_before_declared { |
| 71 | // The referenced-before-declared case only applies to FunctionDeclarations |
| 72 | if let CompiledFunctionNode::FunctionDeclaration(compiled) = rewrite.compiled_fn { |
| 73 | insert_additional_function_declaration( |
| 74 | &mut program.body, |
| 75 | rewrite.original_index, |
| 76 | compiled, |
| 77 | context, |
| 78 | &gating_imported_name, |
| 79 | )?; |
| 80 | } else { |
| 81 | return Err(CompilerDiagnostic::new( |
| 82 | ErrorCategory::Invariant, |
| 83 | "Expected compiled node type to match input type: \ |
| 84 | got non-FunctionDeclaration but expected FunctionDeclaration", |
| 85 | None, |
| 86 | )); |
| 87 | } |
| 88 | } else { |
| 89 | let original_stmt = program.body[rewrite.original_index].clone(); |
| 90 | let original_fn = extract_function_node_from_stmt(&original_stmt)?; |
| 91 | |
| 92 | let gating_expression = |
| 93 | build_gating_expression(rewrite.compiled_fn, original_fn, &gating_imported_name); |
| 94 | |
| 95 | // Determine how to rewrite based on context |
| 96 | if !rewrite.is_export_default { |
| 97 | if let Some(fn_name) = get_fn_decl_name(&original_stmt) { |
| 98 | // Convert function declaration to: const fnName = gating() ? compiled : original |
| 99 | let var_decl = Statement::VariableDeclaration(VariableDeclaration { |
| 100 | base: BaseNode::default(), |
| 101 | declarations: vec![VariableDeclarator { |
| 102 | base: BaseNode::default(), |
| 103 | id: PatternLike::Identifier(make_identifier(&fn_name)), |
| 104 | init: Some(Box::new(gating_expression)), |
| 105 | definite: None, |
| 106 | }], |
| 107 | kind: VariableDeclarationKind::Const, |
| 108 | declare: None, |
| 109 | }); |
| 110 | program.body[rewrite.original_index] = var_decl; |
| 111 | } else { |
| 112 | // Replace with the conditional expression directly (e.g. arrow/expression) |
| 113 | let expr_stmt = Statement::ExpressionStatement(ExpressionStatement { |
| 114 | base: BaseNode::default(), |
| 115 | expression: Box::new(gating_expression), |
| 116 | }); |
| 117 | program.body[rewrite.original_index] = expr_stmt; |
| 118 | } |
| 119 | } else { |
| 120 | // ExportDefaultDeclaration case |
| 121 | if let Some(fn_name) = get_fn_decl_name_from_export_default(&original_stmt) { |
| 122 | // Named export default function: replace with const + re-export |
| 123 | // const fnName = gating() ? compiled : original; |
| 124 | // export default fnName; |
| 125 | let var_decl = Statement::VariableDeclaration(VariableDeclaration { |
| 126 | base: BaseNode::default(), |
| 127 | declarations: vec![VariableDeclarator { |
| 128 | base: BaseNode::default(), |
| 129 | id: PatternLike::Identifier(make_identifier(&fn_name)), |
| 130 | init: Some(Box::new(gating_expression)), |
| 131 | definite: None, |
| 132 | }], |
| 133 | kind: VariableDeclarationKind::Const, |
| 134 | declare: None, |
| 135 | }); |
| 136 | let re_export = Statement::ExportDefaultDeclaration( |
| 137 | react_compiler_ast::declarations::ExportDefaultDeclaration { |
| 138 | base: BaseNode::default(), |
| 139 | declaration: Box::new( |
| 140 | react_compiler_ast::declarations::ExportDefaultDecl::Expression( |
| 141 | Box::new(Expression::Identifier(make_identifier(&fn_name))), |
| 142 | ), |
| 143 | ), |
| 144 | export_kind: None, |
| 145 | }, |
| 146 | ); |
| 147 | // Replace the original statement with the var decl, then insert re-export after |
| 148 | program.body[rewrite.original_index] = var_decl; |
| 149 | program.body.insert(rewrite.original_index + 1, re_export); |
| 150 | } else { |
| 151 | // Anonymous export default or arrow: replace the declaration content |
| 152 | // with the conditional expression |
| 153 | let export_default = Statement::ExportDefaultDeclaration( |
| 154 | react_compiler_ast::declarations::ExportDefaultDeclaration { |
| 155 | base: BaseNode::default(), |
| 156 | declaration: Box::new( |
| 157 | react_compiler_ast::declarations::ExportDefaultDecl::Expression( |
| 158 | Box::new(gating_expression), |
| 159 | ), |
| 160 | ), |
| 161 | export_kind: None, |
| 162 | }, |
| 163 | ); |
| 164 | program.body[rewrite.original_index] = export_default; |
| 165 | } |
| 166 | } |
| 167 | } |
| 168 | } |
| 169 | Ok(()) |
| 170 | } |
| 171 | |
| 172 | /// Gating rewrite for function declarations which are referenced before their |
| 173 | /// declaration site. |
| 174 | /// |
| 175 | /// ```js |
| 176 | /// // original |
| 177 | /// export default React.memo(Foo); |
| 178 | /// function Foo() { ... } |
| 179 | /// |
| 180 | /// // React compiler optimized + gated |
| 181 | /// import {gating} from 'myGating'; |
| 182 | /// export default React.memo(Foo); |
| 183 | /// const gating_result = gating(); // <- inserted |
| 184 | /// function Foo_optimized() {} // <- inserted |
| 185 | /// function Foo_unoptimized() {} // <- renamed from Foo |
| 186 | /// function Foo() { // <- inserted, hoistable by JS engines |
| 187 | /// if (gating_result) return Foo_optimized(); |
| 188 | /// else return Foo_unoptimized(); |
| 189 | /// } |
| 190 | /// ``` |
| 191 | fn insert_additional_function_declaration( |
| 192 | body: &mut Vec<Statement>, |
| 193 | original_index: usize, |
| 194 | mut compiled: FunctionDeclaration, |
| 195 | context: &mut ProgramContext, |
| 196 | gating_function_identifier_name: &str, |
| 197 | ) -> Result<(), CompilerDiagnostic> { |
| 198 | // Extract the original function declaration from body |
| 199 | let original_fn = match &body[original_index] { |
| 200 | Statement::FunctionDeclaration(fd) => fd.clone(), |
| 201 | Statement::ExportNamedDeclaration(end) => { |
| 202 | if let Some(decl) = &end.declaration { |
| 203 | if let react_compiler_ast::declarations::Declaration::FunctionDeclaration(fd) = |
| 204 | decl.as_ref() |
| 205 | { |
| 206 | fd.clone() |
| 207 | } else { |
| 208 | return Err(CompilerDiagnostic::new( |
| 209 | ErrorCategory::Invariant, |
| 210 | "Expected function declaration in export", |
| 211 | None, |
| 212 | )); |
| 213 | } |
| 214 | } else { |
| 215 | return Err(CompilerDiagnostic::new( |
| 216 | ErrorCategory::Invariant, |
| 217 | "Expected declaration in export", |
| 218 | None, |
| 219 | )); |
| 220 | } |
| 221 | } |
| 222 | _ => { |
| 223 | return Err(CompilerDiagnostic::new( |
| 224 | ErrorCategory::Invariant, |
| 225 | "Expected function declaration at original_index", |
| 226 | None, |
| 227 | )); |
| 228 | } |
| 229 | }; |
| 230 | |
| 231 | let original_fn_name = original_fn |
| 232 | .id |
| 233 | .as_ref() |
| 234 | .expect("Expected function declaration referenced elsewhere to have a named identifier"); |
| 235 | let compiled_id = compiled |
| 236 | .id |
| 237 | .as_ref() |
| 238 | .expect("Expected compiled function declaration to have a named identifier"); |
| 239 | assert_eq!( |
| 240 | original_fn.params.len(), |
| 241 | compiled.params.len(), |
| 242 | "Expected compiled function to have the same number of parameters as source" |
| 243 | ); |
| 244 | |
| 245 | let _ = compiled_id; // used above for the assert |
| 246 | |
| 247 | // Generate unique names |
| 248 | let gating_condition_name = |
| 249 | context.new_uid(&format!("{}_result", gating_function_identifier_name)); |
| 250 | let unoptimized_fn_name = context.new_uid(&format!("{}_unoptimized", original_fn_name.name)); |
| 251 | let optimized_fn_name = context.new_uid(&format!("{}_optimized", original_fn_name.name)); |
| 252 | |
| 253 | // Step 1: rename existing functions |
| 254 | compiled.id = Some(make_identifier(&optimized_fn_name)); |
| 255 | |
| 256 | // Rename the original function in-place to *_unoptimized |
| 257 | rename_fn_decl_at(body, original_index, &unoptimized_fn_name)?; |
| 258 | |
| 259 | // Step 2: build new params and args for the dispatcher function |
| 260 | let mut new_params: Vec<PatternLike> = Vec::new(); |
| 261 | let mut new_args_optimized: Vec<Expression> = Vec::new(); |
| 262 | let mut new_args_unoptimized: Vec<Expression> = Vec::new(); |
| 263 | |
| 264 | for (i, param) in original_fn.params.iter().enumerate() { |
| 265 | let arg_name = format!("arg{}", i); |
| 266 | match param { |
| 267 | PatternLike::RestElement(_) => { |
| 268 | new_params.push(PatternLike::RestElement( |
| 269 | react_compiler_ast::patterns::RestElement { |
| 270 | base: BaseNode::default(), |
| 271 | argument: Box::new(PatternLike::Identifier(make_identifier(&arg_name))), |
| 272 | type_annotation: None, |
| 273 | decorators: None, |
| 274 | }, |
| 275 | )); |
| 276 | new_args_optimized.push(Expression::SpreadElement(SpreadElement { |
| 277 | base: BaseNode::default(), |
| 278 | argument: Box::new(Expression::Identifier(make_identifier(&arg_name))), |
| 279 | })); |
| 280 | new_args_unoptimized.push(Expression::SpreadElement(SpreadElement { |
| 281 | base: BaseNode::default(), |
| 282 | argument: Box::new(Expression::Identifier(make_identifier(&arg_name))), |
| 283 | })); |
| 284 | } |
| 285 | _ => { |
| 286 | new_params.push(PatternLike::Identifier(make_identifier(&arg_name))); |
| 287 | new_args_optimized.push(Expression::Identifier(make_identifier(&arg_name))); |
| 288 | new_args_unoptimized.push(Expression::Identifier(make_identifier(&arg_name))); |
| 289 | } |
| 290 | } |
| 291 | } |
| 292 | |
| 293 | // Build the dispatcher function: |
| 294 | // function Foo(...args) { |
| 295 | // if (gating_result) return Foo_optimized(...args); |
| 296 | // else return Foo_unoptimized(...args); |
| 297 | // } |
| 298 | let dispatcher_fn = Statement::FunctionDeclaration(FunctionDeclaration { |
| 299 | base: BaseNode::default(), |
| 300 | id: Some(make_identifier(&original_fn_name.name)), |
| 301 | params: new_params, |
| 302 | body: BlockStatement { |
| 303 | base: BaseNode::default(), |
| 304 | body: vec![Statement::IfStatement(IfStatement { |
| 305 | base: BaseNode::default(), |
| 306 | test: Box::new(Expression::Identifier(make_identifier( |
| 307 | &gating_condition_name, |
| 308 | ))), |
| 309 | consequent: Box::new(Statement::ReturnStatement(ReturnStatement { |
| 310 | base: BaseNode::default(), |
| 311 | argument: Some(Box::new(Expression::CallExpression(CallExpression { |
| 312 | base: BaseNode::default(), |
| 313 | callee: Box::new(Expression::Identifier(make_identifier( |
| 314 | &optimized_fn_name, |
| 315 | ))), |
| 316 | arguments: new_args_optimized, |
| 317 | type_parameters: None, |
| 318 | type_arguments: None, |
| 319 | optional: None, |
| 320 | }))), |
| 321 | })), |
| 322 | alternate: Some(Box::new(Statement::ReturnStatement(ReturnStatement { |
| 323 | base: BaseNode::default(), |
| 324 | argument: Some(Box::new(Expression::CallExpression(CallExpression { |
| 325 | base: BaseNode::default(), |
| 326 | callee: Box::new(Expression::Identifier(make_identifier( |
| 327 | &unoptimized_fn_name, |
| 328 | ))), |
| 329 | arguments: new_args_unoptimized, |
| 330 | type_parameters: None, |
| 331 | type_arguments: None, |
| 332 | optional: None, |
| 333 | }))), |
| 334 | }))), |
| 335 | })], |
| 336 | directives: vec![], |
| 337 | }, |
| 338 | generator: false, |
| 339 | is_async: false, |
| 340 | declare: None, |
| 341 | return_type: None, |
| 342 | type_parameters: None, |
| 343 | predicate: None, |
| 344 | component_declaration: false, |
| 345 | hook_declaration: false, |
| 346 | }); |
| 347 | |
| 348 | // Build: const gating_result = gating(); |
| 349 | let gating_const = Statement::VariableDeclaration(VariableDeclaration { |
| 350 | base: BaseNode::default(), |
| 351 | declarations: vec![VariableDeclarator { |
| 352 | base: BaseNode::default(), |
| 353 | id: PatternLike::Identifier(make_identifier(&gating_condition_name)), |
| 354 | init: Some(Box::new(Expression::CallExpression(CallExpression { |
| 355 | base: BaseNode::default(), |
| 356 | callee: Box::new(Expression::Identifier(make_identifier( |
| 357 | gating_function_identifier_name, |
| 358 | ))), |
| 359 | arguments: vec![], |
| 360 | type_parameters: None, |
| 361 | type_arguments: None, |
| 362 | optional: None, |
| 363 | }))), |
| 364 | definite: None, |
| 365 | }], |
| 366 | kind: VariableDeclarationKind::Const, |
| 367 | declare: None, |
| 368 | }); |
| 369 | |
| 370 | // Build: the compiled (optimized) function declaration |
| 371 | let compiled_stmt = Statement::FunctionDeclaration(compiled); |
| 372 | |
| 373 | // Insert statements. In the TS version: |
| 374 | // fnPath.insertBefore(gating_const) |
| 375 | // fnPath.insertBefore(compiled) |
| 376 | // fnPath.insertAfter(dispatcher_fn) |
| 377 | // |
| 378 | // This means the final order is: |
| 379 | // [before original_index]: gating_const |
| 380 | // [before original_index]: compiled (optimized fn) |
| 381 | // [at original_index]: original fn (renamed to *_unoptimized) |
| 382 | // [after original_index]: dispatcher fn |
| 383 | // |
| 384 | // We insert in order: first the ones before, then the one after. |
| 385 | // Insert before original_index: gating_const, compiled |
| 386 | body.insert(original_index, compiled_stmt); |
| 387 | body.insert(original_index, gating_const); |
| 388 | // The original (now renamed) fn is now at original_index + 2 |
| 389 | // Insert dispatcher after it |
| 390 | body.insert(original_index + 3, dispatcher_fn); |
| 391 | Ok(()) |
| 392 | } |
| 393 | |
| 394 | /// Build a gating conditional expression: |
| 395 | /// `gating_fn() ? build_fn_expr(compiled) : build_fn_expr(original)` |
| 396 | fn build_gating_expression( |
| 397 | compiled: CompiledFunctionNode, |
| 398 | original: CompiledFunctionNode, |
| 399 | gating_name: &str, |
| 400 | ) -> Expression { |
| 401 | Expression::ConditionalExpression(ConditionalExpression { |
| 402 | base: BaseNode::default(), |
| 403 | test: Box::new(Expression::CallExpression(CallExpression { |
| 404 | base: BaseNode::default(), |
| 405 | callee: Box::new(Expression::Identifier(make_identifier(gating_name))), |
| 406 | arguments: vec![], |
| 407 | type_parameters: None, |
| 408 | type_arguments: None, |
| 409 | optional: None, |
| 410 | })), |
| 411 | consequent: Box::new(build_function_expression(compiled)), |
| 412 | alternate: Box::new(build_function_expression(original)), |
| 413 | }) |
| 414 | } |
| 415 | |
| 416 | /// Convert a compiled function node to an expression. |
| 417 | /// Function declarations are converted to function expressions; |
| 418 | /// arrow functions and function expressions are returned as-is. |
| 419 | fn build_function_expression(node: CompiledFunctionNode) -> Expression { |
| 420 | match node { |
| 421 | CompiledFunctionNode::ArrowFunctionExpression(arrow) => { |
| 422 | Expression::ArrowFunctionExpression(arrow) |
| 423 | } |
| 424 | CompiledFunctionNode::FunctionExpression(func_expr) => { |
| 425 | Expression::FunctionExpression(func_expr) |
| 426 | } |
| 427 | CompiledFunctionNode::FunctionDeclaration(func_decl) => { |
| 428 | // Convert FunctionDeclaration to FunctionExpression |
| 429 | Expression::FunctionExpression(FunctionExpression { |
| 430 | base: func_decl.base, |
| 431 | params: func_decl.params, |
| 432 | body: func_decl.body, |
| 433 | id: func_decl.id, |
| 434 | generator: func_decl.generator, |
| 435 | is_async: func_decl.is_async, |
| 436 | return_type: func_decl.return_type, |
| 437 | type_parameters: func_decl.type_parameters, |
| 438 | predicate: func_decl.predicate, |
| 439 | }) |
| 440 | } |
| 441 | } |
| 442 | } |
| 443 | |
| 444 | /// Helper to create a simple Identifier with the given name and default BaseNode. |
| 445 | fn make_identifier(name: &str) -> Identifier { |
| 446 | Identifier { |
| 447 | base: BaseNode::default(), |
| 448 | name: name.to_string(), |
| 449 | type_annotation: None, |
| 450 | optional: None, |
| 451 | decorators: None, |
| 452 | } |
| 453 | } |
| 454 | |
| 455 | /// Extract the function name from a top-level Statement if it is a |
| 456 | /// FunctionDeclaration with an id. |
| 457 | fn get_fn_decl_name(stmt: &Statement) -> Option<String> { |
| 458 | match stmt { |
| 459 | Statement::FunctionDeclaration(fd) => fd.id.as_ref().map(|id| id.name.clone()), |
| 460 | _ => None, |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | /// Extract the function name from an ExportDefaultDeclaration's declaration, |
| 465 | /// if it is a named FunctionDeclaration. |
| 466 | fn get_fn_decl_name_from_export_default(stmt: &Statement) -> Option<String> { |
| 467 | match stmt { |
| 468 | Statement::ExportDefaultDeclaration(ed) => match ed.declaration.as_ref() { |
| 469 | react_compiler_ast::declarations::ExportDefaultDecl::FunctionDeclaration(fd) => { |
| 470 | fd.id.as_ref().map(|id| id.name.clone()) |
| 471 | } |
| 472 | _ => None, |
| 473 | }, |
| 474 | _ => None, |
| 475 | } |
| 476 | } |
| 477 | |
| 478 | /// Extract a CompiledFunctionNode from a statement (for building the |
| 479 | /// "original" side of the gating expression). |
| 480 | fn extract_function_node_from_stmt( |
| 481 | stmt: &Statement, |
| 482 | ) -> Result<CompiledFunctionNode, CompilerDiagnostic> { |
| 483 | match stmt { |
| 484 | Statement::FunctionDeclaration(fd) => { |
| 485 | Ok(CompiledFunctionNode::FunctionDeclaration(fd.clone())) |
| 486 | } |
| 487 | Statement::ExpressionStatement(es) => match es.expression.as_ref() { |
| 488 | Expression::ArrowFunctionExpression(arrow) => { |
| 489 | Ok(CompiledFunctionNode::ArrowFunctionExpression(arrow.clone())) |
| 490 | } |
| 491 | Expression::FunctionExpression(fe) => { |
| 492 | Ok(CompiledFunctionNode::FunctionExpression(fe.clone())) |
| 493 | } |
| 494 | _ => Err(CompilerDiagnostic::new( |
| 495 | ErrorCategory::Invariant, |
| 496 | "Expected function expression in expression statement for gating", |
| 497 | None, |
| 498 | )), |
| 499 | }, |
| 500 | Statement::ExportDefaultDeclaration(ed) => match ed.declaration.as_ref() { |
| 501 | react_compiler_ast::declarations::ExportDefaultDecl::FunctionDeclaration(fd) => { |
| 502 | Ok(CompiledFunctionNode::FunctionDeclaration(fd.clone())) |
| 503 | } |
| 504 | react_compiler_ast::declarations::ExportDefaultDecl::Expression(expr) => { |
| 505 | match expr.as_ref() { |
| 506 | Expression::ArrowFunctionExpression(arrow) => { |
| 507 | Ok(CompiledFunctionNode::ArrowFunctionExpression(arrow.clone())) |
| 508 | } |
| 509 | Expression::FunctionExpression(fe) => { |
| 510 | Ok(CompiledFunctionNode::FunctionExpression(fe.clone())) |
| 511 | } |
| 512 | _ => Err(CompilerDiagnostic::new( |
| 513 | ErrorCategory::Invariant, |
| 514 | "Expected function expression in export default for gating", |
| 515 | None, |
| 516 | )), |
| 517 | } |
| 518 | } |
| 519 | _ => Err(CompilerDiagnostic::new( |
| 520 | ErrorCategory::Invariant, |
| 521 | "Expected function in export default declaration for gating", |
| 522 | None, |
| 523 | )), |
| 524 | }, |
| 525 | Statement::VariableDeclaration(vd) => { |
| 526 | let init = vd.declarations[0] |
| 527 | .init |
| 528 | .as_ref() |
| 529 | .expect("Expected variable declarator to have an init for gating"); |
| 530 | match init.as_ref() { |
| 531 | Expression::ArrowFunctionExpression(arrow) => { |
| 532 | Ok(CompiledFunctionNode::ArrowFunctionExpression(arrow.clone())) |
| 533 | } |
| 534 | Expression::FunctionExpression(fe) => { |
| 535 | Ok(CompiledFunctionNode::FunctionExpression(fe.clone())) |
| 536 | } |
| 537 | _ => Err(CompilerDiagnostic::new( |
| 538 | ErrorCategory::Invariant, |
| 539 | "Expected function expression in variable declaration for gating", |
| 540 | None, |
| 541 | )), |
| 542 | } |
| 543 | } |
| 544 | _ => Err(CompilerDiagnostic::new( |
| 545 | ErrorCategory::Invariant, |
| 546 | "Unexpected statement type for gating rewrite", |
| 547 | None, |
| 548 | )), |
| 549 | } |
| 550 | } |
| 551 | |
| 552 | /// Rename the function declaration at `body[index]` in place. |
| 553 | /// Handles both bare FunctionDeclaration and ExportNamedDeclaration wrapping one. |
| 554 | fn rename_fn_decl_at( |
| 555 | body: &mut [Statement], |
| 556 | index: usize, |
| 557 | new_name: &str, |
| 558 | ) -> Result<(), CompilerDiagnostic> { |
| 559 | match &mut body[index] { |
| 560 | Statement::FunctionDeclaration(fd) => { |
| 561 | fd.id = Some(make_identifier(new_name)); |
| 562 | } |
| 563 | Statement::ExportNamedDeclaration(end) => { |
| 564 | if let Some(decl) = &mut end.declaration { |
| 565 | if let react_compiler_ast::declarations::Declaration::FunctionDeclaration(fd) = |
| 566 | decl.as_mut() |
| 567 | { |
| 568 | fd.id = Some(make_identifier(new_name)); |
| 569 | } |
| 570 | } |
| 571 | } |
| 572 | _ => { |
| 573 | return Err(CompilerDiagnostic::new( |
| 574 | ErrorCategory::Invariant, |
| 575 | "Expected function declaration to rename", |
| 576 | None, |
| 577 | )); |
| 578 | } |
| 579 | } |
| 580 | Ok(()) |
| 581 | } |