Support stop and skip traversal APIs

Summary: Our SimpleTraversal API already supports stopping and skipping traversal, this diff exposes it to the main hermes-transform traversal logic.

Reviewed By: bradzacher

Differential Revision: D41325007

fbshipit-source-id: 0c283713feeed4e2d761553ac356a0e939df14fc

tools/hermes-parser/js/hermes-parser/src/traverse/SimpleTraverser.js

@@ -77,6 +77,7 @@ export class SimpleTraverser {
       if (ex === SimpleTraverserSkip) {
         return;
       }
+      this._setErrorContext(ex, node);
       throw ex;
     }
 
@@ -100,10 +101,26 @@ export class SimpleTraverser {
       if (ex === SimpleTraverserSkip) {
         return;
       }
+      this._setErrorContext(ex, node);
       throw ex;
     }
   }
 
+  /**
+   * Set useful contextual information onto the error object.
+   * @param ex The error object.
+   * @param node The current node.
+   * @private
+   */
+  _setErrorContext(ex: Error, node: ESNode): void {
+    // $FlowFixMe[prop-missing]
+    ex.currentNode = {
+      type: node.type,
+      range: node.range,
+      loc: node.loc,
+    };
+  }
+
   /**
    * Traverse the given AST tree.
    * @param node The root node to traverse.

tools/hermes-parser/js/hermes-transform/__tests__/traverse/traverse-test.js

@@ -34,6 +34,49 @@ describe('traverse', () => {
     expect(visitedNodes).toEqual(['Program', 'VariableDeclaration', 'Literal']);
   });
 
+  it('stops traversal', () => {
+    const code = 'const x = 1, y = 2;';
+    const {ast, scopeManager} = parseForESLint(code);
+
+    const visitedNodes = [];
+    traverse(code, ast, scopeManager, context => ({
+      '*'(node) {
+        visitedNodes.push(node.type);
+      },
+      VariableDeclarator() {
+        context.stopTraversal();
+      },
+    }));
+
+    expect(visitedNodes).toEqual([
+      'Program',
+      'VariableDeclaration',
+      'VariableDeclarator',
+    ]);
+  });
+
+  it('skips traversal', () => {
+    const code = 'const x = 1, y = 2;';
+    const {ast, scopeManager} = parseForESLint(code);
+
+    const visitedNodes = [];
+    traverse(code, ast, scopeManager, context => ({
+      '*'(node) {
+        visitedNodes.push(node.type);
+      },
+      VariableDeclarator() {
+        context.skipTraversal();
+      },
+    }));
+
+    expect(visitedNodes).toEqual([
+      'Program',
+      'VariableDeclaration',
+      'VariableDeclarator',
+      'VariableDeclarator',
+    ]);
+  });
+
   it('visits the AST in the correct order - traversed as defined by the visitor keys', () => {
     const code = 'const x = 1;';
     const {ast, scopeManager} = parseForESLint(code);

tools/hermes-parser/js/hermes-transform/src/transform/transformAST.js

@@ -28,7 +28,7 @@ import {performReplaceNodeMutation} from './mutations/ReplaceNode';
 import {performReplaceStatementWithManyMutation} from './mutations/ReplaceStatementWithMany';
 import type {ParseResult} from './parse';
 
-export type transformASTResult = {
+export type TransformASTResult = {
   ast: Program,
   astWasMutated: boolean,
   mutatedCode: string,
@@ -37,7 +37,7 @@ export type transformASTResult = {
 export function transformAST(
   {ast, scopeManager, code}: ParseResult,
   visitors: TransformVisitor,
-): transformASTResult {
+): TransformASTResult {
   // traverse the AST and colllect the mutations
   const transformContext = getTransformContext();
   traverseWithContext(

tools/hermes-parser/js/hermes-transform/src/traverse/traverse.js

@@ -55,6 +55,17 @@ export type TraversalContextBase = $ReadOnly<{
    * (where 56:44 represents L56, Col44)
    */
   buildSimpleCodeFrame: (node: ESNode, message: string) => string,
+  /**
+   * Can be called at any point during the traversal to immediately stop traversal
+   * entirely.
+   */
+  stopTraversal: () => void,
+  /**
+   * Can be called within the traversal "enter" function to prevent the traverser
+   * from traversing the node any further, essentially culling the remainder of the
+   * AST branch from traversal.
+   */
+  skipTraversal: () => void,
 }>;
 export type TraversalContext<T> = $ReadOnly<{
   ...TraversalContextBase,
@@ -77,19 +88,10 @@ export function traverseWithContext<T = TraversalContextBase>(
   visitor: Visitor<T>,
 ): void {
   const emitter = new SafeEmitter();
-  const nodeQueue: Array<{isEntering: boolean, node: ESNode}> = [];
 
   let currentNode: ESNode = ast;
-
-  // build up the traversal queue
-  SimpleTraverser.traverse(ast, {
-    enter(node) {
-      nodeQueue.push({isEntering: true, node});
-    },
-    leave(node) {
-      nodeQueue.push({isEntering: false, node});
-    },
-  });
+  let shouldSkipTraversal = false;
+  let shouldStopTraversal = false;
 
   const getScope = (givenNode: ESNode = currentNode) => {
     // On Program node, get the outermost scope to avoid return Node.js special function scope or ES modules scope.
@@ -153,6 +155,14 @@ export function traverseWithContext<T = TraversalContextBase>(
     },
 
     getScope,
+
+    stopTraversal: () => {
+      shouldStopTraversal = true;
+    },
+
+    skipTraversal: () => {
+      shouldSkipTraversal = true;
+    },
   });
 
   const traversalContext: TraversalContext<T> = Object.freeze({
@@ -174,19 +184,30 @@ export function traverseWithContext<T = TraversalContextBase>(
   });
 
   const eventGenerator = new NodeEventGenerator(emitter);
-  nodeQueue.forEach(traversalInfo => {
-    currentNode = traversalInfo.node;
-
-    try {
-      if (traversalInfo.isEntering) {
-        eventGenerator.enterNode(currentNode);
-      } else {
-        eventGenerator.leaveNode(currentNode);
-      }
-    } catch (err) {
-      err.currentNode = currentNode;
-      throw err;
+
+  function checkTraversalFlags(): void {
+    if (shouldStopTraversal) {
+      // No need to reset the flag since we won't enter any more nodes.
+      throw SimpleTraverser.Break;
+    }
+
+    if (shouldSkipTraversal) {
+      shouldSkipTraversal = false;
+      throw SimpleTraverser.Skip;
     }
+  }
+
+  SimpleTraverser.traverse(ast, {
+    enter(node) {
+      currentNode = node;
+      eventGenerator.enterNode(node);
+      checkTraversalFlags();
+    },
+    leave(node) {
+      currentNode = node;
+      eventGenerator.leaveNode(node);
+      checkTraversalFlags();
+    },
   });
 }