Range Analysis: Compute range of Binary with >> operator

This patch computes the range of a `Binary` with `>>`
(shift-right) operator when the `lhs` is a non-negative integer.

Bug: 348701956
Test: tint_unittests
Change-Id: I42d0e4ff542a6feb7b6e1f929c6fffa8338b2bed
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/245776
Commit-Queue: Jiawei Shao <[email protected]>
Reviewed-by: James Price <[email protected]>

Author: Jiawei-Shao

src/tint/lang/core/ir/analysis/integer_range_analysis.cc

@@ -284,6 +284,9 @@ struct IntegerRangeAnalysisImpl {
             case BinaryOp::kShiftLeft:
                 return ComputeAndCacheIntegerRangeForBinaryShiftLeft(binary, range_lhs, range_rhs);
 
+            case BinaryOp::kShiftRight:
+                return ComputeAndCacheIntegerRangeForBinaryShiftRight(binary, range_lhs, range_rhs);
+
             default:
                 return nullptr;
         }
@@ -1095,6 +1098,45 @@ struct IntegerRangeAnalysisImpl {
         }
     }
 
+    const IntegerRangeInfo* ComputeAndCacheIntegerRangeForBinaryShiftRight(
+        const Binary* binary,
+        const IntegerRangeInfo* lhs,
+        const IntegerRangeInfo* rhs) {
+        auto rhs_u32 = std::get<IntegerRangeInfo::UnsignedIntegerRange>(rhs->range);
+
+        // rhs_u32 must be less than the bit width of i32 and u32 (32):
+        // rhs_u32.min_bound <= rhs_u32.max_bound < 32
+        if (rhs_u32.max_bound >= 32) {
+            return nullptr;
+        }
+
+        if (std::holds_alternative<IntegerRangeInfo::SignedIntegerRange>(lhs->range)) {
+            auto lhs_i32 = std::get<IntegerRangeInfo::SignedIntegerRange>(lhs->range);
+
+            // Currently we require `lhs` must be non-negative.
+            // 0 <= lhs.min_bound <= lhs.max_bound
+            if (lhs_i32.min_bound < 0) {
+                return nullptr;
+            }
+
+            // [min1, max1] >> [min2, max2] => [min1 >> max2, max1 >> min2]
+            int64_t min_bound = lhs_i32.min_bound >> rhs_u32.max_bound;
+            int64_t max_bound = lhs_i32.max_bound >> rhs_u32.min_bound;
+
+            auto result =
+                integer_binary_range_info_map_.Add(binary, IntegerRangeInfo(min_bound, max_bound));
+            return &result.value;
+        } else {
+            auto lhs_u32 = std::get<IntegerRangeInfo::UnsignedIntegerRange>(lhs->range);
+
+            uint64_t min_bound = lhs_u32.min_bound >> rhs_u32.max_bound;
+            uint64_t max_bound = lhs_u32.max_bound >> rhs_u32.min_bound;
+            auto result =
+                integer_binary_range_info_map_.Add(binary, IntegerRangeInfo(min_bound, max_bound));
+            return &result.value;
+        }
+    }
+
     Hashmap<const FunctionParam*, Vector<IntegerRangeInfo, 3>, 4>
         integer_function_param_range_info_map_;
     Hashmap<const Var*, IntegerRangeInfo, 8> integer_var_range_info_map_;