Factor fold-conts, fold-local-conts into cps.scm

* module/language/cps.scm (fold-conts, fold-local-conts): Consolidate
  these helpers into (language cps).

* module/language/cps/arities.scm:
* module/language/cps/compile-rtl.scm:
* module/language/cps/reify-primitives.scm: Adapt to use the common
  definitions.

4 files changed, 86 insertions, 119 deletions

diff --git a/module/language/cps.scm b/module/language/cps.scm
index 562d07068..d8c30a390 100644
--- a/module/language/cps.scm
+++ b/module/language/cps.scm
@@ -69,9 +69,6 @@
             make-$var make-$void make-$const make-$prim
             make-$call make-$primcall make-$values make-$prompt
 
-            parse-cps
-            unparse-cps
-
             ;; Building macros.
             let-gensyms
             build-cps-term
@@ -79,7 +76,13 @@
             build-cps-cont
             rewrite-cps-term
             rewrite-cps-call
-            rewrite-cps-cont))
+            rewrite-cps-cont
+
+            parse-cps
+            unparse-cps
+
+            fold-conts
+            fold-local-conts))
 
 ;; FIXME: Use SRFI-99, when Guile adds it.
 (define-syntax define-record-type*
@@ -342,3 +345,57 @@
     (_
      (error "unexpected cps" exp))))
 
+(define (fold-conts proc seed fun)
+  (define (cont-folder cont seed)
+    (match cont
+      (($ $cont k src ($ $kargs names syms body))
+       (term-folder body (proc cont seed)))
+
+      (($ $cont k src ($ $kentry arity body))
+       (cont-folder body (proc cont seed)))
+
+      (($ $cont)
+       (proc cont seed))))
+
+  (define (fun-folder fun seed)
+    (match fun
+      (($ $fun meta self free entries)
+       (fold cont-folder seed entries))))
+
+  (define (term-folder term seed)
+    (match term
+      (($ $letk conts body)
+       (fold cont-folder (term-folder body seed) conts))
+
+      (($ $continue k exp)
+       (match exp
+         (($ $fun) (fun-folder exp seed))
+         (_ seed)))
+
+      (($ $letrec names syms funs body)
+       (fold fun-folder funs (term-folder body seed)))))
+
+  (fun-folder fun seed))
+
+(define (fold-local-conts proc seed cont)
+  (define (cont-folder cont seed)
+    (match cont
+      (($ $cont k src ($ $kargs names syms body))
+       (term-folder body (proc cont seed)))
+
+      (($ $cont k src ($ $kentry arity body))
+       (cont-folder body (proc cont seed)))
+
+      (($ $cont)
+       (proc cont seed))))
+
+  (define (term-folder term seed)
+    (match term
+      (($ $letk conts body)
+       (fold cont-folder (term-folder body seed) conts))
+
+      (($ $continue) seed)
+
+      (($ $letrec names syms funs body) (term-folder body seed))))
+
+  (cont-folder cont seed))
diff --git a/module/language/cps/arities.scm b/module/language/cps/arities.scm
index bfae9b573..b89eef640 100644
--- a/module/language/cps/arities.scm
+++ b/module/language/cps/arities.scm
@@ -29,44 +29,6 @@
   #:use-module (language cps primitives)
   #:export (fix-arities))
 
-(define (fold-conts proc seed term)
-  (match term
-    (($ $fun meta self free entries)
-     (fold (lambda (exp seed)
-             (fold-conts proc seed exp))
-           seed
-           entries))
-    
-    (($ $letrec names syms funs body)
-     (fold-conts proc
-                 (fold (lambda (exp seed)
-                         (fold-conts proc seed exp))
-                       seed
-                       funs)
-                 body))
-
-    (($ $letk conts body)
-     (fold-conts proc
-                 (fold (lambda (exp seed)
-                         (fold-conts proc seed exp))
-                       seed
-                       conts)
-                 body))
-
-    (($ $cont sym src ($ $kargs names syms body))
-     (fold-conts proc (proc term seed) body))
-
-    (($ $cont sym src ($ $kentry arity body))
-     (fold-conts proc (proc term seed) body))
-
-    (($ $cont)
-     (proc term seed))
-
-    (($ $continue k exp)
-     (match exp
-       (($ $fun) (fold-conts proc seed exp))
-       (_ seed)))))
-
 (define (lookup-cont conts k)
   (and (not (eq? k 'ktail))
        (let lp ((conts conts))
@@ -76,14 +38,14 @@
               (($ $cont (? (cut eq? <> k))) cont)
               (else (lp conts))))))))
 
