wasm2js: optimize away casts going into a suitable store (#2069)

4 files changed, 128 insertions, 15 deletions

diff --git a/src/tools/wasm2js.cpp b/src/tools/wasm2js.cpp
index 3d5f484da..6e4656878 100644
--- a/src/tools/wasm2js.cpp
+++ b/src/tools/wasm2js.cpp
@@ -129,7 +129,7 @@ static void traversePost(Ref node, std::function<void(Ref)> visit) {
 }
 
 static void optimizeJS(Ref ast) {
-  // helpers
+  // Helpers
 
   auto isOrZero = [](Ref node) {
     return node->isArray() && !node->empty() && node[0] == BINARY &&
@@ -141,6 +141,11 @@ static void optimizeJS(Ref ast) {
            node[1] == PLUS;
   };
 
+  auto isFround = [](Ref node) {
+    return node->isArray() && !node->empty() && node[0] == cashew::CALL &&
+           node[1] == MATH_FROUND;
+  };
+
   auto isBitwise = [](Ref node) {
     if (node->isArray() && !node->empty() && node[0] == BINARY) {
       auto op = node[1];
@@ -150,6 +155,60 @@ static void optimizeJS(Ref ast) {
     return false;
   };
 
+  auto isConstantAnd = [](Ref node, int num) {
+    return node->isArray() && !node->empty() && node[0] == BINARY &&
+           node[1] == AND && node[3]->isNumber() && node[3]->getNumber() == num;
+  };
+
+  auto removeOrZero = [&](Ref node) {
+    while (isOrZero(node)) {
+      node = node[2];
+    }
+    return node;
+  };
+
+  auto removePlus = [&](Ref node) {
+    while (isPlus(node)) {
+      node = node[2];
+    }
+    return node;
+  };
+
+  auto removePlusAndFround = [&](Ref node) {
+    while (1) {
+      if (isFround(node)) {
+        node = node[2][0];
+      } else if (isPlus(node)) {
+        node = node[2];
+      } else {
+        break;
+      }
+    }
+    return node;
+  };
+
+  auto getHeapFromAccess = [](Ref node) { return node[1]->getIString(); };
+
+  auto isIntegerHeap = [](IString heap) {
+    return heap == HEAP8 || heap == HEAPU8 || heap == HEAP16 ||
+           heap == HEAPU16 || heap == HEAP32 || heap == HEAPU32;
+  };
+
+  auto isFloatHeap = [](IString heap) {
+    return heap == HEAPF32 || heap == HEAPF64;
+  };
+
+  auto isHeapAccess = [&](Ref node) {
+    if (node->isArray() && !node->empty() && node[0] == SUB &&
+        node[1]->isString()) {
+      auto heap = getHeapFromAccess(node);
+      return isIntegerHeap(heap) || isFloatHeap(heap);
+    }
+    return false;
+  };
+
+  // Optimizations
+
   // x >> 0  =>  x | 0
   traversePost(ast, [](Ref node) {
     if (node->isArray() && !node->empty() && node[0] == BINARY &&
@@ -175,17 +234,46 @@ static void optimizeJS(Ref ast) {
     }
     // x | 0 going into a bitwise op => skip the | 0
     else if (isBitwise(node)) {
-      while (isOrZero(node[2])) {
-        node[2] = node[2][2];
-      }
-      while (isOrZero(node[3])) {
-        node[3] = node[3][2];
-      }
+      node[2] = removeOrZero(node[2]);
+      node[3] = removeOrZero(node[3]);
     }
     // +(+x) => +x
     else if (isPlus(node)) {
-      while (isPlus(node[2])) {
-        node[2] = node[2][2];
+      node[2] = removePlus(node[2]);
+    }
+    // +(+x) => +x
+    else if (isFround(node)) {
+      node[2] = removePlusAndFround(node[2]);
+    }
+    // Assignment into a heap coerces.
+    else if (node->isAssign()) {
+      auto assign = node->asAssign();
+      auto target = assign->target();
+      if (isHeapAccess(target)) {
+        auto heap = getHeapFromAccess(target);
+        if (isIntegerHeap(heap)) {
+          if (heap == HEAP8 || heap == HEAPU8) {
+            while (isOrZero(assign->value()) ||
+                   isConstantAnd(assign->value(), 255)) {
+              assign->value() = assign->value()[2];
+            }
+          } else if (heap == HEAP16 || heap == HEAPU16) {
+            while (isOrZero(assign->value()) ||
+                   isConstantAnd(assign->value(), 65535)) {
+              assign->value() = assign->value()[2];
+            }
+          } else {
+            assert(heap == HEAP32 || heap == HEAPU32);
+            assign->value() = removeOrZero(assign->value());
+          }
+        } else {
+          assert(isFloatHeap(heap));
+          if (heap == HEAPF32) {
+            assign->value() = removePlusAndFround(assign->value());
+          } else {
+            assign->value() = removePlus(assign->value());
+          }
+        }
       }
     }
   });
diff --git a/test/wasm2js/emscripten.2asm.js b/test/wasm2js/emscripten.2asm.js
index 0b1d7ee67..52828f46e 100644
--- a/test/wasm2js/emscripten.2asm.js
+++ b/test/wasm2js/emscripten.2asm.js
@@ -48,6 +48,11 @@ function asmFunc(global, env, buffer) {
   HEAP8[128 | 0];
   HEAPU16[128 >> 1];
   HEAP16[128 >> 1];
+  HEAP32[16 >> 2] = 1 + 2;
+  HEAPF32[16 >> 2] = Math_fround(3.0) + Math_fround(4.0);
+  HEAPF64[16 >> 3] = 5.0 + 6.0;
+  HEAP8[16 | 0] = 7 + 8;
+  HEAP16[16 >> 1] = 9 + 10;
  }
  
  function __growWasmMemory($0) {
diff --git a/test/wasm2js/emscripten.wast b/test/wasm2js/emscripten.wast
index 7305f0dc9..8f797f8fc 100644
--- a/test/wasm2js/emscripten.wast
+++ b/test/wasm2js/emscripten.wast
@@ -52,6 +52,26 @@
     (i32.const 128)
    )
   )
+  (i32.store
+   (i32.const 16)
+   (i32.add (i32.const 1) (i32.const 2))
+  )
+  (f32.store
+   (i32.const 16)
+   (f32.add (f32.const 3) (f32.const 4))
+  )
+  (f64.store
+   (i32.const 16)
+   (f64.add (f64.const 5) (f64.const 6))
+  )
+  (i32.store8
+   (i32.const 16)
+   (i32.add (i32.const 7) (i32.const 8))
+  )
+  (i32.store16
+   (i32.const 16)
+   (i32.add (i32.const 9) (i32.const 10))
+  )
  )
  (func $__growWasmMemory (param $0 i32) (result i32)
   (grow_memory
diff --git a/test/wasm2js/unaligned.2asm.js b/test/wasm2js/unaligned.2asm.js
index 51e4bf269..6996fd1b9 100644
--- a/test/wasm2js/unaligned.2asm.js
+++ b/test/wasm2js/unaligned.2asm.js
@@ -84,19 +84,19 @@ function asmFunc(global, env, buffer) {
  
  function $4() {
   var wasm2js_i32$0 = 0, wasm2js_i32$1 = 0;
-  (wasm2js_i32$0 = 0, wasm2js_i32$1 = 0), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1 & 255, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8 & 255), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16 & 255), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24 & 255;
+  (wasm2js_i32$0 = 0, wasm2js_i32$1 = 0), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24;
  }
  
  function $5() {
   var i64toi32_i32$1 = 0, wasm2js_i32$0 = 0, wasm2js_i32$1 = 0;
   i64toi32_i32$1 = 0;
-  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = 0), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1 & 255, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8 & 255), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16 & 255), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24 & 255;
-  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = 0), ((HEAP8[wasm2js_i32$0 + 4 | 0] = wasm2js_i32$1 & 255, HEAP8[wasm2js_i32$0 + 5 | 0] = wasm2js_i32$1 >>> 8 & 255), HEAP8[wasm2js_i32$0 + 6 | 0] = wasm2js_i32$1 >>> 16 & 255), HEAP8[wasm2js_i32$0 + 7 | 0] = wasm2js_i32$1 >>> 24 & 255;
+  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = 0), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24;
+  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = 0), ((HEAP8[wasm2js_i32$0 + 4 | 0] = wasm2js_i32$1, HEAP8[wasm2js_i32$0 + 5 | 0] = wasm2js_i32$1 >>> 8), HEAP8[wasm2js_i32$0 + 6 | 0] = wasm2js_i32$1 >>> 16), HEAP8[wasm2js_i32$0 + 7 | 0] = wasm2js_i32$1 >>> 24;
  }
  
  function $6() {
   var wasm2js_i32$0 = 0, wasm2js_i32$1 = 0;
-  (wasm2js_i32$0 = 0, wasm2js_i32$1 = (wasm2js_scratch_store_f32(Math_fround(0.0)), wasm2js_scratch_load_i32(0))), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1 & 255, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8 & 255), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16 & 255), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24 & 255;
+  (wasm2js_i32$0 = 0, wasm2js_i32$1 = (wasm2js_scratch_store_f32(Math_fround(0.0)), wasm2js_scratch_load_i32(0))), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24;
  }
  
  function $7() {
@@ -104,8 +104,8 @@ function asmFunc(global, env, buffer) {
   wasm2js_scratch_store_f64(0.0);
   i64toi32_i32$0 = wasm2js_scratch_load_i32(1 | 0) | 0;
   i64toi32_i32$1 = 0;
-  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = wasm2js_scratch_load_i32(0 | 0) | 0), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1 & 255, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8 & 255), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16 & 255), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24 & 255;
-  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = i64toi32_i32$0), ((HEAP8[wasm2js_i32$0 + 4 | 0] = wasm2js_i32$1 & 255, HEAP8[wasm2js_i32$0 + 5 | 0] = wasm2js_i32$1 >>> 8 & 255), HEAP8[wasm2js_i32$0 + 6 | 0] = wasm2js_i32$1 >>> 16 & 255), HEAP8[wasm2js_i32$0 + 7 | 0] = wasm2js_i32$1 >>> 24 & 255;
+  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = wasm2js_scratch_load_i32(0 | 0) | 0), ((HEAP8[wasm2js_i32$0 | 0] = wasm2js_i32$1, HEAP8[wasm2js_i32$0 + 1 | 0] = wasm2js_i32$1 >>> 8), HEAP8[wasm2js_i32$0 + 2 | 0] = wasm2js_i32$1 >>> 16), HEAP8[wasm2js_i32$0 + 3 | 0] = wasm2js_i32$1 >>> 24;
+  (wasm2js_i32$0 = i64toi32_i32$1, wasm2js_i32$1 = i64toi32_i32$0), ((HEAP8[wasm2js_i32$0 + 4 | 0] = wasm2js_i32$1, HEAP8[wasm2js_i32$0 + 5 | 0] = wasm2js_i32$1 >>> 8), HEAP8[wasm2js_i32$0 + 6 | 0] = wasm2js_i32$1 >>> 16), HEAP8[wasm2js_i32$0 + 7 | 0] = wasm2js_i32$1 >>> 24;
  }
  
  function legalstub$1() {