BlockScopingTransformations.cpp

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include "SemanticValidator.h"

#include "hermes/AST/RecursiveVisitor.h"
#include "hermes/Support/InternalIdentifierMaker.h"
#include "hermes/Support/StringTable.h"
#include "llvh/ADT/SetVector.h"
#include "llvh/Support/SaveAndRestore.h"

namespace hermes {
namespace sem {
namespace {

/// Implements an AST traversal that performs transformations that simplify
/// block scoping downstream.
class BlockScopingTransformations {
 public:
  explicit BlockScopingTransformations(Context &astContext)
      : astContext_(astContext),
        internalIDs_(astContext.getStringTable()),
        identAssign_(astContext.getIdentifier("=").getUnderlyingPointer()),
        identExclaim_(astContext.getIdentifier("!").getUnderlyingPointer()),
        identLet_(astContext.getIdentifier("let").getUnderlyingPointer()),
        identUndefined_(
            astContext.getIdentifier("undefined").getUnderlyingPointer()),
        identVar_(astContext.getIdentifier("var").getUnderlyingPointer()) {}

  void visit(Node *node) {
    visitESTreeChildren(*this, node);
  }

  ESTree::VisitResult visit(ForInStatementNode *forInStmt) {
    return visitAndRewriteForInOf(forInStmt, forInStmt);
  }

  ESTree::VisitResult visit(ForOfStatementNode *forOfStmt) {
    return visitAndRewriteForInOf(forOfStmt, forOfStmt);
  }

  ESTree::VisitResult visit(ForStatementNode *forStmt) {
    return visitAndRewriteFor(forStmt, forStmt);
  }

  /// Intercepts LabeledStatements
  ///
  ///    label : [label:]+ for-in
  ///    label : [label:]+ for-of
  ///    label : [label:]+ for(const/let
  ///
  /// and applies the relevant transformation.
  ESTree::VisitResult visit(LabeledStatementNode *labeledStmt) {
    LabeledStatementNode *firstLabel = labeledStmt;
    LabeledStatementNode *lastLabel = labeledStmt;

    while (auto *bodyLabel =
               llvh::dyn_cast<LabeledStatementNode>(lastLabel->_body)) {
      lastLabel = bodyLabel;
    }

    Node *body = lastLabel->_body;
    switch (body->getKind()) {
      case NodeKind::ForInStatement:
      case NodeKind::ForOfStatement:
        return visitAndRewriteForInOf(body, firstLabel);

      case NodeKind::ForStatement:
        return visitAndRewriteFor(
            llvh::cast<ForStatementNode>(body), firstLabel);

      default:
        visitESTreeNode(*this, body, lastLabel);
        return ESTree::Unmodified;
    }
  }

  bool incRecursionDepth(Node *n) {
    if (LLVM_UNLIKELY(recursionDepth_ == 0)) {
      return false;
    }
    --recursionDepth_;
    if (LLVM_UNLIKELY(recursionDepth_ == 0)) {
      return false;
    }
    return true;
  }

  void decRecursionDepth() {
    assert(
        recursionDepth_ < MAX_RECURSION_DEPTH &&
        "recursionDepth_ cannot go negative");
    if (LLVM_LIKELY(recursionDepth_ != 0))
      ++recursionDepth_;
  }

 private:
  Context &astContext_;
  InternalIdentifierMaker internalIDs_;

