1/**2 * Copyright (c) Meta Platforms, Inc. and affiliates.3 *4 * This source code is licensed under the MIT license found in the5 * LICENSE file in the root directory of this source tree.6 */78import {NodePath} from '@babel/traverse';9import * as t from '@babel/types';10import {CompilerDiagnostic, ErrorCategory} from '..';11import {CodegenFunction} from '../ReactiveScopes';12import {Environment} from '../HIR/Environment';1314/**15 * IMPORTANT: This validation is only intended for use in unit tests.16 * It is not intended for use in production.17 *18 * This validation is used to ensure that the generated AST has proper source locations19 * for "important" original nodes.20 *21 * There's one big gotcha with this validation: it only works if the "important" original nodes22 * are not optimized away by the compiler.23 *24 * When that scenario happens, we should just update the fixture to not include a node that has no25 * corresponding node in the generated AST due to being completely removed during compilation.26 */2728/**29 * Some common node types that are important for coverage tracking.30 * Based on istanbul-lib-instrument + some other common nodes we expect to be present in the generated AST.31 *32 * Note: For VariableDeclaration, VariableDeclarator, and Identifier, we enforce stricter validation33 * that requires both the source location AND node type to match in the generated AST. This ensures34 * that variable declarations maintain their structural integrity through compilation.35 */36const IMPORTANT_INSTRUMENTED_TYPES = new Set([37  'ArrowFunctionExpression',38  'AssignmentPattern',39  'ObjectMethod',40  'ExpressionStatement',41  'BreakStatement',42  'ContinueStatement',43  'ReturnStatement',44  'ThrowStatement',45  'TryStatement',46  'VariableDeclarator',47  'IfStatement',48  'ForStatement',49  'ForInStatement',50  'ForOfStatement',51  'WhileStatement',52  'DoWhileStatement',53  'SwitchStatement',54  'SwitchCase',55  'WithStatement',56  'FunctionDeclaration',57  'FunctionExpression',58  'LabeledStatement',59  'ConditionalExpression',60  'LogicalExpression',6162  /**63   * Note: these aren't important for coverage tracking,64   * but we still want to track them to ensure we aren't regressing them when65   * we fix the source location tracking for other nodes.66   */67  'VariableDeclaration',68  'Identifier',69]);7071/**72 * Check if a node is a manual memoization call that the compiler optimizes away.73 * These include useMemo and useCallback calls, which are intentionally removed74 * by the DropManualMemoization pass.75 */76function isManualMemoization(node: t.Node): boolean {77  // Check if this is a useMemo/useCallback call expression78  if (t.isCallExpression(node)) {79    const callee = node.callee;80    if (t.isIdentifier(callee)) {81      return callee.name === 'useMemo' || callee.name === 'useCallback';82    }83    if (84      t.isMemberExpression(callee) &&85      t.isIdentifier(callee.property) &&86      t.isIdentifier(callee.object)87    ) {88      return (89        callee.object.name === 'React' &&90        (callee.property.name === 'useMemo' ||91          callee.property.name === 'useCallback')92      );93    }94  }9596  return false;97}9899/**100 * Create a location key for comparison. We compare by line/column/source,101 * not by object identity.102 */103function locationKey(loc: t.SourceLocation): string {104  return `${loc.start.line}:${loc.start.column}-${loc.end.line}:${loc.end.column}`;105}106107/**108 * Validates that important source locations from the original code are preserved109 * in the generated AST. This ensures that Istanbul coverage instrumentation can110 * properly map back to the original source code.111 *112 * The validator:113 * 1. Collects locations from "important" nodes in the original AST (those that114 *    Istanbul instruments for coverage tracking)115 * 2. Exempts known compiler optimizations (useMemo/useCallback removal)116 * 3. Verifies that all important locations appear somewhere in the generated AST117 *118 * Missing locations can cause Istanbul to fail to track coverage for certain119 * code paths, leading to inaccurate coverage reports.120 */121export function validateSourceLocations(122  func: NodePath<123    t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression124  >,125  generatedAst: CodegenFunction,126  env: Environment,127): void {128  /*129   * Step 1: Collect important locations from the original source130   * Note: Multiple node types can share the same location (e.g. VariableDeclarator and Identifier)131   */132  const importantOriginalLocations = new Map<133    string,134    {loc: t.SourceLocation; nodeTypes: Set<string>}135  >();136137  func.traverse({138    enter(path) {139      const node = path.node;140141      // Only track node types that Istanbul instruments142      if (!IMPORTANT_INSTRUMENTED_TYPES.has(node.type)) {143        return;144      }145146      // Skip manual memoization that the compiler intentionally removes147      if (isManualMemoization(node)) {148        return;149      }150151      /*152       * Skip return statements inside arrow functions that will be simplified to expression body.153       * The compiler transforms `() => { return expr }` to `() => expr` in CodegenReactiveFunction154       */155      if (t.isReturnStatement(node) && node.argument != null) {156        const parentBody = path.parentPath;157        const parentFunc = parentBody?.parentPath;158        if (159          parentBody?.isBlockStatement() &&160          parentFunc?.isArrowFunctionExpression() &&161          parentBody.node.body.length === 1 &&162          parentBody.node.directives.length === 0163        ) {164          return;165        }166      }167168      // Collect the location if it exists169      if (node.loc) {170        const key = locationKey(node.loc);171        const existing = importantOriginalLocations.get(key);172        if (existing) {173          existing.nodeTypes.add(node.type);174        } else {175          importantOriginalLocations.set(key, {176            loc: node.loc,177            nodeTypes: new Set([node.type]),178          });179        }180      }181    },182  });183184  // Step 2: Collect all locations from the generated AST with their node types185  const generatedLocations = new Map<string, Set<string>>();186187  function collectGeneratedLocations(node: t.Node): void {188    if (node.loc) {189      const key = locationKey(node.loc);190      const nodeTypes = generatedLocations.get(key);191      if (nodeTypes) {192        nodeTypes.add(node.type);193      } else {194        generatedLocations.set(key, new Set([node.type]));195      }196    }197198    // Use Babel's VISITOR_KEYS to traverse only actual node properties199    const keys = t.VISITOR_KEYS[node.type as keyof typeof t.VISITOR_KEYS];200201    if (!keys) {202      return;203    }204205    for (const key of keys) {206      const value = (node as any)[key];207208      if (Array.isArray(value)) {209        for (const item of value) {210          if (t.isNode(item)) {211            collectGeneratedLocations(item);212          }213        }214      } else if (t.isNode(value)) {215        collectGeneratedLocations(value);216      }217    }218  }219220  // Collect from main function body221  collectGeneratedLocations(generatedAst.body);222223  // Collect from outlined functions224  for (const outlined of generatedAst.outlined) {225    collectGeneratedLocations(outlined.fn.body);226  }227228  /*229   * Step 3: Validate that all important locations are preserved230   * For certain node types, also validate that the node type matches231   */232  const strictNodeTypes = new Set([233    'VariableDeclaration',234    'VariableDeclarator',235    'Identifier',236  ]);237238  const reportMissingLocation = (239    loc: t.SourceLocation,240    nodeType: string,241  ): void => {242    env.recordError(243      CompilerDiagnostic.create({244        category: ErrorCategory.Todo,245        reason: 'Important source location missing in generated code',246        description:247          `Source location for ${nodeType} is missing in the generated output. This can cause coverage instrumentation ` +248          `to fail to track this code properly, resulting in inaccurate coverage reports.`,249      }).withDetails({250        kind: 'error',251        loc,252        message: null,253      }),254    );255  };256257  const reportWrongNodeType = (258    loc: t.SourceLocation,259    expectedType: string,260    actualTypes: Set<string>,261  ): void => {262    env.recordError(263      CompilerDiagnostic.create({264        category: ErrorCategory.Todo,265        reason:266          'Important source location has wrong node type in generated code',267        description:268          `Source location for ${expectedType} exists in the generated output but with wrong node type(s): ${Array.from(actualTypes).join(', ')}. ` +269          `This can cause coverage instrumentation to fail to track this code properly, resulting in inaccurate coverage reports.`,270      }).withDetails({271        kind: 'error',272        loc,273        message: null,274      }),275    );276  };277278  for (const [key, {loc, nodeTypes}] of importantOriginalLocations) {279    const generatedNodeTypes = generatedLocations.get(key);280281    if (!generatedNodeTypes) {282      // Location is completely missing283      reportMissingLocation(loc, Array.from(nodeTypes).join(', '));284    } else {285      // Location exists, check each node type286      for (const nodeType of nodeTypes) {287        if (288          strictNodeTypes.has(nodeType) &&289          !generatedNodeTypes.has(nodeType)290        ) {291          /*292           * For strict node types, the specific node type must be present293           * Check if any generated node type is also an important original node type294           */295          const hasValidNodeType = Array.from(generatedNodeTypes).some(296            genType => nodeTypes.has(genType),297          );298299          if (hasValidNodeType) {300            // At least one generated node type is valid (also in original), so this is just missing301            reportMissingLocation(loc, nodeType);302          } else {303            // None of the generated node types are in original - this is wrong node type304            reportWrongNodeType(loc, nodeType, generatedNodeTypes);305          }306        }307      }308    }309  }310}