From a0329d01095d6ddaa42449ec18a4fb2bc83db16e Mon Sep 17 00:00:00 2001
From: Andy Wingo <wingo@pobox.com>
Date: Fri, 11 Apr 2014 14:01:27 +0200
Subject: [PATCH] Root higher-order CPS term is always $kfun $cont

* module/language/cps/arities.scm:
* module/language/cps/closure-conversion.scm:
* module/language/cps/compile-bytecode.scm:
* module/language/cps/constructors.scm:
* module/language/cps/contification.scm:
* module/language/cps/cse.scm:
* module/language/cps/dce.scm:
* module/language/cps/elide-values.scm:
* module/language/cps/prune-bailouts.scm:
* module/language/cps/prune-top-level-scopes.scm:
* module/language/cps/renumber.scm:
* module/language/cps/self-references.scm:
* module/language/cps/simplify.scm:
* module/language/cps/specialize-primcalls.scm:
* module/language/tree-il/compile-cps.scm: Adapt to produce and consume
  raw $kfun $cont instances.

* .dir-locals.el: Update $letrec indentation.
---
 .dir-locals.el                                |   1 +
 module/language/cps/arities.scm               |  24 +--
 module/language/cps/closure-conversion.scm    |  14 +-
 module/language/cps/compile-bytecode.scm      |   3 +-
 module/language/cps/constructors.scm          |  20 +-
 module/language/cps/contification.scm         |   8 +-
 module/language/cps/cse.scm                   |  27 +--
 module/language/cps/dce.scm                   |  36 ++--
 module/language/cps/elide-values.scm          |  22 +-
 module/language/cps/prune-bailouts.scm        |  27 +--
 .../language/cps/prune-top-level-scopes.scm   |   6 +-
 module/language/cps/renumber.scm              | 196 +++++++++---------
 module/language/cps/self-references.scm       |  13 +-
 module/language/cps/simplify.scm              |  15 +-
 module/language/cps/specialize-primcalls.scm  |   6 +-
 module/language/tree-il/compile-cps.scm       |  17 +-
 16 files changed, 212 insertions(+), 223 deletions(-)

diff --git a/.dir-locals.el b/.dir-locals.el
index d3cee5c80..0a2a26623 100644
--- a/.dir-locals.el
+++ b/.dir-locals.el
@@ -27,6 +27,7 @@
      (eval . (put '$continue           'scheme-indent-function 2))
      (eval . (put '$kargs              'scheme-indent-function 2))
      (eval . (put '$kfun               'scheme-indent-function 4))
+     (eval . (put '$letrec             'scheme-indent-function 3))
      (eval . (put '$kclause            'scheme-indent-function 1))
      (eval . (put '$fun                'scheme-indent-function 1))))
  (emacs-lisp-mode . ((indent-tabs-mode . nil)))
diff --git a/module/language/cps/arities.scm b/module/language/cps/arities.scm
index d0491fc5c..c18955809 100644
--- a/module/language/cps/arities.scm
+++ b/module/language/cps/arities.scm
@@ -32,7 +32,7 @@
   #:use-module (language cps primitives)
   #:export (fix-arities))
 
-(define (fix-clause-arities clause dfg)
+(define (fix-arities* clause dfg)
   (let ((ktail (match clause
                  (($ $cont _
                      ($ $kfun src meta _ ($ $cont ktail) _)) ktail))))
@@ -41,8 +41,12 @@
         (($ $letk conts body)
          ($letk ,(map visit-cont conts) ,(visit-term body)))
         (($ $letrec names syms funs body)
-         ($letrec names syms (map (cut fix-arities* <> dfg) funs)
-                  ,(visit-term body)))
+         ($letrec names syms (map (lambda (fun)
+                                    (rewrite-cps-exp fun
+                                      (($ $fun free body)
+                                       ($fun free ,(fix-arities* body dfg)))))
+                                  funs)
+           ,(visit-term body)))
         (($ $continue k src exp)
          ,(visit-exp k src exp))))
 
@@ -135,8 +139,9 @@
              ($ $prim)
              ($ $values (_)))
          ,(adapt-exp 1 k src exp))
-        (($ $fun)
-         ,(adapt-exp 1 k src (fix-arities* exp dfg)))
+        (($ $fun free body)
+         ,(adapt-exp 1 k src (build-cps-exp
+                               ($fun free ,(fix-arities* body dfg)))))
         ((or ($ $call) ($ $callk))
          ;; In general, calls have unknown return arity.  For that
          ;; reason every non-tail call has a $kreceive continuation to
@@ -185,14 +190,7 @@
       (($ $cont sym ($ $kfun src meta self tail clause))
        (sym ($kfun src meta self ,tail ,(and clause (visit-cont clause))))))))
 
