@@ -4688,13 +4688,23 @@ divmod_shift(PyObject *shiftby, Py_ssize_t *wordshift, digit *remshift)
4688
4688
return 0 ;
4689
4689
}
4690
4690
4691
+ /* Inner function for both long_rshift and _PyLong_Rshift, shifting an
4692
+ integer right by PyLong_SHIFT*wordshift + remshift bits.
4693
+ wordshift should be nonnegative. */
4694
+
4691
4695
static PyObject *
4692
4696
long_rshift1 (PyLongObject * a , Py_ssize_t wordshift , digit remshift )
4693
4697
{
4694
4698
PyLongObject * z = NULL ;
4695
- Py_ssize_t newsize , hishift , i , j ;
4699
+ Py_ssize_t newsize , hishift , size_a ;
4696
4700
twodigits accum ;
4701
+ int a_negative ;
4702
+
4703
+ /* Total number of bits shifted must be nonnegative. */
4704
+ assert (wordshift >= 0 );
4705
+ assert (remshift < PyLong_SHIFT );
4697
4706
4707
+ /* Fast path for small a. */
4698
4708
if (IS_MEDIUM_VALUE (a )) {
4699
4709
stwodigits m , x ;
4700
4710
digit shift ;
@@ -4704,37 +4714,67 @@ long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
4704
4714
return _PyLong_FromSTwoDigits (x );
4705
4715
}
4706
4716
4707
- if (Py_SIZE (a ) < 0 ) {
4708
- /* Right shifting negative numbers is harder */
4709
- PyLongObject * a1 , * a2 ;
4710
- a1 = (PyLongObject * ) long_invert (a );
4711
- if (a1 == NULL )
4712
- return NULL ;
4713
- a2 = (PyLongObject * ) long_rshift1 (a1 , wordshift , remshift );
4714
- Py_DECREF (a1 );
4715
- if (a2 == NULL )
4716
- return NULL ;
4717
- z = (PyLongObject * ) long_invert (a2 );
4718
- Py_DECREF (a2 );
4717
+ a_negative = Py_SIZE (a ) < 0 ;
4718
+ size_a = Py_ABS (Py_SIZE (a ));
4719
+
4720
+ if (a_negative ) {
4721
+ /* For negative 'a', adjust so that 0 < remshift <= PyLong_SHIFT,
4722
+ while keeping PyLong_SHIFT*wordshift + remshift the same. This
4723
+ ensures that 'newsize' is computed correctly below. */
4724
+ if (remshift == 0 ) {
4725
+ if (wordshift == 0 ) {
4726
+ /* Can only happen if the original shift was 0. */
4727
+ return long_long ((PyObject * )a );
4728
+ }
4729
+ remshift = PyLong_SHIFT ;
4730
+ -- wordshift ;
4731
+ }
4719
4732
}
4720
- else {
4721
- newsize = Py_SIZE (a ) - wordshift ;
4722
- if (newsize <= 0 )
4723
- return PyLong_FromLong (0 );
4724
- hishift = PyLong_SHIFT - remshift ;
4725
- z = _PyLong_New (newsize );
4726
- if (z == NULL )
4727
- return NULL ;
4728
- j = wordshift ;
4729
- accum = a -> ob_digit [j ++ ] >> remshift ;
4730
- for (i = 0 ; j < Py_SIZE (a ); i ++ , j ++ ) {
4731
- accum |= (twodigits )a -> ob_digit [j ] << hishift ;
4732
- z -> ob_digit [i ] = (digit )(accum & PyLong_MASK );
4733
- accum >>= PyLong_SHIFT ;
4733
+
4734
+ assert (wordshift >= 0 );
4735
+ newsize = size_a - wordshift ;
4736
+ if (newsize <= 0 ) {
4737
+ /* Shifting all the bits of 'a' out gives either -1 or 0. */
4738
+ return PyLong_FromLong (- a_negative );
4739
+ }
4740
+ z = _PyLong_New (newsize );
4741
+ if (z == NULL ) {
4742
+ return NULL ;
4743
+ }
4744
+ hishift = PyLong_SHIFT - remshift ;
4745
+
4746
+ accum = a -> ob_digit [wordshift ];
4747
+ if (a_negative ) {
4748
+ /*
4749
+ For a positive integer a and nonnegative shift, we have:
4750
+
4751
+ (-a) >> shift == -((a + 2**shift - 1) >> shift).
4752
+
4753
+ In the addition `a + (2**shift - 1)`, the low `wordshift` digits of
4754
+ `2**shift - 1` all have value `PyLong_MASK`, so we get a carry out
4755
+ from the bottom `wordshift` digits when at least one of the least
4756
+ significant `wordshift` digits of `a` is nonzero. Digit `wordshift`
4757
+ of `2**shift - 1` has value `PyLong_MASK >> hishift`.
4758
+ */
4759
+ Py_SET_SIZE (z , - newsize );
4760
+
4761
+ digit sticky = 0 ;
4762
+ for (Py_ssize_t j = 0 ; j < wordshift ; j ++ ) {
4763
+ sticky |= a -> ob_digit [j ];
4734
4764
}
4735
- z -> ob_digit [i ] = (digit )accum ;
4736
- z = maybe_small_long (long_normalize (z ));
4765
+ accum += (PyLong_MASK >> hishift ) + (digit )(sticky != 0 );
4737
4766
}
4767
+
4768
+ accum >>= remshift ;
4769
+ for (Py_ssize_t i = 0 , j = wordshift + 1 ; j < size_a ; i ++ , j ++ ) {
4770
+ accum += (twodigits )a -> ob_digit [j ] << hishift ;
4771
+ z -> ob_digit [i ] = (digit )(accum & PyLong_MASK );
4772
+ accum >>= PyLong_SHIFT ;
4773
+ }
4774
+ assert (accum <= PyLong_MASK );
4775
+ z -> ob_digit [newsize - 1 ] = (digit )accum ;
4776
+
4777
+ z = maybe_small_long (long_normalize (z ));
4738
4778
return (PyObject * )z ;
4739
4779
}
4740
4780
0 commit comments