Group reference types in binary format. (#4774)

Grouping all references together makes it easier for baseline compilers to
zero out memory (as the zeroing out may be different for MVP types vs.
references).

This puts all references together, either at the start or the end. As a
heuristic for that we see if the first local is a reference. As the optimizer
will sort locals by frequency, this ensures that the most-frequent local
stays in index 0.

Fixes #4773. See more details there

5 files changed, 182 insertions, 7 deletions

diff --git a/src/wasm/wasm-stack.cpp b/src/wasm/wasm-stack.cpp
index e3740ecee..9c086bb47 100644
--- a/src/wasm/wasm-stack.cpp
+++ b/src/wasm/wasm-stack.cpp
@@ -2327,6 +2327,29 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() {
     }
   }
   countScratchLocals();
+
+  if (parent.getModule()->features.hasReferenceTypes()) {
+    // Sort local types in a way that keeps all MVP types together and all
+    // reference types together. E.g. it is helpful to avoid a block of i32s in
+    // between blocks of different reference types, since clearing out reference
+    // types may require different work.
+    //
+    // See https://github.com/WebAssembly/binaryen/issues/4773
+    //
+    // In order to decide whether to put MVP types or reference types first,
+    // look at the type of the first local. In an optimized binary we will have
+    // sorted the locals by frequency of uses, so this way we'll keep the most
+    // commonly-used local at the top, which should work well in many cases.
+    bool refsFirst = !localTypes.empty() && localTypes[0].isRef();
+    std::stable_sort(localTypes.begin(), localTypes.end(), [&](Type a, Type b) {
+      if (refsFirst) {
+        return a.isRef() && !b.isRef();
+      } else {
+        return !a.isRef() && b.isRef();
+      }
+    });
+  }
+
   std::unordered_map<Type, size_t> currLocalsByType;
   for (Index i = func->getVarIndexBase(); i < func->getNumLocals(); i++) {
     Index j = 0;
diff --git a/test/typed-function-references.wast b/test/typed-function-references.wast
index 1057d4487..40eae6ebd 100644
--- a/test/typed-function-references.wast
+++ b/test/typed-function-references.wast
@@ -41,4 +41,45 @@
       )
     )
   )
+  (func $ref-types-first
+    ;; 6 reference types and 3 MVP types. The binary format should emit all the
+    ;; reference types first since a reference type appears first. In addition,
+    ;; types should be emitted in blocks there, that is, locals of identical
+    ;; types should be adjacent.
+    (local $r1 (ref null $mixed_results))
+    (local $r2 (ref null $mixed_results))
+    (local $i1 i32)
+    (local $r3 anyref)
+    (local $i2 i64)
+    (local $r4 anyref)
+    (local $i3 i64)
+    (local $r5 anyref)
+    (local $r6 funcref)
+  )
+  (func $mvp-types-first
+    ;; Reversed from before, now an MVP type appears first, so they should all
+    ;; be before reference types in the binary format.
+    (local $i1 i32) ;; only this local was moved up.
+    (local $r1 (ref null $mixed_results))
+    (local $r2 (ref null $mixed_results))
+    (local $r3 anyref)
+    (local $i2 i64)
+    (local $r4 anyref)
+    (local $i3 i64)
+    (local $r5 anyref)
+    (local $r6 funcref)
+  )
+  (func $mvp-types-first-param (param $r0 (ref null $mixed_results))
+    ;; As before, but now there is a reference type *parameter*. We should
+    ;; ignore that and sort as in the last function.
+    (local $i1 i32) ;; only this local was moved up.
+    (local $r1 (ref null $mixed_results))
+    (local $r2 (ref null $mixed_results))
+    (local $r3 anyref)
+    (local $i2 i64)
+    (local $r4 anyref)
+    (local $i3 i64)
+    (local $r5 anyref)
+    (local $r6 funcref)
+  )
 )
