Make local.tee's type its local's type

According to the current spec, `local.tee`'s return type should be the
same as its local's type. (Discussions on whether we should change this
rule is going on in WebAssembly/reference-types#55, but here I will
assume this spec does not change. If this changes, we should change many
parts of Binaryen transformation anyway...)

But currently in Binaryen `local.tee`'s type is computed from its
value's type. This didn't make any difference in the MVP, but after we
have subtype relationship in WebAssembly#2451, this can become a problem. For
example:
```
(func $test (result funcref) (local $0 anyref)
  (local.tee $0
    (ref.func $test)
  )
)
```
This shouldn't validate in the spec, but this will pass Binaryen
validation with the current `local.tee` implementation.

This makes `local.tee`'s type computed from the local's type, and makes
`LocalSet::makeTee` get a type parameter, to which we should pass the
its corresponding local's type. We don't embed the local type in the
class `LocalSet` because it may increase memory size.

This also fixes the type of `local.get` to be the local type where
`local.get` and `local.set` pair is created from `local.tee`.

Author: aheejin

src/asm2wasm.h

@@ -1907,7 +1907,7 @@ Function* Asm2WasmBuilder::processFunction(Ref ast) {
         auto ret = allocator.alloc<LocalSet>();
         ret->index = function->getLocalIndex(assign->target());
         ret->value = process(assign->value());
-        ret->setTee(false);
+        ret->makeSet();
         ret->finalize();
         return ret;
       }
@@ -2160,7 +2160,7 @@ Function* Asm2WasmBuilder::processFunction(Ref ast) {
             auto set = allocator.alloc<LocalSet>();
             set->index = function->getLocalIndex(I32_TEMP);
             set->value = value;
-            set->setTee(false);
+            set->makeSet();
             set->finalize();
             auto get = [&]() {
               auto ret = allocator.alloc<LocalGet>();
@@ -2266,7 +2266,7 @@ Function* Asm2WasmBuilder::processFunction(Ref ast) {
                 view.bytes,
                 0,
                 processUnshifted(ast[2][1], view.bytes),
-                builder.makeLocalTee(temp, process(ast[2][2])),
+                builder.makeLocalTee(temp, process(ast[2][2]), type),
                 type),
               builder.makeLocalGet(temp, type));
           } else if (name == Atomics_exchange) {

src/binaryen-c.cpp

@@ -1261,22 +1261,23 @@ BinaryenExpressionRef BinaryenLocalSet(BinaryenModuleRef module,
 
   ret->index = index;
   ret->value = (Expression*)value;
-  ret->setTee(false);
+  ret->makeSet();
   ret->finalize();
   return static_cast<Expression*>(ret);
 }
 BinaryenExpressionRef BinaryenLocalTee(BinaryenModuleRef module,
                                        BinaryenIndex index,
-                                       BinaryenExpressionRef value) {
+                                       BinaryenExpressionRef value,
+                                       BinaryenType type) {
   auto* ret = ((Module*)module)->allocator.alloc<LocalSet>();
 
   if (tracing) {
-    traceExpression(ret, "BinaryenLocalTee", index, value);
+    traceExpression(ret, "BinaryenLocalTee", index, value, type);
   }
 
   ret->index = index;
   ret->value = (Expression*)value;
-  ret->setTee(true);
+  ret->makeTee(Type(type));
   ret->finalize();
   return static_cast<Expression*>(ret);
 }

src/binaryen-c.h

@@ -665,8 +665,10 @@ BINARYEN_API BinaryenExpressionRef BinaryenLocalGet(BinaryenModuleRef module,
                                                     BinaryenType type);
 BINARYEN_API BinaryenExpressionRef BinaryenLocalSet(
   BinaryenModuleRef module, BinaryenIndex index, BinaryenExpressionRef value);
-BINARYEN_API BinaryenExpressionRef BinaryenLocalTee(
-  BinaryenModuleRef module, BinaryenIndex index, BinaryenExpressionRef value);
+BINARYEN_API BinaryenExpressionRef BinaryenLocalTee(BinaryenModuleRef module,
+                                                    BinaryenIndex index,
+                                                    BinaryenExpressionRef value,
+                                                    BinaryenType type);
 BINARYEN_API BinaryenExpressionRef BinaryenGlobalGet(BinaryenModuleRef module,
                                                      const char* name,
                                                      BinaryenType type);

src/ir/ExpressionManipulator.cpp

@@ -91,7 +91,7 @@ flexibleCopy(Expression* original, Module& wasm, CustomCopier custom) {
     }
     Expression* visitLocalSet(LocalSet* curr) {
       if (curr->isTee()) {
-        return builder.makeLocalTee(curr->index, copy(curr->value));
+        return builder.makeLocalTee(curr->index, copy(curr->value), curr->type);
       } else {
         return builder.makeLocalSet(curr->index, copy(curr->value));
       }

src/ir/localize.h

@@ -36,7 +36,7 @@ struct Localizer {
       index = set->index;
     } else {
       index = Builder::addVar(func, expr->type);
-      expr = Builder(*wasm).makeLocalTee(index, expr);
+      expr = Builder(*wasm).makeLocalTee(index, expr, expr->type);
     }
   }
 };

