Fix BitAnd+Add carry bug and ShrU arithmetic shift bug

- Disable mask & (x + c) rules: bits-disjoint ignores carry propagation
  (e.g. mask=4, c=2, x=2: (2+2)&4=4 but 2&4=0)
- Fix ShrU constant folding: use u32-shr primitive for logical shift
  instead of egglog's >> which is i64 arithmetic shift
- Fix Shl constant folding: use u32-shl for proper 32-bit semantics

Author: LegNeato

rust/spirv-tools-opt/src/egglog_opt/mod.rs

@@ -998,6 +998,14 @@ pub fn create_spirv_egraph() -> Result<EGraph, EgglogOptError> {
     );
     add_primitive!(&mut egraph, "u32-div" = |a: i64, b: i64| -?> i64 { u32_div(a, b) });
     add_primitive!(&mut egraph, "u32-mod" = |a: i64, b: i64| -?> i64 { u32_mod(a, b) });
+    add_primitive!(
+        &mut egraph,
+        "u32-shr" = |a: i64, b: i64| -> i64 { u32_shr(a, b) }
+    );
+    add_primitive!(
+        &mut egraph,
+        "u32-shl" = |a: i64, b: i64| -> i64 { u32_shl(a, b) }
+    );
 
     // NaN-aware float comparison primitives (FOrd* returns 0 if NaN, FUnord* returns 1 if NaN)
     add_primitive!(

rust/spirv-tools-opt/src/egglog_opt/primitives/bitwise.rs

@@ -253,6 +253,22 @@ pub fn u32_mod(a: i64, b: i64) -> Option<i64> {
     }
 }
 
+/// Logical (unsigned) right shift: cast to u32, shift, sign-extend back to i64.
+/// egglog's native `>>` is arithmetic (sign-extending), which is wrong for ShrU.
+pub fn u32_shr(a: i64, b: i64) -> i64 {
+    let a = a as u32;
+    let b = (b as u32) & 31; // mask shift amount to 0-31
+    (a >> b) as i32 as i64
+}
+
+/// Logical left shift: cast to u32, shift, sign-extend back to i64.
+/// egglog's native `<<` on i64 can produce results wider than 32 bits.
+pub fn u32_shl(a: i64, b: i64) -> i64 {
+    let a = a as u32;
+    let b = (b as u32) & 31;
+    (a << b) as i32 as i64
+}
+
 // =============================================================================
 // Type conversion primitives (cross-type: F <-> i64)
 // =============================================================================

rust/spirv-tools-opt/src/egglog_opt/primitives/mod.rs

@@ -10,7 +10,7 @@ pub(super) use bitwise::{
     funord_ge, funord_gt, funord_le, funord_lt, funord_ne, has_exact_recip, int_to_float_signed,
     int_to_float_unsigned, is_float_one64, is_float_zero64, is_pow2, log2_pow2, mask_superset,
     popcount, sdiv32, shl_clears_mask, shr_clears_mask, smod, srem32, u32_div, u32_ge, u32_gt,
-    u32_le, u32_lt, u32_max, u32_min, u32_mod,
+    u32_le, u32_lt, u32_max, u32_min, u32_mod, u32_shl, u32_shr,
 };
 
 #[allow(unused_imports)]

rust/spirv-tools-opt/src/rules/constant_folding.egg

@@ -20,9 +20,9 @@
 (rule ((= e (BitNot (Const a))))
       ((union e (Const (not-i64 a)))))
 (rule ((= e (Shl (Const a) (Const b))))
-      ((union e (Const (<< a b)))))
+      ((union e (Const (u32-shl a b)))))
 (rule ((= e (ShrU (Const a) (Const b))))
-      ((union e (Const (>> a b)))))
+      ((union e (Const (u32-shr a b)))))
 
 ; Signed right shift of constants
 ; Note: egglog's >> is arithmetic shift for i64

rust/spirv-tools-opt/src/rules/primitives.egg

@@ -20,21 +20,17 @@
 ; Mask-Aware Optimizations
 ; =============================================================================
 
-; mask & (x + const) = mask & x when const has no bits in mask
-(rule ((= e (BitAnd (Const mask) (Add x (Const c))))
-       (bits-disjoint mask c))
-      ((union e (BitAnd (Const mask) x))))
-(rule ((= e (BitAnd (Add x (Const c)) (Const mask)))
-       (bits-disjoint mask c))
-      ((union e (BitAnd x (Const mask)))))
-
-; mask & (x - const) = mask & x when const has no bits in mask
-(rule ((= e (BitAnd (Const mask) (Sub x (Const c))))
-       (bits-disjoint mask c))
-      ((union e (BitAnd (Const mask) x))))
-(rule ((= e (BitAnd (Sub x (Const c)) (Const mask)))
-       (bits-disjoint mask c))
-      ((union e (BitAnd x (Const mask)))))
+; DISABLED: mask & (x + const) = mask & x — incorrect due to carry propagation.
+; bits-disjoint(mask, c) only checks mask & c == 0, but adding c causes carries
+; that can affect bits in mask. E.g. mask=4, c=2, x=2: (2+2)&4=4 but 2&4=0.
+; (rule ((= e (BitAnd (Const mask) (Add x (Const c)))) (bits-disjoint mask c))
+;       ((union e (BitAnd (Const mask) x))))
+; (rule ((= e (BitAnd (Add x (Const c)) (Const mask))) (bits-disjoint mask c))
+;       ((union e (BitAnd x (Const mask)))))
+; (rule ((= e (BitAnd (Const mask) (Sub x (Const c)))) (bits-disjoint mask c))
+;       ((union e (BitAnd (Const mask) x))))
+; (rule ((= e (BitAnd (Sub x (Const c)) (Const mask))) (bits-disjoint mask c))
+;       ((union e (BitAnd x (Const mask)))))
 
 ; mask & (x | const) = mask & x when const has no bits in mask
 (rule ((= e (BitAnd (Const mask) (BitOr x (Const c))))