From f4a90589410725c899311b52546f7e2958b4f15c Mon Sep 17 00:00:00 2001
From: Matthew Flatt <mflatt@racket-lang.org>
Date: Mon, 26 Apr 2021 20:03:55 -0600
Subject: [PATCH] stream-{ref,tail,take}: avoid retaining stream for potential
 error

Although retaining the original stream argument to `stream-ref`,
`stream-tail`, or `stream-take` can enable a better error message if
the stream runs out of elements too soon, it can also interfere with
the intended memory use of a stream.

Closes #2870
---
 .../racket-test-core/tests/racket/stream.rktl | 23 +++++++++++++++++++
 racket/collects/racket/stream.rkt             | 17 ++++++++++----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/pkgs/racket-test-core/tests/racket/stream.rktl b/pkgs/racket-test-core/tests/racket/stream.rktl
index 6da79501aa..6b8bd2d8dc 100644
--- a/pkgs/racket-test-core/tests/racket/stream.rktl
+++ b/pkgs/racket-test-core/tests/racket/stream.rktl
@@ -194,4 +194,27 @@
                       (stream-lazy
                        (stream-lazy '(1)))))
 
+;; Make sure certain operations that could encounter a too-short stream don't
+;; retain the original stream just in case of the error:
+(unless (eq? 'cgc (system-type 'gc))
+  (let ([check (lambda (op)
+                 (define s (stream-cons
+                            1
+                            (stream-cons
+                             2
+                             (begin
+                               (collect-garbage)
+                               (let ([v (weak-box-value wb)])
+                                 (stream-cons
+                                  3
+                                  (stream-cons
+                                   v
+                                   empty)))))))
+                 (define wb (make-weak-box s))
+                 (test #f 'check-stream-no-retain (op s 3)))])
+    (check stream-ref)
+    (check (lambda (s n) (stream-first (stream-tail s n))))
+    (check (lambda (s n) (stream-ref (stream-take s (add1 n)) n)))))
+
+
 (report-errs)
diff --git a/racket/collects/racket/stream.rkt b/racket/collects/racket/stream.rkt
index 5baeda37c9..3632342c0a 100644
--- a/racket/collects/racket/stream.rkt
+++ b/racket/collects/racket/stream.rkt
@@ -103,7 +103,14 @@
       (raise-arguments-error 'stream-ref
                              "stream ended before index"
                              "index" i
-                             "stream" st)]
+                             ;; Why `"stream" st` is omitted:
+                             ;; including `st` in the error message
+                             ;; means that it has to be kept live;
+                             ;; that's not so great for a stream, where
+                             ;; lazy construction could otherwise allow
+                             ;; a element to be reached without consuming
+                             ;; proportional memory
+                             #;"stream" #;st)]
      [(zero? n)
       (stream-first s)]
      [else
@@ -120,12 +127,11 @@
       (raise-arguments-error 'stream-tail
                              "stream ended before index"
                              "index" i
-                             "stream" st)]
+                             ;; See "Why `"stream" st` is omitted" above
+                             #;"stream" #;st)]
      [else
       (loop (sub1 n) (stream-rest s))])))
 
-
-
 (define (stream-take st i)
   (unless (stream? st) (raise-argument-error 'stream-take "stream?" st))
   (unless (exact-nonnegative-integer? i)
@@ -137,7 +143,8 @@
       (raise-arguments-error 'stream-take
                              "stream ended before index"
                              "index" i
-                             "stream" st)]
+                             ;; See "Why `"stream" st` is omitted" above
+                             #;"stream" #;st)]
      [else
       (make-do-stream (lambda () #f)
                       (lambda () (stream-first s))