  /// Rewrites
  ///
  ///                   |---- left ---|
  ///   [label:]* for (let/const/var x in/of right) body
  ///
  /// into
  ///
  ///   wrapperBlock: {
  ///     let temp;
  ///     [label:]* for (temp in/of right) {
  ///       let/const/var x = temp;
  ///       body;
  ///     }
  ///     break wrapperBlock;
  ///     let x;  // for TDZ (e.g., right: { x })
  ///   }
  ESTree::VisitResult visitAndRewriteForInOf(Node *forInOfStmt, Node *current) {
    visitESTreeChildren(*this, forInOfStmt);

    Node **left{};
    VariableDeclarationNode *decl;
    Node **body;

    if (auto *forIn = llvh::dyn_cast<ForInStatementNode>(forInOfStmt)) {
      left = &forIn->_left;
      body = &forIn->_body;
    } else if (auto *forOf = llvh::dyn_cast<ForOfStatementNode>(forInOfStmt)) {
      left = &forOf->_left;
      body = &forOf->_body;
    } else {
      assert(false && "If not for-in or for-of, then what?");
    }

    decl = llvh::dyn_cast_or_null<VariableDeclarationNode>(*left);

    if (!decl) {
      // forInOfStmt is not in the form
      //
      //   for (const/let/var
      //
      // so nothing to rewrite.
      return ESTree::Unmodified;
    }

    // The parser ensures there's a single declaration.
    assert(
        decl->_declarations.size() == 1 &&
        "for-in/of(let/const/var should have a single declaration");

    auto *varDecl =
        llvh::cast<VariableDeclaratorNode>(&decl->_declarations.front());

    if (varDecl->_init != nullptr) {
      // This is a semantic error (unless ES2023 B.3.5 is implemented), so don't
      // transform the AST.
      return ESTree::Unmodified;
    }

    UniqueString *temp = internalIDs_.next("forInOf").getUnderlyingPointer();

    // The outer block is:
    //
    // wrapperBlock: {
    //   let temp;
    //   [label:]* for (temp in right) ...
    //   break wrapperBlock;
    //   let x; // for TDZ
    // }
    UniqueString *wrapperBlockLabel =
        internalIDs_.next("wrapperBlock").getUnderlyingPointer();
    BlockStatementNode *wrapperBlock = makeBlock();

    VariableDeclarationNode *tempDecl = makeVarDeclaration(identLet_);
    tempDecl->_declarations.push_back(*makeVarDeclarator(temp));
    wrapperBlock->_body.push_back(*tempDecl);

    wrapperBlock->_body.push_back(*current);

    if (decl->_kind != identVar_) {
      // forInOfStmt is in the for
      //
      //   for (const/let ...
      //
      // Thus emit a fake "let x" declaration after the loop to preserve TDZ
      // semantics in case right is, e.g., right: { x }.
      VariableDeclarationNode *xTDZDecl = makeVarDeclaration(identLet_);
      llvh::SetVector<UniqueString *> ids;
      collectAllIDs(varDecl->_id, ids);
      for (UniqueString *id : ids) {
        xTDZDecl->_declarations.push_back(*makeVarDeclarator(id));
      }
      wrapperBlock->_body.push_back(*makeBreak(wrapperBlockLabel));
      wrapperBlock->_body.push_back(*xTDZDecl);
    }

    // Patching forInOfStmt in-place so that it no longer has forDeclarations.
    *left = makeIdentifier(temp);

    // Now add an initializer (= temp) to decl.
    varDecl->_init = makeIdentifier(temp);

    // Build the new for-in/of body:
    //
    // {
    //   const/let/var x = temp;
    //   body;
    // }
    BlockStatementNode *newBody = makeBlock();
    newBody->_body.push_back(*decl);
    newBody->_body.push_back(**body);
    *body = newBody;

    return makeLabel(wrapperBlockLabel, wrapperBlock);
  }

