python · JukkaL · Oct 6, 2020 · Aug 9, 2020 · Sep 26, 2020 · Sep 26, 2020
diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py
@@ -558,33 +558,35 @@ def matching_primitive_op(self,
     def binary_op(self,
                   lreg: Value,
                   rreg: Value,
-                  expr_op: str,
+                  op: str,
                   line: int) -> Value:
-        # special case tuple comparison here so that nested tuples can be supported
-        if (isinstance(lreg.type, RTuple) and isinstance(rreg.type, RTuple)
-                and expr_op in ('==', '!=')):
-            return self.compare_tuples(lreg, rreg, expr_op, line)
-        # Special case == and != when we can resolve the method call statically.
-        value = None
-        if expr_op in ('==', '!='):
-            value = self.translate_eq_cmp(lreg, rreg, expr_op, line)
-        if value is not None:
-            return value
-
-        # Special case 'is' and 'is not'
-        if expr_op in ('is', 'is not'):
-            return self.translate_is_op(lreg, rreg, expr_op, line)
-
-        if (is_str_rprimitive(lreg.type) and is_str_rprimitive(rreg.type)
-                and expr_op in ('==', '!=')):
-            return self.compare_strings(lreg, rreg, expr_op, line)
-
-        if is_tagged(lreg.type) and is_tagged(rreg.type) and expr_op in int_comparison_op_mapping:
-            return self.compare_tagged(lreg, rreg, expr_op, line)
-
-        call_c_ops_candidates = c_binary_ops.get(expr_op, [])
+        ltype = lreg.type
+        rtype = rreg.type
+
+        # Special case tuple comparison here so that nested tuples can be supported
+        if isinstance(ltype, RTuple) and isinstance(rtype, RTuple) and op in ('==', '!='):
+            return self.compare_tuples(lreg, rreg, op, line)
+
+        # Special case == and != when we can resolve the method call statically
+        if op in ('==', '!='):
+            value = self.translate_eq_cmp(lreg, rreg, op, line)
+            if value is not None:
+                return value
+
+        # Special case various ops
+        if op in ('is', 'is not'):
+            return self.translate_is_op(lreg, rreg, op, line)
+        if is_str_rprimitive(ltype) and is_str_rprimitive(rtype) and op in ('==', '!='):
+            return self.compare_strings(lreg, rreg, op, line)
+        if is_tagged(ltype) and is_tagged(rtype) and op in int_comparison_op_mapping:
+            return self.compare_tagged(lreg, rreg, op, line)
+        if is_bool_rprimitive(ltype) and is_bool_rprimitive(rtype) and op in (
+                '&', '&=', '|', '|=', '^', '^='):
+            return self.bool_bitwise_op(lreg, rreg, op[0], line)
+
+        call_c_ops_candidates = c_binary_ops.get(op, [])
         target = self.matching_call_c(call_c_ops_candidates, [lreg, rreg], line)
-        assert target, 'Unsupported binary operation: %s' % expr_op
+        assert target, 'Unsupported binary operation: %s' % op
         return target
 
     def check_tagged_short_int(self, val: Value, line: int) -> Value:
@@ -710,6 +712,17 @@ def compare_tuples(self,
         self.goto_and_activate(out)
         return result
 
+    def bool_bitwise_op(self, lreg: Value, rreg: Value, op: str, line: int) -> Value:
+        if op == '&':
+            code = BinaryIntOp.AND
+        elif op == '|':
+            code = BinaryIntOp.OR
+        elif op == '^':
+            code = BinaryIntOp.XOR
+        else:
+            assert False, op
+        return self.add(BinaryIntOp(bool_rprimitive, lreg, rreg, code, line))
+
     def unary_not(self,
                   value: Value,
                   line: int) -> Value:

diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
@@ -129,11 +129,17 @@ void CPyTagged_IncRef(CPyTagged x);
 void CPyTagged_DecRef(CPyTagged x);
 void CPyTagged_XDecRef(CPyTagged x);
 CPyTagged CPyTagged_Negate(CPyTagged num);
+CPyTagged CPyTagged_Invert(CPyTagged num);
 CPyTagged CPyTagged_Add(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_Subtract(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_Multiply(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_FloorDivide(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_Remainder(CPyTagged left, CPyTagged right);
+CPyTagged CPyTagged_And(CPyTagged left, CPyTagged right);
+CPyTagged CPyTagged_Or(CPyTagged left, CPyTagged right);
+CPyTagged CPyTagged_Xor(CPyTagged left, CPyTagged right);
+CPyTagged CPyTagged_Rshift(CPyTagged left, CPyTagged right);
+CPyTagged CPyTagged_Lshift(CPyTagged left, CPyTagged right);
 bool CPyTagged_IsEq_(CPyTagged left, CPyTagged right);
 bool CPyTagged_IsLt_(CPyTagged left, CPyTagged right);
 PyObject *CPyTagged_Str(CPyTagged n);

diff --git a/mypyc/lib-rt/int_ops.c b/mypyc/lib-rt/int_ops.c
@@ -273,3 +273,211 @@ PyObject *CPyLong_FromFloat(PyObject *o) {
 PyObject *CPyBool_Str(bool b) {
     return PyObject_Str(b ? Py_True : Py_False);
 }
+
+static void CPyLong_NormalizeUnsigned(PyLongObject *v) {
+    Py_ssize_t i = v->ob_base.ob_size;
+    while (i > 0 && v->ob_digit[i - 1] == 0)
+        i--;
+    v->ob_base.ob_size = i;
+}
+
+// Bitwise op '&', '|' or '^' using the generic (slow) API
+static CPyTagged GenericBitwiseOp(CPyTagged a, CPyTagged b, char op) {
+    PyObject *aobj = CPyTagged_AsObject(a);
+    PyObject *bobj = CPyTagged_AsObject(b);
+    PyObject *r;
+    if (op == '&') {
+        r = PyNumber_And(aobj, bobj);
+    } else if (op == '|') {
+        r = PyNumber_Or(aobj, bobj);
+    } else {
+        r = PyNumber_Xor(aobj, bobj);
+    }
+    if (unlikely(r == NULL)) {
+        CPyError_OutOfMemory();
+    }
+    Py_DECREF(aobj);
+    Py_DECREF(bobj);
+    return CPyTagged_StealFromObject(r);
+}
+
+// Return pointer to digits of a PyLong object. If it's a short
+// integer, place digits in the buffer buf instead to avoid memory
+// allocation (it's assumed to be big enough). Return the number of
+// digits in *size. *size is negative if the integer is negative.
+static digit *GetIntDigits(CPyTagged n, Py_ssize_t *size, digit *buf) {
+    if (CPyTagged_CheckShort(n)) {
+        Py_ssize_t val = CPyTagged_ShortAsSsize_t(n);
+        bool neg = val < 0;
+        int len = 1;
+        if (neg) {
+            val = -val;
+        }
+        buf[0] = val & PyLong_MASK;
+        if (val > PyLong_MASK) {
+            val >>= PyLong_SHIFT;
+            buf[1] = val & PyLong_MASK;
+            if (val > PyLong_MASK) {
+                buf[2] = val >> PyLong_SHIFT;
+                len = 3;
+            } else {
+                len = 2;
+            }
+        }
+        *size = neg ? -len : len;
+        return buf;
+    } else {
+        PyLongObject *obj = (PyLongObject *)CPyTagged_LongAsObject(n);
+        *size = obj->ob_base.ob_size;
+        return obj->ob_digit;
+    }
+}
+
+// Shared implementation of bitwise '&', '|' and '^' (specified by op) for at least
+// one long operand. This is somewhat optimized for performance.
+static CPyTagged BitwiseLongOp(CPyTagged a, CPyTagged b, char op) {
+    // Directly access the digits, as there is no fast C API function for this.
+    digit abuf[3];
+    digit bbuf[3];
+    Py_ssize_t asize;
+    Py_ssize_t bsize;
+    digit *adigits = GetIntDigits(a, &asize, abuf);
+    digit *bdigits = GetIntDigits(b, &bsize, bbuf);
+
+    PyLongObject *r;
+    if (unlikely(asize < 0 || bsize < 0)) {
+        // Negative operand. This is slower, but bitwise ops on them are pretty rare.
+        return GenericBitwiseOp(a, b, op);
+    }
+    // Optimized implementation for two non-negative integers.
+    // Swap a and b as needed to ensure a is no longer than b.
+    if (asize > bsize) {
+        digit *tmp = adigits;
+        adigits = bdigits;
+        bdigits = tmp;
+        Py_ssize_t tmp_size = asize;
+        asize = bsize;
+        bsize = tmp_size;
+    }
+    r = _PyLong_New(op == '&' ? asize : bsize);
+    if (unlikely(r == NULL)) {
+        CPyError_OutOfMemory();
+    }
+    Py_ssize_t i;
+    if (op == '&') {
+        for (i = 0; i < asize; i++) {
+            r->ob_digit[i] = adigits[i] & bdigits[i];
+        }
+    } else {
+        if (op == '|') {
+            for (i = 0; i < asize; i++) {
+                r->ob_digit[i] = adigits[i] | bdigits[i];
+            }
+        } else {
+            for (i = 0; i < asize; i++) {
+                r->ob_digit[i] = adigits[i] ^ bdigits[i];
+            }
+        }
+        for (; i < bsize; i++) {
+            r->ob_digit[i] = bdigits[i];
+        }
+    }
+    CPyLong_NormalizeUnsigned(r);
+    return CPyTagged_StealFromObject((PyObject *)r);
+}
+
+// Bitwise '&'
+CPyTagged CPyTagged_And(CPyTagged left, CPyTagged right) {
+    if (likely(CPyTagged_CheckShort(left) && CPyTagged_CheckShort(right))) {
+        return left & right;
+    }
+    return BitwiseLongOp(left, right, '&');
+}
+
+// Bitwise '|'
+CPyTagged CPyTagged_Or(CPyTagged left, CPyTagged right) {
+    if (likely(CPyTagged_CheckShort(left) && CPyTagged_CheckShort(right))) {
+        return left | right;
+    }
+    return BitwiseLongOp(left, right, '|');
+}
+
+// Bitwise '^'
+CPyTagged CPyTagged_Xor(CPyTagged left, CPyTagged right) {
+    if (likely(CPyTagged_CheckShort(left) && CPyTagged_CheckShort(right))) {
+        return left ^ right;
+    }
+    return BitwiseLongOp(left, right, '^');
+}
+
+// Bitwise '~'
+CPyTagged CPyTagged_Invert(CPyTagged num) {
+    if (likely(CPyTagged_CheckShort(num) && num != CPY_TAGGED_ABS_MIN)) {
+        return ~num & ~CPY_INT_TAG;
+    } else {
+        PyObject *obj = CPyTagged_AsObject(num);
+        PyObject *result = PyNumber_Invert(obj);
+        if (unlikely(result == NULL)) {
+            CPyError_OutOfMemory();
+        }
+        Py_DECREF(obj);
+        return CPyTagged_StealFromObject(result);
+    }
+}
+
+// Bitwise '>>'
+CPyTagged CPyTagged_Rshift(CPyTagged left, CPyTagged right) {
+    if (likely(CPyTagged_CheckShort(left)
+               && CPyTagged_CheckShort(right)
+               && (Py_ssize_t)right >= 0)) {
+        CPyTagged count = CPyTagged_ShortAsSsize_t(right);
+        if (unlikely(count >= CPY_INT_BITS)) {
+            if ((Py_ssize_t)left >= 0) {
+                return 0;
+            } else {
+                return CPyTagged_ShortFromInt(-1);
+            }
+        }
+        return ((Py_ssize_t)left >> count) & ~CPY_INT_TAG;
+    } else {
+        // Long integer or negative shift -- use generic op
+        PyObject *lobj = CPyTagged_AsObject(left);
+        PyObject *robj = CPyTagged_AsObject(right);
+        PyObject *result = PyNumber_Rshift(lobj, robj);
+        Py_DECREF(lobj);
+        Py_DECREF(robj);
+        if (result == NULL) {
+            // Propagate error (could be negative shift count)
+            return CPY_INT_TAG;
+        }
+        return CPyTagged_StealFromObject(result);
+    }
+}
+
+static inline bool IsShortLshiftOverflow(Py_ssize_t short_int, Py_ssize_t shift) {
+    return ((Py_ssize_t)(short_int << shift) >> shift) != short_int;
+}
+
+// Bitwise '<<'
+CPyTagged CPyTagged_Lshift(CPyTagged left, CPyTagged right) {
+    if (likely(CPyTagged_CheckShort(left)
+               && CPyTagged_CheckShort(right)
+               && (Py_ssize_t)right >= 0
+               && right < CPY_INT_BITS * 2)) {
+        CPyTagged shift = CPyTagged_ShortAsSsize_t(right);
+        if (!IsShortLshiftOverflow(left, shift))
+            // Short integers, no overflow
+            return left << shift;
+    }
+    // Long integer or out of range shift -- use generic op
+    PyObject *lobj = CPyTagged_AsObject(left);
+    PyObject *robj = CPyTagged_AsObject(right);
+    PyObject *result = PyNumber_Lshift(lobj, robj);
+    Py_DECREF(lobj);
+    Py_DECREF(robj);
+    if (result == NULL) {
+        // Propagate error (could be negative shift count)
+        return CPY_INT_TAG;
+    }
+    return CPyTagged_StealFromObject(result);
+}
diff --git a/mypyc/primitives/int_ops.py b/mypyc/primitives/int_ops.py
@@ -84,19 +84,30 @@ def int_binary_op(name: str, c_function_name: str,
 int_binary_op('+', 'CPyTagged_Add')
 int_binary_op('-', 'CPyTagged_Subtract')
 int_binary_op('*', 'CPyTagged_Multiply')
+int_binary_op('&', 'CPyTagged_And')
+int_binary_op('|', 'CPyTagged_Or')
+int_binary_op('^', 'CPyTagged_Xor')
 # Divide and remainder we honestly propagate errors from because they
 # can raise ZeroDivisionError
 int_binary_op('//', 'CPyTagged_FloorDivide', error_kind=ERR_MAGIC)
 int_binary_op('%', 'CPyTagged_Remainder', error_kind=ERR_MAGIC)
+# Negative shift counts raise an exception
+int_binary_op('>>', 'CPyTagged_Rshift', error_kind=ERR_MAGIC)
+int_binary_op('<<', 'CPyTagged_Lshift', error_kind=ERR_MAGIC)
 
 # This should work because assignment operators are parsed differently
 # and the code in irbuild that handles it does the assignment
 # regardless of whether or not the operator works in place anyway.
 int_binary_op('+=', 'CPyTagged_Add')
 int_binary_op('-=', 'CPyTagged_Subtract')
 int_binary_op('*=', 'CPyTagged_Multiply')
+int_binary_op('&=', 'CPyTagged_And')
+int_binary_op('|=', 'CPyTagged_Or')
+int_binary_op('^=', 'CPyTagged_Xor')
 int_binary_op('//=', 'CPyTagged_FloorDivide', error_kind=ERR_MAGIC)
 int_binary_op('%=', 'CPyTagged_Remainder', error_kind=ERR_MAGIC)
+int_binary_op('>>=', 'CPyTagged_Rshift', error_kind=ERR_MAGIC)
+int_binary_op('<<=', 'CPyTagged_Lshift', error_kind=ERR_MAGIC)
 
 
 def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
@@ -108,6 +119,7 @@ def int_unary_op(name: str, c_function_name: str) -> CFunctionDescription:
 
 
 int_neg_op = int_unary_op('-', 'CPyTagged_Negate')
+int_invert_op = int_unary_op('~', 'CPyTagged_Invert')
 
 # integer comparsion operation implementation related:
 

diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -37,6 +37,12 @@ def __floordiv__(self, x: int) -> int: pass
     def __mod__(self, x: int) -> int: pass
     def __neg__(self) -> int: pass
     def __pos__(self) -> int: pass
+    def __invert__(self) -> int: pass
+    def __and__(self, n: int) -> int: pass
+    def __or__(self, n: int) -> int: pass
+    def __xor__(self, n: int) -> int: pass
+    def __lshift__(self, x: int) -> int: pass
+    def __rshift__(self, x: int) -> int: pass
     def __eq__(self, n: object) -> bool: pass
     def __ne__(self, n: object) -> bool: pass
     def __lt__(self, n: int) -> bool: pass
@@ -89,9 +95,20 @@ def __eq__(self, x:object) -> bool:pass
     def __ne__(self, x: object) -> bool: pass
     def join(self, x: Iterable[object]) -> bytes: pass
 
-class bool:
+class bool(int):
     def __init__(self, o: object = ...) -> None: ...
-
+    @overload
+    def __and__(self, n: bool) -> bool: ...
+    @overload
+    def __and__(self, n: int) -> int: ...
+    @overload
+    def __or__(self, n: bool) -> bool: ...
+    @overload
+    def __or__(self, n: int) -> int: ...
+    @overload
+    def __xor__(self, n: bool) -> bool: ...
+    @overload
+    def __xor__(self, n: int) -> int: ...
 
 class tuple(Generic[T_co], Sequence[T_co], Iterable[T_co]):
     def __init__(self, i: Iterable[T_co]) -> None: pass
@@ -231,6 +248,7 @@ def enumerate(x: Iterable[T]) -> Iterator[Tuple[int, T]]: ...
 def zip(x: Iterable[T], y: Iterable[S]) -> Iterator[Tuple[T, S]]: ...
 @overload
 def zip(x: Iterable[T], y: Iterable[S], z: Iterable[V]) -> Iterator[Tuple[T, S, V]]: ...
+def eval(e: str) -> Any: ...
 
 # Dummy definitions.
 class classmethod: pass