1 files changed, 22 insertions, 1 deletions
diff --git a/src/ir/type-updating.cpp b/src/ir/type-updating.cpp
index d909bffb4..658fabc83 100644
--- a/src/ir/type-updating.cpp
+++ b/src/ir/type-updating.cpp
@@ -39,6 +39,7 @@ void handleNonDefaultableLocals(Function* func, Module& wasm) {
   if (!hasNonNullable) {
     return;
   }
+
   // Rewrite the local.gets.
   Builder builder(wasm);
   for (auto** getp : FindAllPointers<LocalGet>(func->body).list) {
@@ -56,8 +57,28 @@ void handleNonDefaultableLocals(Function* func, Module& wasm) {
     }
   }
 
+  // Update tees, whose type must match the local (if the wasm spec changes for
+  // the type to be that of the value, then this can be removed).
+  for (auto** setp : FindAllPointers<LocalSet>(func->body).list) {
+    auto* set = (*setp)->cast<LocalSet>();
+    if (!func->isVar(set->index)) {
+      // We do not need to process params, which can legally be non-nullable.
+      continue;
+    }
+    // Non-tees do not change, and unreachable tees can be ignored here as their
+    // type is unreachable anyhow.
+    if (!set->isTee() || set->type == Type::unreachable) {
+      continue;
+    }
+    auto type = func->getLocalType(set->index);
+    if (type.isRef() && !type.isNullable()) {
+      set->type = Type(type.getHeapType(), Nullable);
+      *setp = builder.makeRefAs(RefAsNonNull, set);
+    }
+  }
+
   // Rewrite the types of the function's vars (which we can do now, after we
-  // are done using them to know which local.gets to fix).
+  // are done using them to know which local.gets etc to fix).
   for (auto& type : func->vars) {
     if (type.isRef() && !type.isNullable()) {
       type = Type(type.getHeapType(), Nullable);