  /// Rewrites
  ///
  ///   [label:]* for (const/let x = init; test; update) body
  ///
  /// into
  ///
  ///  {
  ///    const/let x = init;
  ///    let temp_x = x;
  ///    let first = true;
  ///    undefined;
  ///    for (;;) {
  ///      const/let x = temp_x;
  ///      if (first) {
  ///        first = false;
  ///      } else {
  ///        update;
  ///      }
  ///      if (!test) break;
  ///      control = true;
  ///      [label:]* for (; control; control = false, temp_x = x) {
  ///        body
  ///      }
  ///      if (control) break;
  ///    }
  ///  }
  ///
  /// which preserves the semantics required by ES2023 14.7.4.3:
  ///
  ///  * each iteration happens in a new environment in which the
  ///    forDeclarations are copied to.
  ///  * update executes in the next iteration's context.
  ///
  /// while preserving the original statement's completion value.
  ESTree::VisitResult visitAndRewriteFor(
      ForStatementNode *forStmt,
      Node *current) {
    visitESTreeChildren(*this, forStmt);

    auto *init =
        llvh::dyn_cast_or_null<VariableDeclarationNode>(forStmt->_init);
    if (!init || init->_kind == identVar_) {
      // forStmt is
      //
      //   for (var ...; ...; ...)
      //   for (expr; ...; ...)
      //
      // which doesn't need to be rewritten.
      return ESTree::Unmodified;
    }

    BlockStatementNode *wrapperBlock = makeBlock();

    BlockStatementNode *outerBlock = makeBlock(); // outer for body.

    Node *test = forStmt->_test;
    Node *update = forStmt->_update;

    // forStmt becomes the inner loop in the transformed AST, and its init,
    // test, and update expressions are used by the outer loop. Thus, "break"
    // the link between forStmt and those components early to prevent unintended
    // usage.
    forStmt->_init = nullptr;
    forStmt->_test = nullptr;
    forStmt->_update = nullptr;

    // The transformation introduces one temporary variable for each declaration
    // in init, and then "copies" between then at specific points:
    //
    //   * tempsDecl declares the temp variables (one per declaration in init),
    //     and each temporary is initialized with the values introduced by init.
    //
    //   * initFromTemps re-declares all variables in init inside the outer
    //     loop; this is a const/let declaration (based on the original init
    //     kind), and each variable is initialize with its corresponding
    //     temporary.
    //
    //   * newUpdate copies the values from the declarations introduced by
    //     initFromTemps back into the temporaries; this means that next time
    //     the outer loop executes initFromTemps will use the updated values.
    VariableDeclarationNode *tempsDecl = makeVarDeclaration(identLet_);
    VariableDeclarationNode *initFromTemps = makeVarDeclaration(init->_kind);
    SequenceExpressionNode *newUpdate = makeSequenceExpression();

    llvh::DenseMap<UniqueString *, UniqueString *> tempIds;
    for (Node &n : init->_declarations) {
      auto *decl = llvh::cast<VariableDeclaratorNode>(&n);
      traverseForLexicalDecl(
          decl->_id, tempIds, tempsDecl, initFromTemps, newUpdate);
    }

    // Creating the wrapper block:
    //
    // {
    //   const/let x = init;
    //   let temp_x = x, first = true, first;
    //   undefined;
    //   outer : for (;;) outerBlock
    // }
    wrapperBlock->_body.push_back(*init);
    wrapperBlock->_body.push_back(*tempsDecl);
    wrapperBlock->_body.push_back(
        *toStatement(makeIdentifier(identUndefined_)));
    wrapperBlock->_body.push_back(
        *makeFor(nullptr, nullptr, nullptr, outerBlock));

    // Creating outerBlock:
    // {
    //   const/let x = temp_x;
    //   if (first)                //
    //     first = false;          // only when
    //   else                      //   update != nullptr
    //     update;                 //
    //   if (!cond);               // only when
    //     break;                  //   cond != nullptr
    //   control = true;
    //   [label:]* for (; control; control = 0, temp_x = x) body;
    //   if (control) break;
    // }
    outerBlock->_body.push_back(*initFromTemps);

    if (update) {
      // if (first)
      //   first = false;
      // else
      //   update;
      UniqueString *first = internalIDs_.next("first").getUnderlyingPointer();
      tempsDecl->_declarations.push_back(
          *makeVarDeclarator(first, makeBooleanLiteral(true)));

      IdentifierNode *firstIsTrue = makeIdentifier(first);
      Node *firstEqFalse = toStatement(
          makeAssignment(makeIdentifier(first), makeBooleanLiteral(false)));
      outerBlock->_body.push_back(
          *makeIf(firstIsTrue, firstEqFalse, toStatement(update)));
    }

    if (test) {
      // if (!cond) break;
      outerBlock->_body.push_back(*makeIf(makeNot(test), makeBreak()));
    }

    // control = true (also add the variable to tempsDecl in the wrapperBlock)
    UniqueString *control =
        internalIDs_.next("forControl").getUnderlyingPointer();
    tempsDecl->_declarations.push_back(*makeVarDeclarator(control));
    outerBlock->_body.push_back(*toStatement(
        makeAssignment(makeIdentifier(control), makeBooleanLiteral(true))));

    // Add control = false to new update expression list.
    newUpdate->_expressions.push_front(
        *makeAssignment(makeIdentifier(control), makeBooleanLiteral(false)));

    // Rewrite forStatement into
    //
    //   for (; control; control = 0, temp_x = x)
    forStmt->_test = makeIdentifier(control);
    forStmt->_update = newUpdate;
    outerBlock->_body.push_back(*current);

    // if (control) break;
    outerBlock->_body.push_back(*makeIf(makeIdentifier(control), makeBreak()));

    return wrapperBlock;
  }

