py: Fix rshift and not of zero/one edge cases in mpz.

Addresses issue #1027.

2 files changed, 29 insertions, 2 deletions

diff --git a/py/mpz.c b/py/mpz.c
index fc1109fdd..6213b4a57 100644
--- a/py/mpz.c
+++ b/py/mpz.c
@@ -871,11 +871,17 @@ void mpz_not_inpl(mpz_t *dest, const mpz_t *z) {
     if (dest != z) {
         mpz_set(dest, z);
     }
-    if (dest->neg) {
+    if (dest->len == 0) {
+        mpz_need_dig(dest, 1);
+        dest->dig[0] = 1;
+        dest->len = 1;
+        dest->neg = 1;
+    } else if (dest->neg) {
         dest->neg = 0;
         mpz_dig_t k = 1;
         dest->len = mpn_sub(dest->dig, dest->dig, dest->len, &k, 1);
     } else {
+        mpz_need_dig(dest, dest->len + 1);
         mpz_dig_t k = 1;
         dest->len = mpn_add(dest->dig, dest->dig, dest->len, &k, 1);
         dest->neg = 1;
@@ -924,7 +930,14 @@ void mpz_shr_inpl(mpz_t *dest, const mpz_t *lhs, mp_int_t rhs) {
                 round_up = 1;
             }
             if (round_up) {
-                dest->len = mpn_add(dest->dig, dest->dig, dest->len, &round_up, 1);
+                if (dest->len == 0) {
+                    // dest == 0, so need to add 1 by hand (answer will be -1)
+                    dest->dig[0] = 1;
+                    dest->len = 1;
+                } else {
+                    // dest > 0, so can use mpn_add to add 1
+                    dest->len = mpn_add(dest->dig, dest->dig, dest->len, &round_up, 1);
+                }
             }
         }
     }
diff --git a/tests/basics/int_big_zeroone.py b/tests/basics/int_big_zeroone.py
new file mode 100644
index 000000000..e467714bd
--- /dev/null
+++ b/tests/basics/int_big_zeroone.py
@@ -0,0 +1,14 @@
+# test [0,-0,1,-1] edge cases of bignum
+
+long_zero = (2**64) >> 65
+long_neg_zero = -long_zero
+long_one = long_zero + 1
+long_neg_one = -long_one
+
+cases = [long_zero, long_neg_zero, long_one, long_neg_one]
+
+print(cases)
+print([-c for c in cases])
+print([~c for c in cases])
+print([c >> 1 for c in cases])
+print([c << 1 for c in cases])