module/language/cps/contification.scm

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