  /// The maximum AST nesting level. Once we reach it, we report an error and
  /// stop.
  static constexpr unsigned MAX_RECURSION_DEPTH =
#if defined(HERMES_LIMIT_STACK_DEPTH) || defined(_MSC_VER)
      512
#else
      1024
#endif
      ;

  /// MAX_RECURSION_DEPTH minus the current AST nesting level. Once it reaches
  /// 0 stop transforming it.
  unsigned recursionDepth_ = MAX_RECURSION_DEPTH;

  UniqueString *const identAssign_;
  UniqueString *const identExclaim_;
  UniqueString *const identLet_;
  UniqueString *const identUndefined_;
  UniqueString *const identVar_;

  /// Collect all IDs appearing in \p node and store them in \p ids. These IDs
  /// appear in a forInOf declaration.
  void collectAllIDs(Node *node, llvh::SetVector<UniqueString *> &ids) {
    switch (node->getKind()) {
      case NodeKind::Empty: {
        break;
      }
      case NodeKind::AssignmentPattern: {
        auto *assignment = llvh::cast<AssignmentPatternNode>(node);
        collectAllIDs(assignment->_left, ids);
        break;
      }
      case NodeKind::RestElement: {
        auto *rest = llvh::cast<RestElementNode>(node);
        collectAllIDs(rest->_argument, ids);
        break;
      }
      case NodeKind::Property: {
        auto *prop = llvh::cast<PropertyNode>(node);
        collectAllIDs(prop->_value, ids);
        break;
      }
      case NodeKind::ObjectPattern: {
        auto *objectPattern = llvh::cast<ObjectPatternNode>(node);
        for (Node &prop : objectPattern->_properties) {
          collectAllIDs(&prop, ids);
        }
        break;
      }
      case NodeKind::ArrayPattern: {
        auto *arrayPattern = llvh::cast<ArrayPatternNode>(node);
        for (Node &element : arrayPattern->_elements) {
          collectAllIDs(&element, ids);
        }
        break;
      }
      case NodeKind::Identifier: {
        auto *id = llvh::cast<IdentifierNode>(node);
        ids.insert(id->_name);
        break;
      }
      default:
        llvm_unreachable("unhandled node in collectAllIDs");
    }
  }

  /// Traverses \p node, which is a ForLexicalDeclaration in a
  ///
  ///   for (const/let ... ; ... ; ...) ...
  ///
  /// statement, and emits VariableDeclarators and AssignmentExpressions used
  /// during for rewrite.
  void traverseForLexicalDecl(
      Node *node,
      llvh::DenseMap<UniqueString *, UniqueString *> &tempIds,
      VariableDeclarationNode *tempsDecl,
      VariableDeclarationNode *initFromTemps,
      SequenceExpressionNode *newUpdate) {
    switch (node->getKind()) {
      case NodeKind::Empty: {
        break;
      }
      case NodeKind::AssignmentPattern: {
        auto *assignment = llvh::cast<AssignmentPatternNode>(node);
        traverseForLexicalDecl(
            assignment->_left, tempIds, tempsDecl, initFromTemps, newUpdate);
        break;
      }
      case NodeKind::RestElement: {
        auto *rest = llvh::dyn_cast<RestElementNode>(node);
        traverseForLexicalDecl(
            rest->_argument, tempIds, tempsDecl, initFromTemps, newUpdate);
        break;
      }
      case NodeKind::Property: {
        auto *prop = llvh::cast<PropertyNode>(node);
        traverseForLexicalDecl(
            prop->_value, tempIds, tempsDecl, initFromTemps, newUpdate);
        break;
      }
      case NodeKind::ObjectPattern: {
        auto *objectPattern = llvh::cast<ObjectPatternNode>(node);
        for (Node &prop : objectPattern->_properties) {
          traverseForLexicalDecl(
              &prop, tempIds, tempsDecl, initFromTemps, newUpdate);
        }
        break;
      }
      case NodeKind::ArrayPattern: {
        auto *arrayPattern = llvh::cast<ArrayPatternNode>(node);
        for (Node &elt : arrayPattern->_elements) {
          traverseForLexicalDecl(
              &elt, tempIds, tempsDecl, initFromTemps, newUpdate);
        }
        break;
      }
      case NodeKind::Identifier: {
        auto *id = llvh::cast<IdentifierNode>(node);
        auto res = tempIds.insert(std::make_pair(id->_name, nullptr));
        if (res.second) {
          res.first->second =
              internalIDs_.next("forDecl").getUnderlyingPointer();
        }

        UniqueString *temp = res.first->second;

        // let temp_x = x (in wrapper block)
        tempsDecl->_declarations.push_back(
            *makeVarDeclarator(temp, makeIdentifier(id->_name)));

        // const/let x = temp_x (in outer loop)
        initFromTemps->_declarations.push_back(
            *makeVarDeclarator(id->_name, makeIdentifier(temp)));

        // temp_x = x; (inner loop update)
        newUpdate->_expressions.push_back(
            *makeAssignment(makeIdentifier(temp), makeIdentifier(id->_name)));

        break;
      }
      default:
        llvm_unreachable("unhandled node in duplicateID");
    }
  }