-(define (fix-arities* fun dfg)
-  (rewrite-cps-exp fun
-    (($ $fun free body)
-     ($fun free ,(fix-clause-arities body dfg)))))
-
 (define (fix-arities fun)
-  (let ((dfg (match fun
-               (($ $fun free body)
-                (compute-dfg body)))))
+  (let ((dfg (compute-dfg fun)))
     (with-fresh-name-state-from-dfg dfg
       (fix-arities* fun dfg))))
diff --git a/module/language/cps/closure-conversion.scm b/module/language/cps/closure-conversion.scm
index 151448ea0..08e511d4e 100644
--- a/module/language/cps/closure-conversion.scm
+++ b/module/language/cps/closure-conversion.scm
@@ -272,13 +272,11 @@ convert functions to flat closures."
     (($ $cont sym ($ $kfun src meta self tail clause))
      (sym ($kfun src meta self ,tail ,(and clause (visit-cont clause)))))))
 
-(define (convert-closures exp)
+(define (convert-closures fun)
   "Convert free reference in @var{exp} to primcalls to @code{free-ref},
 and allocate and initialize flat closures."
-  (match exp
-    (($ $fun () body)
-     (with-fresh-name-state body
-       (receive (body free) (cc body #f '())
-         (unless (null? free)
-           (error "Expected no free vars in toplevel thunk" exp body free))
-         (convert-to-indices body free))))))
+  (with-fresh-name-state fun
+    (receive (body free) (cc fun #f '())
+      (unless (null? free)
+        (error "Expected no free vars in toplevel thunk" fun body free))
+      (convert-to-indices body free))))
diff --git a/module/language/cps/compile-bytecode.scm b/module/language/cps/compile-bytecode.scm
index 20414a7b0..eb873d398 100644
--- a/module/language/cps/compile-bytecode.scm
+++ b/module/language/cps/compile-bytecode.scm
@@ -506,8 +506,7 @@
          (exp (optimize exp opts))
          (exp (convert-closures exp))
          (exp (reify-primitives exp))
-         (exp (match (renumber (build-cps-exp ($fun '() ,exp)))
-                (($ $fun free body) body)))
+         (exp (renumber exp))
          (asm (make-assembler)))
     (visit-funs (lambda (fun)
                   (compile-fun fun asm))
diff --git a/module/language/cps/constructors.scm b/module/language/cps/constructors.scm
index e4ab6a9d0..16de82569 100644
--- a/module/language/cps/constructors.scm
+++ b/module/language/cps/constructors.scm
@@ -47,8 +47,8 @@
        ($letk ,(map visit-cont conts)
          ,(visit-term body)))
       (($ $letrec names syms funs body)
-       ($letrec names syms (map inline-constructors* funs)
-                ,(visit-term body)))
+       ($letrec names syms (map visit-fun funs)
+         ,(visit-term body)))
       (($ $continue k src ($ $primcall 'list args))
        ,(let-fresh (kvalues) (val)
           (build-cps-term
@@ -90,16 +90,16 @@
                 ($continue kalloc src
                   ($primcall 'make-vector (len init))))))))
       (($ $continue k src (and fun ($ $fun)))
-       ($continue k src ,(inline-constructors* fun)))
+       ($continue k src ,(visit-fun fun)))
       (($ $continue)
        ,term)))
+  (define (visit-fun fun)
+    (rewrite-cps-exp fun
+      (($ $fun free body)
+       ($fun free ,(inline-constructors* body)))))
 
-  (rewrite-cps-exp fun
-    (($ $fun free body)
-     ($fun free ,(visit-cont body)))))
+  (visit-cont fun))
 
 (define (inline-constructors fun)
-  (match fun
-    (($ $fun free body)
-     (with-fresh-name-state body
-       (inline-constructors* fun)))))
+  (with-fresh-name-state fun
+    (inline-constructors* fun)))
diff --git a/module/language/cps/contification.scm b/module/language/cps/contification.scm
index 267a4d65c..dc832c338 100644
--- a/module/language/cps/contification.scm
+++ b/module/language/cps/contification.scm
@@ -39,9 +39,7 @@
   #:export (contify))
 
 (define (compute-contification fun)
-  (let* ((dfg (match fun
-                (($ $fun free body)
-                 (compute-dfg body))))
+  (let* ((dfg (compute-dfg fun))
          (scope-table (make-hash-table))
          (call-substs '())
          (cont-substs '())
@@ -294,7 +292,7 @@
                 (visit-fun exp)))
            (_ #t)))))
 
-    (visit-fun fun)
+    (visit-cont fun)
     (values call-substs cont-substs fun-elisions cont-splices)))
 
 (define (apply-contification fun call-substs cont-substs fun-elisions cont-splices)
@@ -401,7 +399,7 @@
            (or (contify-call src proc args)
                (continue k src exp)))
           (_ (continue k src exp)))))))
-  (visit-fun fun))
+  (visit-cont fun))
 
 (define (contify fun)
   (call-with-values (lambda () (compute-contification fun))
diff --git a/module/language/cps/cse.scm b/module/language/cps/cse.scm
index 236254648..2ecf40ca5 100644
--- a/module/language/cps/cse.scm
+++ b/module/language/cps/cse.scm
@@ -229,7 +229,7 @@ be that both true and false proofs are available."
 
 (define (compute-label-and-var-ranges fun)
   (match fun
-    (($ $fun free (and body ($ $cont kfun ($ $kfun src meta self))))
+    (($ $cont kfun ($ $kfun src meta self))
      ((make-cont-folder #f min-label label-count min-var var-count)
       (lambda (k cont min-label label-count min-var var-count)
         (let ((min-label (min k min-label))
@@ -250,7 +250,7 @@ be that both true and false proofs are available."
              (values min-label label-count (min self min-var) (1+ var-count)))
             (_
              (values min-label label-count min-var var-count)))))
-      body kfun 0 self 0))))
+      fun kfun 0 self 0))))
 
 (define (compute-idoms dfg min-label label-count)
   (define (label->idx label) (- label min-label))
@@ -458,8 +458,10 @@ be that both true and false proofs are available."
 
     (define (visit-exp* k src exp)
       (match exp
-        ((and fun ($ $fun))
-         (build-cps-term ($continue k src ,(cse fun dfg))))
+        (($ $fun free body)
+         (build-cps-term
+           ($continue k src
+             ($fun (map subst-var free) ,(cse body dfg)))))
         (_
          (cond
           ((vector-ref equiv-labels (label->idx label))
@@ -501,8 +503,13 @@ be that both true and false proofs are available."
       (($ $letk conts body)
        ,(visit-term body label))
       (($ $letrec names syms funs body)
-       ($letrec names syms (map (lambda (fun) (cse fun dfg)) funs)
-                ,(visit-term body label)))
+       ($letrec names syms
+                (map (lambda (fun)
+                       (rewrite-cps-exp fun
+                         (($ $fun free body)
+                          ($fun (map subst-var free) ,(cse body dfg)))))
+                     funs)
+         ,(visit-term body label)))
       (($ $continue k src exp)
        ,(let ((conts (append-map visit-dom-conts
                                  (vector-ref doms (label->idx label)))))
@@ -511,9 +518,7 @@ be that both true and false proofs are available."
               (build-cps-term
                 ($letk ,conts ,(visit-exp* k src exp))))))))
 
-  (rewrite-cps-exp fun
-    (($ $fun free body)
-     ($fun (map subst-var free) ,(visit-fun-cont body)))))
+  (visit-fun-cont fun))
 
 (define (cse fun dfg)
   (call-with-values (lambda () (compute-equivalent-subexpressions fun dfg))
@@ -525,6 +530,4 @@ be that both true and false proofs are available."
 (define (eliminate-common-subexpressions fun)
   (call-with-values (lambda () (renumber fun))
     (lambda (fun nlabels nvars)
-      (match fun
-        (($ $fun free body)
-         (cse fun (compute-dfg body)))))))
+      (cse fun (compute-dfg fun)))))
diff --git a/module/language/cps/dce.scm b/module/language/cps/dce.scm
index 2ef26100e..9105c4eaa 100644
--- a/module/language/cps/dce.scm
+++ b/module/language/cps/dce.scm
@@ -78,9 +78,7 @@
 
 (define (compute-live-code fun)
   (let* ((fun-data-table (make-hash-table))
-         (dfg (match fun
-                (($ $fun free body)
-                 (compute-dfg body #:global? #t))))
+         (dfg (compute-dfg fun #:global? #t))
          (live-vars (make-bitvector (dfg-var-count dfg) #f))
          (changed? #f))
     (define (mark-live! var)
@@ -92,12 +90,10 @@
     (define (ensure-fun-data fun)
       (or (hashq-ref fun-data-table fun)
           (call-with-values (lambda ()
-                              (match fun
-                                (($ $fun free body)
-                                 ((make-cont-folder #f label-count max-label)
-                                  (lambda (k cont label-count max-label)
-                                    (values (1+ label-count) (max k max-label)))
-                                  body 0 -1))))
+                              ((make-cont-folder #f label-count max-label)
+                               (lambda (k cont label-count max-label)
+                                 (values (1+ label-count) (max k max-label)))
+                               fun 0 -1))
             (lambda (label-count max-label)
               (let* ((min-label (- (1+ max-label) label-count))
                      (effects (compute-effects dfg min-label label-count))
@@ -133,7 +129,9 @@
                        (lp body)
                        (for-each (lambda (sym fun)
                                    (when (value-live? sym)
-                                     (visit-fun fun)))
+                                     (match fun
+                                       (($ $fun free body)
+                                        (visit-fun body)))))
                                  syms funs))
                       (($ $continue k src exp)
                        (unless (bitvector-ref live-conts n)
@@ -144,8 +142,8 @@
                          (match exp
                            ((or ($ $void) ($ $const) ($ $prim))
                             #f)
-                           ((and fun ($ $fun))
-                            (visit-fun fun))
+                           (($ $fun free body)
+                            (visit-fun body))
                            (($ $prompt escape? tag handler)
                             (mark-live! tag))
                            (($ $call proc args)
@@ -248,7 +246,12 @@
               (match (filter-map
                       (lambda (name sym fun)
                         (and (value-live? sym)
-                             (list name sym (visit-fun fun))))
+                             (match fun
+                               (($ $fun free body)
+                                (list name
+                                      sym
+                                      (build-cps-exp
+                                        ($fun free ,(visit-fun body))))))))
                       names syms funs)
                 (() body)
                 (((names syms funs) ...)
@@ -266,7 +269,8 @@
            (($ $continue k src exp)
             (if (bitvector-ref live-conts (label->idx term-k))
                 (rewrite-cps-term exp
-                  (($ $fun) ($continue k src ,(visit-fun exp)))
+                  (($ $fun free body)
+                   ($continue k src ($fun free ,(visit-fun body))))
                   (_
                    ,(match (vector-ref defs (label->idx term-k))
                       ((or #f ((? value-live?) ...))
@@ -278,9 +282,7 @@
                            ($letk (,(make-adaptor adapt k syms))
                              ($continue adapt src ,exp))))))))
                 (build-cps-term ($continue k src ($values ())))))))
-       (rewrite-cps-exp fun
-         (($ $fun free body)
-          ($fun free ,(visit-cont body)))))))
+       (visit-cont fun))))
   (visit-fun fun))
 
 (define (eliminate-dead-code fun)
diff --git a/module/language/cps/elide-values.scm b/module/language/cps/elide-values.scm
index c86702574..6823debbc 100644
--- a/module/language/cps/elide-values.scm
+++ b/module/language/cps/elide-values.scm
@@ -53,8 +53,8 @@
        ($letk ,(map visit-cont conts)
          ,(visit-term body)))
       (($ $letrec names syms funs body)
-       ($letrec names syms (map (cut elide-values* <> conts) funs)
-                ,(visit-term body)))
+       ($letrec names syms (map visit-fun funs)
+         ,(visit-term body)))
       (($ $continue k src ($ $primcall 'values vals))
        ,(rewrite-cps-term (vector-ref conts k)
           (($ $ktail)
@@ -94,17 +94,17 @@
                   (build-cps-term
                     ($continue k src ($values vals))))))))
       (($ $continue k src (and fun ($ $fun)))
-       ($continue k src ,(elide-values* fun conts)))
+       ($continue k src ,(visit-fun fun)))
       (($ $continue)
        ,term)))
+  (define (visit-fun fun)
+    (rewrite-cps-exp fun
+      (($ $fun free cont)
+       ($fun free ,(visit-cont cont)))))
 
-  (rewrite-cps-exp fun
-    (($ $fun free body)
-     ($fun free ,(visit-cont body)))))
+  (visit-cont fun))
 
 (define (elide-values fun)
-  (match fun
-    (($ $fun free funk)
-     (with-fresh-name-state funk
-       (let ((conts (build-cont-table funk)))
-         (elide-values* fun conts))))))
+  (with-fresh-name-state fun
+    (let ((conts (build-cont-table fun)))
+      (elide-values* fun conts))))
diff --git a/module/language/cps/prune-bailouts.scm b/module/language/cps/prune-bailouts.scm
index b241781a9..3ba28d900 100644
--- a/module/language/cps/prune-bailouts.scm
+++ b/module/language/cps/prune-bailouts.scm
@@ -61,7 +61,7 @@
   (define (visit-term term ktail)
     (rewrite-cps-term term
       (($ $letrec names vars funs body)
-       ($letrec names vars (map prune-bailouts* funs)
+       ($letrec names vars (map visit-fun funs)
                 ,(visit-term body ktail)))
       (($ $letk conts body)
        ($letk ,(map (lambda (cont) (visit-cont cont ktail)) conts)
@@ -71,7 +71,7 @@
 
   (define (visit-exp k src exp ktail)
     (rewrite-cps-term exp
-      (($ $fun) ($continue k src ,(prune-bailouts* exp)))
+      (($ $fun) ($continue k src ,(visit-fun exp)))
       (($ $primcall (and name (or 'error 'scm-error 'throw)) args)
        ,(if (eq? k ktail)
             (build-cps-term ($continue k src ,exp))
@@ -86,16 +86,17 @@
                   ,(primitive-ref name kprim src))))))
       (_ ($continue k src ,exp))))
 
-  (rewrite-cps-exp fun
-    (($ $fun free
-        ($ $cont kfun
-           ($ $kfun src meta self ($ $cont ktail ($ $ktail)) clause)))
-     ($fun free
-           (kfun ($kfun src meta self (ktail ($ktail))
-                     ,(and clause (visit-cont clause ktail))))))))
+  (define (visit-fun fun)
+    (rewrite-cps-exp fun
+      (($ $fun free body)
+       ($fun free ,(prune-bailouts* body)))))
+
+  (rewrite-cps-cont fun
+    (($ $cont kfun
+        ($ $kfun src meta self ($ $cont ktail ($ $ktail)) clause))
+     (kfun ($kfun src meta self (ktail ($ktail))
+             ,(and clause (visit-cont clause ktail)))))))
 
 (define (prune-bailouts fun)
-  (match fun
-    (($ $fun free body)
-     (with-fresh-name-state body
-       (prune-bailouts* fun)))))
+  (with-fresh-name-state fun
+    (prune-bailouts* fun)))
diff --git a/module/language/cps/prune-top-level-scopes.scm b/module/language/cps/prune-top-level-scopes.scm
index 8f6c0247c..2330d31c4 100644
--- a/module/language/cps/prune-top-level-scopes.scm
+++ b/module/language/cps/prune-top-level-scopes.scm
@@ -85,7 +85,7 @@
         (($ $fun free body)
          (visit-cont body))))
 
-    (visit-fun fun)
+    (visit-cont fun)
     scope-var->used?))
 
 (define (prune-top-level-scopes fun)
@@ -114,6 +114,4 @@
          ($continue k src ($primcall 'values ())))
         (($ $continue)
          ,term)))
-    (rewrite-cps-exp fun
-      (($ $fun free body)
-       ($fun free ,(visit-cont body))))))
+    (visit-cont fun)))
diff --git a/module/language/cps/renumber.scm b/module/language/cps/renumber.scm
index 217d6b0d7..0621ec92e 100644
--- a/module/language/cps/renumber.scm
+++ b/module/language/cps/renumber.scm
@@ -74,10 +74,7 @@
             (lp (1+ n) next))))))
 
 (define (compute-new-labels-and-vars fun)
-  (call-with-values (lambda ()
-                      (match fun
-                        (($ $fun free body)
-                         (compute-max-label-and-var body))))
+  (call-with-values (lambda () (compute-max-label-and-var fun))
     (lambda (max-label max-var)
       (let ((labels (make-vector (1+ max-label) #f))
             (next-label 0)
@@ -113,9 +110,7 @@
               (($ $letrec names syms funs body)
                (visit-term body))
               (($ $continue k src _) #f)))
-          (match fun
-            (($ $fun free body)
-             (visit-cont body))))
+          (visit-cont fun))
 
         (define (compute-names-in-fun fun)
           (define queue '())
@@ -162,108 +157,111 @@
               (($ $letrec names syms funs body)
                (when reachable?
                  (for-each rename! syms)
-                 (set! queue (fold cons queue funs)))
+                 (set! queue (fold (lambda (fun queue)
+                                     (match fun
+                                       (($ $fun free body)
+                                        (cons body queue))))
+                                   queue
+                                   funs)))
                (visit-term body reachable?))
-              (($ $continue k src (and fun ($ $fun)))
+              (($ $continue k src ($ $fun free body))
                (when reachable?
-                 (set! queue (cons fun queue))))
+                 (set! queue (cons body queue))))
               (($ $continue) #f)))
 
           (collect-conts fun)
           (match fun
-            (($ $fun free (and entry ($ $cont kfun)))
+            (($ $cont kfun)
              (set! next-label (sort-conts kfun labels next-label))
-             (visit-cont entry)
+             (visit-cont fun)
              (for-each compute-names-in-fun (reverse queue)))))
 
         (compute-names-in-fun fun)
         (values labels vars next-label next-var)))))
 
 (define (renumber fun)
-  (match fun
-    (($ $fun free cont)
-     (call-with-values (lambda () (compute-new-labels-and-vars fun))
-       (lambda (labels vars nlabels nvars)
-         (define (relabel label) (vector-ref labels label))
-         (define (rename var) (vector-ref vars var))
-         (define (rename-kw-arity arity)
-           (match arity
-             (($ $arity req opt rest kw aok?)
-              (make-$arity req opt rest
-                           (map (match-lambda
-                                 ((kw kw-name kw-var)
-                                  (list kw kw-name (rename kw-var))))
-                                kw)
-                           aok?))))
-         (define (must-visit-cont cont)
-           (or (visit-cont cont)
-               (error "internal error -- failed to visit cont")))
-         (define (visit-conts conts)
-           (match conts
-             (() '())
-             ((cont . conts)
-              (cond
-               ((visit-cont cont)
-                => (lambda (cont)
-                     (cons cont (visit-conts conts))))
-               (else (visit-conts conts))))))
-         (define (visit-cont cont)
-           (match cont
-             (($ $cont label cont)
-              (let ((label (relabel label)))
-                (and
-                 label
-                 (rewrite-cps-cont cont
-                   (($ $kargs names vars body)
-                    (label ($kargs names (map rename vars) ,(visit-term body))))
-                   (($ $kfun src meta self tail clause)
-                    (label
-                     ($kfun src meta (rename self) ,(must-visit-cont tail)
-                       ,(and clause (must-visit-cont clause)))))
-                   (($ $ktail)
-                    (label ($ktail)))
-                   (($ $kclause arity body alternate)
-                    (label
-                     ($kclause ,(rename-kw-arity arity) ,(must-visit-cont body)
-                               ,(and alternate (must-visit-cont alternate)))))
-                   (($ $kreceive ($ $arity req () rest () #f) kargs)
-                    (label ($kreceive req rest (relabel kargs))))
-                   (($ $kif kt kf)
-                    (label ($kif (relabel kt) (relabel kf))))))))))
-         (define (visit-term term)
-           (rewrite-cps-term term
-             (($ $letk conts body)
-              ,(match (visit-conts conts)
-                 (() (visit-term body))
-                 (conts (build-cps-term ($letk ,conts ,(visit-term body))))))
-             (($ $letrec names vars funs body)
-              ($letrec names (map rename vars) (map visit-fun funs)
-                       ,(visit-term body)))
-             (($ $continue k src exp)
-              ($continue (relabel k) src ,(visit-exp exp)))))
-         (define (visit-exp exp)
-           (match exp
-             ((or ($ $void) ($ $const) ($ $prim))
-              exp)
-             (($ $fun)
-              (visit-fun exp))
-             (($ $values args)
-              (let ((args (map rename args)))
-                (build-cps-exp ($values args))))
-             (($ $call proc args)
-              (let ((args (map rename args)))
-                (build-cps-exp ($call (rename proc) args))))
-             (($ $callk k proc args)
-              (let ((args (map rename args)))
-                (build-cps-exp ($callk (relabel k) (rename proc) args))))
-             (($ $primcall name args)
-              (let ((args (map rename args)))
-                (build-cps-exp ($primcall name args))))
-             (($ $prompt escape? tag handler)
-              (build-cps-exp
-                ($prompt escape? (rename tag) (relabel handler))))))
-         (define (visit-fun fun)
-           (rewrite-cps-exp fun
-             (($ $fun free body)
-              ($fun (map rename free) ,(must-visit-cont body)))))
-         (values (visit-fun fun) nlabels nvars))))))
+  (call-with-values (lambda () (compute-new-labels-and-vars fun))
+    (lambda (labels vars nlabels nvars)
+      (define (relabel label) (vector-ref labels label))
+      (define (rename var) (vector-ref vars var))
+      (define (rename-kw-arity arity)
+        (match arity
+          (($ $arity req opt rest kw aok?)
+           (make-$arity req opt rest
+                        (map (match-lambda
+                              ((kw kw-name kw-var)
+                               (list kw kw-name (rename kw-var))))
+                             kw)
+                        aok?))))
+      (define (must-visit-cont cont)
+        (or (visit-cont cont)
+            (error "internal error -- failed to visit cont")))
+      (define (visit-conts conts)
+        (match conts
+          (() '())
+          ((cont . conts)
+           (cond
+            ((visit-cont cont)
+             => (lambda (cont)
+                  (cons cont (visit-conts conts))))
+            (else (visit-conts conts))))))
+      (define (visit-cont cont)
+        (match cont
+          (($ $cont label cont)
+           (let ((label (relabel label)))
+             (and
+              label
+              (rewrite-cps-cont cont
+                (($ $kargs names vars body)
+                 (label ($kargs names (map rename vars) ,(visit-term body))))
+                (($ $kfun src meta self tail clause)
+                 (label
+                  ($kfun src meta (rename self) ,(must-visit-cont tail)
+                    ,(and clause (must-visit-cont clause)))))
+                (($ $ktail)
+                 (label ($ktail)))
+                (($ $kclause arity body alternate)
+                 (label
+                  ($kclause ,(rename-kw-arity arity) ,(must-visit-cont body)
+                            ,(and alternate (must-visit-cont alternate)))))
+                (($ $kreceive ($ $arity req () rest () #f) kargs)
+                 (label ($kreceive req rest (relabel kargs))))
+                (($ $kif kt kf)
+                 (label ($kif (relabel kt) (relabel kf))))))))))
+      (define (visit-term term)
+        (rewrite-cps-term term
+          (($ $letk conts body)
+           ,(match (visit-conts conts)
+              (() (visit-term body))
+              (conts (build-cps-term ($letk ,conts ,(visit-term body))))))
+          (($ $letrec names vars funs body)
+           ($letrec names (map rename vars) (map visit-fun funs)
+                    ,(visit-term body)))
+          (($ $continue k src exp)
+           ($continue (relabel k) src ,(visit-exp exp)))))
+      (define (visit-exp exp)
+        (match exp
+          ((or ($ $void) ($ $const) ($ $prim))
+           exp)
+          (($ $fun)
+           (visit-fun exp))
+          (($ $values args)
+           (let ((args (map rename args)))
+             (build-cps-exp ($values args))))
+          (($ $call proc args)
+           (let ((args (map rename args)))
+             (build-cps-exp ($call (rename proc) args))))
+          (($ $callk k proc args)
+           (let ((args (map rename args)))
+             (build-cps-exp ($callk (relabel k) (rename proc) args))))
+          (($ $primcall name args)
+           (let ((args (map rename args)))
+             (build-cps-exp ($primcall name args))))
+          (($ $prompt escape? tag handler)
+           (build-cps-exp
+             ($prompt escape? (rename tag) (relabel handler))))))
+      (define (visit-fun fun)
+        (rewrite-cps-exp fun
+          (($ $fun free body)
+           ($fun (map rename free) ,(must-visit-cont body)))))
+      (values (must-visit-cont fun) nlabels nvars))))
diff --git a/module/language/cps/self-references.scm b/module/language/cps/self-references.scm
index 7a2c57d34..69113208f 100644
--- a/module/language/cps/self-references.scm
+++ b/module/language/cps/self-references.scm
@@ -47,7 +47,7 @@
     (rewrite-cps-term term
       (($ $letrec names vars funs body)
        ($letrec names vars (map visit-recursive-fun funs vars)
-                ,(visit-term body)))
+         ,(visit-term body)))
       (($ $letk conts body)
        ($letk ,(map visit-cont conts)
          ,(visit-term body)))
@@ -57,7 +57,8 @@
   (define (visit-exp exp)
     (rewrite-cps-exp exp
       ((or ($ $void) ($ $const) ($ $prim)) ,exp)
-      (($ $fun) ,(resolve-self-references exp env))
+      (($ $fun free body)
+       ($fun free ,(resolve-self-references body env)))
       (($ $call proc args)
        ($call (subst proc) ,(map subst args)))
       (($ $callk k proc args)
@@ -70,10 +71,8 @@
        ($prompt escape? (subst tag) handler))))
 
   (define (visit-recursive-fun fun var)
-    (match fun
+    (rewrite-cps-exp fun
       (($ $fun free (and cont ($ $cont _ ($ $kfun src meta self))))
-       (resolve-self-references fun (acons var self env)))))
+       ($fun free ,(resolve-self-references cont (acons var self env))))))
 
-  (rewrite-cps-exp fun
-    (($ $fun free cont)
-     ($fun (map subst free) ,(visit-cont cont)))))
+  (visit-cont fun))
diff --git a/module/language/cps/simplify.scm b/module/language/cps/simplify.scm
index 3d09f63fc..0dd98e24f 100644
--- a/module/language/cps/simplify.scm
+++ b/module/language/cps/simplify.scm
@@ -64,14 +64,12 @@
       (match fun
         (($ $fun free body)
          (visit-cont body))))
-    (visit-fun fun)
+    (visit-cont fun)
     table))
 
 (define (eta-reduce fun)
   (let ((table (compute-eta-reductions fun))
-        (dfg (match fun
-               (($ $fun free body)
-                (compute-dfg body)))))
+        (dfg (compute-dfg fun)))
     (define (reduce* k scope values?)
       (match (hashq-ref table k)
         (#f k)
@@ -119,7 +117,7 @@
       (rewrite-cps-exp fun
         (($ $fun free body)
          ($fun free ,(visit-cont body #f)))))
-    (visit-fun fun)))
+    (visit-cont fun #f)))
 
 (define (compute-beta-reductions fun)
   ;; A continuation's body can be inlined in place of a $values
@@ -127,8 +125,7 @@
   ;; inlined if it is used only once, and not recursively.
   (let ((var-table (make-hash-table))
         (k-table (make-hash-table))
-        (dfg (match fun
-               (($ $fun free body) (compute-dfg body)))))
+        (dfg (compute-dfg fun)))
     (define (visit-cont cont)
       (match cont
         (($ $cont sym ($ $kargs names syms body))
@@ -171,7 +168,7 @@
       (match fun
         (($ $fun free body)
          (visit-cont body))))
-    (visit-fun fun)
+    (visit-cont fun)
     (values var-table k-table)))
 
 (define (beta-reduce fun)
@@ -235,7 +232,7 @@
       (rewrite-cps-exp fun
         (($ $fun free body)
          ($fun (map subst free) ,(must-visit-cont body)))))
-    (visit-fun fun)))
+    (must-visit-cont fun)))
 
 (define (simplify fun)
   ;; Renumbering prunes continuations that are made unreachable by
diff --git a/module/language/cps/specialize-primcalls.scm b/module/language/cps/specialize-primcalls.scm
index d58c85329..e03eb6222 100644
--- a/module/language/cps/specialize-primcalls.scm
+++ b/module/language/cps/specialize-primcalls.scm
@@ -31,9 +31,7 @@
   #:export (specialize-primcalls))
 
 (define (specialize-primcalls fun)
-  (let ((dfg (match fun
-               (($ $fun free body)
-                (compute-dfg body #:global? #t)))))
+  (let ((dfg (compute-dfg fun #:global? #t)))
     (with-fresh-name-state-from-dfg dfg
       (define (immediate-u8? sym)
         (call-with-values (lambda () (find-constant-value sym dfg))
@@ -113,4 +111,4 @@
           (($ $fun free body)
            ($fun free ,(visit-cont body)))))
 
-      (visit-fun fun))))
+      (visit-cont fun))))
diff --git a/module/language/tree-il/compile-cps.scm b/module/language/tree-il/compile-cps.scm
index e2b4fb3b0..96f27cd44 100644
--- a/module/language/tree-il/compile-cps.scm
+++ b/module/language/tree-il/compile-cps.scm
@@ -603,15 +603,14 @@ integer."
                  (scope-counter 0))
     (let ((src (tree-il-src exp)))
       (let-fresh (kinit ktail kclause kbody) (init)
-        (build-cps-exp
-          ($fun '()
-            (kinit ($kfun src '() init (ktail ($ktail))
-                     (kclause
-                      ($kclause ('() '() #f '() #f)
-                        (kbody ($kargs () ()
-                                 ,(convert exp ktail
-                                           (build-subst exp))))
-                        ,#f))))))))))
+        (build-cps-cont
+          (kinit ($kfun src '() init (ktail ($ktail))
+                   (kclause
+                    ($kclause ('() '() #f '() #f)
+                      (kbody ($kargs () ()
+                               ,(convert exp ktail
+                                         (build-subst exp))))
+                      ,#f)))))))))
 
 (define *comp-module* (make-fluid))
 
-- 
2.20.1