[llvm] [WebAssembly] Optimizes [zext](and 63|32) for shl, srl and sra with selectShiftMask (PR #170572)

[Luke Lau via llvm-commits]

================
@@ -548,6 +557,49 @@ bool WebAssemblyDAGToDAGISel::SelectAddrOperands64(SDValue Op, SDValue &Offset,
   return SelectAddrOperands(MVT::i64, WebAssembly::CONST_I64, Op, Offset, Addr);
 }
 
+bool WebAssemblyDAGToDAGISel::selectShiftMask(SDValue N, unsigned ShiftWidth,
+                                              SDValue &ShAmt) {
+
+  ShAmt = N;
+
+  bool HasZext = false;
+  // Peek through zext.
+  if (ShAmt->getOpcode() == ISD::ZERO_EXTEND) {
+    ShAmt = ShAmt.getOperand(0);
+    HasZext = true;
+  }
+
+  if (ShAmt.getOpcode() == ISD::AND &&
+      isa<ConstantSDNode>(ShAmt.getOperand(1))) {
+    const APInt &AndMask = ShAmt.getConstantOperandAPInt(1);
+
+    // Since the max shift amount is a power of 2 we can subtract 1 to make a
+    // mask that covers the bits needed to represent all shift amounts.
+    assert(isPowerOf2_32(ShiftWidth) && "Unexpected max shift amount!");
+    APInt ShMask(AndMask.getBitWidth(), ShiftWidth - 1);
+
+    if (ShMask.isSubsetOf(AndMask)) {
+      ShAmt = ShAmt.getOperand(0);
+    } else {
+      // TODO: port computeKnownBits from riscv in another PR about rotr and
+      // rotl
+      return false;
+    }
+
+    // Only reinstate zext if it's i32 -> i64, WebAssembly would have legalize
+    // i16 to i32 in the dag otherwise.
+    if (HasZext && ShiftWidth == 64) {
+      ShAmt = SDValue(CurDAG->getMachineNode(WebAssembly::I64_EXTEND_U_I32,
+                                             SDLoc(N), MVT::i64, ShAmt),
+                      0);
+    }
+    return true;
+  }
+
+  // TODO: Port rest of riscv if applicable
----------------
lukel97 wrote:

```suggestion
  // TODO: Port rest of RISCVDAGToDAGISel::selectShiftMask if applicable
```

https://github.com/llvm/llvm-project/pull/170572