  //****** Helpers for creating AST nodes ******//

  /// \return (IdentifierNode <<name>>)
  IdentifierNode *makeIdentifier(UniqueString *name) {
    return new (astContext_) IdentifierNode(name, nullptr, false);
  }

  /// \return (VariableDeclarationNode <<kind>> <<{}>>)
  VariableDeclarationNode *makeVarDeclaration(UniqueString *kind) {
    return new (astContext_) VariableDeclarationNode(kind, {});
  }

  /// \return (VariableDeclaratorNode <<init>> (Identifier <<name>>))
  VariableDeclaratorNode *makeVarDeclarator(
      UniqueString *name,
      Node *init = nullptr) {
    return new (astContext_) VariableDeclaratorNode(init, makeIdentifier(name));
  }

  /// \return (AssignmentExpressionNode <<"=">> <<dst>> <<src>>)
  AssignmentExpressionNode *makeAssignment(Node *dst, Node *src) {
    return new (astContext_) AssignmentExpressionNode(identAssign_, dst, src);
  }

  /// \return <<n>> if n is a StatementNode;
  ///         (ExpressionStatementNode <<n>> <<>>) otherwise
  StatementNode *toStatement(Node *n) {
    if (auto *stmt = llvh::dyn_cast<StatementNode>(n)) {
      return stmt;
    }
    return new (astContext_) ExpressionStatementNode(n, nullptr);
  }

  /// \return (UnaryExpressionNode "!" <<n>>)
  UnaryExpressionNode *makeNot(Node *n) {
    return new (astContext_) UnaryExpressionNode(identExclaim_, n, false);
  }

  /// \return (LabeledStatementNode (Identifier <<label>>) <<body>>)
  LabeledStatementNode *makeLabel(UniqueString *label, Node *body) {
    return new (astContext_) LabeledStatementNode(makeIdentifier(label), body);
  }

  /// \return (BreakStatementNode (Identifier <<label>>)?)
  BreakStatementNode *makeBreak(UniqueString *label = nullptr) {
    Node *labelNode{};
    if (label) {
      labelNode = makeIdentifier(label);
    }
    return new (astContext_) BreakStatementNode(labelNode);
  }

  /// \return (BooleanLiteralNode <<value>>)
  BooleanLiteralNode *makeBooleanLiteral(bool value) {
    return new (astContext_) BooleanLiteralNode(value);
  }

  /// \return (BlockStatementNode <<{}>>)
  BlockStatementNode *makeBlock() {
    return new (astContext_) BlockStatementNode({});
  }

  /// \return (SequenceExpressionNode <<{}>>)
  SequenceExpressionNode *makeSequenceExpression() {
    return new (astContext_) SequenceExpressionNode({});
  }

  /// \return (IfStatementNode <<test>> <<consequent>> <<alternate>>?)
  IfStatementNode *
  makeIf(Node *test, Node *consequent, Node *alternate = nullptr) {
    return new (astContext_) IfStatementNode(test, consequent, alternate);
  }

  /// \return (ForStatementNode <<init>>? <<test>>? <<update>>? <<body>>)
  ForStatementNode *makeFor(Node *init, Node *test, Node *update, Node *body) {
    return new (astContext_) ForStatementNode(init, test, update, body);
  }
};

} // namespace

void canonicalizeForBlockScoping(Context &astContext, Node *root) {
  BlockScopingTransformations BST{astContext};
  visitESTreeNode(BST, root);
}
} // namespace sem
} // namespace hermes