src/js/binaryen.js-post.js

@@ -541,8 +541,8 @@ function wrapModule(module, self) {
     'set': function(index, value) {
       return Module['_BinaryenLocalSet'](module, index, value);
     },
-    'tee': function(index, value) {
-      return Module['_BinaryenLocalTee'](module, index, value);
+    'tee': function(index, value, type) {
+      return Module['_BinaryenLocalTee'](module, index, value, type);
     }
   }
 

src/passes/Flatten.cpp

@@ -172,9 +172,10 @@ struct Flatten
             replaceCurrent(set->value); // trivial, no set happens
           } else {
             // use a set in a prelude + a get
-            set->setTee(false);
+            set->makeSet();
             ourPreludes.push_back(set);
-            replaceCurrent(builder.makeLocalGet(set->index, set->value->type));
+            replaceCurrent(builder.makeLocalGet(
+              set->index, getFunction()->getLocalType(set->index)));
           }
         }
       } else if (auto* br = curr->dynCast<Break>()) {

src/passes/LocalCSE.cpp

@@ -184,8 +184,9 @@ struct LocalCSE : public WalkerPass<LinearExecutionWalker<LocalCSE>> {
         if (iter != usables.end()) {
           // already exists in the table, this is good to reuse
           auto& info = iter->second;
+          Type localType = getFunction()->getLocalType(info.index);
           set->value =
-            Builder(*getModule()).makeLocalGet(info.index, value->type);
+            Builder(*getModule()).makeLocalGet(info.index, localType);
           anotherPass = true;
         } else {
           // not in table, add this, maybe we can help others later

src/passes/MergeLocals.cpp

@@ -88,7 +88,7 @@ struct MergeLocals
     if (auto* get = curr->value->dynCast<LocalGet>()) {
       if (get->index != curr->index) {
         Builder builder(*getModule());
-        auto* trivial = builder.makeLocalTee(get->index, get);
+        auto* trivial = builder.makeLocalTee(get->index, get, get->type);
         curr->value = trivial;
         copies.push_back(curr);
       }

src/passes/RemoveUnusedBrs.cpp

@@ -287,7 +287,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         Expression* z;
         replaceCurrent(
           z = builder.makeIf(
-            builder.makeLocalTee(temp, curr->condition),
+            builder.makeLocalTee(temp, curr->condition, i32),
             builder.makeIf(builder.makeBinary(EqInt32,
                                               builder.makeLocalGet(temp, i32),
                                               builder.makeConst(Literal(int32_t(
@@ -1074,7 +1074,7 @@ struct RemoveUnusedBrs : public WalkerPass<PostWalker<RemoveUnusedBrs>> {
         iff->finalize();
         Expression* replacement = iff;
         if (tee) {
-          set->setTee(false);
+          set->makeSet();
           // We need a block too.
           replacement = builder.makeSequence(iff,
                                              get // reuse the get

src/passes/SSAify.cpp

@@ -154,7 +154,7 @@ struct SSAify : public Pass {
         if (set) {
           // a set exists, just add a tee of its value
           auto* value = set->value;
-          auto* tee = builder.makeLocalTee(new_, value);
+          auto* tee = builder.makeLocalTee(new_, value, get->type);
           set->value = tee;
           // the value may have been something we tracked the location
           // of. if so, update that, since we moved it into the tee

src/passes/SimplifyLocals.cpp

@@ -256,7 +256,7 @@ struct SimplifyLocals
       } else {
         this->replaceCurrent(set);
         assert(!set->isTee());
-        set->setTee(true);
+        set->makeTee(this->getFunction()->getLocalType(set->index));
       }
       // reuse the local.get that is dying
       *found->second.item = curr;
@@ -271,7 +271,7 @@ struct SimplifyLocals
     auto* set = curr->value->dynCast<LocalSet>();
     if (set) {
       assert(set->isTee());
-      set->setTee(false);
+      set->makeSet();
       this->replaceCurrent(set);
     }
   }
@@ -559,7 +559,7 @@ struct SimplifyLocals
       auto* set = (*breakLocalSetPointer)->template cast<LocalSet>();
       if (br->condition) {
         br->value = set;
-        set->setTee(true);
+        set->makeTee(this->getFunction()->getLocalType(set->index));
         *breakLocalSetPointer =
           this->getModule()->allocator.template alloc<Nop>();
         // in addition, as this is a conditional br that now has a value, it now
@@ -728,7 +728,8 @@ struct SimplifyLocals
     ifTrueBlock->finalize();
     assert(ifTrueBlock->type != none);
     // Update the ifFalse side.
-    iff->ifFalse = builder.makeLocalGet(set->index, set->value->type);
+    iff->ifFalse = builder.makeLocalGet(
+      set->index, this->getFunction()->getLocalType(set->index));
     iff->finalize(); // update type
     // Update the get count.
     getCounter.num[set->index]++;

src/passes/Untee.cpp

@@ -41,9 +41,10 @@ struct Untee : public WalkerPass<PostWalker<Untee>> {
       } else {
         // a normal tee. replace with set and get
         Builder builder(*getModule());
-        replaceCurrent(builder.makeSequence(
-          curr, builder.makeLocalGet(curr->index, curr->value->type)));
-        curr->setTee(false);
+        LocalGet* get = builder.makeLocalGet(
+          curr->index, getFunction()->getLocalType(curr->index));
+        replaceCurrent(builder.makeSequence(curr, get));
+        curr->makeSet();
       }
     }
   }

src/passes/Vacuum.cpp

@@ -346,7 +346,7 @@ struct Vacuum : public WalkerPass<ExpressionStackWalker<Vacuum>> {
     // a drop of a tee is a set
     if (auto* set = curr->value->dynCast<LocalSet>()) {
       assert(set->isTee());
-      set->setTee(false);
+      set->makeSet();
       replaceCurrent(set);
       return;
     }

src/tools/fuzzing.h

@@ -1282,7 +1282,7 @@ class TranslateToFuzzReader {
     }
     auto* value = make(valueType);
     if (tee) {
-      return builder.makeLocalTee(pick(locals), value);
+      return builder.makeLocalTee(pick(locals), value, valueType);
     } else {
       return builder.makeLocalSet(pick(locals), value);
     }

src/wasm-builder.h

@@ -241,14 +241,16 @@ class Builder {
     auto* ret = allocator.alloc<LocalSet>();
     ret->index = index;
     ret->value = value;
+    ret->makeSet();
     ret->finalize();
     return ret;
   }
-  LocalSet* makeLocalTee(Index index, Expression* value) {
+  LocalSet* makeLocalTee(Index index, Expression* value, Type type) {
     auto* ret = allocator.alloc<LocalSet>();
     ret->index = index;
     ret->value = value;
-    ret->setTee(true);
+    ret->makeTee(type);
+    ret->finalize();
     return ret;
   }
   GlobalGet* makeGlobalGet(Name name, Type type) {

src/wasm.h

@@ -725,8 +725,9 @@ class LocalSet : public SpecificExpression<Expression::LocalSetId> {
   Index index;
   Expression* value;
 
-  bool isTee();
-  void setTee(bool is);
+  bool isTee() const;
+  void makeTee(Type type);
+  void makeSet();
 };
 
 class GlobalGet : public SpecificExpression<Expression::GlobalGetId> {

src/wasm/wasm-binary.cpp

@@ -2527,8 +2527,11 @@ void WasmBinaryBuilder::visitLocalSet(LocalSet* curr, uint8_t code) {
     throwError("bad local.set index");
   }
   curr->value = popNonVoidExpression();
-  curr->type = curr->value->type;
-  curr->setTee(code == BinaryConsts::LocalTee);
+  if (code == BinaryConsts::LocalTee) {
+    curr->makeTee(currFunction->getLocalType(curr->index));
+  } else {
+    curr->makeSet();
+  }
   curr->finalize();
 }
 

src/wasm/wasm-emscripten.cpp

@@ -113,7 +113,7 @@ inline Expression* stackBoundsCheck(Builder& builder,
   auto check =
     builder.makeIf(builder.makeBinary(
                      BinaryOp::LtUInt32,
-                     builder.makeLocalTee(newSP, value),
+                     builder.makeLocalTee(newSP, value, stackPointer->type),
                      builder.makeGlobalGet(stackLimit->name, stackLimit->type)),
                    builder.makeCall(handler, {}, none));
   // (global.set $__stack_pointer (local.get $newSP))
@@ -173,7 +173,7 @@ void EmscriptenGlueGenerator::generateStackAllocFunction() {
   const static uint32_t bitMask = bitAlignment - 1;
   Const* subConst = builder.makeConst(Literal(~bitMask));
   Binary* maskedSub = builder.makeBinary(AndInt32, sub, subConst);
-  LocalSet* teeStackLocal = builder.makeLocalTee(1, maskedSub);
+  LocalSet* teeStackLocal = builder.makeLocalTee(1, maskedSub, i32);
   Expression* storeStack = generateStoreStackPointer(function, teeStackLocal);
 
   Block* block = builder.makeBlock();

src/wasm/wasm-s-parser.cpp

@@ -1040,7 +1040,7 @@ Expression* SExpressionWasmBuilder::makeLocalTee(Element& s) {
   auto ret = allocator.alloc<LocalSet>();
   ret->index = getLocalIndex(*s[1]);
   ret->value = parseExpression(s[2]);
-  ret->setTee(true);
+  ret->makeTee(currFunction->getLocalType(ret->index));
   ret->finalize();
   return ret;
 }
@@ -1049,7 +1049,7 @@ Expression* SExpressionWasmBuilder::makeLocalSet(Element& s) {
   auto ret = allocator.alloc<LocalSet>();
   ret->index = getLocalIndex(*s[1]);
   ret->value = parseExpression(s[2]);
-  ret->setTee(false);
+  ret->makeSet();
   ret->finalize();
   return ret;
 }

src/wasm/wasm-validator.cpp

@@ -720,7 +720,13 @@ void FunctionValidator::visitLocalSet(LocalSet* curr) {
   if (curr->value->type != unreachable) {
     if (curr->type != none) { // tee is ok anyhow
       shouldBeEqualOrFirstIsUnreachable(
-        curr->value->type, curr->type, curr, "local.set type must be correct");
+        curr->value->type,
+        curr->type,
+        curr,
+        "local.set's value type must be correct");
+      shouldBeTrue(curr->type == getFunction()->getLocalType(curr->index),
+                   curr,
+                   "local.set's tpye must be correct");
     }
     shouldBeEqual(getFunction()->getLocalType(curr->index),
                   curr->value->type,

src/wasm/wasm.cpp

@@ -469,24 +469,23 @@ bool FunctionType::operator==(FunctionType& b) {
 }
 bool FunctionType::operator!=(FunctionType& b) { return !(*this == b); }
 
-bool LocalSet::isTee() { return type != none; }
+bool LocalSet::isTee() const { return type != none; }
 
-void LocalSet::setTee(bool is) {
-  if (is) {
-    type = value->type;
-  } else {
-    type = none;
-  }
+// Changes to local.tee. The type of the local should be given.
+void LocalSet::makeTee(Type type_) {
+  type = type_;
+  finalize(); // type may need to be unreachable
+}
+
+// Changes to local.set.
+void LocalSet::makeSet() {
+  type = none;
   finalize(); // type may need to be unreachable
 }
 
 void LocalSet::finalize() {
   if (value->type == unreachable) {
     type = unreachable;
-  } else if (isTee()) {
-    type = value->type;
-  } else {
-    type = none;
   }
 }
 

test/binaryen.js/kitchen-sink.js

@@ -469,7 +469,7 @@ function test_core() {
     ),
     module.drop(module.local.get(0, Binaryen.i32)),
     module.local.set(0, makeInt32(101)),
-    module.drop(module.local.tee(0, makeInt32(102))),
+    module.drop(module.local.tee(0, makeInt32(102), Binaryen.i32)),
     module.i32.load(0, 0, makeInt32(1)),
     module.i64.load16_s(2, 1, makeInt32(8)),
     module.f32.load(0, 0, makeInt32(2)),

test/example/c-api-kitchen-sink.c

@@ -642,7 +642,9 @@ void test_core() {
         module, makeInt32(module, 2449), callOperands4b, 4, "iiIfF")),
     BinaryenDrop(module, BinaryenLocalGet(module, 0, BinaryenTypeInt32())),
     BinaryenLocalSet(module, 0, makeInt32(module, 101)),
-    BinaryenDrop(module, BinaryenLocalTee(module, 0, makeInt32(module, 102))),
+    BinaryenDrop(
+      module,
+      BinaryenLocalTee(module, 0, makeInt32(module, 102), BinaryenTypeInt32())),
     BinaryenLoad(module, 4, 0, 0, 0, BinaryenTypeInt32(), makeInt32(module, 1)),
     BinaryenLoad(module, 2, 1, 2, 1, BinaryenTypeInt64(), makeInt32(module, 8)),
     BinaryenLoad(

test/example/c-api-kitchen-sink.txt

@@ -3835,7 +3835,7 @@ int main() {
   expressions[747] = BinaryenConst(the_module, BinaryenLiteralInt32(101));
   expressions[748] = BinaryenLocalSet(the_module, 0, expressions[747]);
   expressions[749] = BinaryenConst(the_module, BinaryenLiteralInt32(102));
-  expressions[750] = BinaryenLocalTee(the_module, 0, expressions[749]);
+  expressions[750] = BinaryenLocalTee(the_module, 0, expressions[749], 2);
   expressions[751] = BinaryenDrop(the_module, expressions[750]);
   expressions[752] = BinaryenConst(the_module, BinaryenLiteralInt32(1));
   expressions[753] = BinaryenLoad(the_module, 4, 0, 0, 0, 2, expressions[752]);