[pypy-commit] pypy bigint-with-int-ops: Support add and sub ops

[stian]

Author: stian
Branch: bigint-with-int-ops
Changeset: r66022:ec59bbb3636b
Date: 2013-08-08 20:35 +0200
http://bitbucket.org/pypy/pypy/changeset/ec59bbb3636b/

Log:	Support add and sub ops

diff --git a/pypy/objspace/std/longobject.py b/pypy/objspace/std/longobject.py
--- a/pypy/objspace/std/longobject.py
+++ b/pypy/objspace/std/longobject.py
@@ -192,9 +192,15 @@
 def add__Long_Long(space, w_long1, w_long2):
     return W_LongObject(w_long1.num.add(w_long2.num))
 
+def add__Long_Int(space, w_long1, w_int2):
+    return W_LongObject(w_long1.num.int_add(w_int2.intval))
+
 def sub__Long_Long(space, w_long1, w_long2):
     return W_LongObject(w_long1.num.sub(w_long2.num))
 
+def sub__Long_Int(space, w_long1, w_int2):
+    return W_LongObject(w_long1.num.int_sub(w_int2.intval))
+
 def mul__Long_Long(space, w_long1, w_long2):
     return W_LongObject(w_long1.num.mul(w_long2.num))
 
diff --git a/rpython/rlib/rbigint.py b/rpython/rlib/rbigint.py
--- a/rpython/rlib/rbigint.py
+++ b/rpython/rlib/rbigint.py
@@ -44,6 +44,23 @@
 MASK = int((1 << SHIFT) - 1)
 FLOAT_MULTIPLIER = float(1 << SHIFT)
 
+# For BIGINT and INT mix.
+#
+# The VALID range of an int is different than a valid range of a bigint of length one.
+# -1 << LONG_BIT is actually TWO digits, because they are stored without the sign.
+if SHIFT == LONG_BIT - 1:
+    MIN_INT_VALUE = -1 << SHIFT
+    def int_in_valid_range(x):
+        if x == MIN_INT_VALUE:
+            return False
+        return True
+else:
+    # Means we don't have INT128 on 64bit.
+    def int_in_valid_range(x):
+        if x > MASK or x < -MASK:
+            return False
+        return True
+
 # Debugging digit array access.
 #
 # False == no checking at all
@@ -551,6 +568,25 @@
         return result
 
     @jit.elidable
+    def int_add(self, other):
+        if not int_in_valid_range(other):
+            # Fallback to long.
+            return self.add(rbigint.fromint(other))
+        elif self.sign == 0:
+            return rbigint.fromint(other)
+        elif other == 0:
+            return self
+
+        sign = -1 if other < 0 else 1
+        if self.sign == sign:
+            result = _x_int_add(self, other)
+        else:
+            result = _x_int_sub(self, other)
+            result.sign *= -1
+        result.sign *= sign
+        return result
+
+    @jit.elidable
     def sub(self, other):
         if other.sign == 0:
             return self
@@ -564,6 +600,22 @@
         return result
 
     @jit.elidable
+    def int_sub(self, other):
+        if not int_in_valid_range(other):
+            # Fallback to long.
+            return self.sub(rbigint.fromint(other))
+        elif other == 0:
+            return self
+        elif self.sign == 0:
+            return rbigint.fromint(-other)
+        elif self.sign == (-1 if other < 0 else 1):
+            result = _x_int_sub(self, other)
+        else:
+            result = _x_int_add(self, other)
+        result.sign *= self.sign
+        return result
+
+    @jit.elidable
     def mul(self, b):
         asize = self.numdigits()
         bsize = b.numdigits()
@@ -1129,6 +1181,25 @@
     z._normalize()
     return z
 
+def _x_int_add(a, b):
+    """ Add the absolute values of one bigint and one integer. """
+    size_a = a.numdigits()
+
+    z = rbigint([NULLDIGIT] * (size_a + 1), 1)
+    i = UDIGIT_TYPE(1)
+    carry = a.udigit(0) + abs(b)
+    z.setdigit(0, carry)
+    carry >>= SHIFT
+
+    while i < size_a:
+        carry += a.udigit(i)
+        z.setdigit(i, carry)
+        carry >>= SHIFT
+        i += 1
+    z.setdigit(i, carry)
+    z._normalize()
+    return z
+
 def _x_sub(a, b):
     """ Subtract the absolute values of two integers. """
 
@@ -1175,6 +1246,42 @@
     z._normalize()
     return z
 
+def _x_int_sub(a, b):
+    """ Subtract the absolute values of two integers. """
+
+    size_a = a.numdigits()
+
+    bdigit = abs(b)
+
+    if size_a == 1:
+        # Find highest digit where a and b differ:
+        adigit = a.digit(0)
+
+        if adigit == bdigit:
+            return NULLRBIGINT
+    
+        return rbigint.fromint(adigit - bdigit)
+
+    z = rbigint([NULLDIGIT] * size_a, 1, size_a)
+    i = _load_unsigned_digit(1)
+    # The following assumes unsigned arithmetic
+    # works modulo 2**N for some N>SHIFT.
+    borrow = a.udigit(0) - bdigit
+    z.setdigit(0, borrow)
+    borrow >>= SHIFT
+    #borrow &= 1 # Keep only one sign bit
+
+    while i < size_a:
+        borrow = a.udigit(i) - borrow
+        z.setdigit(i, borrow)
+        borrow >>= SHIFT
+        #borrow &= 1
+        i += 1
+
+    assert borrow == 0
+    z._normalize()
+    return z
+
 # A neat little table of power of twos.
 ptwotable = {}
 for x in range(SHIFT-1):
diff --git a/rpython/rlib/test/test_rbigint.py b/rpython/rlib/test/test_rbigint.py
--- a/rpython/rlib/test/test_rbigint.py
+++ b/rpython/rlib/test/test_rbigint.py
@@ -237,6 +237,16 @@
                 result = f1.add(f2)
                 assert result.tolong() == x * i + y * j
 
+    def test_int_add(self):
+        x = 123456789123456789000000L
+        y = 9999
+        for i in [-1, 1]:
+            for j in [-1, 1]:
+                f1 = rbigint.fromlong(x * i)
+                f2 = y * j
+                result = f1.int_add(f2)
+                assert result.tolong() == x * i + y * j
+
     def test_sub(self):
         x = 12378959520302182384345L
         y = 88961284756491823819191823L
@@ -247,6 +257,16 @@
                 result = f1.sub(f2)
                 assert result.tolong() == x * i - y * j
 
+    def test_int_sub(self):
+        x = 12378959520302182384345L
+        y = 8888
+        for i in [-1, 1]:
+            for j in [-1, 1]:
+                f1 = rbigint.fromlong(x * i)
+                f2 = y * j
+                result = f1.int_sub(f2)
+                assert result.tolong() == x * i - y * j
+
     def test_subzz(self):
         w_l0 = rbigint.fromint(0)
         assert w_l0.sub(w_l0).tolong() == 0