diff --git a/test/typed-function-references.wast.from-wast b/test/typed-function-references.wast.from-wast
index cb9dd5671..a695e826d 100644
--- a/test/typed-function-references.wast.from-wast
+++ b/test/typed-function-references.wast.from-wast
@@ -1,4 +1,5 @@
 (module
+ (type $mixed_results (func (result anyref f32 anyref f32)))
  (type $none_=>_none (func))
  (type $i32-i32 (func (param i32) (result i32)))
  (type $=>eqref (func (result eqref)))
@@ -8,7 +9,7 @@
  (type $f64_=>_ref_null<_->_eqref> (func (param f64) (result (ref null $=>eqref))))
  (type $=>anyref (func (result anyref)))
  (type $none_=>_i32_ref?|$mixed_results|_f64 (func (result i32 (ref null $mixed_results) f64)))
- (type $mixed_results (func (result anyref f32 anyref f32)))
+ (type $ref?|$mixed_results|_=>_none (func (param (ref null $mixed_results))))
  (elem declare func $call-ref $call-ref-more)
  (func $call-ref
   (call_ref
@@ -62,4 +63,40 @@
    )
   )
  )
+ (func $ref-types-first
+  (local $r1 (ref null $mixed_results))
+  (local $r2 (ref null $mixed_results))
+  (local $i1 i32)
+  (local $r3 anyref)
+  (local $i2 i64)
+  (local $r4 anyref)
+  (local $i3 i64)
+  (local $r5 anyref)
+  (local $r6 funcref)
+  (nop)
+ )
+ (func $mvp-types-first
+  (local $i1 i32)
+  (local $r1 (ref null $mixed_results))
+  (local $r2 (ref null $mixed_results))
+  (local $r3 anyref)
+  (local $i2 i64)
+  (local $r4 anyref)
+  (local $i3 i64)
+  (local $r5 anyref)
+  (local $r6 funcref)
+  (nop)
+ )
+ (func $mvp-types-first-param (param $r0 (ref null $mixed_results))
+  (local $i1 i32)
+  (local $r1 (ref null $mixed_results))
+  (local $r2 (ref null $mixed_results))
+  (local $r3 anyref)
+  (local $i2 i64)
+  (local $r4 anyref)
+  (local $i3 i64)
+  (local $r5 anyref)
+  (local $r6 funcref)
+  (nop)
+ )
 )
diff --git a/test/typed-function-references.wast.fromBinary b/test/typed-function-references.wast.fromBinary
index e3839f33d..fefe37405 100644
--- a/test/typed-function-references.wast.fromBinary
+++ b/test/typed-function-references.wast.fromBinary
@@ -1,7 +1,7 @@
 (module
+ (type $mixed_results (func (result anyref f32 anyref f32)))
  (type $none_=>_none (func))
  (type $i32-i32 (func (param i32) (result i32)))
- (type $mixed_results (func (result anyref f32 anyref f32)))
  (type $=>eqref (func (result eqref)))
  (type $ref|$i32-i32|_=>_i32 (func (param (ref $i32-i32)) (result i32)))
  (type $ref?|$i32-i32|_=>_i32 (func (param (ref null $i32-i32)) (result i32)))
@@ -9,6 +9,7 @@
  (type $f64_=>_ref_null<_->_eqref> (func (param f64) (result (ref null $=>eqref))))
  (type $=>anyref (func (result anyref)))
  (type $none_=>_i32_ref?|$mixed_results|_f64 (func (result i32 (ref null $mixed_results) f64)))
+ (type $ref?|$mixed_results|_=>_none (func (param (ref null $mixed_results))))
  (elem declare func $call-ref $call-ref-more)
  (func $call-ref
   (call_ref
@@ -53,8 +54,8 @@
  )
  (func $type-only-in-tuple-local
   (local $x i32)
-  (local $1 (ref null $=>anyref))
-  (local $2 f64)
+  (local $1 f64)
+  (local $2 (ref null $=>anyref))
   (nop)
  )
  (func $type-only-in-tuple-block
@@ -92,5 +93,41 @@
    )
   )
  )
