[type-profile] Handle returns correctly.

src/interpreter/bytecode-generator.cc

Index: src/interpreter/bytecode-generator.cc
diff --git a/src/interpreter/bytecode-generator.cc b/src/interpreter/bytecode-generator.cc
index 1bf1a3351e42ba00c27f4277072b9c86e8891d23..12b07339bad0772a3f7082d591fc54cfef25e4da 100644
--- a/src/interpreter/bytecode-generator.cc
+++ b/src/interpreter/bytecode-generator.cc
@@ -1092,13 +1092,6 @@ void BytecodeGenerator::VisitBreakStatement(BreakStatement* stmt) {
 void BytecodeGenerator::VisitReturnStatement(ReturnStatement* stmt) {
   builder()->SetStatementPosition(stmt);
   VisitForAccumulatorValue(stmt->expression());
-
-  if (stmt->HasTypeProfileSlot()) {
-    FeedbackSlot collect_type_feedback_slot = stmt->TypeProfileSlot();
-    builder()->CollectTypeProfile(stmt->position(),
-                                  feedback_index(collect_type_feedback_slot));
-  }
-
   if (stmt->is_async_return()) {
     execution_control()->AsyncReturnAccumulator();
   } else {
@@ -2030,6 +2023,11 @@ void BytecodeGenerator::BuildReturn() {
     builder()->StoreAccumulatorInRegister(result).CallRuntime(
         Runtime::kTraceExit, result);
   }
+  if (!info()->literal()->TypeProfileSlot().IsInvalid()) {
+    builder()->CollectTypeProfile(
+        info()->literal()->position(),
+        feedback_index(info()->literal()->TypeProfileSlot()));
+  }
   builder()->Return();
 }
 
@@ -2339,16 +2337,6 @@ void BytecodeGenerator::VisitAssignment(Assignment* expr) {
       break;
     }
   }
-
-  // Value is in accumulator.
-  // TODO(franzih): Collect type profile once we can handle more than just
-  // return statements.
-  if (false && expr->HasTypeProfileSlot()) {
-    FeedbackSlot collect_type_feedback_slot = expr->TypeProfileSlot();
-
-    builder()->CollectTypeProfile(expr->position(),
-                                  feedback_index(collect_type_feedback_slot));
-  }
 }
 
 void BytecodeGenerator::VisitYield(Yield* expr) {