Enable visitor in rewriter to replace statements.

This is in preparation of implementing ES6 completion semantics and
depends on #1362333002.

R=rossberg
BUG=

Review URL: https://codereview.chromium.org/1362363002

Cr-Commit-Position: refs/heads/master@{#31041}

src/ast.h

@@ -640,6 +640,7 @@ class IterationStatement : public BreakableStatement {
   IterationStatement* AsIterationStatement() final { return this; }
 
   Statement* body() const { return body_; }
+  void set_body(Statement* s) { body_ = s; }
 
   static int num_ids() { return parent_num_ids() + 1; }
   BailoutId OsrEntryId() const { return BailoutId(local_id(0)); }
@@ -988,6 +989,7 @@ class WithStatement final : public Statement {
   Scope* scope() { return scope_; }
   Expression* expression() const { return expression_; }
   Statement* statement() const { return statement_; }
+  void set_statement(Statement* s) { statement_ = s; }
 
   void set_base_id(int id) { base_id_ = id; }
   static int num_ids() { return parent_num_ids() + 1; }
@@ -1092,6 +1094,9 @@ class IfStatement final : public Statement {
   Statement* then_statement() const { return then_statement_; }
   Statement* else_statement() const { return else_statement_; }
 
+  void set_then_statement(Statement* s) { then_statement_ = s; }
+  void set_else_statement(Statement* s) { else_statement_ = s; }
+
   bool IsJump() const override {
     return HasThenStatement() && then_statement()->IsJump()
         && HasElseStatement() && else_statement()->IsJump();
@@ -1131,6 +1136,7 @@ class IfStatement final : public Statement {
 class TryStatement : public Statement {
  public:
   Block* try_block() const { return try_block_; }
+  void set_try_block(Block* b) { try_block_ = b; }
 
   void set_base_id(int id) { base_id_ = id; }
   static int num_ids() { return parent_num_ids() + 1; }
@@ -1163,6 +1169,7 @@ class TryCatchStatement final : public TryStatement {
   Scope* scope() { return scope_; }
   Variable* variable() { return variable_; }
   Block* catch_block() const { return catch_block_; }
+  void set_catch_block(Block* b) { catch_block_ = b; }
 
  protected:
   TryCatchStatement(Zone* zone, Block* try_block, Scope* scope,
@@ -1184,6 +1191,7 @@ class TryFinallyStatement final : public TryStatement {
   DECLARE_NODE_TYPE(TryFinallyStatement)
 
   Block* finally_block() const { return finally_block_; }
+  void set_finally_block(Block* b) { finally_block_ = b; }
 
  protected:
   TryFinallyStatement(Zone* zone, Block* try_block, Block* finally_block,

src/rewriter.cc

@@ -17,6 +17,7 @@ class Processor: public AstVisitor {
             AstValueFactory* ast_value_factory)
       : result_(result),
         result_assigned_(false),
+        replacement_(nullptr),
         is_set_(false),
         factory_(ast_value_factory) {
     InitializeAstVisitor(isolate, ast_value_factory->zone());
@@ -38,6 +39,10 @@ class Processor: public AstVisitor {
   // there was ever an assignment to result_.
   bool result_assigned_;
 
+  // When visiting a node, we "return" a replacement for that node in
+  // [replacement_].  In many cases this will just be the original node.
+  Statement* replacement_;
+
   // To avoid storing to .result all the time, we eliminate some of
   // the stores by keeping track of whether or not we're sure .result
   // will be overwritten anyway. This is a bit more tricky than what I
@@ -67,6 +72,7 @@ class Processor: public AstVisitor {
 void Processor::Process(ZoneList<Statement*>* statements) {
   for (int i = statements->length() - 1; i >= 0; --i) {
     Visit(statements->at(i));
+    statements->Set(i, replacement_);
   }
 }
 
@@ -81,6 +87,7 @@ void Processor::VisitBlock(Block* node) {
   // returns 'undefined'. To obtain the same behavior with v8, we need
   // to prevent rewriting in that case.
   if (!node->ignore_completion_value()) Process(node->statements());
+  replacement_ = node;
 }
 
 
@@ -90,6 +97,7 @@ void Processor::VisitExpressionStatement(ExpressionStatement* node) {
     node->set_expression(SetResult(node->expression()));
     is_set_ = true;
   }
+  replacement_ = node;
 }
 
 
@@ -97,10 +105,13 @@ void Processor::VisitIfStatement(IfStatement* node) {
   // Rewrite both branches.
   bool set_after = is_set_;
   Visit(node->then_statement());
+  node->set_then_statement(replacement_);
   bool set_in_then = is_set_;
   is_set_ = set_after;
   Visit(node->else_statement());
+  node->set_else_statement(replacement_);
   is_set_ = is_set_ && set_in_then;
+  replacement_ = node;
 }
 
 
@@ -109,7 +120,9 @@ void Processor::VisitIterationStatement(IterationStatement* node) {
   bool set_after = is_set_;
   is_set_ = false;  // We are in a loop, so we can't rely on [set_after].
   Visit(node->body());
+  node->set_body(replacement_);
   is_set_ = is_set_ && set_after;
+  replacement_ = node;
 }
 
 
@@ -142,17 +155,23 @@ void Processor::VisitTryCatchStatement(TryCatchStatement* node) {
   // Rewrite both try and catch block.
   bool set_after = is_set_;
   Visit(node->try_block());
+  node->set_try_block(static_cast<Block*>(replacement_));
   bool set_in_try = is_set_;
   is_set_ = set_after;
   Visit(node->catch_block());
+  node->set_catch_block(static_cast<Block*>(replacement_));
   is_set_ = is_set_ && set_in_try;
+  replacement_ = node;
 }
 
 
 void Processor::VisitTryFinallyStatement(TryFinallyStatement* node) {
   // Rewrite both try and finally block (in reverse order).
   Visit(node->finally_block());
-  Visit(node->try_block());
+  node->set_finally_block(replacement_->AsBlock());
+  Visit(node->try_block());  // Exception will not be caught.
+  node->set_try_block(replacement_->AsBlock());
+  replacement_ = node;
 }
 
 
@@ -165,36 +184,51 @@ void Processor::VisitSwitchStatement(SwitchStatement* node) {
     Process(clause->statements());
   }
   is_set_ = is_set_ && set_after;
+  replacement_ = node;
 }
 
 
 void Processor::VisitContinueStatement(ContinueStatement* node) {
   is_set_ = false;
+  replacement_ = node;
 }
 
 
 void Processor::VisitBreakStatement(BreakStatement* node) {
   is_set_ = false;
+  replacement_ = node;
 }
 
 
 void Processor::VisitWithStatement(WithStatement* node) {
   Visit(node->statement());
+  node->set_statement(replacement_);
+  replacement_ = node;
 }
 
 
 void Processor::VisitSloppyBlockFunctionStatement(
     SloppyBlockFunctionStatement* node) {
   Visit(node->statement());
+  node->set_statement(replacement_);
+  replacement_ = node;
+}
+
+
+void Processor::VisitEmptyStatement(EmptyStatement* node) {
+  replacement_ = node;
 }
 
 
-void Processor::VisitReturnStatement(ReturnStatement* node) { is_set_ = true; }
+void Processor::VisitReturnStatement(ReturnStatement* node) {
+  is_set_ = true;
+  replacement_ = node;
+}
 
 
-// Do nothing:
-void Processor::VisitEmptyStatement(EmptyStatement* node) {}
-void Processor::VisitDebuggerStatement(DebuggerStatement* node) {}
+void Processor::VisitDebuggerStatement(DebuggerStatement* node) {
+  replacement_ = node;
+}
 
 
 // Expressions are never visited.