[pypy-commit] pypy math-improvements: Make rshift invert (in most cases) in place, this makes a huge speedup for rshift with negative numbers as it avoids two extra copies, also make an rqshift for the power of twos

[stian]

Author: stian
Branch: math-improvements
Changeset: r92929:f30c2f38b0b5
Date: 2017-11-04 19:18 +0100
http://bitbucket.org/pypy/pypy/changeset/f30c2f38b0b5/

Log:	Make rshift invert (in most cases) in place, this makes a huge
	speedup for rshift with negative numbers as it avoids two extra
	copies, also make an rqshift for the power of twos

diff --git a/rpython/rlib/rbigint.py b/rpython/rlib/rbigint.py
--- a/rpython/rlib/rbigint.py
+++ b/rpython/rlib/rbigint.py
@@ -787,7 +787,7 @@
             if digit == 1:
                 return rbigint(self._digits[:self.numdigits()], 1, self.numdigits())
             elif digit and digit & (digit - 1) == 0:
-                return self.rshift(ptwotable[digit])
+                return self.rqshift(ptwotable[digit])
 
         div, mod = _divrem(self, other)
         if mod.sign * other.sign == -1:
@@ -816,7 +816,7 @@
             if digit == 1:
                 return self
             elif digit & (digit - 1) == 0:
-                return self.rshift(ptwotable[digit])
+                return self.rqshift(ptwotable[digit])
             
         div, mod = _divrem1(self, digit)
 
@@ -1267,31 +1267,85 @@
             raise ValueError("negative shift count")
         elif int_other == 0:
             return self
+        invert = False
         if self.sign == -1 and not dont_invert:
-            a = self.invert().rshift(int_other)
-            return a.invert()
+            first = self.digit(0)
+            if first == 0:
+                a = self.invert().rshift(int_other)
+                return a.invert()
+            invert = True
 
         wordshift = int_other / SHIFT
+        loshift = int_other % SHIFT
         newsize = self.numdigits() - wordshift
         if newsize <= 0:
-            return NULLRBIGINT
-
-        loshift = int_other % SHIFT
+            if invert:
+                return ONENEGATIVERBIGINT
+            else:
+                return NULLRBIGINT
+
+        
         hishift = SHIFT - loshift
         z = rbigint([NULLDIGIT] * newsize, self.sign, newsize)
         i = 0
         while i < newsize:
-            newdigit = (self.udigit(wordshift) >> loshift)
+            digit = self.udigit(wordshift)
+            if i == 0 and invert and wordshift == 0:
+                digit -= 1
+            newdigit = (digit >> loshift)
             if i+1 < newsize:
                 newdigit |= (self.udigit(wordshift+1) << hishift)
             z.setdigit(i, newdigit)
             i += 1
             wordshift += 1
+        if invert:
+            z.setdigit(0, z.digit(0)+1)
         z._normalize()
         return z
     rshift._always_inline_ = 'try' # It's so fast that it's always benefitial.
 
     @jit.elidable
+    def rqshift(self, int_other):
+        wordshift = int_other / SHIFT
+        loshift = int_other % SHIFT
+        newsize = self.numdigits() - wordshift
+
+        invert = False
+        if self.sign == -1:
+            first = self.digit(0)
+            if first == 0:
+                a = self.invert().rqshift(int_other)
+                return a.invert()
+            invert = True
+            
+        if newsize <= 0:
+            if invert:
+                return ONENEGATIVERBIGINT
+            else:
+                return NULLRBIGINT
+                
+        
+        hishift = SHIFT - loshift
+        z = rbigint([NULLDIGIT] * newsize, self.sign, newsize)
+        i = 0
+        inverted = False
+        while i < newsize:
+            digit = self.udigit(wordshift)
+            if invert and i == 0 and wordshift == 0:
+                digit -= 1
+            newdigit = (digit >> loshift)
+            if i+1 < newsize:
+                newdigit |= (self.udigit(wordshift+1) << hishift)
+            z.setdigit(i, newdigit)
+            i += 1
+            wordshift += 1
+        if invert:
+            z.setdigit(0, z.digit(0)+1)      
+        z._normalize()
+        return z
+    rshift._always_inline_ = 'try' # It's so fast that it's always benefitial.
+    
+    @jit.elidable
     def abs_rshift_and_mask(self, bigshiftcount, mask):
         assert isinstance(bigshiftcount, r_ulonglong)
         assert mask >= 0
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
@@ -598,6 +598,33 @@
                     res3 = f1.abs_rshift_and_mask(r_ulonglong(y), mask)
                     assert res3 == (abs(x) >> y) & mask
 
+    def test_qshift(self):
+        for x in range(10):
+            for y in range(1, 161, 16):
+                num = (x << y) + x
+                f1 = rbigint.fromlong(num)
+                nf1 = rbigint.fromlong(-num)
+                
+                for z in range(1, 31):
+                    res1 = f1.lqshift(z).tolong() 
+                    res2 = f1.rqshift(z).tolong() 
+                    res3 = nf1.lqshift(z).tolong() 
+                    res4 = nf1.rqshift(z).tolong() 
+                    
+                    assert res1 == num << z
+                    assert res2 == num >> z
+                    assert res3 == -num << z
+                    assert res4 == -num >> z
+                    
+        # Large digit
+        for x in range((1 << SHIFT) - 10, (1 << SHIFT) + 10):
+            f1 = rbigint.fromlong(x)
+            nf1 = rbigint.fromlong(-x)
+            assert f1.rqshift(SHIFT).tolong() == x >> SHIFT 
+            assert nf1.rqshift(SHIFT).tolong() == -x >> SHIFT
+            assert f1.rqshift(SHIFT+1).tolong() == x >> (SHIFT+1)
+            assert nf1.rqshift(SHIFT+1).tolong() == -x >> (SHIFT+1)
+                    
     def test_from_list_n_bits(self):
         for x in ([3L ** 30L, 5L ** 20L, 7 ** 300] +
                   [1L << i for i in range(130)] +