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 special7// hoisting pattern is used (see `insert_additional_function_declaration`).8//9// Ported from `Entrypoint/Gating.ts`.1011use react_compiler_ast::common::BaseNode;12use react_compiler_ast::expressions::*;13use react_compiler_ast::patterns::PatternLike;14use react_compiler_ast::statements::*;15use react_compiler_diagnostics::CompilerDiagnostic;16use react_compiler_diagnostics::ErrorCategory;1718use super::imports::ProgramContext;19use super::plugin_options::GatingConfig;2021/// A compiled function node, can be any function type.22#[derive(Debug, Clone)]23pub enum CompiledFunctionNode {24    FunctionDeclaration(FunctionDeclaration),25    FunctionExpression(FunctionExpression),26    ArrowFunctionExpression(ArrowFunctionExpression),27}2829/// Represents a compiled function that needs gating.30/// In the Rust version, we work with indices into the program body31/// rather than Babel paths.32pub struct GatingRewrite {33    /// Index in program.body where the original function is34    pub original_index: usize,35    /// The compiled function AST node36    pub compiled_fn: CompiledFunctionNode,37    /// The gating config38    pub gating: GatingConfig,39    /// Whether the function is referenced before its declaration at top level40    pub referenced_before_declared: bool,41    /// Whether the parent statement is an ExportDefaultDeclaration42    pub is_export_default: bool,43}4445/// 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 reverse50/// index order to maintain validity of earlier indices.51pub 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 insertions57    // at higher indices don't invalidate lower indices.58    rewrites.sort_by(|a, b| b.original_index.cmp(&a.original_index));5960    for rewrite in rewrites {61        let gating_imported_name = context62            .add_import_specifier(63                &rewrite.gating.source,64                &rewrite.gating.import_specifier_name,65                None,66            )67            .name68            .clone();6970        if rewrite.referenced_before_declared {71            // The referenced-before-declared case only applies to FunctionDeclarations72            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)?;9192            let gating_expression =93                build_gating_expression(rewrite.compiled_fn, original_fn, &gating_imported_name);9495            // Determine how to rewrite based on context96            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 : original99                    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 case121                if let Some(fn_name) = get_fn_decl_name_from_export_default(&original_stmt) {122                    // Named export default function: replace with const + re-export123                    //   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 after148                    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 content152                    // with the conditional expression153                    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}171172/// Gating rewrite for function declarations which are referenced before their173/// declaration site.174///175/// ```js176/// // original177/// export default React.memo(Foo);178/// function Foo() { ... }179///180/// // React compiler optimized + gated181/// import {gating} from 'myGating';182/// export default React.memo(Foo);183/// const gating_result = gating();  // <- inserted184/// function Foo_optimized() {}      // <- inserted185/// function Foo_unoptimized() {}    // <- renamed from Foo186/// function Foo() {                 // <- inserted, hoistable by JS engines187///   if (gating_result) return Foo_optimized();188///   else return Foo_unoptimized();189/// }190/// ```191fn 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 body199    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    };230231    let original_fn_name = original_fn232        .id233        .as_ref()234        .expect("Expected function declaration referenced elsewhere to have a named identifier");235    let compiled_id = compiled236        .id237        .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    );244245    let _ = compiled_id; // used above for the assert246247    // Generate unique names248    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));252253    // Step 1: rename existing functions254    compiled.id = Some(make_identifier(&optimized_fn_name));255256    // Rename the original function in-place to *_unoptimized257    rename_fn_decl_at(body, original_index, &unoptimized_fn_name)?;258259    // Step 2: build new params and args for the dispatcher function260    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();263264    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    }292293    // 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    });347348    // 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    });369370    // Build: the compiled (optimized) function declaration371    let compiled_stmt = Statement::FunctionDeclaration(compiled);372373    // 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_const380    //   [before original_index]: compiled (optimized fn)381    //   [at original_index]:     original fn (renamed to *_unoptimized)382    //   [after original_index]:  dispatcher fn383    //384    // We insert in order: first the ones before, then the one after.385    // Insert before original_index: gating_const, compiled386    body.insert(original_index, compiled_stmt);387    body.insert(original_index, gating_const);388    // The original (now renamed) fn is now at original_index + 2389    // Insert dispatcher after it390    body.insert(original_index + 3, dispatcher_fn);391    Ok(())392}393394/// Build a gating conditional expression:395/// `gating_fn() ? build_fn_expr(compiled) : build_fn_expr(original)`396fn 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}415416/// 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.419fn 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 FunctionExpression429            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}443444/// Helper to create a simple Identifier with the given name and default BaseNode.445fn 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}454455/// Extract the function name from a top-level Statement if it is a456/// FunctionDeclaration with an id.457fn 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}463464/// Extract the function name from an ExportDefaultDeclaration's declaration,465/// if it is a named FunctionDeclaration.466fn 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}477478/// Extract a CompiledFunctionNode from a statement (for building the479/// "original" side of the gating expression).480fn 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                .init528                .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}551552/// Rename the function declaration at `body[index]` in place.553/// Handles both bare FunctionDeclaration and ExportNamedDeclaration wrapping one.554fn 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}