schemify: adjust simple? for multiple returns

The check for whether an expression is "simple" at the schemify layer
didn't take into account the expected arity.

racket/src/cs/primitive/kernel.ss

@@ -685,7 +685,7 @@
   [pseudo-random-generator-vector? (known-procedure/no-prompt 2)]
   [pseudo-random-generator? (known-procedure/pure/folding 2)]
   [quotient (known-procedure/folding 4)]
-  [quotient/remainder (known-procedure/no-prompt 4)]
+  [quotient/remainder (known-procedure 4)] ; no `no-prompt` due to multiple return values
   [raise (known-procedure 6)]
   [raise-argument-error (known-procedure -8)]
   [raise-arguments-error (known-procedure -4)]

racket/src/schemify/mutated.rkt

@@ -143,7 +143,8 @@
             (find-mutated! rhs (unwrap-list ids))
             (define new-maybe-cc? (or maybe-cc?
                                       (not (simple? rhs prim-knowns knowns imports mutated simples
-                                                    #:pure? #f))))
+                                                    #:pure? #f
+                                                    #:result-arity (length ids)))))
             ;; Each `id` in `ids` is now ready (but might also hold a delay):
             (for ([id (in-wrap-list ids)])
               (let ([u-id (unwrap id)])

racket/src/schemify/schemify.rkt

@@ -343,7 +343,8 @@
               (finish-wrapped-definition (list id) rhs)])]
           [`(define-values ,ids ,rhs)
            (cond
-             [(simple? #:pure? #f rhs prim-knowns knowns imports mutated simples)
+             [(simple? #:pure? #f rhs prim-knowns knowns imports mutated simples
+                       #:result-arity (length ids))
               (match rhs
                 [`(values ,rhss ...)
                  ;; Flatten `(define-values (id ...) (values rhs ...))` to
@@ -378,7 +379,8 @@
               (finish-definition ids (append set-vars accum-exprs) null)]
              [`,_
               (cond
-                [(simple? #:pure? #f schemified prim-knowns knowns imports mutated simples)
+                [(simple? #:pure? #f schemified prim-knowns knowns imports mutated simples
+                          #:result-arity #f)
                  (loop (cdr l) mut-l (cons schemified accum-exprs) accum-ids knowns)]
                 [else
                  ;; In case `schemified` triggers an error, sync exported variables
@@ -581,7 +583,7 @@
               (authentic-valued? key knowns prim-knowns imports mutated))
             (cond
               [(and authentic-key?
-                    (simple? s-body prim-knowns knowns imports mutated simples))
+                    (simple? s-body prim-knowns knowns imports mutated simples #:result-arity #f))
                `(begin ,(ensure-single-valued s-key knowns prim-knowns imports mutated)
                        ,(ensure-single-valued s-val knowns prim-knowns imports mutated)
                        ,s-body)]

racket/src/schemify/simple.rkt

@@ -12,54 +12,79 @@
 ;; of evaluation isn't detectable. This function receives both
 ;; schemified and non-schemified expressions.
 (define (simple? e prim-knowns knowns imports mutated simples
-                 #:pure? [pure? #t])
+                 #:pure? [pure? #t]
-  (let simple? ([e e])
+                 #:result-arity [result-arity 1])
+  (let simple? ([e e] [result-arity result-arity])
     (define-syntax-rule (cached expr)
-      (let* ([c (hash-ref simples e '(unknown . unknown))]
+      (let* ([c (hash-ref simples e #(unknown unknown 1))]
-             [r (if pure? (car c) (cdr c))])
+             [r (vector-ref c (if pure? 0 1))]
-        (if (eq? 'unknown r)
+             [arity-match? (eqv? result-arity (vector-ref c 2))])
+        (if (or (eq? 'unknown r)
+                (not arity-match?))
             (let ([r expr])
-              (hash-set! simples e (if pure? (cons r (cdr c)) (cons (car c) r)))
+              (hash-set! simples e (if pure?
+                                       (vector r
+                                               (if arity-match? (vector-ref c 1) 'unknown)
+                                               result-arity)
+                                       (vector (if arity-match? (vector-ref c 0) 'unknown)
+                                               r
+                                               result-arity)))
               r)
             r)))
+    (define (returns n)
+      (or (not result-arity)
+          (eqv? n result-arity)))
     (match e
-      [`(lambda . ,_) #t]
+      [`(lambda . ,_) (returns 1)]
-      [`(case-lambda . ,_) #t]
+      [`(case-lambda . ,_) (returns 1)]
-      [`(quote . ,_) #t]
+      [`(quote . ,_) (returns 1)]
-      [`(#%variable-reference . ,_) #t]
+      [`(#%variable-reference . ,_) (returns 1)]
-      [`(let-values ([,_ ,rhss] ...) ,body)
+      [`(let-values ([,idss ,rhss] ...) ,body)
        (cached
-        (and (for/and ([rhs (in-list rhss)])
+        (and (for/and ([ids (in-list idss)]
-               (simple? rhs))
+                       [rhs (in-list rhss)])
-             (simple? body)))]
+               (simple? rhs (length ids)))
+             (simple? body result-arity)))]
       [`(let ([,_ ,rhss] ...) ,body)
        (cached
         (and (for/and ([rhs (in-list rhss)])
-               (simple? rhs))
+               (simple? rhs 1))
-             (simple? body)))]
+             (simple? body result-arity)))]
       [`(letrec-values ([(,idss ...) ,rhss] ...) ,body)
        (cached
-        (and (for/and ([rhs (in-list rhss)])
+        (and (for/and ([ids (in-list idss)]
-               (simple? rhs))
+                       [rhs (in-list rhss)])
-             (simple? body)))]
+               (simple? rhs (length ids)))
+             (simple? body result-arity)))]
       [`(letrec* ([,ids ,rhss] ...) ,body)
        (cached
         (and (for/and ([rhs (in-list rhss)])
-               (simple? rhs))
+               (simple? rhs 1))
-             (simple? body)))]
+             (simple? body result-arity)))]
       [`(begin ,es ...)
        #:guard (not pure?)
        (cached
+        (let loop ([es es])
+          (cond
+            [(null? (cdr es))
+             (simple? (car es) result-arity)]
+            [else
+             (and (simple? (car es) #f)
+                  (loop (cdr es)))])))]
+      [`(begin0 ,e0 ,es ...)
+       (cached
+        (and (simple? e0 result-arity)
              (for/and ([e (in-list es)])
-          (simple? e)))]
+               (simple? e #f))))]
       [`(set! ,_ ,e)
        #:guard (not pure?)
-       (simple? e)]
+       (simple? e 1)
+       (returns 1)]
       [`(values ,es ...)
-       #:guard (not pure?)
        (cached
+        (and (returns (length es))
              (for/and ([e (in-list es)])
-          (simple? e)))]
+               (simple? e 1))))]
       [`(,proc . ,args)
        (cached
         (let ([proc (unwrap proc)])
@@ -67,21 +92,24 @@
                (let ([v (or (hash-ref-either knowns imports proc)
                             (hash-ref prim-knowns proc #f))])
                  (and (if pure?
-                          (known-procedure/pure? v)
+                          (and (known-procedure/pure? v)
-                          (known-procedure/no-prompt? v))
+                               (returns 1))
+                          (and (known-procedure/no-prompt? v)
+                               (eqv? result-arity #f)))
                       (bitwise-bit-set? (known-procedure-arity-mask v) (length args))))
                (simple-mutated-state? (hash-ref mutated proc #f))
                (for/and ([arg (in-list args)])
-                 (simple? arg)))))]
+                 (simple? arg 1)))))]
       [`,_
        (let ([e (unwrap e)])
+         (and (returns 1)
               (or (and (symbol? e)
                        (simple-mutated-state? (hash-ref mutated e #f)))
                   (integer? e)
                   (boolean? e)
                   (string? e)
                   (bytes? e)
-             (regexp? e)))])))
+                  (regexp? e))))])))
 
 (define (simple/can-copy? e prim-knowns knowns imports mutated)
   (match e