Starlark: specialize StarlarkInt.shift* for longs

Closes #12623. PiperOrigin-RevId: 345953110
bazelbuild · Dec 6, 2020 · f8e45f4 · f8e45f4
1 parent 46468e4
commit f8e45f4
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/src/main/java/net/starlark/java/eval/StarlarkInt.java b/src/main/java/net/starlark/java/eval/StarlarkInt.java
@@ -570,12 +570,14 @@ public static StarlarkInt shiftRight(StarlarkInt x, StarlarkInt y) throws EvalEx
     if (yi < 0) {
       throw Starlark.errorf("negative shift count: %d", yi);
     }
-    if (x instanceof Int32) {
-      long xl = ((Int32) x).v;
-      if (yi >= Integer.SIZE) {
+    try {
+      long xl = x.toLongFast();
+      if (yi >= Long.SIZE) {
         return xl < 0 ? StarlarkInt.of(-1) : ZERO;
       }
       return StarlarkInt.of(xl >> yi);
+    } catch (Overflow unused) {
+      /* fall through */
     }
 
     BigInteger xbig = x.toBigInteger();
@@ -591,13 +593,15 @@ public static StarlarkInt shiftLeft(StarlarkInt x, StarlarkInt y) throws EvalExc
     } else if (yi >= 512) {
       throw Starlark.errorf("shift count too large: %d", yi);
     }
-    if (x instanceof Int32) {
-      long xl = ((Int32) x).v;
+    try {
+      long xl = x.toLongFast();
       long z = xl << yi; // only uses low 6 bits of yi
       if ((z >> yi) == xl && yi < 64) {
         return StarlarkInt.of(z);
       }
       /* overflow */
+    } catch (Overflow unused) {
+      /* fall through */
     }
 
     BigInteger xbig = x.toBigInteger();

diff --git a/src/test/java/net/starlark/java/eval/testdata/int.star b/src/test/java/net/starlark/java/eval/testdata/int.star
@@ -223,6 +223,12 @@ assert_eq(~maxlong, minlong)
 1 % 0  ### integer modulo by zero
 
 ---
+# Not using << to define constants because we are testing <<
+maxint = 0x7fffffff
+maxlong = 0x7fffffffffffffff
+minint = -0x80000000
+minlong = -0x8000000000000000
+
 assert_eq(1 | 2, 3)
 assert_eq(3 | 6, 7)
 assert_eq(7 | 0, 7)
@@ -243,9 +249,29 @@ assert_eq(~2147483647, -2147483647 - 1);
 assert_eq(1 << 2, 4)
 assert_eq(7 << 0, 7)
 assert_eq(-1 << 31, -2147483647 - 1)
+assert_eq(1 << 31, maxint + 1)
+assert_eq(1 << 32, (maxint + 1) * 2)
+assert_eq(1 << 63, maxlong + 1)
+assert_eq(1 << 64, (maxlong + 1) * 2)
+assert_eq(-1 << 31, minint)
+assert_eq(-1 << 32, minint * 2)
+assert_eq(-1 << 63, minlong)
+assert_eq(-1 << 64, minlong * 2)
 assert_eq(2 >> 1, 1)
 assert_eq(7 >> 0, 7)
 assert_eq(0 >> 0, 0)
+assert_eq(minint >> 9999, -1)
+assert_eq(minlong >> 9999, -1)
+assert_eq(maxint >> 9999, 0)
+assert_eq(maxlong >> 9999, 0)
+assert_eq(minint >> 31, -1)
+assert_eq(minint >> 30, -2)
+assert_eq(minlong >> 63, -1)
+assert_eq(minlong >> 62, -2)
+assert_eq(maxint >> 31, 0)
+assert_eq(maxint >> 30, 1)
+assert_eq(maxlong >> 63, 0)
+assert_eq(maxlong >> 62, 1)
 assert_eq(1000 >> 100, 0)
 assert_eq(-10 >> 1000, -1)
 assert_eq(1 << 500 >> 499, 2)