+ (func $ref-types-first
+  (local $r1 (ref null $mixed_results))
+  (local $r2 (ref null $mixed_results))
+  (local $r3 anyref)
+  (local $r4 anyref)
+  (local $r5 anyref)
+  (local $r6 funcref)
+  (local $i1 i32)
+  (local $i2 i64)
+  (local $i3 i64)
+  (nop)
+ )
+ (func $mvp-types-first
+  (local $i1 i32)
+  (local $i2 i64)
+  (local $i3 i64)
+  (local $r1 (ref null $mixed_results))
+  (local $r2 (ref null $mixed_results))
+  (local $r3 anyref)
+  (local $r4 anyref)
+  (local $r5 anyref)
+  (local $r6 funcref)
+  (nop)
+ )
+ (func $mvp-types-first-param (param $r0 (ref null $mixed_results))
+  (local $i1 i32)
+  (local $i2 i64)
+  (local $i3 i64)
+  (local $r1 (ref null $mixed_results))
+  (local $r2 (ref null $mixed_results))
+  (local $r3 anyref)
+  (local $r4 anyref)
+  (local $r5 anyref)
+  (local $r6 funcref)
+  (nop)
+ )
 )
 
diff --git a/test/typed-function-references.wast.fromBinary.noDebugInfo b/test/typed-function-references.wast.fromBinary.noDebugInfo
index cf462ffba..9e68efcc1 100644
--- a/test/typed-function-references.wast.fromBinary.noDebugInfo
+++ b/test/typed-function-references.wast.fromBinary.noDebugInfo
@@ -1,7 +1,7 @@
 (module
+ (type $none_=>_anyref_f32_anyref_f32 (func (result anyref f32 anyref f32)))
  (type $none_=>_none (func))
  (type $i32_=>_i32 (func (param i32) (result i32)))
- (type $none_=>_anyref_f32_anyref_f32 (func (result anyref f32 anyref f32)))
  (type $none_=>_eqref (func (result eqref)))
  (type $ref|i32_->_i32|_=>_i32 (func (param (ref $i32_=>_i32)) (result i32)))
  (type $ref?|i32_->_i32|_=>_i32 (func (param (ref null $i32_=>_i32)) (result i32)))
@@ -9,6 +9,7 @@
  (type $f64_=>_ref?|none_->_eqref| (func (param f64) (result (ref null $none_=>_eqref))))
  (type $none_=>_anyref (func (result anyref)))
  (type $none_=>_i32_ref?|none_->_anyref_f32_anyref_f32|_f64 (func (result i32 (ref null $none_=>_anyref_f32_anyref_f32) f64)))
+ (type $ref?|none_->_anyref_f32_anyref_f32|_=>_none (func (param (ref null $none_=>_anyref_f32_anyref_f32))))
  (elem declare func $0 $2)
  (func $0
   (call_ref
@@ -53,8 +54,8 @@
  )
  (func $7
   (local $0 i32)
-  (local $1 (ref null $none_=>_anyref))
-  (local $2 f64)
+  (local $1 f64)
+  (local $2 (ref null $none_=>_anyref))
   (nop)
  )
  (func $8
@@ -92,5 +93,41 @@
    )
   )
  )
+ (func $9
+  (local $0 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $1 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $2 anyref)
+  (local $3 anyref)
+  (local $4 anyref)
+  (local $5 funcref)
+  (local $6 i32)
+  (local $7 i64)
+  (local $8 i64)
+  (nop)
+ )
+ (func $10
+  (local $0 i32)
+  (local $1 i64)
+  (local $2 i64)
+  (local $3 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $4 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $5 anyref)
+  (local $6 anyref)
+  (local $7 anyref)
+  (local $8 funcref)
+  (nop)
+ )
+ (func $11 (param $0 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $1 i32)
+  (local $2 i64)
+  (local $3 i64)
+  (local $4 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $5 (ref null $none_=>_anyref_f32_anyref_f32))
+  (local $6 anyref)
+  (local $7 anyref)
+  (local $8 anyref)
+  (local $9 funcref)
+  (nop)
+ )
 )