Merge pull request #2152 from Shaikh-Ubaid/unsigned_bit_operators

Support unsigned bitshift operators

Author: certik

integration_tests/CMakeLists.txt

@@ -566,6 +566,7 @@ RUN(NAME test_unary_op_03    LABELS cpython llvm c wasm) # unary bitinvert
 RUN(NAME test_unary_op_04    LABELS cpython llvm c) # unary bitinvert
 RUN(NAME test_unary_op_05    LABELS cpython llvm c) # unsigned unary minus, plus
 RUN(NAME test_unary_op_06    LABELS cpython llvm c) # unsigned unary bitnot
+RUN(NAME test_unsigned_01    LABELS cpython llvm c) # unsigned bitshift left, right
 RUN(NAME test_bool_binop     LABELS cpython llvm c)
 RUN(NAME test_issue_518      LABELS cpython llvm c NOFAST)
 RUN(NAME structs_01          LABELS cpython llvm c)

integration_tests/test_unsigned_01.py

@@ -0,0 +1,37 @@
+from lpython import u8, u16, u32, u64
+
+def f():
+
+    h: u8
+    h = u8(5)
+    print(h << u8(4), h << u8(7), h >> u8(4), h >> u8(7))
+    assert h << u8(4) == u8(80)
+    assert h << u8(7) == u8(128)
+    assert h >> u8(4) == u8(0)
+    assert h >> u8(7) == u8(0)
+
+    i: u16
+    i = u16(67)
+    print(i << u16(4), i << u16(7), i >> u16(4), i >> u16(7))
+    assert i << u16(4) == u16(1072)
+    assert i << u16(7) == u16(8576)
+    assert i >> u16(4) == u16(4)
+    assert i >> u16(7) == u16(0)
+
+    j: u32
+    j = u32(25)
+    print(j << u32(4), j << u32(7), j >> u32(4), j >> u32(7))
+    assert j << u32(4) == u32(400)
+    assert j << u32(7) == u32(3200)
+    assert j >> u32(4) == u32(1)
+    assert j >> u32(7) == u32(0)
+
+    k: u64
+    k = u64(100000000000123)
+    print(k << u64(4), k << u64(7), k >> u64(4), k >> u64(7))
+    assert k << u64(4) == u64(1600000000001968)
+    assert k << u64(7) == u64(12800000000015744)
+    assert k >> u64(4) == u64(6250000000007)
+    assert k >> u64(7) == u64(781250000000)
+
+f()

src/libasr/codegen/asr_to_c_cpp.h

@@ -2190,6 +2190,8 @@ PyMODINIT_FUNC PyInit_lpython_module_)" + fn_name + R"((void) {
 
     void visit_UnsignedIntegerBinOp(const ASR::UnsignedIntegerBinOp_t &x) {
         handle_BinOp(x);
+        int kind = ASRUtils::extract_kind_from_ttype_t(x.m_type);
+        src = "(uint" + std::to_string(kind * 8) + "_t)(" + src + ")";
     }
 
     void visit_RealBinOp(const ASR::RealBinOp_t &x) {

src/runtime/lpython/lpython.py

@@ -139,6 +139,18 @@ def __ge__(self, other):
         else:
             raise TypeError("Unsupported operand type")
 
+    def __lshift__(self, other):
+        if isinstance(other, self.__class__):
+            return UnsignedInteger(self.bit_width, self.value << other.value)
+        else:
+            raise TypeError("Unsupported operand type")
+
+    def __rshift__(self, other):
+        if isinstance(other, self.__class__):
+            return UnsignedInteger(self.bit_width, self.value >> other.value)
+        else:
+            raise TypeError("Unsupported operand type")
+
     # conversion to integer
     def __int__(self):
         return self.value