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