1 files changed, 22 insertions, 17 deletions

diff --git a/src/passes/PickLoadSigns.cpp b/src/passes/PickLoadSigns.cpp
index 587314209..693fd4a68 100644
--- a/src/passes/PickLoadSigns.cpp
+++ b/src/passes/PickLoadSigns.cpp
@@ -20,10 +20,11 @@
 
 namespace wasm {
 
+//
 // Adjust load signedness based on usage. If a load only has uses that sign or
 // unsign it anyhow, then it could be either, and picking the popular one can
-// help remove the most sign/unsign operations
-// unsigned, then it could be either
+// help remove the most sign/unsign operations.
+//
 
 struct PickLoadSigns : public WalkerPass<ExpressionStackWalker<PickLoadSigns>> {
   bool isFunctionParallel() override { return true; }
@@ -59,13 +60,20 @@ struct PickLoadSigns : public WalkerPass<ExpressionStackWalker<PickLoadSigns>> {
   }
 
   void visitLocalGet(LocalGet* curr) {
-    // this is a use. check from the context what it is, signed or unsigned,
-    // etc.
+    // This is a use. check from the context what it is, signed or unsigned,
+    // etc. Sign- and zero-extends may have two levels of nesting (x << K >> K)
+    // or one (x & K), so check both.
     auto& usage = usages[curr->index];
     usage.totalUsages++;
-    if (expressionStack.size() >= 2) {
-      auto* parent = expressionStack[expressionStack.size() - 2];
-      if (Properties::getZeroExtValue(parent)) {
+    for (Index i = 2; i <= 3; i++) {
+      if (expressionStack.size() < i) {
+        // We do not have that many expressions above us. (And if we don't have
+        // 2 then we definitely don't have 3.)
+        break;
+      }
+
+      auto* parent = expressionStack[expressionStack.size() - i];
+      if (Properties::getZeroExtValue(parent) == curr) {
         auto bits = Properties::getZeroExtBits(parent);
         if (usage.unsignedUsages == 0) {
           usage.unsignedBits = bits;
@@ -73,17 +81,14 @@ struct PickLoadSigns : public WalkerPass<ExpressionStackWalker<PickLoadSigns>> {
           usage.unsignedBits = 0;
         }
         usage.unsignedUsages++;
-      } else if (expressionStack.size() >= 3) {
-        auto* grandparent = expressionStack[expressionStack.size() - 3];
-        if (Properties::getSignExtValue(grandparent)) {
-          auto bits = Properties::getSignExtBits(grandparent);
-          if (usage.signedUsages == 0) {
-            usage.signedBits = bits;
-          } else if (usage.signedBits != bits) {
-            usage.signedBits = 0;
-          }
-          usage.signedUsages++;
+      } else if (Properties::getSignExtValue(parent) == curr) {
+        auto bits = Properties::getSignExtBits(parent);
+        if (usage.signedUsages == 0) {
+          usage.signedBits = bits;
+        } else if (usage.signedBits != bits) {
+          usage.signedBits = 0;
         }
+        usage.signedUsages++;
       }
     }
   }