[llvm] [SCEV] Generalize A + zext(-A + B) fold to A + zext(C + X) (PR #209160)

[Antonio Frighetto via llvm-commits]

================
@@ -2809,17 +2809,27 @@ const SCEV *ScalarEvolution::getAddExpr(SmallVectorImpl<SCEVUse> &Ops,
       }
     }
 
-    // Try to push the constant operand into a ZExt: A + zext (-A + B) -> zext
-    // (B), if trunc (A) + -A + B  does not unsigned-wrap.
-    const SCEVAddExpr *InnerAdd;
-    if (match(B, m_scev_ZExt(m_scev_Add(InnerAdd)))) {
-      const SCEV *NarrowA = getTruncateExpr(A, InnerAdd->getType());
-      if (NarrowA == getNegativeSCEV(InnerAdd->getOperand(0)) &&
-          getZeroExtendExpr(NarrowA, B->getType()) == A &&
-          hasFlags(StrengthenNoWrapFlags(this, scAddExpr, {NarrowA, InnerAdd},
-                                         SCEV::FlagAnyWrap),
+    // Push a negative constant addend out of a ZExt when the inner add is
+    // provably non-negative in the narrow type:
+    //
+    //   A + zext(C + X)  ->  WideAC + zext(X)   [WideAC = A + sext(C)]
+    //
+    // Require A to be a constant so that `A + sext(C)` folds into a single wide
+    // constant, actually simplifying the expression.
+    const SCEVAddExpr *Add;
+    if (isa<SCEVConstant>(A) && match(B, m_scev_ZExt(m_scev_Add(Add)))) {
+      const auto *CInner = dyn_cast<SCEVConstant>(Add->getOperand(0));
+      if (CInner && CInner->getAPInt().isNegative() &&
+          // NUW on `(-C) + (C + X) = X` proves that `C + X` did not wrap
+          // below zero, so `zext(C + X) == sext(C) + zext(X)`.
+          hasFlags(StrengthenNoWrapFlags(
+                       this, scAddExpr, {getConstant(-CInner->getAPInt()), Add},
+                       SCEV::FlagAnyWrap),
                    SCEV::FlagNUW)) {
-        return getZeroExtendExpr(getAddExpr(NarrowA, InnerAdd), B->getType());
+        SmallVector<SCEVUse, 4> XOps(drop_begin(Add->operands()));
+        const SCEV *X = getAddExpr(XOps);
----------------
antoniofrighetto wrote:

```suggestion
        const SCEV *XAdd = getAddExpr(XOps);
```

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