Guard out-of-range shift counts in integer opcodes

Author: uwezkhan

numexpr/interp_body.cpp

@@ -266,8 +266,8 @@
         case OP_POW_III: VEC_ARG2(i_dest = (i2 < 0) ? (1 / i1) : (int)pow((double)i1, i2));
         case OP_MOD_III: VEC_ARG2(i_dest = i2 == 0 ? 0 :((i1 % i2) + i2) % i2);
         case OP_FLOORDIV_III: VEC_ARG2(i_dest = i2 ? (i1 / i2) - ((i1 % i2 != 0) && (i1 < 0 != i2 < 0)) : 0);
-        case OP_LSHIFT_III: VEC_ARG2(i_dest = i1 << i2);
-        case OP_RSHIFT_III: VEC_ARG2(i_dest = i1 >> i2);
+        case OP_LSHIFT_III: VEC_ARG2(i_dest = (unsigned int)i2 < 32 ? i1 << i2 : 0);
+        case OP_RSHIFT_III: VEC_ARG2(i_dest = i1 >> ((unsigned int)i2 < 32 ? i2 : 31));
 
         case OP_WHERE_IBII: VEC_ARG3(i_dest = b1 ? i2 : i3);
         //Bitwise ops
@@ -292,8 +292,8 @@
 #endif
         case OP_MOD_LLL: VEC_ARG2(l_dest = l2 == 0 ? 0 :((l1 % l2) + l2) % l2);
         case OP_FLOORDIV_LLL: VEC_ARG2(l_dest = l2 ? (l1 / l2) - ((l1 % l2 != 0) && (l1 < 0 != l2 < 0)): 0);
-        case OP_LSHIFT_LLL: VEC_ARG2(l_dest = l1 << l2);
-        case OP_RSHIFT_LLL: VEC_ARG2(l_dest = l1 >> l2);
+        case OP_LSHIFT_LLL: VEC_ARG2(l_dest = (unsigned long long)l2 < 64 ? l1 << l2 : 0);
+        case OP_RSHIFT_LLL: VEC_ARG2(l_dest = l1 >> ((unsigned long long)l2 < 64 ? l2 : 63));
 
         case OP_WHERE_LBLL: VEC_ARG3(l_dest = b1 ? l2 : l3);
         //Bitwise ops

numexpr/tests/test_numexpr.py

@@ -457,6 +457,17 @@ def test_right_shift(self):
         x = arange(10, dtype='i4')
         assert_array_equal(evaluate("x>>2"), x >> 2)
 
+    def test_shift_out_of_range(self):
+        # Shift counts that are negative or >= the operand width are
+        # undefined behavior in C. Match NumPy, which treats them as a
+        # full shift (0 for <<, sign fill for >>).
+        for dtype in ('i4', 'i8'):
+            x = array([5, -5, 0], dtype=dtype)
+            for count in (-1, 64, 200):
+                y = array([count] * len(x), dtype=dtype)
+                assert_array_equal(evaluate("x << y"), x << y)
+                assert_array_equal(evaluate("x >> y"), x >> y)
+
     # PyTables uses __nonzero__ among ExpressionNode objects internally
     # so this should be commented out for the moment.  See #24.
     def test_boolean_operator(self):