module/language/cps/contification.scm

   1 ;;; Continuation-passing style (CPS) intermediate language (IL)
   2
   3 ;; Copyright (C) 2013, 2014 Free Software Foundation, Inc.
   4
   5 ;;;; This library is free software; you can redistribute it and/or
   6 ;;;; modify it under the terms of the GNU Lesser General Public
   7 ;;;; License as published by the Free Software Foundation; either
   8 ;;;; version 3 of the License, or (at your option) any later version.
   9 ;;;;
  10 ;;;; This library is distributed in the hope that it will be useful,
  11 ;;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 ;;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  13 ;;;; Lesser General Public License for more details.
  14 ;;;;
  15 ;;;; You should have received a copy of the GNU Lesser General Public
  16 ;;;; License along with this library; if not, write to the Free Software
  17 ;;;; Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  18
  19 ;;; Commentary:
  20 ;;;
  21 ;;; Contification is a pass that turns $fun instances into $cont
  22 ;;; instances if all calls to the $fun return to the same continuation.
  23 ;;; This is a more rigorous variant of our old "fixpoint labels
  24 ;;; allocation" optimization.
  25 ;;;
  26 ;;; See Kennedy's "Compiling with Continuations, Continued", and Fluet
  27 ;;; and Weeks's "Contification using Dominators".
  28 ;;;
  29 ;;; Code:
  30
  31 (define-module (language cps contification)
  32   #:use-module (ice-9 match)
  33   #:use-module ((srfi srfi-1) #:select (concatenate filter-map))
  34   #:use-module (srfi srfi-26)
  35   #:use-module (language cps)
  36   #:use-module (language cps dfg)
  37   #:use-module (language cps primitives)
  38   #:use-module (language bytecode)
  39   #:export (contify))
  40
  41 (define (compute-contification fun)
  42   (let* ((dfg (compute-dfg fun))
  43          (scope-table (make-hash-table))
  44          (call-substs '())
  45          (cont-substs '())
  46          (fun-elisions '())
  47          (cont-splices (make-hash-table)))
  48     (define (subst-call! sym arities body-ks)
  49       (set! call-substs (acons sym (map cons arities body-ks) call-substs)))
  50     (define (subst-return! old-tail new-tail)
  51       (set! cont-substs (acons old-tail new-tail cont-substs)))
  52     (define (elide-function! k cont)
  53       (set! fun-elisions (acons k cont fun-elisions)))
  54     (define (splice-conts! scope conts)
  55       (for-each (match-lambda
  56                  (($ $cont k) (hashq-set! scope-table k scope)))
  57                 conts)
  58       (hashq-set! cont-splices scope
  59                   (append conts (hashq-ref cont-splices scope '()))))
  60
  61     (define (lookup-return-cont k)
  62       (match (assq-ref cont-substs k)
  63         (#f k)
  64         (k (lookup-return-cont k))))
  65
  66     ;; If K is a continuation that binds one variable, and it has only
  67     ;; one predecessor, return that variable.
  68     (define (bound-symbol k)
  69       (match (lookup-cont k dfg)
  70         (($ $kargs (_) (sym))
  71          (match (lookup-predecessors k dfg)
  72            ((_)
  73             ;; K has one predecessor, the one that defined SYM.
  74             sym)
  75            (_ #f)))
  76         (_ #f)))
  77
  78     (define (extract-arities clause)
  79       (match clause
  80         (($ $cont _ ($ $kclause arity body alternate))
  81          (cons arity (extract-arities alternate)))
  82         (#f '())))
  83     (define (extract-bodies clause)
  84       (match clause
  85         (($ $cont _ ($ $kclause arity body alternate))
  86          (cons body (extract-bodies alternate)))
  87         (#f '())))
  88
  89     (define (contify-fun term-k sym self tail arities bodies)
  90       (contify-funs term-k
  91                     (list sym) (list self) (list tail)
  92                     (list arities) (list bodies)))
  93
  94     ;; Given a set of mutually recursive functions bound to local
  95     ;; variables SYMS, with self symbols SELFS, tail continuations
  96     ;; TAILS, arities ARITIES, and bodies BODIES, all bound in TERM-K,
  97     ;; contify them if we can prove that they all return to the same
  98     ;; continuation.  Returns a true value on success, and false
  99     ;; otherwise.
 100     (define (contify-funs term-k syms selfs tails arities bodies)
 101       (define (unused? sym)
 102         (null? (lookup-uses sym dfg)))
 103
 104       ;; Are the given args compatible with any of the arities?
 105       (define (applicable? proc args)
 106         (let lp ((arities (assq-ref (map cons syms arities) proc)))
 107           (match arities
 108             ((($ $arity req () #f () #f) . arities)
 109              (or (= (length args) (length req))
 110                  (lp arities)))
 111             ;; If we reached the end of the arities, fail.  Also fail if
 112             ;; the next arity in the list has optional, keyword, or rest
 113             ;; arguments.
 114             (_ #f))))
 115
 116       ;; If the use of PROC in continuation USE is a call to PROC that
 117       ;; is compatible with one of the procedure's arities, return the
 118       ;; target continuation.  Otherwise return #f.
 119       (define (call-target use proc)
 120         (match (find-call (lookup-cont use dfg))
 121           (($ $continue k src ($ $call proc* args))
 122            (and (eq? proc proc*) (not (memq proc args)) (applicable? proc args)
 123                 ;; Converge more quickly by resolving already-contified
 124                 ;; call targets.
 125                 (lookup-return-cont k)))
 126           (_ #f)))
 127
 128       ;; If this set of functions is always called with one
 129       ;; continuation, not counting tail calls between the functions,
 130       ;; return that continuation.
 131       (define (find-common-continuation)
 132         (let visit-syms ((syms syms) (k #f))
 133           (match syms
 134             (() k)
 135             ((sym . syms)
 136              (let visit-uses ((uses (lookup-uses sym dfg)) (k k))
 137                (match uses
 138                  (() (visit-syms syms k))
 139                  ((use . uses)
 140                   (and=> (call-target use sym)
 141                          (lambda (k*)
 142                            (cond
 143                             ((memq k* tails) (visit-uses uses k))
 144                             ((not k) (visit-uses uses k*))
 145                             ((eq? k k*) (visit-uses uses k))
 146                             (else #f)))))))))))
 147
 148       ;; Given that the functions are called with the common
 149       ;; continuation K, determine the scope at which to contify the
 150       ;; functions.  If K is in scope in the term, we go ahead and
 151       ;; contify them there.  Otherwise the scope is inside the letrec
 152       ;; body, and so choose the scope in which the continuation is
 153       ;; defined, whose free variables are a superset of the free
 154       ;; variables of the functions.
 155       ;;
 156       ;; There is some slight trickiness here.  Call-target already uses
 157       ;; the information we compute within this pass.  Previous
 158       ;; contifications may cause functions to be contified not at their
 159       ;; point of definition but at their point of non-recursive use.
 160       ;; That will cause the scope nesting to change.  (It may
 161       ;; effectively push a function deeper down the tree -- the second
 162       ;; case above, a call within the letrec body.)  What if we contify
 163       ;; to the tail of a previously contified function?  We have to
 164       ;; track what the new scope tree will be when asking whether K
 165       ;; will be bound in TERM-K's scope, not the scope tree that
 166       ;; existed when we started the pass.
 167       ;;
 168       ;; FIXME: Does this choose the right scope for contified let-bound
 169       ;; functions?
 170       (define (find-contification-scope k)
 171         (define (scope-contains? scope k)
 172           (let ((k-scope (or (hashq-ref scope-table k)
 173                              (let ((k-scope (lookup-block-scope k dfg)))
 174                                (hashq-set! scope-table k k-scope)
 175                                k-scope))))
 176             (or (eq? scope k-scope)
 177                 (and k-scope (scope-contains? scope k-scope)))))
 178
 179         ;; Find the scope of K.
 180         (define (continuation-scope k)
 181           (or (hashq-ref scope-table k)
 182               (let ((scope (lookup-block-scope k dfg)))
 183                 (hashq-set! scope-table k scope)
 184                 scope)))
 185
 186         (let ((k-scope (continuation-scope k)))
 187           (if (scope-contains? k-scope term-k)
 188               term-k
 189               (match (lookup-cont k-scope dfg)
 190                 (($ $kentry self tail clause)
 191                  ;; K is the tail of some function.  If that function
 192                  ;; has just one clause, return that clause.  Otherwise
 193                  ;; bail.
 194                  (match clause
 195                    (($ $cont _ ($ $kclause arity ($ $cont kargs) #f))
 196                     kargs)
 197                    (_ #f)))
 198                 (_ k-scope)))))
 199
 200       ;; We are going to contify.  Mark all SYMs for replacement in
 201       ;; calls, and mark the tail continuations for replacement by K.
 202       ;; Arrange for the continuations to be spliced into SCOPE.
 203       (define (enqueue-contification! k scope)
 204         (for-each (lambda (sym tail arities bodies)
 205                     (match bodies
 206                       ((($ $cont body-k) ...)
 207                        (subst-call! sym arities body-k)))
 208                     (subst-return! tail k))
 209                   syms tails arities bodies)
 210         (splice-conts! scope (concatenate bodies))
 211         #t)
 212
 213       ;; "Call me maybe"
 214       (and (and-map unused? selfs)
 215            (and=> (find-common-continuation)
 216                   (lambda (k)
 217                     (and=> (find-contification-scope k)
 218                            (cut enqueue-contification! k <>))))))
 219
 220     (define (visit-fun term)
 221       (match term
 222         (($ $fun src meta free body)
 223          (visit-cont body))))
 224     (define (visit-cont cont)
 225       (match cont
 226         (($ $cont sym ($ $kargs _ _ body))
 227          (visit-term body sym))
 228         (($ $cont sym ($ $kentry self tail clause))
 229          (when clause (visit-cont clause)))
 230         (($ $cont sym ($ $kclause arity body alternate))
 231          (visit-cont body)
 232          (when alternate (visit-cont alternate)))
 233         (($ $cont)
 234          #t)))
 235     (define (visit-term term term-k)
 236       (match term
 237         (($ $letk conts body)
 238          (for-each visit-cont conts)
 239          (visit-term body term-k))
 240         (($ $letrec names syms funs body)
 241          (define (split-components nsf)
 242            ;; FIXME: Compute strongly-connected components.  Currently
 243            ;; we just put non-recursive functions in their own
 244            ;; components, and lump everything else in the remaining
 245            ;; component.
 246            (define (recursive? k)
 247              (or-map (cut variable-free-in? <> k dfg) syms))
 248            (let lp ((nsf nsf) (rec '()))
 249              (match nsf
 250                (()
 251                 (if (null? rec)
 252                     '()
 253                     (list rec)))
 254                (((and elt (n s ($ $fun src meta free ($ $cont kentry))))
 255                  . nsf)
 256                 (if (recursive? kentry)
 257                     (lp nsf (cons elt rec))
 258                     (cons (list elt) (lp nsf rec)))))))
 259          (define (extract-arities+bodies clauses)
 260            (values (map extract-arities clauses)
 261                    (map extract-bodies clauses)))
 262          (define (visit-component component)
 263            (match component
 264              (((name sym fun) ...)
 265               (match fun
 266                 ((($ $fun src meta free
 267                      ($ $cont fun-k
 268                         ($ $kentry self ($ $cont tail-k ($ $ktail)) clause)))
 269                   ...)
 270                  (call-with-values (lambda () (extract-arities+bodies clause))
 271                    (lambda (arities bodies)
 272                      (if (contify-funs term-k sym self tail-k arities bodies)
 273                          (for-each (cut for-each visit-cont <>) bodies)
 274                          (for-each visit-fun fun)))))))))
 275          (visit-term body term-k)
 276          (for-each visit-component
 277                    (split-components (map list names syms funs))))
 278         (($ $continue k src exp)
 279          (match exp
 280            (($ $fun src meta free
 281                ($ $cont fun-k
 282                   ($ $kentry self ($ $cont tail-k ($ $ktail)) clause)))
 283             (if (and=> (bound-symbol k)
 284                        (lambda (sym)
 285                          (contify-fun term-k sym self tail-k
 286                                       (extract-arities clause)
 287                                       (extract-bodies clause))))
 288                 (begin
 289                   (elide-function! k (lookup-cont k dfg))
 290                   (for-each visit-cont (extract-bodies clause)))
 291                 (visit-fun exp)))
 292            (_ #t)))))
 293
 294     (visit-fun fun)
 295     (values call-substs cont-substs fun-elisions cont-splices)))
 296
 297 (define (apply-contification fun call-substs cont-substs fun-elisions cont-splices)
 298   (define (contify-call src proc args)
 299     (and=> (assq-ref call-substs proc)
 300            (lambda (clauses)
 301              (let lp ((clauses clauses))
 302                (match clauses
 303                  (() (error "invalid contification"))
 304                  (((($ $arity req () #f () #f) . k) . clauses)
 305                   (if (= (length req) (length args))
 306                       (build-cps-term
 307                         ($continue k src
 308                           ($values args)))
 309                       (lp clauses)))
 310                  ((_ . clauses) (lp clauses)))))))
 311   (define (continue k src exp)
 312     (define (lookup-return-cont k)
 313       (match (assq-ref cont-substs k)
 314         (#f k)
 315         (k (lookup-return-cont k))))
 316     (let ((k* (lookup-return-cont k)))
 317       ;; We are contifying this return.  It must be a call or a
 318       ;; primcall to values, return, or return-values.
 319       (if (eq? k k*)
 320           (build-cps-term ($continue k src ,exp))
 321           (rewrite-cps-term exp
 322             (($ $primcall 'return (val))
 323              ($continue k* src ($primcall 'values (val))))
 324             (($ $values vals)
 325              ($continue k* src ($primcall 'values vals)))
 326             (_ ($continue k* src ,exp))))))
 327   (define (splice-continuations term-k term)
 328     (match (hashq-ref cont-splices term-k)
 329       (#f term)
 330       ((cont ...)
 331        (let lp ((term term))
 332          (rewrite-cps-term term
 333            (($ $letrec names syms funs body)
 334             ($letrec names syms funs ,(lp body)))
 335            (($ $letk conts* body)
 336             ($letk ,(append conts* (filter-map visit-cont cont))
 337               ,body))
 338            (body
 339             ($letk ,(filter-map visit-cont cont)
 340               ,body)))))))
 341   (define (visit-fun term)
 342     (rewrite-cps-exp term
 343       (($ $fun src meta free body)
 344        ($fun src meta free ,(visit-cont body)))))
 345   (define (visit-cont cont)
 346     (rewrite-cps-cont cont
 347       (($ $cont (? (cut assq <> fun-elisions)))
 348        ;; This cont gets inlined in place of the $fun.
 349        ,#f)
 350       (($ $cont sym ($ $kargs names syms body))
 351        (sym ($kargs names syms ,(visit-term body sym))))
 352       (($ $cont sym ($ $kentry self tail clause))
 353        (sym ($kentry self ,tail ,(and clause (visit-cont clause)))))
 354       (($ $cont sym ($ $kclause arity body alternate))
 355        (sym ($kclause ,arity ,(visit-cont body)
 356                       ,(and alternate (visit-cont alternate)))))
 357       (($ $cont)
 358        ,cont)))
 359   (define (visit-term term term-k)
 360     (match term
 361       (($ $letk conts body)
 362        ;; Visit the body first, so we rewrite depth-first.
 363        (let lp ((body (visit-term body term-k)))
 364          ;; Because we attach contified functions on a particular
 365          ;; term-k, and one term-k can correspond to an arbitrarily
 366          ;; nested sequence of $letrec and $letk instances, normalize
 367          ;; so that all continuations are bound by one $letk --
 368          ;; guaranteeing that they are in the same scope.
 369          (rewrite-cps-term body
 370            (($ $letrec names syms funs body)
 371             ($letrec names syms funs ,(lp body)))
 372            (($ $letk conts* body)
 373             ($letk ,(append conts* (filter-map visit-cont conts))
 374               ,body))
 375            (body
 376             ($letk ,(filter-map visit-cont conts)
 377               ,body)))))
 378       (($ $letrec names syms funs body)
 379        (rewrite-cps-term (filter (match-lambda
 380                                   ((n s f) (not (assq s call-substs))))
 381                                  (map list names syms funs))
 382          (((names syms funs) ...)
 383           ($letrec names syms (map visit-fun funs)
 384                    ,(visit-term body term-k)))))
 385       (($ $continue k src exp)
 386        (splice-continuations
 387         term-k
 388         (match exp
 389           (($ $fun)
 390            (cond
 391             ((assq-ref fun-elisions k)
 392              => (match-lambda
 393                  (($ $kargs (_) (_) body)
 394                   (visit-term body k))))
 395             (else
 396              (continue k src (visit-fun exp)))))
 397           (($ $call proc args)
 398            (or (contify-call src proc args)
 399                (continue k src exp)))
 400           (_ (continue k src exp)))))))
 401   (visit-fun fun))
 402
 403 (define (contify fun)
 404   (call-with-values (lambda () (compute-contification fun))
 405     (lambda (call-substs cont-substs fun-elisions cont-splices)
 406       (if (null? call-substs)
 407           fun
 408           ;; Iterate to fixed point.
 409           (contify
 410            (apply-contification fun call-substs cont-substs fun-elisions cont-splices))))))