wasm-interp: Handle refs_ correctly

include/wabt/opcode.def

@@ -235,6 +235,9 @@ WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xe3, InterpData, "data", "")
 WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xe4, InterpDropKeep, "drop_keep", "")
 WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xe5, InterpCatchDrop, "catch_drop", "")
 WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xe6, InterpAdjustFrameForReturnCall, "adjust_frame_for_return_call", "")
+WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xe7, InterpGlobalGetRef, "global.get.ref", "")
+WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xe9, InterpLocalGetRef, "local.get.ref", "")
+WABT_OPCODE(___,  ___,  ___,  ___,  0,  0,    0xea, InterpMarkRef, "mark_ref", "")
 
 /* Saturating float-to-int opcodes (--enable-saturating-float-to-int) */
 WABT_OPCODE(I32,  F32,  ___,  ___,  0,  0xfc, 0x00, I32TruncSatF32S, "i32.trunc_sat_f32_s", "")

src/interp/binary-reader-interp.cc

@@ -875,6 +875,11 @@ Result BinaryReaderInterp::OnLocalDecl(Index decl_index,
 Result BinaryReaderInterp::EndLocalDecls() {
   if (local_count_ != 0) {
     istream_.Emit(Opcode::InterpAlloca, local_count_);
+    for (Index i = 0; i < local_count_; i++) {
+      if (func_->GetLocalType(func_->type.params.size() + i).IsRef()) {
+        istream_.Emit(Opcode::InterpMarkRef, local_count_ - i);
+      }
+    }
   }
   // Continuation of the implicit func label, used for exception handling. (See
   // BeginFunctionBody.)
@@ -1274,7 +1279,13 @@ Result BinaryReaderInterp::OnV128ConstExpr(v128 value_bits) {
 Result BinaryReaderInterp::OnGlobalGetExpr(Index global_index) {
   CHECK_RESULT(
       validator_.OnGlobalGet(GetLocation(), Var(global_index, GetLocation())));
-  istream_.Emit(Opcode::GlobalGet, global_index);
+
+  Type type = global_types_.at(global_index).type;
+  if (type.IsRef()) {
+    istream_.Emit(Opcode::InterpGlobalGetRef, global_index);
+  } else {
+    istream_.Emit(Opcode::GlobalGet, global_index);
+  }
   return Result::Ok;
 }
 
@@ -1297,7 +1308,13 @@ Result BinaryReaderInterp::OnLocalGetExpr(Index local_index) {
   Index translated_local_index = TranslateLocalIndex(local_index);
   CHECK_RESULT(
       validator_.OnLocalGet(GetLocation(), Var(local_index, GetLocation())));
-  istream_.Emit(Opcode::LocalGet, translated_local_index);
+
+  Type type = func_->GetLocalType(local_index);
+  if (type.IsRef()) {
+    istream_.Emit(Opcode::InterpLocalGetRef, translated_local_index);
+  } else {
+    istream_.Emit(Opcode::LocalGet, translated_local_index);
+  }
   return Result::Ok;
 }
 
@@ -1306,13 +1323,15 @@ Result BinaryReaderInterp::OnLocalSetExpr(Index local_index) {
   Index translated_local_index = TranslateLocalIndex(local_index);
   CHECK_RESULT(
       validator_.OnLocalSet(GetLocation(), Var(local_index, GetLocation())));
+
   istream_.Emit(Opcode::LocalSet, translated_local_index);
   return Result::Ok;
 }
 
 Result BinaryReaderInterp::OnLocalTeeExpr(Index local_index) {
   CHECK_RESULT(
       validator_.OnLocalTee(GetLocation(), Var(local_index, GetLocation())));
+
   istream_.Emit(Opcode::LocalTee, TranslateLocalIndex(local_index));
   return Result::Ok;
 }

src/interp/interp.cc

@@ -1131,7 +1131,7 @@ T WABT_VECTORCALL Thread::Pop() {
 }
 
 Value Thread::Pop() {
-  if (!refs_.empty() && refs_.back() >= values_.size()) {
+  if (!refs_.empty() && refs_.back() >= values_.size() - 1) {
     refs_.pop_back();
   }
   auto value = values_.back();
@@ -1254,16 +1254,24 @@ RunResult Thread::StepInternal(Trap::Ptr* out_trap) {
       break;
 
     case O::Select: {
-      // TODO: need to mark whether this is a ref.
       auto cond = Pop<u32>();
+      auto ref = false;
+      // check if either is a ref
+      ref |= !refs_.empty() && refs_.back() == values_.size();
       Value false_ = Pop();
+      ref |= !refs_.empty() && refs_.back() == values_.size();
       Value true_ = Pop();
+      if (ref) {
+        refs_.push_back(values_.size());
+      }
       Push(cond ? true_ : false_);
       break;
     }
 
+    case O::InterpLocalGetRef:
+      refs_.push_back(values_.size());
+      [[fallthrough]];
     case O::LocalGet:
-      // TODO: need to mark whether this is a ref.
       Push(Pick(instr.imm_u32));
       break;
 
@@ -1277,8 +1285,14 @@ RunResult Thread::StepInternal(Trap::Ptr* out_trap) {
       Pick(instr.imm_u32) = Pick(1);
       break;
 
+    case O::InterpMarkRef:
+      refs_.push_back(values_.size() - instr.imm_u32);
+      break;
+
+    case O::InterpGlobalGetRef:
+      refs_.push_back(values_.size());
+      [[fallthrough]];
     case O::GlobalGet: {
-      // TODO: need to mark whether this is a ref.
       Global::Ptr global{store_, inst_->globals()[instr.imm_u32]};
       Push(global->Get());
       break;
@@ -1478,9 +1492,7 @@ RunResult Thread::StepInternal(Trap::Ptr* out_trap) {
 
     case O::InterpAlloca:
       values_.resize(values_.size() + instr.imm_u32);
-      // refs_ doesn't need to be updated; We may be allocating space for
-      // references, but they will be initialized to null, so it is OK if we
-      // don't mark them.
+      // refs_ will be marked in InterpMarkRef.
       break;
 
     case O::InterpBrUnless:
@@ -1499,13 +1511,23 @@ RunResult Thread::StepInternal(Trap::Ptr* out_trap) {
       auto drop = instr.imm_u32x2.fst;
       auto keep = instr.imm_u32x2.snd;
       // Shift kept refs down.
-      for (auto iter = refs_.rbegin(); iter != refs_.rend(); ++iter) {
+      auto iter = refs_.rbegin();
+      for (; iter != refs_.rend(); ++iter) {
         if (*iter >= values_.size() - keep) {
           *iter -= drop;
         } else {
           break;
         }
       }
+      // Find dropped refs.
+      auto drop_iter = iter;
+      for (; drop_iter != refs_.rend(); ++drop_iter) {
+        if (*iter < values_.size() - keep - drop) {
+          break;
+        }
+      }
+      // Erase dropped refs.
+      refs_.erase(drop_iter.base(), iter.base());
       std::move(values_.end() - keep, values_.end(),
                 values_.end() - drop - keep);
       values_.resize(values_.size() - drop);

src/interp/istream.cc

@@ -515,6 +515,9 @@ Instr Istream::Read(Offset* offset) const {
 
     case Opcode::GlobalGet:
     case Opcode::LocalGet:
+    case Opcode::InterpLocalGetRef:
+    case Opcode::InterpGlobalGetRef:
+    case Opcode::InterpMarkRef:
     case Opcode::MemorySize:
     case Opcode::TableSize:
     case Opcode::DataDrop: