wasm2js: more js optimization (#2050)

* Emit ints as signed, so -1 isn't a big unsigned number.
* x - -c (where c is a constant) is larger than x + c in js (but not wasm)
* +(+x) => +x
* Avoid unnecessary coercions on calls, return, load, etc. - we just need coercions when entering or exiting "wasm" (not internally), and on actual operations that need them.

4 files changed, 119 insertions, 34 deletions

diff --git a/src/binaryen-c.cpp b/src/binaryen-c.cpp
index aa94d1d30..87dd631d3 100644
--- a/src/binaryen-c.cpp
+++ b/src/binaryen-c.cpp
@@ -2469,7 +2469,7 @@ void BinaryenModulePrintAsmjs(BinaryenModuleRef module) {
 
   Module* wasm = (Module*)module;
   Wasm2JSBuilder::Flags flags;
-  Wasm2JSBuilder wasm2js(flags);
+  Wasm2JSBuilder wasm2js(flags, globalPassOptions);
   Ref asmjs = wasm2js.processWasm(wasm);
   JSPrinter jser(true, true, asmjs);
   Output out("", Flags::Text, Flags::Release); // stdout
diff --git a/src/emscripten-optimizer/simple_ast.h b/src/emscripten-optimizer/simple_ast.h
index 3d8508a15..30ef42711 100644
--- a/src/emscripten-optimizer/simple_ast.h
+++ b/src/emscripten-optimizer/simple_ast.h
@@ -1437,6 +1437,9 @@ public:
   static Ref makeInt(uint32_t num) {
     return makeDouble(double(num));
   }
+  static Ref makeInt(int32_t num) {
+    return makeDouble(double(num));
+  }
   static Ref makeNum(double num) {
     return makeDouble(num);
   }
diff --git a/src/tools/wasm2js.cpp b/src/tools/wasm2js.cpp
index f55f2586b..34ed2d9cb 100644
--- a/src/tools/wasm2js.cpp
+++ b/src/tools/wasm2js.cpp
@@ -24,6 +24,7 @@
 #include "wasm-s-parser.h"
 #include "wasm2js.h"
 #include "optimization-options.h"
+#include "pass.h"
 
 using namespace cashew;
 using namespace wasm;
@@ -32,6 +33,33 @@ using namespace wasm;
 
 namespace {
 
+static void optimizeWasm(Module& wasm, PassOptions options) {
+  // Perform various optimizations that will be good for JS, but would not be great
+  // for wasm in general
+  struct OptimizeForJS : public WalkerPass<PostWalker<OptimizeForJS>> {
+    bool isFunctionParallel() override { return true; }
+
+    Pass* create() override { return new OptimizeForJS; }
+
+    void visitBinary(Binary* curr) {
+      // x - -c (where c is a constant) is larger than x + c, in js (but not
+      // necessarily in wasm, where LEBs prefer negatives).
+      if (curr->op == SubInt32) {
+        if (auto* c = curr->right->dynCast<Const>()) {
+          if (c->value.geti32() < 0) {
+            curr->op = AddInt32;
+            c->value = c->value.neg();
+          }
+        }
+      }
+    }
+  };
+
+  PassRunner runner(&wasm, options);
+  runner.add<OptimizeForJS>();
+  runner.run();
+}
+
 template<typename T>
 static void printJS(Ref ast, T& output) {
   JSPrinter jser(true, true, ast);
@@ -41,10 +69,15 @@ static void printJS(Ref ast, T& output) {
 
 static void optimizeJS(Ref ast) {
   // helpers
+
   auto isOrZero = [](Ref node) {
     return node->isArray() && node->size() > 0 && node[0] == BINARY && node[1] == OR && node[3]->isNumber() && node[3]->getNumber() == 0;
   };
 
+  auto isPlus = [](Ref node) {
+    return node->isArray() && node->size() > 0 && node[0] == UNARY_PREFIX && node[1] == PLUS;
+  };
+
   auto isBitwise = [](Ref node) {
     if (node->isArray() && node->size() > 0 && node[0] == BINARY) {
       auto op = node[1];
@@ -84,13 +117,22 @@ static void optimizeJS(Ref ast) {
         node[3] = node[3][2];
       }
     }
+    // +(+x) => +x
+    else if (isPlus(node)) {
+      while (isPlus(node[2])) {
+        node[2] = node[2][2];
+      }
+    }
   });
 }
 
-static void emitWasm(Module& wasm, Output& output, Wasm2JSBuilder::Flags flags, Name name, bool optimize=false) {
-  Wasm2JSBuilder wasm2js(flags);
+static void emitWasm(Module& wasm, Output& output, Wasm2JSBuilder::Flags flags, PassOptions options, Name name) {
+  if (options.optimizeLevel > 0) {
+    optimizeWasm(wasm, options);
+  }
+  Wasm2JSBuilder wasm2js(flags, options);
   auto js = wasm2js.processWasm(&wasm, name);
-  if (optimize) {
+  if (options.optimizeLevel >= 2) {
     optimizeJS(js);
   }
   Wasm2JSGlue glue(wasm, output, flags, name);
@@ -104,7 +146,8 @@ public:
   AssertionEmitter(Element& root,
                    SExpressionWasmBuilder& sexpBuilder,
                    Output& out,
-                   Wasm2JSBuilder::Flags flags) : root(root), sexpBuilder(sexpBuilder), out(out), flags(flags) {}
+                   Wasm2JSBuilder::Flags flags,
+                   PassOptions options) : root(root), sexpBuilder(sexpBuilder), out(out), flags(flags), options(options) {}
 
   void emit();
 
@@ -113,6 +156,7 @@ private:
   SExpressionWasmBuilder& sexpBuilder;
   Output& out;
   Wasm2JSBuilder::Flags flags;
+  PassOptions options;
   Module tempAllocationModule;
 
   Ref emitAssertReturnFunc(Builder& wasmBuilder,
@@ -131,7 +175,7 @@ private:
   void fixCalls(Ref asmjs, Name asmModule);
 
   Ref processFunction(Function* func) {
-    Wasm2JSBuilder sub(flags);
+    Wasm2JSBuilder sub(flags, options);
     return sub.processStandaloneFunction(&tempAllocationModule, func);
   }
 
@@ -370,7 +414,7 @@ void AssertionEmitter::emit() {
       asmModule = Name(moduleNameS.str().c_str());
       Module wasm;
       SExpressionWasmBuilder builder(wasm, e);
-      emitWasm(wasm, out, flags, funcName);
+      emitWasm(wasm, out, flags, options, funcName);
       continue;
     }
     if (!isAssertHandled(e)) {
@@ -492,9 +536,9 @@ int main(int argc, const char *argv[]) {
   if (options.debug) std::cerr << "j-printing..." << std::endl;
   Output output(options.extra["output"], Flags::Text, options.debug ? Flags::Debug : Flags::Release);
   if (!binaryInput && options.extra["asserts"] == "1") {
-    AssertionEmitter(*root, *sexprBuilder, output, flags).emit();
+    AssertionEmitter(*root, *sexprBuilder, output, flags, options.passOptions).emit();
   } else {
-    emitWasm(wasm, output, flags, "asmFunc", options.passOptions.optimizeLevel > 0);
+    emitWasm(wasm, output, flags, options.passOptions, "asmFunc");
   }
 
   if (options.debug) std::cerr << "done." << std::endl;
diff --git a/src/wasm2js.h b/src/wasm2js.h
index 7bc356bea..84c65b189 100644
--- a/src/wasm2js.h
+++ b/src/wasm2js.h
@@ -126,7 +126,7 @@ public:
     bool emscripten = false;
   };
 
-  Wasm2JSBuilder(Flags f) : flags(f) {}
+  Wasm2JSBuilder(Flags f, PassOptions options) : flags(f), options(options) {}
 
   Ref processWasm(Module* wasm, Name funcName = ASM_FUNC);
   Ref processFunction(Module* wasm, Function* func, bool standalone=false);
@@ -136,7 +136,7 @@ public:
 
   // The second pass on an expression: process it fully, generating
   // JS
-  Ref processFunctionBody(Module* m, Function* func);
+  Ref processFunctionBody(Module* m, Function* func, bool standalone);
 
   // Get a temp var.
   IString getTemp(Type type, Function* func) {
@@ -225,6 +225,7 @@ public:
 
 private:
   Flags flags;
+  PassOptions options;
 
   // How many temp vars we need
   std::vector<size_t> temps; // type => num temps
@@ -238,6 +239,11 @@ private:
 
   size_t tableSize;
 
+  // If a function is callable from outside, we'll need to cast the inputs
+  // and our return value. Otherwise, internally, casts are only needed
+  // on operations.
+  std::unordered_set<Name> functionsCallableFromOutside;
+
   void addBasics(Ref ast);
   void addFunctionImport(Ref ast, Function* import);
   void addGlobalImport(Ref ast, Global* import);
@@ -252,6 +258,18 @@ private:
 };
 
 Ref Wasm2JSBuilder::processWasm(Module* wasm, Name funcName) {
+  // Scan the wasm for important things.
+  for (auto& exp : wasm->exports) {
+    if (exp->kind == ExternalKind::Function) {
+      functionsCallableFromOutside.insert(exp->value);
+    }
+  }
+  for (auto& segment : wasm->table.segments) {
+    for (auto name : segment.data) {
+      functionsCallableFromOutside.insert(name);
+    }
+  }
+
   // Ensure the scratch memory helpers.
   // If later on they aren't needed, we'll clean them up.
   ABI::wasm2js::ensureScratchMemoryHelpers(wasm);
@@ -645,26 +663,29 @@ Ref Wasm2JSBuilder::processFunction(Module* m, Function* func, bool standaloneFu
   temps.resize(std::max(i32, std::max(f32, f64)) + 1);
   temps[i32] = temps[f32] = temps[f64] = 0;
   // arguments
+  bool needCoercions = options.optimizeLevel == 0 || standaloneFunction || functionsCallableFromOutside.count(func->name);
   for (Index i = 0; i < func->getNumParams(); i++) {
     IString name = fromName(func->getLocalNameOrGeneric(i), NameScope::Local);
     ValueBuilder::appendArgumentToFunction(ret, name);
-    ret[3]->push_back(
-      ValueBuilder::makeStatement(
-        ValueBuilder::makeBinary(
-          ValueBuilder::makeName(name), SET,
-          makeAsmCoercion(
-            ValueBuilder::makeName(name),
-            wasmToAsmType(func->getLocalType(i))
+    if (needCoercions) {
+      ret[3]->push_back(
+        ValueBuilder::makeStatement(
+          ValueBuilder::makeBinary(
+            ValueBuilder::makeName(name), SET,
+            makeAsmCoercion(
+              ValueBuilder::makeName(name),
+              wasmToAsmType(func->getLocalType(i))
+            )
           )
         )
-      )
-    );
+      );
+    }
   }
   Ref theVar = ValueBuilder::makeVar();
   size_t theVarIndex = ret[3]->size();
   ret[3]->push_back(theVar);
   // body
-  flattenAppend(ret, processFunctionBody(m, func));
+  flattenAppend(ret, processFunctionBody(m, func, standaloneFunction));
   // vars, including new temp vars
   for (Index i = func->getVarIndexBase(); i < func->getNumLocals(); i++) {
     ValueBuilder::appendToVar(
@@ -683,15 +704,17 @@ Ref Wasm2JSBuilder::processFunction(Module* m, Function* func, bool standaloneFu
   return ret;
 }
 
-Ref Wasm2JSBuilder::processFunctionBody(Module* m, Function* func) {
+Ref Wasm2JSBuilder::processFunctionBody(Module* m, Function* func, bool standaloneFunction) {
   struct ExpressionProcessor : public Visitor<ExpressionProcessor, Ref> {
     Wasm2JSBuilder* parent;
     IString result; // TODO: remove
     Function* func;
     Module* module;
+    bool standaloneFunction;
     MixedArena allocator;
-    ExpressionProcessor(Wasm2JSBuilder* parent, Module* m, Function* func)
-      : parent(parent), func(func), module(m) {}
+
+    ExpressionProcessor(Wasm2JSBuilder* parent, Module* m, Function* func, bool standaloneFunction)
+      : parent(parent), func(func), module(m), standaloneFunction(standaloneFunction) {}
 
     // A scoped temporary variable.
     struct ScopedTemp {
@@ -847,12 +870,19 @@ Ref Wasm2JSBuilder::processFunctionBody(Module* m, Function* func) {
 
     Ref visitCall(Call* curr) {
       Ref theCall = ValueBuilder::makeCall(fromName(curr->target, NameScope::Top));
+      // For wasm => wasm calls, we don't need coercions. TODO: even imports might be safe?
+      bool needCoercions = parent->options.optimizeLevel == 0 || standaloneFunction || module->getFunction(curr->target)->imported();
       for (auto operand : curr->operands) {
-        theCall[2]->push_back(
-          makeAsmCoercion(visit(operand, EXPRESSION_RESULT),
-                          wasmToAsmType(operand->type)));
+        auto value = visit(operand, EXPRESSION_RESULT);
+        if (needCoercions) {
+          value = makeAsmCoercion(value, wasmToAsmType(operand->type));
+        }
+        theCall[2]->push_back(value);
+      }
+      if (needCoercions) {
+        theCall = makeAsmCoercion(theCall, wasmToAsmType(curr->type));
       }
-      return makeAsmCoercion(theCall, wasmToAsmType(curr->type));
+      return theCall;
     }
 
     Ref visitCallIndirect(CallIndirect* curr) {
@@ -992,7 +1022,14 @@ Ref Wasm2JSBuilder::processFunctionBody(Module* m, Function* func) {
           abort();
         }
       }
-      return makeAsmCoercion(ret, wasmToAsmType(curr->type));
+      // Coercions are not actually needed, as if the user reads beyond valid memory, it's
+      // undefined behavior anyhow, and so we don't care much about slowness of undefined
+      // values etc.
+      bool needCoercions = parent->options.optimizeLevel == 0 || standaloneFunction;
+      if (needCoercions) {
+        ret = makeAsmCoercion(ret, wasmToAsmType(curr->type));
+      }
+      return ret;
     }
 
     Ref visitStore(Store* curr) {
@@ -1521,10 +1558,11 @@ Ref Wasm2JSBuilder::processFunctionBody(Module* m, Function* func) {
       if (!curr->value) {
         return ValueBuilder::makeReturn(Ref());
       }
-      Ref val = makeAsmCoercion(
-        visit(curr->value, EXPRESSION_RESULT),
-        wasmToAsmType(curr->value->type)
-      );
+      Ref val = visit(curr->value, EXPRESSION_RESULT);
+      bool needCoercion = parent->options.optimizeLevel == 0 || standaloneFunction || parent->functionsCallableFromOutside.count(func->name);
+      if (needCoercion) {
+        val = makeAsmCoercion(val, wasmToAsmType(curr->value->type));
+      }
       return ValueBuilder::makeReturn(val);
     }
 
@@ -1564,7 +1602,7 @@ Ref Wasm2JSBuilder::processFunctionBody(Module* m, Function* func) {
     }
   };
 
-  return ExpressionProcessor(this, m, func).visit(func->body, NO_RESULT);
+  return ExpressionProcessor(this, m, func, standaloneFunction).visit(func->body, NO_RESULT);
 }
 
 void Wasm2JSBuilder::addMemoryGrowthFuncs(Ref ast, Module* wasm) {