JIT: Transform arithmetic using distributive property

src/coreclr/jit/compiler.h

@@ -3904,6 +3904,7 @@ class Compiler
 
     GenTree* gtFoldExpr(GenTree* tree);
     GenTree* gtFoldExprConst(GenTree* tree);
+    GenTree* gtFoldDistributiveArithmetic(GenTree* tree);
     GenTree* gtFoldIndirConst(GenTreeIndir* indir);
     GenTree* gtFoldExprSpecial(GenTree* tree);
     GenTree* gtFoldExprSpecialFloating(GenTree* tree);

src/coreclr/jit/gentree.cpp

@@ -15041,6 +15041,10 @@ GenTree* Compiler::gtFoldExpr(GenTree* tree)
 
             return gtFoldExprCompare(tree);
         }
+        else if (tree->OperIs(GT_AND, GT_OR, GT_XOR, GT_ADD, GT_SUB))
+        {
+            return gtFoldDistributiveArithmetic(tree);
+        }
     }
 
     /* Return the original node (folded/bashed or not) */
@@ -18010,6 +18014,80 @@ GenTree* Compiler::gtFoldExprConst(GenTree* tree)
     return tree;
 }
 
+//------------------------------------------------------------------------
+// gtFoldDistributiveArithmetic: Optimizes distributive Arithmetic.
+//
+// Arguments:
+//   tree - the unchecked GT_AND/GT_OR/GT_XOR/GT_ADD/GT_SUB tree to optimize.
+//
+// Return Value:
+//   The unchanged tree or optimized tree with oper GT_MUL/GT_OR/GT_AND.
+//
+GenTree* Compiler::gtFoldDistributiveArithmetic(GenTree* tree)
+{
+    assert(tree->OperIs(GT_AND, GT_OR, GT_XOR, GT_ADD, GT_SUB));
+
+    if (opts.OptimizationDisabled())
+    {
+        return tree;
+    }
+
+    if (tree->gtOverflowEx() || !varTypeIsIntegralOrI(tree))
+    {
+        return tree;
+    }
+
+    if ((tree->gtFlags & (GTF_PERSISTENT_SIDE_EFFECTS | GTF_ORDER_SIDEEFF)) != 0)
+    {
+        return tree;
+    }
+
+    GenTree* op1 = tree->gtGetOp1();
+    GenTree* op2 = tree->gtGetOp2();
+
+    auto isLeftDistributive = [](genTreeOps op1, genTreeOps op2) {
+        // op1 is left distributive over op2 iff:
+        // "A op1 (B op2 C)" <==> "(A op1 B) op2 (A op1 C)"
+        switch (op1)
+        {
+            case GT_AND:
+                return op2 == GT_OR || op2 == GT_XOR || op2 == GT_AND;
+
+            case GT_OR:
+                return op2 == GT_AND || op2 == GT_OR;
+
+            case GT_MUL:
+                return op2 == GT_ADD || op2 == GT_SUB;
+
+            default:
+                return false;
+        }
+    };
+
+    if ((op1->OperGet() == op2->OperGet()) && isLeftDistributive(op1->OperGet(), tree->OperGet()))
+    {
+        if (op1->gtGetOp1()->OperIs(GT_LCL_VAR) && op2->gtGetOp1()->OperIs(GT_LCL_VAR))
+        {
+            bool sameLcl =
+                op1->gtGetOp1()->AsLclVarCommon()->GetLclNum() == op2->gtGetOp1()->AsLclVarCommon()->GetLclNum();
+            if (sameLcl)
+            {
+                tree->AsOp()->gtOp1 = op1->gtGetOp1();
+                tree->AsOp()->gtOp2 =
+                    gtFoldExpr(gtNewOperNode(tree->OperGet(), tree->TypeGet(), op1->gtGetOp2(), op2->gtGetOp2()));
+                tree->SetOper(op1->OperGet(), GenTree::PRESERVE_VN);
+
+                if (fgGlobalMorph)
+                {
+                    fgMorphTreeDone(tree->gtGetOp2());
+                }
+            }
+        }
+    }
+
+    return tree;
+}
+
 //------------------------------------------------------------------------
 // gtFoldIndirConst: Attempt to fold an "IND(addr)" expression to a constant.
 //

src/coreclr/jit/morph.cpp

@@ -7853,6 +7853,8 @@ GenTree* Compiler::fgMorphSmpOp(GenTree* tree, MorphAddrContext* mac, bool* optA
                 goto CM_OVF_OP;
             }
 
+            // TODO: Call fgOptimizeDistributiveArithmetic + fgOptimizeCommutativeArithmetic
+
             if (!fgGlobalMorph)
             {
                 break;
@@ -10303,7 +10305,7 @@ GenTree* Compiler::fgOptimizeHWIntrinsicAssociative(GenTreeHWIntrinsic* tree)
 #endif // FEATURE_HW_INTRINSICS
 
 //------------------------------------------------------------------------
-// fgOptimizeCommutativeArithmetic: Optimizes commutative operations.
+// fgOptimizeCommutativeArithmetic: Optimizes commutative Arithmetic.
 //
 // Arguments:
 //   tree - the unchecked GT_ADD/GT_MUL/GT_OR/GT_XOR/GT_AND tree to optimize.

src/coreclr/jit/optimizebools.cpp

@@ -1198,6 +1198,12 @@ void OptBoolsDsc::optOptimizeBoolsUpdateTrees()
 
     GenTree* cmpOp1 = m_foldOp == GT_NONE ? m_c1 : m_compiler->gtNewOperNode(m_foldOp, m_foldType, m_c1, m_c2);
 
+    // There may be new opportunities for distributive arithmetic optimization
+    if (m_foldOp != GT_NONE)
+    {
+        cmpOp1 = m_compiler->gtFoldExpr(cmpOp1);
+    }
+
     GenTree* t1Comp = m_testInfo1.compTree;
     t1Comp->SetOper(m_cmpOp);
     t1Comp->AsOp()->gtOp1         = cmpOp1;