Heap2Local: Optimize packed fields (#6480)

Previously we did not optimize a struct or an array with a packed field. As a result a
single packed field in a struct prevented the entire struct from being localized,
which this fixes. This is also useful for arrays as packed arrays are common (e.g. for
string data).

3 files changed, 261 insertions, 41 deletions

diff --git a/src/passes/Heap2Local.cpp b/src/passes/Heap2Local.cpp
index a9465b1d7..3b0410a77 100644
--- a/src/passes/Heap2Local.cpp
+++ b/src/passes/Heap2Local.cpp
@@ -150,6 +150,7 @@
 // This optimization focuses on such cases.
 //
 
+#include "ir/bits.h"
 #include "ir/branch-utils.h"
 #include "ir/find_all.h"
 #include "ir/local-graph.h"
@@ -648,6 +649,21 @@ struct Struct2Local : PostWalker<Struct2Local> {
     curr->finalize();
   }
 
+  // Add a mask for packed fields. We add masks on sets rather than on gets
+  // because gets tend to be more numerous both in code appearances and in
+  // runtime execution. As a result of masking on sets, the value in the local
+  // is always the masked value (which is also nice for debugging,
+  // incidentally).
+  Expression* addMask(Expression* value, const Field& field) {
+    if (!field.isPacked()) {
+      return value;
+    }
+
+    auto mask = Bits::lowBitMask(field.getByteSize() * 8);
+    return builder.makeBinary(
+      AndInt32, value, builder.makeConst(int32_t(mask)));
+  }
+
   void visitStructNew(StructNew* curr) {
     if (curr != allocation) {
       return;
@@ -693,9 +709,10 @@ struct Struct2Local : PostWalker<Struct2Local> {
 
       // Copy them to the normal ones.
       for (Index i = 0; i < tempIndexes.size(); i++) {
-        contents.push_back(builder.makeLocalSet(
-          localIndexes[i],
-          builder.makeLocalGet(tempIndexes[i], fields[i].type)));
+        auto* value = builder.makeLocalGet(tempIndexes[i], fields[i].type);
+        // Add a mask on the values we write.
+        contents.push_back(
+          builder.makeLocalSet(localIndexes[i], addMask(value, fields[i])));
       }
 
       // TODO Check if the nondefault case does not increase code size in some
@@ -704,8 +721,11 @@ struct Struct2Local : PostWalker<Struct2Local> {
       //      defaults.
     } else {
       // Set the default values.
+      //
       // Note that we must assign the defaults because we might be in a loop,
       // that is, there might be a previous value.
+      //
+      // Note there is no need to mask as these are zeros anyhow.
       for (Index i = 0; i < localIndexes.size(); i++) {
         contents.push_back(builder.makeLocalSet(
           localIndexes[i],
@@ -766,7 +786,8 @@ struct Struct2Local : PostWalker<Struct2Local> {
     // write the data to the local instead of the heap allocation.
     replaceCurrent(builder.makeSequence(
       builder.makeDrop(curr->ref),
-      builder.makeLocalSet(localIndexes[curr->index], curr->value)));
+      builder.makeLocalSet(localIndexes[curr->index],
+                           addMask(curr->value, fields[curr->index]))));
   }
 
   void visitStructGet(StructGet* curr) {
@@ -1111,14 +1132,7 @@ struct Heap2Local {
   }
 
   bool canHandleAsLocal(const Field& field) {
-    if (!TypeUpdating::canHandleAsLocal(field.type)) {
-      return false;
-    }
-    if (field.isPacked()) {
-      // TODO: support packed fields by adding coercions/truncations.
-      return false;
-    }
-    return true;
+    return TypeUpdating::canHandleAsLocal(field.type);
   }
 
   bool canHandleAsLocals(Type type) {
diff --git a/test/lit/passes/heap2local.wast b/test/lit/passes/heap2local.wast
index 6bdfa0be8..96fef4701 100644
--- a/test/lit/passes/heap2local.wast
+++ b/test/lit/passes/heap2local.wast
@@ -19,8 +19,10 @@
 
   ;; CHECK:      (type $6 (func (result anyref)))
 
-  ;; CHECK:      (type $struct.packed (struct (field (mut i8))))
-  (type $struct.packed (struct (field (mut i8))))
+  ;; CHECK:      (type $7 (func (param i32) (result f64)))
+
+  ;; CHECK:      (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
+  (type $struct.packed (struct (field (mut i8)) (field (mut i32))))
 
   (type $struct.nondefaultable (struct (field (ref $struct.A))))
 
@@ -28,8 +30,6 @@
 
   (type $struct.nonnullable (struct (field (ref $struct.A))))
 
-  ;; CHECK:      (type $8 (func (param i32) (result f64)))
-
   ;; CHECK:      (type $9 (func (param (ref null $struct.recursive))))
 
   ;; CHECK:      (type $10 (func (param (ref $struct.A))))
@@ -175,19 +175,114 @@
   )
 
   ;; CHECK:      (func $packed (type $1)
+  ;; CHECK-NEXT:  (local $temp (ref $struct.packed))
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (local $4 i32)
+  ;; CHECK-NEXT:  (local $5 i32)
+  ;; CHECK-NEXT:  (local $6 i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (i32.const 1337)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $4
+  ;; CHECK-NEXT:     (i32.const 1338)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (i32.and
+  ;; CHECK-NEXT:      (local.get $3)
+  ;; CHECK-NEXT:      (i32.const 255)
+  ;; CHECK-NEXT:     )
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (local.get $4)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $1
+  ;; CHECK-NEXT:    (i32.and
+  ;; CHECK-NEXT:     (i32.const 99998)
+  ;; CHECK-NEXT:     (i32.const 255)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
   ;; CHECK-NEXT:  (drop
-  ;; CHECK-NEXT:   (struct.get_u $struct.packed 0
-  ;; CHECK-NEXT:    (struct.new_default $struct.packed)
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $1)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $2
+  ;; CHECK-NEXT:    (i32.const 99999)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result i32)
+  ;; CHECK-NEXT:    (drop
+  ;; CHECK-NEXT:     (ref.null none)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.get $2)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $5
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $6
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
   ;; CHECK-NEXT:   )
   ;; CHECK-NEXT:  )
   ;; CHECK-NEXT: )
   (func $packed
-    ;; We do not optimize packed structs yet.
+    (local $temp (ref $struct.packed))
+    ;; Packed fields require masking of irrelevant bits, which we apply on the
+    ;; sets.
+    (local.set $temp
+      (struct.new $struct.packed
+        (i32.const 1337)
+        (i32.const 1338)
+      )
+    )
+    (struct.set $struct.packed 0
+      (local.get $temp)
+      (i32.const 99998)
+    )
     (drop
       (struct.get $struct.packed 0
-        (struct.new_default $struct.packed)
+        (local.get $temp)
       )
     )
+    ;; Unpacked fields in the same struct do not need anything.
+    (struct.set $struct.packed 1
+      (local.get $temp)
+      (i32.const 99999)
+    )
+    (drop
+      (struct.get $struct.packed 1
+        (local.get $temp)
+      )
+    )
+    ;; When using struct.new_default we do not need any masking, as the values
+    ;; written are 0 anyhow.
+    (local.set $temp
+      (struct.new_default $struct.packed)
+    )
   )
 
   ;; CHECK:      (func $with-init-values (type $1)
@@ -647,7 +742,7 @@
     )
   )
 
-  ;; CHECK:      (func $local-copies-conditional (type $8) (param $x i32) (result f64)
+  ;; CHECK:      (func $local-copies-conditional (type $7) (param $x i32) (result f64)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local $2 i32)
   ;; CHECK-NEXT:  (local $3 f64)
@@ -741,7 +836,7 @@
     )
   )
 
-  ;; CHECK:      (func $non-exclusive-get (type $8) (param $x i32) (result f64)
+  ;; CHECK:      (func $non-exclusive-get (type $7) (param $x i32) (result f64)
   ;; CHECK-NEXT:  (local $ref (ref null $struct.A))
   ;; CHECK-NEXT:  (local.set $ref
   ;; CHECK-NEXT:   (struct.new_default $struct.A)
@@ -3311,3 +3406,124 @@
     )
   )
 )
+
+;; Packed arrays.
+(module
+  ;; CHECK:      (type $0 (func (result i32)))
+
+  ;; CHECK:      (type $array8 (array (mut i8)))
+  (type $array8 (array (mut i8)))
+
+  ;; CHECK:      (type $array16 (array (mut i16)))
+  (type $array16 (array (mut i16)))
+
+  ;; CHECK:      (func $array8 (type $0) (result i32)
+  ;; CHECK-NEXT:  (local $temp (ref $array8))
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $2
+  ;; CHECK-NEXT:    (i32.and
+  ;; CHECK-NEXT:     (i32.const 1337)
+  ;; CHECK-NEXT:     (i32.const 255)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block (result i32)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.get $2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $array8 (result i32)
+    (local $temp (ref $array8))
+    (local.set $temp
+      (array.new_default $array8
+        (i32.const 3)
+      )
+    )
+    (array.set $array8
+      (local.get $temp)
+      (i32.const 1)
+      (i32.const 1337)
+    )
+    (array.get $array8
+      (local.get $temp)
+      (i32.const 1)
+    )
+  )
+
+  ;; CHECK:      (func $array16 (type $0) (result i32)
+  ;; CHECK-NEXT:  (local $temp (ref $array16))
+  ;; CHECK-NEXT:  (local $1 i32)
+  ;; CHECK-NEXT:  (local $2 i32)
+  ;; CHECK-NEXT:  (local $3 i32)
+  ;; CHECK-NEXT:  (drop
+  ;; CHECK-NEXT:   (block (result nullref)
+  ;; CHECK-NEXT:    (local.set $1
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $2
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (local.set $3
+  ;; CHECK-NEXT:     (i32.const 0)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.set $2
+  ;; CHECK-NEXT:    (i32.and
+  ;; CHECK-NEXT:     (i32.const 1337)
+  ;; CHECK-NEXT:     (i32.const 65535)
+  ;; CHECK-NEXT:    )
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (block (result i32)
+  ;; CHECK-NEXT:   (drop
+  ;; CHECK-NEXT:    (ref.null none)
+  ;; CHECK-NEXT:   )
+  ;; CHECK-NEXT:   (local.get $2)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $array16 (result i32)
+    (local $temp (ref $array16))
+    (local.set $temp
+      (array.new_default $array16
+        (i32.const 3)
+      )
+    )
+    (array.set $array16
+      (local.get $temp)
+      (i32.const 1)
+      (i32.const 1337)
+    )
+    (array.get $array16
+      (local.get $temp)
+      (i32.const 1)
+    )
+  )
+)
diff --git a/test/passes/Oz_fuzz-exec_all-features.txt b/test/passes/Oz_fuzz-exec_all-features.txt
index 20cbca924..c46085be1 100644
--- a/test/passes/Oz_fuzz-exec_all-features.txt
+++ b/test/passes/Oz_fuzz-exec_all-features.txt
@@ -269,37 +269,27 @@
   )
  )
  (func $array.new_fixed (type $void_func) (; has Stack IR ;)
-  (local $0 (ref $bytes))
+  (local $0 i32)
+  (local $1 i32)
   (local.set $0
-   (array.new_fixed $bytes 2
-    (i32.const 42)
-    (i32.const 50)
-   )
+   (i32.const 42)
+  )
+  (local.set $1
+   (i32.const 50)
   )
   (call $log
    (i32.const 2)
   )
   (call $log
-   (array.get_u $bytes
-    (local.get $0)
-    (i32.const 0)
-   )
+   (i32.const 42)
   )
   (call $log
-   (array.get_u $bytes
-    (local.get $0)
-    (i32.const 1)
-   )
+   (i32.const 50)
   )
  )
  (func $array.new_fixed-packed (type $void_func) (; has Stack IR ;)
   (call $log
-   (array.get_u $bytes
-    (array.new_fixed $bytes 1
-     (i32.const -11512)
-    )
-    (i32.const 0)
-   )
+   (i32.const 8)
   )
  )
  (func $static-casts (type $void_func) (; has Stack IR ;)