1 files changed, 50 insertions, 21 deletions

diff --git a/src/interp/binary-reader-interp.cc b/src/interp/binary-reader-interp.cc
index d847fa8e..2c66348e 100644
--- a/src/interp/binary-reader-interp.cc
+++ b/src/interp/binary-reader-interp.cc
@@ -283,11 +283,18 @@ class BinaryReaderInterp : public BinaryReaderNop {
   wabt::Result AppendFixup(IstreamOffsetVectorVector* fixups_vector,
                            Index index);
   wabt::Result EmitBrOffset(Index depth, IstreamOffset offset);
+  wabt::Result GetDropCount(Index keep_count,
+                            size_t type_stack_limit,
+                            Index* out_drop_count);
   wabt::Result GetBrDropKeepCount(Index depth,
                                   Index* out_drop_count,
                                   Index* out_keep_count);
   wabt::Result GetReturnDropKeepCount(Index* out_drop_count,
                                       Index* out_keep_count);
+  wabt::Result GetReturnCallDropKeepCount(FuncSignature* sig,
+                                          Index keep_extra,
+                                          Index* out_drop_count,
+                                          Index* out_keep_count);
   wabt::Result EmitBr(Index depth, Index drop_count, Index keep_count);
   wabt::Result EmitBrTableOffset(Index depth);
   wabt::Result FixupTopLabel();
@@ -532,19 +539,28 @@ wabt::Result BinaryReaderInterp::EmitBrOffset(Index depth,
   return wabt::Result::Ok;
 }
 
+wabt::Result BinaryReaderInterp::GetDropCount(Index keep_count,
+                                              size_t type_stack_limit,
+                                              Index* out_drop_count) {
+  assert(typechecker_.type_stack_size() >= type_stack_limit);
+  Index type_stack_count = typechecker_.type_stack_size() - type_stack_limit;
+  // The keep_count may be larger than the type_stack_count if the typechecker
+  // is currently unreachable. In that case, it doesn't matter what value we
+  // drop, but 0 is a reasonable choice.
+  *out_drop_count =
+      type_stack_count >= keep_count ? type_stack_count - keep_count : 0;
+  return wabt::Result::Ok;
+}
+
 wabt::Result BinaryReaderInterp::GetBrDropKeepCount(Index depth,
                                                     Index* out_drop_count,
                                                     Index* out_keep_count) {
   TypeChecker::Label* label;
   CHECK_RESULT(typechecker_.GetLabel(depth, &label));
-  *out_keep_count = label->br_types().size();
-  if (typechecker_.IsUnreachable()) {
-    *out_drop_count = 0;
-  } else {
-    *out_drop_count =
-        (typechecker_.type_stack_size() - label->type_stack_limit) -
-        *out_keep_count;
-  }
+  Index keep_count = label->br_types().size();
+  CHECK_RESULT(
+      GetDropCount(keep_count, label->type_stack_limit, out_drop_count));
+  *out_keep_count = keep_count;
   return wabt::Result::Ok;
 }
 
@@ -556,6 +572,18 @@ wabt::Result BinaryReaderInterp::GetReturnDropKeepCount(Index* out_drop_count,
   return wabt::Result::Ok;
 }
 
+wabt::Result BinaryReaderInterp::GetReturnCallDropKeepCount(
+    FuncSignature* sig,
+    Index keep_extra,
+    Index* out_drop_count,
+    Index* out_keep_count) {
+  Index keep_count = static_cast<Index>(sig->param_types.size()) + keep_extra;
+  CHECK_RESULT(GetDropCount(keep_count, 0, out_drop_count));
+  *out_drop_count += current_func_->param_and_local_types.size();
+  *out_keep_count = keep_count;
+  return wabt::Result::Ok;
+}
+
 wabt::Result BinaryReaderInterp::EmitBr(Index depth,
                                         Index drop_count,
                                         Index keep_count) {
@@ -1482,7 +1510,8 @@ wabt::Result BinaryReaderInterp::OnCallExpr(Index func_index) {
   return wabt::Result::Ok;
 }
 
-wabt::Result BinaryReaderInterp::OnCallIndirectExpr(Index sig_index, Index table_index) {
+wabt::Result BinaryReaderInterp::OnCallIndirectExpr(Index sig_index,
+                                                    Index table_index) {
   if (module_->table_index == kInvalidIndex) {
     PrintError("found call_indirect operator, but no table");
     return wabt::Result::Error;
@@ -1500,13 +1529,12 @@ wabt::Result BinaryReaderInterp::OnCallIndirectExpr(Index sig_index, Index table
 wabt::Result BinaryReaderInterp::OnReturnCallExpr(Index func_index) {
   Func* func = GetFuncByModuleIndex(func_index);
   FuncSignature* sig = env_->GetFuncSignature(func->sig_index);
-  CHECK_RESULT(typechecker_.OnReturnCall(sig->param_types, sig->result_types));
 
   Index drop_count, keep_count;
-  CHECK_RESULT(GetReturnDropKeepCount(&drop_count, &keep_count));
-
-  keep_count = static_cast<Index>(sig->param_types.size());
-
+  CHECK_RESULT(GetReturnCallDropKeepCount(sig, 0, &drop_count, &keep_count));
+  // The typechecker must be run after we get the drop/keep counts, since it
+  // will change the type stack.
+  CHECK_RESULT(typechecker_.OnReturnCall(sig->param_types, sig->result_types));
   CHECK_RESULT(EmitDropKeep(drop_count, keep_count));
 
   if (func->is_host) {
@@ -1521,20 +1549,21 @@ wabt::Result BinaryReaderInterp::OnReturnCallExpr(Index func_index) {
   return wabt::Result::Ok;
 }
 
-wabt::Result BinaryReaderInterp::OnReturnCallIndirectExpr(Index sig_index, Index table_index) {
+wabt::Result BinaryReaderInterp::OnReturnCallIndirectExpr(Index sig_index,
+                                                          Index table_index) {
   if (module_->table_index == kInvalidIndex) {
     PrintError("found return_call_indirect operator, but no table");
     return wabt::Result::Error;
   }
   FuncSignature* sig = GetSignatureByModuleIndex(sig_index);
-  CHECK_RESULT(
-    typechecker_.OnReturnCallIndirect(sig->param_types, sig->result_types));
 
   Index drop_count, keep_count;
-  CHECK_RESULT(GetReturnDropKeepCount(&drop_count, &keep_count));
-
-  keep_count = static_cast<Index>(sig->param_types.size()+1); // Include the index of the function
-
+  // +1 to include the index of the function.
+  CHECK_RESULT(GetReturnCallDropKeepCount(sig, +1, &drop_count, &keep_count));
+  // The typechecker must be run after we get the drop/keep counts, since it
+  // changes the type stack.
+  CHECK_RESULT(
+      typechecker_.OnReturnCallIndirect(sig->param_types, sig->result_types));
   CHECK_RESULT(EmitDropKeep(drop_count, keep_count));
 
   CHECK_RESULT(EmitOpcode(Opcode::ReturnCallIndirect));