wasm2js: Support arbitrary temp variable types (#6661)

Previously only basic types were allowed.

Generalizing this to arbitrary types means we use a map instead of a vector,
which is slower, but I can't measure any noticeable difference. Temp vars are
pretty rare, and there are just much slower parts of wasm2js, I think.

4 files changed, 59 insertions, 22 deletions

diff --git a/src/wasm2js.h b/src/wasm2js.h
index 34260547d..321734688 100644
--- a/src/wasm2js.h
+++ b/src/wasm2js.h
@@ -206,12 +206,13 @@ public:
   // Get a temp var.
   IString getTemp(Type type, Function* func) {
     IString ret;
-    TODO_SINGLE_COMPOUND(type);
-    if (frees[type.getBasic()].size() > 0) {
-      ret = frees[type.getBasic()].back();
-      frees[type.getBasic()].pop_back();
+    // TODO: handle tuples
+    assert(!type.isTuple() && "Unexpected tuple type");
+    if (frees[type].size() > 0) {
+      ret = frees[type].back();
+      frees[type].pop_back();
     } else {
-      size_t index = temps[type.getBasic()]++;
+      auto index = temps[type]++;
       ret = IString((std::string("wasm2js_") + type.toString() + "$" +
                      std::to_string(index))
                       .c_str(),
@@ -225,8 +226,9 @@ public:
 
   // Free a temp var.
   void freeTemp(Type type, IString temp) {
-    TODO_SINGLE_COMPOUND(type);
-    frees[type.getBasic()].push_back(temp);
+    // TODO: handle tuples
+    assert(!type.isTuple() && "Unexpected tuple type");
+    frees[type].push_back(temp);
   }
 
   // Generates a mangled name from `name` within the specified scope.
@@ -300,10 +302,10 @@ private:
   Flags flags;
   PassOptions options;
 
-  // How many temp vars we need
-  std::vector<size_t> temps; // type => num temps
-  // Which are currently free to use
-  std::vector<std::vector<IString>> frees; // type => list of free names
+  // How many temp vars we need for each type (type => num).
+  std::unordered_map<Type, Index> temps;
+  // Which temp vars are currently free to use for each type (type => freelist).
+  std::unordered_map<Type, std::vector<IString>> frees;
 
   // Mangled names cache by interned names.
   // Utilizes the usually reused underlying cstring's pointer as the key.
@@ -874,15 +876,15 @@ Ref Wasm2JSBuilder::processFunction(Module* m,
     runner.runOnFunction(func);
   }
 
+  // We process multiple functions from a single Wasm2JSBuilder instance, so
+  // clean up the function-specific local state before each function.
+  frees.clear();
+  temps.clear();
+
   // We will be symbolically referring to all variables in the function, so make
   // sure that everything has a name and it's unique.
   Names::ensureNames(func);
   Ref ret = ValueBuilder::makeFunction(fromName(func->name, NameScope::Top));
-  frees.clear();
-  frees.resize(std::max(Type::i32, std::max(Type::f32, Type::f64)) + 1);
-  temps.clear();
-  temps.resize(std::max(Type::i32, std::max(Type::f32, Type::f64)) + 1);
-  temps[Type::i32] = temps[Type::f32] = temps[Type::f64] = 0;
   // arguments
   bool needCoercions = options.optimizeLevel == 0 || standaloneFunction ||
                        functionsCallableFromOutside.count(func->name);
@@ -915,10 +917,6 @@ Ref Wasm2JSBuilder::processFunction(Module* m,
   if (theVar[1]->size() == 0) {
     ret[3]->splice(theVarIndex, 1);
   }
-  // checks: all temp vars should be free at the end
-  assert(frees[Type::i32].size() == temps[Type::i32]);
-  assert(frees[Type::f32].size() == temps[Type::f32]);
-  assert(frees[Type::f64].size() == temps[Type::f64]);
   return ret;
 }
 
diff --git a/test/wasm2js/refs.2asm.js b/test/wasm2js/refs.2asm.js
index f1d60af6c..f4c08408f 100644
--- a/test/wasm2js/refs.2asm.js
+++ b/test/wasm2js/refs.2asm.js
@@ -50,6 +50,17 @@ function asmFunc(imports) {
   return temp;
  }
  
+ function funcref_temps($0, $1) {
+  $1 = +$1;
+  var $2 = null, $3 = null, wasm2js_funcref$0 = null, wasm2js_funcref$1 = null, wasm2js_i32$0 = 0;
+  $2 = $0;
+  loop : while (1) {
+   $3 = funcref_temps;
+   break loop;
+  };
+  funcref_temps(funcref_temps, +(+((wasm2js_funcref$0 = $2, wasm2js_funcref$1 = $3 || wasm2js_trap(), wasm2js_i32$0 = 0, wasm2js_i32$0 ? wasm2js_funcref$0 : wasm2js_funcref$1) == null | 0)));
+ }
+ 
  return {
   "null_": null_, 
   "is_null": is_null, 
@@ -57,7 +68,8 @@ function asmFunc(imports) {
   "ref_eq": ref_eq, 
   "ref_as": ref_as, 
   "use_global": use_global, 
-  "use_global_ref": use_global_ref
+  "use_global_ref": use_global_ref, 
+  "funcref_temps": funcref_temps
  };
 }
 
@@ -70,3 +82,4 @@ export var ref_eq = retasmFunc.ref_eq;
 export var ref_as = retasmFunc.ref_as;
 export var use_global = retasmFunc.use_global;
 export var use_global_ref = retasmFunc.use_global_ref;
+export var funcref_temps = retasmFunc.funcref_temps;
diff --git a/test/wasm2js/refs.2asm.js.opt b/test/wasm2js/refs.2asm.js.opt
index fa1947215..0071a15b5 100644
--- a/test/wasm2js/refs.2asm.js.opt
+++ b/test/wasm2js/refs.2asm.js.opt
@@ -48,6 +48,11 @@ function asmFunc(imports) {
   return $1;
  }
  
+ function funcref_temps($0, $1) {
+  $1 = +$1;
+  funcref_temps(funcref_temps, 0.0);
+ }
+ 
  return {
   "null_": null_, 
   "is_null": is_null, 
@@ -55,7 +60,8 @@ function asmFunc(imports) {
   "ref_eq": ref_eq, 
   "ref_as": ref_as, 
   "use_global": use_global, 
-  "use_global_ref": use_global_ref
+  "use_global_ref": use_global_ref, 
+  "funcref_temps": funcref_temps
  };
 }
 
@@ -68,3 +74,4 @@ export var ref_eq = retasmFunc.ref_eq;
 export var ref_as = retasmFunc.ref_as;
 export var use_global = retasmFunc.use_global;
 export var use_global_ref = retasmFunc.use_global_ref;
+export var funcref_temps = retasmFunc.funcref_temps;
diff --git a/test/wasm2js/refs.wast b/test/wasm2js/refs.wast
index 244cf2e97..087567e8c 100644
--- a/test/wasm2js/refs.wast
+++ b/test/wasm2js/refs.wast
@@ -62,4 +62,23 @@
     )
     (local.get $temp)
   )
+
+  (func $funcref_temps (export "funcref_temps") (param $0 funcref) (param $1 f64)
+    ;; A deeply-nested expression that ends up requiring multiple function type
+    ;; temp variables.
+    (call $funcref_temps
+      (ref.func $funcref_temps)
+      (f64.convert_i32_s
+        (ref.is_null
+          (select (result funcref)
+            (local.get $0)
+            (loop $loop (result funcref)
+              (ref.func $funcref_temps)
+            )
+            (i32.const 0)
+          )
+        )
+      )
+    )
+  )
 )