[OptimizeInstructions] Handle post-MVP sign extended operations (#3910)

fixes part of #3906

2 files changed, 37 insertions, 3 deletions

diff --git a/src/ir/bits.h b/src/ir/bits.h
index 95481eeab..950864a23 100644
--- a/src/ir/bits.h
+++ b/src/ir/bits.h
@@ -385,9 +385,26 @@ Index getMaxBits(Expression* curr,
       case WrapInt64:
       case ExtendUInt32:
         return std::min(Index(32), getMaxBits(unary->value, localInfoProvider));
+      case ExtendS8Int32: {
+        auto maxBits = getMaxBits(unary->value, localInfoProvider);
+        return maxBits >= 8 ? Index(32) : maxBits;
+      }
+      case ExtendS16Int32: {
+        auto maxBits = getMaxBits(unary->value, localInfoProvider);
+        return maxBits >= 16 ? Index(32) : maxBits;
+      }
+      case ExtendS8Int64: {
+        auto maxBits = getMaxBits(unary->value, localInfoProvider);
+        return maxBits >= 8 ? Index(64) : maxBits;
+      }
+      case ExtendS16Int64: {
+        auto maxBits = getMaxBits(unary->value, localInfoProvider);
+        return maxBits >= 16 ? Index(64) : maxBits;
+      }
+      case ExtendS32Int64:
       case ExtendSInt32: {
         auto maxBits = getMaxBits(unary->value, localInfoProvider);
-        return maxBits == 32 ? Index(64) : maxBits;
+        return maxBits >= 32 ? Index(64) : maxBits;
       }
       default: {
       }
diff --git a/src/ir/properties.h b/src/ir/properties.h
index fdf31f4e2..63dd6c72c 100644
--- a/src/ir/properties.h
+++ b/src/ir/properties.h
@@ -146,6 +146,12 @@ inline Expression* getSignExtValue(Expression* curr) {
   if (curr->type != Type::i32) {
     return nullptr;
   }
+  if (auto* unary = curr->dynCast<Unary>()) {
+    if (unary->op == ExtendS8Int32 || unary->op == ExtendS16Int32) {
+      return unary->value;
+    }
+    return nullptr;
+  }
   using namespace Match;
   int32_t leftShift = 0, rightShift = 0;
   Expression* extended = nullptr;
@@ -162,8 +168,19 @@ inline Expression* getSignExtValue(Expression* curr) {
 // gets the size of the sign-extended value
 inline Index getSignExtBits(Expression* curr) {
   assert(curr->type == Type::i32);
-  auto* rightShift = curr->cast<Binary>()->right;
-  return 32 - Bits::getEffectiveShifts(rightShift);
+  if (auto* unary = curr->dynCast<Unary>()) {
+    switch (unary->op) {
+      case ExtendS8Int32:
+        return 8;
+      case ExtendS16Int32:
+        return 16;
+      default:
+        WASM_UNREACHABLE("invalid unary operation");
+    }
+  } else {
+    auto* rightShift = curr->cast<Binary>()->right;
+    return 32 - Bits::getEffectiveShifts(rightShift);
+  }
 }
 
 // Check if an expression is almost a sign-extend: perhaps the inner shift