From c48afdb16b5b3e6aac60c39ae65196eaf3d2d8b8 Mon Sep 17 00:00:00 2001
From: Matthew Flatt <mflatt@racket-lang.org>
Date: Sun, 26 Jan 2020 08:14:56 -0700
Subject: [PATCH] repairs for `arithmetic-shift`

Produce 0, -1, or out-of-memory for bignum shifts. For large fixnum
shifts, check memory limits.

The repairs are mostly for Racket CS, but traditional Racket
incorrectly reported out-of-memory for 0 shifted by a positive bignum.
---
 .../racket-test-core/tests/racket/number.rktl | 10 +++++--
 racket/src/cs/rumble/number.ss                | 29 ++++++++++++++++---
 racket/src/racket/src/number.c                |  2 ++
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/pkgs/racket-test-core/tests/racket/number.rktl b/pkgs/racket-test-core/tests/racket/number.rktl
index d8846db324..dc50978f9e 100644
--- a/pkgs/racket-test-core/tests/racket/number.rktl
+++ b/pkgs/racket-test-core/tests/racket/number.rktl
@@ -1289,14 +1289,20 @@
   (test 0 arithmetic-shift 42 (+ i (- (expt 2 31))))
   (test 0 arithmetic-shift 42 (+ i (- (expt 2 63)))))
 
+(test 0 arithmetic-shift 0 (expt 2 100))
+(test 0 arithmetic-shift 0 (- (expt 2 100)))
+(test 0 arithmetic-shift 1 (- (expt 2 100)))
+(test 0 arithmetic-shift (expt 2 100) (- (expt 2 100)))
+(test -1 arithmetic-shift -1 (- (expt 2 100)))
+(test -1 arithmetic-shift (- (expt 2 100)) (- (expt 2 100)))
+
 (arity-test arithmetic-shift 2 2)
 (err/rt-test (arithmetic-shift "a" 1))
 (err/rt-test (arithmetic-shift 1 "a"))
 (err/rt-test (arithmetic-shift 1.0 1))
 (err/rt-test (arithmetic-shift 1 1.0))
 (err/rt-test (arithmetic-shift 1 1.0+0.0i))
-(unless (eq? 'chez-scheme (system-type 'vm))
-  (err/rt-test (eval '(arithmetic-shift 1 (expt 2 80))) exn:fail:out-of-memory?))
+(err/rt-test (eval '(arithmetic-shift 1 (expt 2 80))) exn:fail:out-of-memory?)
 
 (test #f bitwise-bit-set? 13 1)
 (test #t bitwise-bit-set? 13 2)
diff --git a/racket/src/cs/rumble/number.ss b/racket/src/cs/rumble/number.ss
index 2208c473c1..c63ee90602 100644
--- a/racket/src/cs/rumble/number.ss
+++ b/racket/src/cs/rumble/number.ss
@@ -36,14 +36,35 @@
                        (#3%fxsra x max-fx-shift)
                        (#3%fxsra x (#3%fx- n)))
                    (if (#3%fx> n max-fx-shift)
-                       (#3%bitwise-arithmetic-shift x n)
+                       (general-arithmetic-shift x n)
                        (let ([m (#3%fxsll x n)])
                          (if (#3%fx= (#3%fxsra m n) x)
                              m
                              (#3%bitwise-arithmetic-shift x n))))))
-             (#2%bitwise-arithmetic-shift x n)))]
-    [(_ expr ...) #'(#2%bitwise-arithmetic-shift expr ...)]
-    [_ #'#2%bitwise-arithmetic-shift]))
+             (general-arithmetic-shift x n)))]
+    [(_ expr ...) #'(general-arithmetic-shift expr ...)]
+    [_ #'general-arithmetic-shift]))
+
+(define general-arithmetic-shift
+  (|#%name|
+   arithmetic-shift
+   (lambda (x n)
+     (cond
+      [(not (exact-integer? x))
+       (#2%bitwise-arithmetic-shift x n)]
+      [(fixnum? n)
+       (unless (or (eqv? x 0) (fx< n 1000))
+         (guard-large-allocation 'arithmetic-shift 'number n 1))
+       (#2%bitwise-arithmetic-shift x n)]
+      [(and (not (eqv? x 0))
+            (bignum? n)
+            (positive? n))
+       (raise (|#%app|
+               exn:fail:out-of-memory
+               "arithmetic-shift: out of memory"
+               (current-continuation-marks)))]
+      [else
+       (#2%bitwise-arithmetic-shift x n)]))))
 
 (define-syntax-rule (define-bitwise op fxop)
   (...
diff --git a/racket/src/racket/src/number.c b/racket/src/racket/src/number.c
index 5a5e5ca3df..661ec35476 100644
--- a/racket/src/racket/src/number.c
+++ b/racket/src/racket/src/number.c
@@ -4281,6 +4281,8 @@ scheme_bitwise_shift(int argc, Scheme_Object *argv[])
 	  return scheme_make_integer(-1);
 	else
 	  return scheme_make_integer(0);
+      } else if (SAME_OBJ(v, scheme_exact_zero)) {
+        return scheme_make_integer(0);
       } else
 	scheme_raise_out_of_memory("arithmetic-shift", NULL);
     } else