-(define (fix-arities fun)
-  (let ((conts (fold-conts cons '() fun)))
+(define (fix-entry-arities entry)
+  (let ((conts (fold-local-conts cons '() entry)))
     (define (visit-term term)
       (rewrite-cps-term term
         (($ $letk conts body)
          ($letk ,(map visit-cont conts) ,(visit-term body)))
         (($ $letrec names syms funs body)
-         ($letrec names syms (map visit-fun funs) ,(visit-term body)))
+         ($letrec names syms (map fix-arities funs) ,(visit-term body)))
         (($ $continue k exp)
          ,(visit-call k exp))))
 
@@ -145,7 +107,7 @@
              ($ $var))
          ,(adapt-call 1 k exp))
         (($ $fun)
-         ,(adapt-call 1 k (visit-fun exp)))
+         ,(adapt-call 1 k (fix-arities exp)))
         (($ $call)
          ;; In general, calls have unknown return arity.  For that
          ;; reason every non-tail call has an implicit adaptor
@@ -176,18 +138,18 @@
         (($ $prompt)
          ($continue k ,exp))))
 
-    (define (visit-fun fun)
-      (rewrite-cps-call fun
-        (($ $fun meta self free entries)
-         ($fun meta self free ,(map visit-cont entries)))))
-
     (define (visit-cont cont)
       (rewrite-cps-cont cont
         (($ $cont sym src ($ $kargs names syms body))
          (sym src ($kargs names syms ,(visit-term body))))
-        (($ $cont sym src ($ $kentry arity body))
-         (sym src ($kentry ,arity ,(visit-cont body))))
         (($ $cont)
          ,cont)))
 
-    (visit-fun fun)))
+    (rewrite-cps-cont entry
+      (($ $cont sym src ($ $kentry arity body))
+       (sym src ($kentry ,arity ,(visit-cont body)))))))
+
+(define (fix-arities fun)
+  (rewrite-cps-call fun
+    (($ $fun meta self free entries)
+     ($fun meta self free ,(map fix-entry-arities entries)))))
diff --git a/module/language/cps/compile-rtl.scm b/module/language/cps/compile-rtl.scm
index c0fc9b8e5..901aaea4c 100644
--- a/module/language/cps/compile-rtl.scm
+++ b/module/language/cps/compile-rtl.scm
@@ -87,28 +87,14 @@
 
     (_ (values))))
 
-(define (fold-conts proc seed exp)
-  (match exp
-    (($ $letk conts body)
-     (fold (lambda (exp seed)
-             (fold-conts proc seed exp))
-           (fold-conts proc seed body)
-           conts))
-
-    (($ $cont k src cont)
-     (fold-conts proc (proc k src cont seed) cont))
-
-    (($ $kargs names syms body)
-     (fold-conts proc seed body))
-
-    (_ seed)))
-
 (define (emit-rtl-sequence exp moves slots nlocals)
-  (define (intern-cont! k src cont table)
-    (hashq-set! table k cont)
-    table)
+  (define (intern-cont! cont table)
+    (match cont
+      (($ $cont k src cont)
+       (hashq-set! table k cont)
+       table)))
 
-  (let* ((cont-table (fold-conts intern-cont! (make-hash-table) exp))
+  (let* ((cont-table (fold-local-conts intern-cont! (make-hash-table) exp))
          (rtl '()))
     (define (slot sym)
       (lookup-slot sym slots))
@@ -310,19 +296,19 @@
         (($ $ktrunc ($ $arity req () rest () #f) k)
          (emit-trunc (length req) (and rest #t) k))))
 
-    (define (collect-exps k src cont tail)
-      (define (find-exp term)
+    (define (collect-exps cont tail)
+      (define (find-exp k src term)
         (match term
           (($ $continue exp-k exp)
            (cons (list k src exp-k exp) tail))
           (($ $letk conts body)
-           (find-exp body))))
+           (find-exp k src body))))
       (match cont
-        (($ $kargs names syms body)
-         (find-exp body))
-        (else tail)))
+        (($ $cont k src ($ $kargs names syms body))
+         (find-exp k src body))
+        (_ tail)))
 
-    (let lp ((exps (reverse (fold-conts collect-exps '() exp))))
+    (let lp ((exps (reverse (fold-local-conts collect-exps '() exp))))
       (match exps
         (() (reverse rtl))
         (((k src exp-k exp) . exps)
diff --git a/module/language/cps/reify-primitives.scm b/module/language/cps/reify-primitives.scm
index a9233916e..fe63b4446 100644
--- a/module/language/cps/reify-primitives.scm
+++ b/module/language/cps/reify-primitives.scm
@@ -36,44 +36,6 @@
 ;; FIXME: Some of these common utilities should be factored elsewhere,
 ;; perhaps (language cps).
 
-(define (fold-conts proc seed term)
-  (match term
-    (($ $fun meta self free entries)
-     (fold (lambda (exp seed)
-             (fold-conts proc seed exp))
-           seed
-           entries))
-
-    (($ $letrec names syms funs body)
-     (fold-conts proc
-                 (fold (lambda (exp seed)
-                         (fold-conts proc seed exp))
-                       seed
-                       funs)
-                 body))
-
-    (($ $letk conts body)
-     (fold-conts proc
-                 (fold (lambda (exp seed)
-                         (fold-conts proc seed exp))
-                       seed
-                       conts)
-                 body))
-
-    (($ $cont sym src ($ $kargs names syms body))
-     (fold-conts proc (proc term seed) body))
-
-    (($ $cont sym src ($ $kentry arity body))
-     (fold-conts proc (proc term seed) body))
-
-    (($ $cont)
-     (proc term seed))
-
-    (($ $continue k exp)
-     (match exp
-       (($ $fun) (fold-conts proc seed exp))
-       (_ seed)))))
-
 (define (lookup-cont table k)
   (cond
    ((vhash-assq k table) => cdr)