| 1 | /** |
| 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | * |
| 4 | * This source code is licensed under the MIT license found in the |
| 5 | * LICENSE file in the root directory of this source tree. |
| 6 | */ |
| 7 | |
| 8 | import type {NodePath} from '@babel/traverse'; |
| 9 | import type * as t from '@babel/types'; |
| 10 | import {CompilerError} from '../CompilerError'; |
| 11 | import {getOrInsertDefault} from '../Utils/utils'; |
| 12 | import {GeneratedSource} from './HIR'; |
| 13 | |
| 14 | type IdentifierInfo = { |
| 15 | reassigned: boolean; |
| 16 | reassignedByInnerFn: boolean; |
| 17 | referencedByInnerFn: boolean; |
| 18 | }; |
| 19 | const DEFAULT_IDENTIFIER_INFO: IdentifierInfo = { |
| 20 | reassigned: false, |
| 21 | reassignedByInnerFn: false, |
| 22 | referencedByInnerFn: false, |
| 23 | }; |
| 24 | |
| 25 | type BabelFunction = |
| 26 | | NodePath<t.FunctionDeclaration> |
| 27 | | NodePath<t.FunctionExpression> |
| 28 | | NodePath<t.ArrowFunctionExpression> |
| 29 | | NodePath<t.ObjectMethod>; |
| 30 | type FindContextIdentifierState = { |
| 31 | currentFn: Array<BabelFunction>; |
| 32 | identifiers: Map<t.Identifier, IdentifierInfo>; |
| 33 | }; |
| 34 | |
| 35 | const withFunctionScope = { |
| 36 | enter: function ( |
| 37 | path: BabelFunction, |
| 38 | state: FindContextIdentifierState, |
| 39 | ): void { |
| 40 | state.currentFn.push(path); |
| 41 | }, |
| 42 | exit: function (_: BabelFunction, state: FindContextIdentifierState): void { |
| 43 | state.currentFn.pop(); |
| 44 | }, |
| 45 | }; |
| 46 | |
| 47 | export function findContextIdentifiers( |
| 48 | func: NodePath<t.Function>, |
| 49 | ): Set<t.Identifier> { |
| 50 | const state: FindContextIdentifierState = { |
| 51 | currentFn: [], |
| 52 | identifiers: new Map(), |
| 53 | }; |
| 54 | |
| 55 | func.traverse<FindContextIdentifierState>( |
| 56 | { |
| 57 | FunctionDeclaration: withFunctionScope, |
| 58 | FunctionExpression: withFunctionScope, |
| 59 | ArrowFunctionExpression: withFunctionScope, |
| 60 | ObjectMethod: withFunctionScope, |
| 61 | AssignmentExpression( |
| 62 | path: NodePath<t.AssignmentExpression>, |
| 63 | state: FindContextIdentifierState, |
| 64 | ): void { |
| 65 | const left = path.get('left'); |
| 66 | if (left.isLVal()) { |
| 67 | const currentFn = state.currentFn.at(-1) ?? null; |
| 68 | handleAssignment(currentFn, state.identifiers, left); |
| 69 | } else { |
| 70 | /** |
| 71 | * OptionalMemberExpressions as the left side of an AssignmentExpression are Stage 1 and |
| 72 | * not supported by React Compiler yet. |
| 73 | */ |
| 74 | CompilerError.throwTodo({ |
| 75 | reason: `Unsupported syntax on the left side of an AssignmentExpression`, |
| 76 | description: `Expected an LVal, got: ${left.type}`, |
| 77 | loc: left.node.loc ?? null, |
| 78 | }); |
| 79 | } |
| 80 | }, |
| 81 | UpdateExpression( |
| 82 | path: NodePath<t.UpdateExpression>, |
| 83 | state: FindContextIdentifierState, |
| 84 | ): void { |
| 85 | const argument = path.get('argument'); |
| 86 | const currentFn = state.currentFn.at(-1) ?? null; |
| 87 | if (argument.isLVal()) { |
| 88 | handleAssignment(currentFn, state.identifiers, argument); |
| 89 | } |
| 90 | }, |
| 91 | Identifier( |
| 92 | path: NodePath<t.Identifier>, |
| 93 | state: FindContextIdentifierState, |
| 94 | ): void { |
| 95 | const currentFn = state.currentFn.at(-1) ?? null; |
| 96 | if (path.isReferencedIdentifier()) { |
| 97 | handleIdentifier(currentFn, state.identifiers, path); |
| 98 | } |
| 99 | }, |
| 100 | }, |
| 101 | state, |
| 102 | ); |
| 103 | |
| 104 | const result = new Set<t.Identifier>(); |
| 105 | for (const [id, info] of state.identifiers.entries()) { |
| 106 | if (info.reassignedByInnerFn) { |
| 107 | result.add(id); |
| 108 | } else if (info.reassigned && info.referencedByInnerFn) { |
| 109 | result.add(id); |
| 110 | } |
| 111 | } |
| 112 | return result; |
| 113 | } |
| 114 | |
| 115 | function handleIdentifier( |
| 116 | currentFn: BabelFunction | null, |
| 117 | identifiers: Map<t.Identifier, IdentifierInfo>, |
| 118 | path: NodePath<t.Identifier>, |
| 119 | ): void { |
| 120 | const name = path.node.name; |
| 121 | const binding = path.scope.getBinding(name); |
| 122 | if (binding == null) { |
| 123 | return; |
| 124 | } |
| 125 | const identifier = getOrInsertDefault(identifiers, binding.identifier, { |
| 126 | ...DEFAULT_IDENTIFIER_INFO, |
| 127 | }); |
| 128 | |
| 129 | if (currentFn != null) { |
| 130 | const bindingAboveLambdaScope = currentFn.scope.parent.getBinding(name); |
| 131 | |
| 132 | if (binding === bindingAboveLambdaScope) { |
| 133 | identifier.referencedByInnerFn = true; |
| 134 | } |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | function handleAssignment( |
| 139 | currentFn: BabelFunction | null, |
| 140 | identifiers: Map<t.Identifier, IdentifierInfo>, |
| 141 | lvalPath: NodePath<t.LVal>, |
| 142 | ): void { |
| 143 | /* |
| 144 | * Find all reassignments to identifiers declared outside of currentFn |
| 145 | * This closely follows destructuring assignment assumptions and logic in BuildHIR |
| 146 | */ |
| 147 | const lvalNode = lvalPath.node; |
| 148 | switch (lvalNode.type) { |
| 149 | case 'Identifier': { |
| 150 | const path = lvalPath as NodePath<t.Identifier>; |
| 151 | const name = path.node.name; |
| 152 | const binding = path.scope.getBinding(name); |
| 153 | if (binding == null) { |
| 154 | break; |
| 155 | } |
| 156 | const state = getOrInsertDefault(identifiers, binding.identifier, { |
| 157 | ...DEFAULT_IDENTIFIER_INFO, |
| 158 | }); |
| 159 | state.reassigned = true; |
| 160 | |
| 161 | if (currentFn != null) { |
| 162 | const bindingAboveLambdaScope = currentFn.scope.parent.getBinding(name); |
| 163 | |
| 164 | if (binding === bindingAboveLambdaScope) { |
| 165 | state.reassignedByInnerFn = true; |
| 166 | } |
| 167 | } |
| 168 | break; |
| 169 | } |
| 170 | case 'ArrayPattern': { |
| 171 | const path = lvalPath as NodePath<t.ArrayPattern>; |
| 172 | for (const element of path.get('elements')) { |
| 173 | if (nonNull(element)) { |
| 174 | handleAssignment(currentFn, identifiers, element); |
| 175 | } |
| 176 | } |
| 177 | break; |
| 178 | } |
| 179 | case 'ObjectPattern': { |
| 180 | const path = lvalPath as NodePath<t.ObjectPattern>; |
| 181 | for (const property of path.get('properties')) { |
| 182 | if (property.isObjectProperty()) { |
| 183 | const valuePath = property.get('value'); |
| 184 | CompilerError.invariant(valuePath.isLVal(), { |
| 185 | reason: `[FindContextIdentifiers] Expected object property value to be an LVal, got: ${valuePath.type}`, |
| 186 | loc: valuePath.node.loc ?? GeneratedSource, |
| 187 | }); |
| 188 | handleAssignment(currentFn, identifiers, valuePath); |
| 189 | } else { |
| 190 | CompilerError.invariant(property.isRestElement(), { |
| 191 | reason: `[FindContextIdentifiers] Invalid assumptions for babel types.`, |
| 192 | loc: property.node.loc ?? GeneratedSource, |
| 193 | }); |
| 194 | handleAssignment(currentFn, identifiers, property); |
| 195 | } |
| 196 | } |
| 197 | break; |
| 198 | } |
| 199 | case 'AssignmentPattern': { |
| 200 | const path = lvalPath as NodePath<t.AssignmentPattern>; |
| 201 | const left = path.get('left'); |
| 202 | handleAssignment(currentFn, identifiers, left); |
| 203 | break; |
| 204 | } |
| 205 | case 'RestElement': { |
| 206 | const path = lvalPath as NodePath<t.RestElement>; |
| 207 | handleAssignment(currentFn, identifiers, path.get('argument')); |
| 208 | break; |
| 209 | } |
| 210 | case 'MemberExpression': { |
| 211 | // Interior mutability (not a reassign) |
| 212 | break; |
| 213 | } |
| 214 | default: { |
| 215 | CompilerError.throwTodo({ |
| 216 | reason: `[FindContextIdentifiers] Cannot handle Object destructuring assignment target ${lvalNode.type}`, |
| 217 | description: null, |
| 218 | loc: lvalNode.loc ?? GeneratedSource, |
| 219 | suggestions: null, |
| 220 | }); |
| 221 | } |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | function nonNull<T extends NonNullable<t.Node>>( |
| 226 | t: NodePath<T | null>, |
| 227 | ): t is NodePath<T> { |
| 228 | return t.node != null; |
| 229 | } |