[PATCH 6/9] Add arity-adapting module

[Andy Wingo <wingo@pobox.com>]

* module/Makefile.am:
* module/language/cps/arities.scm: New module.  Adapts call and return
  arities, especially for primcalls.
---
 module/Makefile.am              |   1 +
 module/language/cps/arities.scm | 152 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 153 insertions(+)
 create mode 100644 module/language/cps/arities.scm

diff --git a/module/Makefile.am b/module/Makefile.am
index 790db8c..bcef4bf 100644
--- a/module/Makefile.am
+++ b/module/Makefile.am
@@ -120,6 +120,7 @@ TREE_IL_LANG_SOURCES =						\
 
 CPS_LANG_SOURCES =						\
   language/cps.scm						\
+  language/cps/arities.scm					\
   language/cps/closure-conversion.scm				\
   language/cps/primitives.scm					\
   language/cps/spec.scm						\
diff --git a/module/language/cps/arities.scm b/module/language/cps/arities.scm
new file mode 100644
index 0000000..b697ec0
--- /dev/null
+++ b/module/language/cps/arities.scm
@@ -0,0 +1,152 @@
+;;; Continuation-passing style (CPS) intermediate language (IL)
+
+;; Copyright (C) 2013 Free Software Foundation, Inc.
+
+;;;; This library is free software; you can redistribute it and/or
+;;;; modify it under the terms of the GNU Lesser General Public
+;;;; License as published by the Free Software Foundation; either
+;;;; version 3 of the License, or (at your option) any later version.
+;;;;
+;;;; This library is distributed in the hope that it will be useful,
+;;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
+;;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+;;;; Lesser General Public License for more details.
+;;;;
+;;;; You should have received a copy of the GNU Lesser General Public
+;;;; License along with this library; if not, write to the Free Software
+;;;; Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+
+;;; Commentary:
+;;;
+;;; A pass to adapt expressions to the arities of their continuations,
+;;; and to rewrite some tail expressions as primcalls to "return".
+;;;
+;;; Code:
+
+(define-module (language cps arities)
+  #:use-module (ice-9 match)
+  #:use-module ((srfi srfi-1) #:select (fold))
+  #:use-module (srfi srfi-26)
+  #:use-module (language cps)
+  #:use-module (language cps dfg)
+  #:use-module (language cps primitives)
+  #:export (fix-arities))
+
+(define (fix-clause-arities clause)
+  (let ((conts (build-local-cont-table clause))
+        (ktail (match clause
+                 (($ $cont _ _ ($ $kentry _ ($ $cont ktail) _)) ktail))))
+    (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 fix-arities funs) ,(visit-term body)))
+        (($ $continue k exp)
+         ,(visit-exp k exp))))
+
+    (define (adapt-exp nvals k exp)
+      (match nvals
+        (0
+         (rewrite-cps-term (lookup-cont k conts)
+           (($ $ktail)
+            ,(let-gensyms (kvoid kunspec unspec)
+               (build-cps-term
+                 ($letk* ((kunspec #f ($kargs (unspec) (unspec)
+                                        ($continue k
+                                          ($primcall 'return (unspec)))))
+                          (kvoid #f ($kargs () ()
+                                      ($continue kunspec ($void)))))
+                   ($continue kvoid ,exp)))))
+           (($ $ktrunc ($ $arity () () #f () #f) kseq)
+            ($continue kseq ,exp))
+           (($ $kargs () () _)
+            ($continue k ,exp))
+           (_
+            ,(let-gensyms (k*)
+               (build-cps-term
+                 ($letk ((k* #f ($kargs () () ($continue k ($void)))))
+                   ($continue k* ,exp)))))))
+        (1
+         (let ((drop-result
+                (lambda (kseq)
+                  (let-gensyms (k* drop)
+                    (build-cps-term
+                      ($letk ((k* #f ($kargs ('drop) (drop)
+                                       ($continue kseq ($values ())))))
+                        ($continue k* ,exp)))))))
+           (rewrite-cps-term (lookup-cont k conts)
+             (($ $ktail)
+              ,(rewrite-cps-term exp
+                 (($var sym)
+                  ($continue ktail ($primcall 'return (sym))))
+                 (_
+                  ,(let-gensyms (k* v)
+                     (build-cps-term
+                       ($letk ((k* #f ($kargs (v) (v)
+                                        ($continue k
+                                          ($primcall 'return (v))))))
+                         ($continue k* ,exp)))))))
+             (($ $ktrunc ($ $arity () () #f () #f) kseq)
+              ,(drop-result kseq))
+             (($ $kargs () () _)
+              ,(drop-result k))
+             (_
+              ($continue k ,exp)))))))
+
+    (define (visit-exp k exp)
+      (rewrite-cps-term exp
+        ((or ($ $void)
+             ($ $const)
+             ($ $prim)
+             ($ $var))
+         ,(adapt-exp 1 k exp))
+        (($ $fun)
+         ,(adapt-exp 1 k (fix-arities exp)))
+        (($ $call)
+         ;; In general, calls have unknown return arity.  For that
+         ;; reason every non-tail call has an implicit adaptor
+         ;; continuation to adapt the return to the target
+         ;; continuation, and we don't need to do any adapting here.
+         ($continue k ,exp))
+        (($ $primcall 'return (arg))
+         ;; Primcalls to return are in tail position.
+         ($continue ktail ,exp))
+        (($ $primcall (? (lambda (name)
+                           (and (not (prim-rtl-instruction name))
+                                (not (branching-primitive? name))))))
+         ($continue k ,exp))
+        (($ $primcall name args)
+         ,(match (prim-arity name)
+            ((out . in)
+             (if (= in (length args))
+                 (adapt-exp out k exp)
+                 (let-gensyms (k* p*)
+                   (build-cps-term
+                     ($letk ((k* #f ($kargs ('prim) (p*)
+                                      ($continue k ($call p* args)))))
+                       ($continue k* ($prim name)))))))))
+        (($ $values)
+         ;; Values nodes are inserted by CPS optimization passes, so
+         ;; we assume they are correct.
+         ($continue k ,exp))
+        (($ $prompt)
+         ($continue k ,exp))))
+
+    (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 ($ $kclause arity body))
+         (sym src ($kclause ,arity ,(visit-cont body))))
+        (($ $cont)
+         ,cont)))
+
+    (rewrite-cps-cont clause
+      (($ $cont sym src ($ $kentry self tail clauses))
+       (sym src ($kentry self ,tail ,(map visit-cont clauses)))))))
+
+(define (fix-arities fun)
+  (rewrite-cps-exp fun
+    (($ $fun meta free body)
+     ($fun meta free ,(fix-clause-arities body)))))
-- 
1.8.3.2