Update tailification pass to handle return-types optimizationwip-tailify

The return-types optimization can produce $callk that continues to
$kargs, not $kreceive, indicating that the number of returned values is
known.  Therefore knowing that a cont is $kargs isn't enough to know
what kind it is: it could be a join or it could be a return
continuation.  Therefore update to determine continuation type based on
predecessors.  Once we have the tails (and know which ones are joins),
extend tailify-tails for known-arity returns.

* module/language/cps/tailify.scm (tailify-tail): Update to detect tail
end via predecessor, and to transform appropriately.
(tailify-tails): Remove the variables bound by the head (if any: can
only apply to join conts and known-arity returns) from the saved set.
Handle $kargs differently depending on whether it is a join cont or a
return cont.
(compute-tails): Identify which heads are joins.
(tailify-function): Pass join-cont set through to tailify-tails.

1 files changed, 149 insertions, 95 deletions

diff --git a/module/language/cps/tailify.scm b/module/language/cps/tailify.scm
index 9840c8cfe..31cf4581d 100644
--- a/module/language/cps/tailify.scm
+++ b/module/language/cps/tailify.scm
@@ -222,40 +222,40 @@ be rewritten to continue to the tail's ktail."
            (with-cps cps
              (build-term ($continue k src ,exp))))
           (else
-           (match (intmap-ref cps k)
-             (($ $kreceive)
+           (match exp
+             ((or ($ $call) ($ $callk) ($ $calli))
               ;; A non-tail-call: push the pending continuation and tail
               ;; call instead.
               ;; (pk 'non-tail-call head k exp)
-              (match exp
-                ((or ($ $call) ($ $callk) ($ $calli))
-                 (call-with-values (lambda ()
-                                     (compute-saved-vars fresh-names k))
-                   (lambda (reprs vars)
-                     ;; (pk 'saved-vars reprs vars)
-                     (with-cps cps
-                       (letk kexp ($kargs () ()
-                                    ($continue local-ktail src ,exp)))
-                       (letv cont)
-                       (letk kcont ($kargs ('cont) (cont)
-                                     ($continue kexp src
-                                       ($primcall 'save
-                                                  (append reprs (list 'ptr))
-                                                  ,(append vars (list cont))))))
-                       (build-term ($continue kcont src
-                                     ($code (intmap-ref entries k))))))))))
-             (($ $kargs names vars)
+              (call-with-values (lambda ()
+                                  (compute-saved-vars fresh-names k))
+                (lambda (reprs vars)
+                  ;; (pk 'saved-vars reprs vars)
+                  (with-cps cps
+                    (letk kexp ($kargs () ()
+                                 ($continue local-ktail src ,exp)))
+                    (letv cont)
+                    (letk kcont ($kargs ('cont) (cont)
+                                  ($continue kexp src
+                                    ($primcall 'save
+                                               (append reprs (list 'ptr))
+                                               ,(append vars (list cont))))))
+                    (build-term ($continue kcont src
+                                  ($code (intmap-ref entries k))))))))
+             (_
               ;; Calling a join continuation.  This is one of those
               ;; cases where it might be nice in CPS to have names for
               ;; phi predecessor values.  Ah well.
-              (let ((vars' (map (lambda (_) (fresh-var)) vars)))
-                (with-cps cps
-                  (letk kvals
-                        ($kargs names vars'
-                          ($continue local-ktail src
-                            ,(compute-join-call (map cons vars vars') k))))
-                  (build-term
-                    ($continue kvals src ,exp))))))))))
+              (match (intmap-ref cps k)
+                (($ $kargs names vars)
+                 (let ((vars' (map (lambda (_) (fresh-var)) vars)))
+                   (with-cps cps
+                     (letk kvals
+                           ($kargs names vars'
+                             ($continue local-ktail src
+                               ,(compute-join-call (map cons vars vars') k))))
+                     (build-term
+                       ($continue kvals src ,exp))))))))))))
       (($ $branch kf kt src op param args)
        (with-cps cps
          (let$ kf (rewrite-branch-target src kf))
@@ -323,11 +323,12 @@ be rewritten to continue to the tail's ktail."
           (setk label ($kargs names vals ,term))))))
    body cps))
 
-(define (tailify-tails cps winds live-in constants reprs tails)
+(define (tailify-tails cps winds live-in constants reprs tails joins)
   "Given that the conts in a function were partitioned into tails in the
-intmap TAILS, mapping tail entries to tail bodies, return a new CPS
-program in which the tails have been split to separate functions in
-which all calls are tail calls.
+intmap TAILS, mapping tail entries to tail bodies, of which the intset
+JOINS indicates join continuations, return a new CPS program in which
+the tails have been split to separate functions in which all calls are
+tail calls.
 
 WINDS associates 'unwind primcalls with the corresponding conts that
 pushes on the dynamic stack.
@@ -349,19 +350,32 @@ REPRS holds the representation of each var."
          (($ $prompt k kh src) src)
          (($ $throw src) src)))))
 
+  ;; Compute the set of vars that we need to save for each head, which
+  ;; excludes the vars bound by the head cont itself.
+  (define heads-live-in
+    (intmap-map
+     (lambda (head body)
+       (let ((live (intmap-ref live-in head)))
+         (match (intmap-ref cps head)
+           (($ $kargs names vars)
+            (fold1 (lambda (var live) (intset-remove live var))
+                   vars live))
+           (_ live))))
+     tails))
+
   ;; For live values that flow into a tail, each tail will need to give
   ;; them unique names.
   (define fresh-names-per-tail
     (intmap-map (lambda (head body)
                   (intset-fold (lambda (var fresh)
                                  (intmap-add fresh var (fresh-var)))
-                               (intmap-ref live-in head)
+                               (intmap-ref heads-live-in head)
                                empty-intmap))
                 tails))
 
   (define (compute-saved-vars head)
     (compute-saved-vars* (intmap-ref fresh-names-per-tail head)
-                         live-in constants reprs head))
+                         heads-live-in constants reprs head))
 
   ;; For a tail whose head in the source program is HEAD, rewrite to be
   ;; a $kfun.  For the "main" tail, no change needed.  For join tails,
@@ -381,7 +395,7 @@ REPRS holds the representation of each var."
                        (letk k ($kargs ('const) (var') ,term))
                        ($ (values (intset-add body k)
                                   (build-term ($continue k #f ,exp))))))
-                   (live-constants live-in constants head)
+                   (live-constants heads-live-in constants head)
                    cps body term))
     (define (restore-saved cps body term)
       (call-with-values (lambda () (compute-saved-vars head))
@@ -395,59 +409,88 @@ REPRS holds the representation of each var."
                 ($ (values (intset-add body krestore)
                            (build-term ($continue krestore #f
                                          ($primcall 'restore reprs ()))))))))))
-    (match (intmap-ref cps head)
-      (($ $kfun)
-       ;; The main entry.
-       (values cps (intmap-add entries head head) tails))
-      (($ $kreceive ($ $arity req () rest () #f) kargs)
-       ;; The continuation of a non-tail call, or a prompt handler.
-       (match (intmap-ref cps kargs)
-         (($ $kargs names vars)
-          (let ((vars' (map (lambda (_) (fresh-var)) vars))
-                (src (cont-source kargs)))
-            (let*-values (((cps body term)
-                           (restore-constants
-                            cps
-                            body
-                            (build-term
-                              ($continue kargs src ($values vars')))))
-                          ((cps body term) (restore-saved cps body term)))
-              (with-cps cps
-                (letk ktail ($ktail))
-                (letk krestore ($kargs names vars' ,term))
-                (letk kclause ($kclause (req '() rest '() #f) krestore #f))
-                (letk kfun ($kfun src '() #f ktail kclause))
-                ($ (values
-                    (intmap-add entries head kfun)
-                    (let ((added (intset kfun kclause krestore ktail))
-                          (removed (intset head)))
-                      (intmap-add (intmap-remove tails head)
-                                  kfun
-                                  (intset-subtract (intset-union body added)
-                                                   removed)))))))))))
-      (($ $kargs names vars term)
-       ;; A join point.
-       (call-with-values (lambda () (compute-saved-vars head))
-         (lambda (reprs vars')
-           (define names'
-             (let ((names (map cons vars names)))
-               (map (lambda (var) (assq-ref names var))
-                    vars')))
-           (define meta `((arg-representations . ,reprs)))
-           (let*-values (((cps body term)
-                          (restore-constants cps body term)))
+    (cond
+     ((intset-ref joins head)
+      ;; A join point.
+      (match (intmap-ref cps head)
+        (($ $kargs names vars term)
+         (call-with-values (lambda () (compute-saved-vars head))
+           (lambda (reprs vars')
+             (define names'
+               (let ((names (map cons vars names)))
+                 (map (lambda (var) (assq-ref names var))
+                      vars')))
+             (define meta `((arg-representations . ,reprs)))
+             (let*-values (((cps body term)
+                            (restore-constants cps body term)))
+               (with-cps cps
+                 (letk ktail ($ktail))
+                 (letk kargs ($kargs names' vars' ,term))
+                 (letk kfun ($kfun (cont-source head) meta #f ktail kargs))
+                 ($ (values
+                     (intmap-add entries head kfun)
+                     (let ((added (intset kfun kargs ktail))
+                           (removed (intset head)))
+                       (intmap-add (intmap-remove tails head)
+                                   kfun
+                                   (intset-subtract (intset-union body added)
+                                                    removed))))))))))))
+     (else
+      (match (intmap-ref cps head)
+        (($ $kfun)
+         ;; The main entry.
+         (values cps (intmap-add entries head head) tails))
+        (($ $kreceive ($ $arity req () rest () #f) kargs)
+         ;; The continuation of a non-tail call, or a prompt handler.
+         ;; In either case we don't know the return arity of the caller
+         ;; so we have to parse the return values count.
+         (match (intmap-ref cps kargs)
+           (($ $kargs names vars)
+            (let ((vars' (map (lambda (_) (fresh-var)) vars))
+                  (src (cont-source kargs)))
+              (let*-values (((cps body term)
+                             (restore-constants
+                              cps
+                              body
+                              (build-term
+                                ($continue kargs src ($values vars')))))
+                            ((cps body term) (restore-saved cps body term)))
+                (with-cps cps
+                  (letk ktail ($ktail))
+                  (letk krestore ($kargs names vars' ,term))
+                  (letk kclause ($kclause (req '() rest '() #f) krestore #f))
+                  (letk kfun ($kfun src '() #f ktail kclause))
+                  ($ (values
+                      (intmap-add entries head kfun)
+                      (let ((added (intset kfun kclause krestore ktail))
+                            (removed (intset head)))
+                        (intmap-add (intmap-remove tails head)
+                                    kfun
+                                    (intset-subtract (intset-union body added)
+                                                     removed)))))))))))
+        (($ $kargs names vars term)
+         ;; The continuation of a known-return-arity call, from the
+         ;; return-types optimization.
+         (let ((vars' (map (lambda (_) (fresh-var)) vars))
+               (src (cont-source head)))
+           (let*-values (((cps body restore-term)
+                          (restore-constants
+                           cps
+                           body
+                           (build-term
+                             ($continue head src ($values vars')))))
+                         ((cps body restore-term)
+                          (restore-saved cps body restore-term)))
              (with-cps cps
                (letk ktail ($ktail))
-               (letk kargs ($kargs names' vars' ,term))
-               (letk kfun ($kfun (cont-source head) meta #f ktail kargs))
+               (letk kentry ($kargs names vars' ,restore-term))
+               (letk kfun ($kfun src '() #f ktail kentry))
                ($ (values
                    (intmap-add entries head kfun)
-                   (let ((added (intset kfun kargs ktail))
-                         (removed (intset head)))
+                   (let ((added (intset kfun kentry ktail)))
                      (intmap-add (intmap-remove tails head)
                                  kfun
-                                 (intset-subtract (intset-union body added)
-                                                  removed))))))))))))
+                                 (intset-union body added)))))))))))))
 
   (define original-ktail
     (match (intmap-ref cps (intmap-next tails))
@@ -462,7 +505,7 @@ REPRS holds the representation of each var."
        (lambda (old-head head cps)
          (define fresh-names (intmap-ref fresh-names-per-tail old-head))
          (define body (intmap-ref tails head))
-         (tailify-tail cps head body fresh-names winds live-in constants
+         (tailify-tail cps head body fresh-names winds heads-live-in constants
                        reprs entries original-ktail))
        entries cps))))
 
@@ -470,12 +513,17 @@ REPRS holds the representation of each var."
   "Compute the set of tails in the function with entry KFUN and body
 BODY. Return as an intset mapping the head label for each tail to its
 body, as an intset."
-  ;; Initially, we start with the requirement that kfun and kreceive
-  ;; labels are split heads.
+  ;; Initially, we start with the requirement that kfun and
+  ;; continuations of non-tail calls are split heads.
   (define (initial-split label splits)
     (match (intmap-ref cps label)
-      ((or ($ $kfun) ($ $kreceive))
+      (($ $kfun)
        (intmap-add splits label label))
+      (($ $kargs names vars
+          ($ $continue k src (or ($ $call) ($ $callk) ($ $calli))))
+       (match (intmap-ref cps k)
+         (($ $ktail) splits)
+         ((or ($ $kargs) ($ $kreceive)) (intmap-add splits k k))))
       (_
        splits)))
   ;; Then we build tails by propagating splits forward in the CFG,
@@ -516,8 +564,9 @@ body, as an intset."
   ;; (pk initial-splits)
   (cond
    ((trivial-intmap initial-splits)
-    ;; There's only one split head, so only one tail.
-    (intmap-add empty-intmap kfun body))
+    ;; There's only one split head, so only one tail, and no joins.
+    (values (intmap-add empty-intmap kfun body)
+            empty-intset))
    (else
     ;; Otherwise, assign each label to a tail, identified by the split
     ;; head, then collect the tails by split head.
@@ -526,11 +575,16 @@ body, as an intset."
                      ;; (pk 'fixpoint splits)
                      (intset-fold compute-split body splits))
                    initial-splits)))
-      (intmap-fold
-       (lambda (label head split-bodies)
-         (intmap-add split-bodies head (intset label) intset-union))
-       splits
-       empty-intmap)))))
+      (values
+       (intmap-fold
+        (lambda (label head split-bodies)
+          (intmap-add split-bodies head (intset label) intset-union))
+        splits
+        empty-intmap)
+       (intset-subtract (intmap-fold (lambda (label head heads)
+                                       (intset-add heads head))
+                                     splits empty-intset)
+                        (intmap-keys initial-splits)))))))
 
 (define (intset-pop set)
   "Return two values: all values in intset SET except the first one, and
@@ -647,7 +701,7 @@ tails in such a way that they enter via a $kfun and leave only via tail
 calls."
   (define succs (compute-successors cps kfun))
   (define preds (invert-graph succs))
-  (define tails (compute-tails kfun body preds cps))
+  (define-values (tails joins) (compute-tails kfun body preds cps))
   ;; (pk 'tails tails)
   (cond
    ((trivial-intmap tails)
@@ -658,7 +712,7 @@ calls."
           (live-in (compute-live-in cps body preds))
           (constants (compute-constants cps preds))
           (reprs (compute-var-representations cps)))
-      (tailify-tails cps winds live-in constants reprs tails)))))
+      (tailify-tails cps winds live-in constants reprs tails joins)))))
 
 (define (tailify cps)
   ;; Renumber so that label order is topological order.