[Wasm GC] Optimize struct stores like stores to memory, ignore unneeded bits (#3680)

When storing to an i8, we can ignore any higher bits, etc.

Adds a getByteSize utility to Field to make this convenient.

5 files changed, 108 insertions, 23 deletions

diff --git a/src/passes/OptimizeInstructions.cpp b/src/passes/OptimizeInstructions.cpp
index c12a48b0d..7f02eb283 100644
--- a/src/passes/OptimizeInstructions.cpp
+++ b/src/passes/OptimizeInstructions.cpp
@@ -933,28 +933,40 @@ struct OptimizeInstructions
       return;
     }
     optimizeMemoryAccess(curr->ptr, curr->offset);
-    if (curr->valueType.isInteger()) {
-      // truncates constant values during stores
-      // (i32|i64).store(8|16|32)(p, C)   ==>
-      //    (i32|i64).store(8|16|32)(p, C & mask)
-      if (auto* c = curr->value->dynCast<Const>()) {
-        if (curr->valueType == Type::i64 && curr->bytes == 4) {
-          c->value = c->value.and_(Literal(uint64_t(0xffffffff)));
-        } else {
-          c->value = c->value.and_(Literal::makeFromInt32(
-            Bits::lowBitMask(curr->bytes * 8), curr->valueType));
-        }
+    optimizeStoredValue(curr->value, curr->bytes);
+    if (auto* unary = curr->value->dynCast<Unary>()) {
+      if (unary->op == WrapInt64) {
+        // instead of wrapping to 32, just store some of the bits in the i64
+        curr->valueType = Type::i64;
+        curr->value = unary->value;
+      }
+    }
+  }
+
+  void optimizeStoredValue(Expression*& value, Index bytes) {
+    if (!value->type.isInteger()) {
+      return;
+    }
+    // truncates constant values during stores
+    // (i32|i64).store(8|16|32)(p, C)   ==>
+    //    (i32|i64).store(8|16|32)(p, C & mask)
+    if (auto* c = value->dynCast<Const>()) {
+      if (value->type == Type::i64 && bytes == 4) {
+        c->value = c->value.and_(Literal(uint64_t(0xffffffff)));
+      } else {
+        c->value = c->value.and_(
+          Literal::makeFromInt32(Bits::lowBitMask(bytes * 8), value->type));
       }
     }
     // stores of fewer bits truncates anyhow
-    if (auto* binary = curr->value->dynCast<Binary>()) {
+    if (auto* binary = value->dynCast<Binary>()) {
       if (binary->op == AndInt32) {
         if (auto* right = binary->right->dynCast<Const>()) {
           if (right->type == Type::i32) {
             auto mask = right->value.geti32();
-            if ((curr->bytes == 1 && mask == 0xff) ||
-                (curr->bytes == 2 && mask == 0xffff)) {
-              curr->value = binary->left;
+            if ((bytes == 1 && mask == 0xff) ||
+                (bytes == 2 && mask == 0xffff)) {
+              value = binary->left;
             }
           }
         }
@@ -962,16 +974,10 @@ struct OptimizeInstructions
         // if sign extending the exact bit size we store, we can skip the
         // extension if extending something bigger, then we just alter bits we
         // don't save anyhow
-        if (Properties::getSignExtBits(binary) >= Index(curr->bytes) * 8) {
-          curr->value = ext;
+        if (Properties::getSignExtBits(binary) >= Index(bytes) * 8) {
+          value = ext;
         }
       }
-    } else if (auto* unary = curr->value->dynCast<Unary>()) {
-      if (unary->op == WrapInt64) {
-        // instead of wrapping to 32, just store some of the bits in the i64
-        curr->valueType = Type::i64;
-        curr->value = unary->value;
-      }
     }
   }
 
@@ -985,6 +991,13 @@ struct OptimizeInstructions
     }
   }
 
+  void visitStructSet(StructSet* curr) {
+    if (curr->ref->type != Type::unreachable && curr->value->type.isInteger()) {
+      const auto& fields = curr->ref->type.getHeapType().getStruct().fields;
+      optimizeStoredValue(curr->value, fields[curr->index].getByteSize());
+    }
+  }
+
   Index getMaxBitsForLocal(LocalGet* get) {
     // check what we know about the local
     return localInfo[get->index].maxBits;
diff --git a/src/wasm-type.h b/src/wasm-type.h
index 7f037d5e7..bb499c7c1 100644
--- a/src/wasm-type.h
+++ b/src/wasm-type.h
@@ -440,6 +440,8 @@ struct Field {
   }
   bool operator!=(const Field& other) const { return !(*this == other); }
   std::string toString() const;
+
+  unsigned getByteSize() const;
 };
 
 typedef std::vector<Field> FieldList;
diff --git a/src/wasm/wasm-type.cpp b/src/wasm/wasm-type.cpp
index 86f95ccd4..dcc385fb8 100644
--- a/src/wasm/wasm-type.cpp
+++ b/src/wasm/wasm-type.cpp
@@ -779,6 +779,21 @@ std::ostream& operator<<(std::ostream& os, Rtt rtt) {
   return TypePrinter(os).print(rtt);
 }
 
+unsigned Field::getByteSize() const {
+  if (type != Type::i32) {
+    return type.getByteSize();
+  }
+  switch (packedType) {
+    case Field::PackedType::i8:
+      return 1;
+    case Field::PackedType::i16:
+      return 2;
+    case Field::PackedType::not_packed:
+      return 4;
+  }
+  WASM_UNREACHABLE("impossible packed type");
+}
+
 namespace {
 
 bool TypeComparator::lessThan(Type a, Type b) {
diff --git a/test/example/cpp-unit.cpp b/test/example/cpp-unit.cpp
index 9552473ce..49893d1d2 100644
--- a/test/example/cpp-unit.cpp
+++ b/test/example/cpp-unit.cpp
@@ -569,11 +569,24 @@ void test_literals() {
   }
 }
 
+void test_field() {
+  // Simple types
+  assert_equal(Field(Type::i32, Immutable).getByteSize(), 4);
+  assert_equal(Field(Type::i64, Immutable).getByteSize(), 8);
+
+  // Packed types
+  assert_equal(Field(Field::PackedType::i8, Immutable).getByteSize(), 1);
+  assert_equal(Field(Field::PackedType::i16, Immutable).getByteSize(), 2);
+  assert_equal(Field(Field::PackedType::not_packed, Immutable).getByteSize(),
+               4);
+}
+
 int main() {
   test_bits();
   test_cost();
   test_effects();
   test_literals();
+  test_field();
 
   if (failsCount > 0) {
     abort();
diff --git a/test/lit/passes/optimize-instructions-gc.wast b/test/lit/passes/optimize-instructions-gc.wast
index 31ecd0f63..6a0341856 100644
--- a/test/lit/passes/optimize-instructions-gc.wast
+++ b/test/lit/passes/optimize-instructions-gc.wast
@@ -3,6 +3,15 @@
 ;; RUN:   | filecheck %s
 
 (module
+  (import "env" "get-i32" (func $get-i32 (result i32)))
+
+  (type $struct (struct
+    (field $i8  (mut i8))
+    (field $i16 (mut i16))
+    (field $i32 (mut i32))
+    (field $i64 (mut i64))
+  ))
+
   ;; These functions test if an `if` with subtyped arms is correctly folded
   ;; 1. if its `ifTrue` and `ifFalse` arms are identical (can fold)
   ;; CHECK:      (func $if-arms-subtype-fold (result anyref)
@@ -30,4 +39,37 @@
       (ref.null func)
     )
   )
+
+  ;; Stored values automatically truncate unneeded bytes.
+  ;; CHECK:      (func $store-trunc (param $x (ref null $struct))
+  ;; CHECK-NEXT:  (struct.set $struct $i8
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (i32.const 35)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.set $struct $i16
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (i32.const 9029)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT:  (struct.set $struct $i8
+  ;; CHECK-NEXT:   (local.get $x)
+  ;; CHECK-NEXT:   (call $get-i32)
+  ;; CHECK-NEXT:  )
+  ;; CHECK-NEXT: )
+  (func $store-trunc (param $x (ref null $struct))
+    (struct.set $struct $i8
+      (local.get $x)
+      (i32.const 0x123) ;; data over 0xff is unnecessary
+    )
+    (struct.set $struct $i16
+      (local.get $x)
+      (i32.const 0x12345) ;; data over 0xffff is unnecessary
+    )
+    (struct.set $struct $i8
+      (local.get $x)
+      (i32.and ;; truncating bits using an and is unnecessary
+       (call $get-i32)
+       (i32.const 0xff)
+      )
+    )
+  )
 )