Remove nested begin0's added by sfs pass

pkgs/racket-test-core/tests/racket/optimize.rktl

@@ -4270,6 +4270,35 @@
                  (eval (read (open-input-bytes (get-output-bytes o2)))))
                exn:fail:read?))
 
+;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
+;; make sure sfs pass doesn't add a nested begin0
+;; to clear the variables used in the first expression
+
+(let ()
+  (define c
+    '(module c racket/base
+       (define z (let ([result (random)])
+                   (begin0 (lambda () result) (newline))))))
+
+  (define o (open-output-bytes))
+
+  (parameterize ([current-namespace (make-base-namespace)])
+    (write (compile c) o))
+
+  (define m (zo-parse (open-input-bytes (get-output-bytes o))))
+
+  ; extract the content of the begin0 expression
+  (define (analyze-beg0 m)
+    (define def-z (car (mod-body (compilation-top-code m))))
+    (define body-z (let-one-body (def-values-rhs def-z)))
+    (define expr-z (car (beg0-seq body-z)))
+    (cond
+      [(lam? expr-z) 'ok]
+      [(beg0? expr-z) 'not-reduced-beg0-in-sfs]
+      [else 'unexpected]))
+
+  (test 'ok (lambda () (analyze-beg0 m))))
+
 ;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; make sure `begin0' propertly propagates "multiple results" flags
 

racket/src/racket/src/sfs.c

@@ -1023,6 +1023,44 @@ static Scheme_Object *bangboxenv_sfs(Scheme_Object *data, SFS_Info *info)
   }
 }
 
+static Scheme_Object *flatten_begin0(Scheme_Object *o)
+{
+  /* At this point, we sometimes have (begin0 (begin0 (begin0  ...) ...)).
+     Flatten those out. */
+  Scheme_Sequence *s = (Scheme_Sequence *)o, *s2;
+  int i, extra = 0;
+
+  o = s->array[0];
+
+  while (SAME_TYPE(SCHEME_TYPE(o), scheme_begin0_sequence_type)) {
+    s2 = (Scheme_Sequence *)o;
+    extra += s2->count - 1;
+    o = s2->array[0];
+  }
+
+  if (extra) {
+    s2 = scheme_malloc_sequence(s->count + extra);
+    s2->so.type = scheme_begin0_sequence_type;
+    s2->count = s->count + extra;
+
+    extra = s2->count -1;
+    o = (Scheme_Object *)s;
+    while (SAME_TYPE(SCHEME_TYPE(o), scheme_begin0_sequence_type)) {
+      s = (Scheme_Sequence *)o;
+      for (i = s->count - 1; i ; i--) {
+        s2->array[extra--] = s->array[i];
+      }
+      o = s->array[i];
+    }
+    s2->array[extra--] = o;
+
+    if (extra != -1) scheme_signal_error("internal error: flatten begin0 failed");
+
+    return (Scheme_Object *)s2;
+  } else
+    return (Scheme_Object *)s;
+}
+
 static Scheme_Object *
 begin0_sfs(Scheme_Object *obj, SFS_Info *info)
 {
@@ -1038,6 +1076,9 @@ begin0_sfs(Scheme_Object *obj, SFS_Info *info)
     ((Scheme_Sequence *)obj)->array[i] = le;
   }
 
+  if (info->pass)
+    obj = flatten_begin0(obj);
+
   return obj;
 }