[Date Prev][Date Next][Thread Prev][Thread Next][Date Index][Thread Index]
[Guile-commits] 05/05: Add tailify pass
|
From: |
Andy Wingo |
|
Subject: |
[Guile-commits] 05/05: Add tailify pass |
|
Date: |
Wed, 9 Jun 2021 10:47:05 -0400 (EDT) |
wingo pushed a commit to branch wip-tailify
in repository guile.
commit 5bb0ffb5584d385e2fae0f1808e173f1e3bb12a6
Author: Andy Wingo <wingo@pobox.com>
AuthorDate: Thu Jun 3 21:37:59 2021 +0200
Add tailify pass
* module/Makefile.am (SOURCES): Add tailify.scm
* module/language/cps/tailify.scm: New file.
---
module/Makefile.am | 1 +
module/language/cps/tailify.scm | 636 ++++++++++++++++++++++++++++++++++++++++
2 files changed, 637 insertions(+)
diff --git a/module/Makefile.am b/module/Makefile.am
index 37786ed..a43640e 100644
--- a/module/Makefile.am
+++ b/module/Makefile.am
@@ -73,6 +73,7 @@ SOURCES = \
language/cps/specialize-numbers.scm \
language/cps/split-rec.scm \
language/cps/switch.scm \
+ language/cps/tailify.scm \
language/cps/type-checks.scm \
language/cps/type-fold.scm \
language/cps/types.scm \
diff --git a/module/language/cps/tailify.scm b/module/language/cps/tailify.scm
new file mode 100644
index 0000000..b5a7477
--- /dev/null
+++ b/module/language/cps/tailify.scm
@@ -0,0 +1,636 @@
+;;; Continuation-passing style (CPS) intermediate language (IL)
+
+;; Copyright (C) 2021 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:
+;;;
+;;; Tailification converts a program so that all calls are tail calls.
+;;; It is a minimal form of global CPS conversion that stack-allocates
+;;; "return continuations" -- minimal in the sense that the only
+;;; additionally residualized continuations are the ones necessary to
+;;; preserve the all-tail-calls property. Notably, loops, conditionals,
+;;; and similar features in the source program are left as is unless
+;;; it's necessary to split them.
+;;;
+;;; The first step of tailification computes the set of "tails" in a
+;;; function. The function entry starts a tail, as does each return
+;;; point from non-tail calls. Join points between different tails
+;;; also start tails.
+;;;
+;;; In the residual program, there are four ways that a continuation
+;;; exits:
+;;;
+;;; - Tail calls in the source program are tail calls in the residual
+;;; program; no change.
+;;;
+;;; - For non-tail calls in the source program, the caller saves the
+;;; state of the continuation (the live variables flowing into the
+;;; continuation) on an explicit stack, and saves the label of the
+;;; continuation. The return continuation will be converted into a
+;;; arity-checking function entry, to handle multi-value returns;
+;;; when it is invoked, it will pop its incoming live variables from
+;;; the continuation stack.
+;;;
+;;; - Terms that continue to a join continuation are converted to
+;;; label calls in tail position, passing the state of the
+;;; continuation as arguments.
+;;;
+;;; - Returning values from a continuation pops the return label from
+;;; the stack and does an indirect tail label call on that label,
+;;; with the given return values.
+;;;
+;;; Additionally, the abort-to-prompt run-time routine may unwind the
+;;; explicit stack and tail-call a handler continuation. If the
+;;; continuation is not escape-only, then the slice of the continuation
+;;; that would be popped off is captured before unwinding. Resuming a
+;;; continuation splats the saved continuation back on the stack and
+;;; returns to the top continuation, just as in the tail return case
+;;; above.
+;;;
+;;; We expect that a tailified program will probably be slower than a
+;;; non-tailified program. However a tailified program has a few
+;;; interesting properties: the stack is packed and only contains live
+;;; data; the stack can be traversed in a portable way, allowing for
+;;; implementation of prompts on systems that don't support them
+;;; natively; and as all calls are tail calls, the whole system can be
+;;; implemented naturally with a driver trampoline on targets that don't
+;;; support tail calls (e.g. JavaScript and WebAssembly).
+;;;
+;;; Code:
+
+(define-module (language cps tailify)
+ #:use-module (srfi srfi-1)
+ #:use-module (ice-9 match)
+ #:use-module (language cps)
+ #:use-module (language cps intmap)
+ #:use-module (language cps intset)
+ #:use-module (language cps graphs)
+ #:use-module (language cps utils)
+ #:use-module (language cps renumber)
+ #:use-module (language cps with-cps)
+ #:export (tailify))
+
+(define (trivial-intmap x)
+ (let ((next (intmap-next x)))
+ (and (eqv? next (intmap-prev x))
+ next)))
+
+(define (live-constants live-in constants head)
+ (intmap-select constants
+ (intmap-intersect (intmap-ref live-in head)
+ (intmap-keys constants))))
+(define (live-vars live-in constants head)
+ (intset-subtract (intmap-ref live-in head)
+ (intmap-keys constants)))
+
+(define (rename-var* fresh-names var)
+ (intmap-ref fresh-names var (lambda (var) var)))
+(define (rename-vars* fresh-names vars)
+ (match vars
+ (() '())
+ ((var . vars)
+ (cons (rename-var* fresh-names var)
+ (rename-vars* fresh-names vars)))))
+
+(define (compute-saved-vars* fresh-names live-in constants reprs k)
+ (intset-fold-right
+ (lambda (var reprs vars)
+ (values (cons (intmap-ref reprs var) reprs)
+ (cons (rename-var* fresh-names var) vars)))
+ (live-vars live-in constants k) '() '()))
+
+(define (tailify-tail cps head body fresh-names winds live-in constants
+ reprs entries original-ktail)
+ "Rewrite the conts with labels in the intset BODY, forming the body of
+the tail which begins at HEAD in the source program. The entry to the
+tail was already rewritten, with ENTRIES containing an intmap of tail
+heads to $kfun labels. WINDS associates 'unwind primcalls with the
+corresponding conts that pushes on the dynamic stack. LIVE-IN indicates
+the variables that are live at tail heads, and CONSTANTS is an intmap
+associating vars known to be constant with their values. REPRS holds
+the representation of each var. ORIGINAL-KTAIL is the tail cont of the
+source function; terms in the tail that continue to ORIGINAL-KTAIL will
+be rewritten to continue to the tail's ktail."
+
+ ;; HEAD will have been given a corresponding entry $kfun by
+ ;; tailify-tails. Here we find the tail-label for the current tail.
+ (define local-ktail
+ (match (intmap-ref cps (intmap-ref entries head))
+ (($ $kfun src meta self ktail kentry)
+ ktail)))
+
+ (define (rename-var var) (rename-var* fresh-names var))
+ (define (rename-vars vars) (rename-vars* fresh-names vars))
+ (define (rename-exp exp)
+ (rewrite-exp exp
+ ((or ($ $const) ($ $prim) ($ $const-fun) ($ $code)) ,exp)
+ (($ $call proc args)
+ ($call (rename-var proc) ,(rename-vars args)))
+ (($ $callk k proc args)
+ ($callk k (and proc (rename-var proc)) ,(rename-vars args)))
+ (($ $primcall name param args)
+ ($primcall name param ,(rename-vars args)))
+ (($ $values args)
+ ($values ,(rename-vars args)))))
+
+ (define (compute-saved-vars fresh-names k)
+ (compute-saved-vars* fresh-names live-in constants reprs k))
+
+ ;; Return a $callk to the join tail with head K. To allow for
+ ;; tail-local names for values bound by K, JOIN-VARS is an alist of
+ ;; mappings to add to FRESH-NAMES.
+ (define (compute-join-call join-vars k)
+ (let ((fresh-names (fold (lambda (pair fresh-names)
+ (match pair
+ ((old . new)
+ (intmap-add fresh-names old new))))
+ fresh-names join-vars)))
+ (call-with-values (lambda () (compute-saved-vars fresh-names k))
+ (lambda (reprs vars)
+ (build-exp
+ ($callk (intmap-ref entries k) #f vars))))))
+
+ ;; A branch target can either be in the current tail, or it starts a
+ ;; join continuation. It can't be $ktail, it can't be $kreceive, and
+ ;; it takes no values, hence we pass () to compute-join-call.
+ (define (rewrite-branch-target cps src k)
+ (cond
+ ((intset-ref body k)
+ (with-cps cps k))
+ (else
+ (when (eqv? k original-ktail) (error "what!!"))
+ (with-cps cps
+ (letk kcall
+ ($kargs () ()
+ ($continue local-ktail src ,(compute-join-call '() k))))
+ kcall))))
+ (define (rewrite-branch-targets cps src k*)
+ (match k*
+ (()
+ (with-cps cps '()))
+ ((k . k*)
+ (with-cps cps
+ (let$ k* (rewrite-branch-targets src k*))
+ (let$ k (rewrite-branch-target src k))
+ (cons k k*)))))
+
+ ;; Rewrite TERM. Generally speaking we just rename variable uses.
+ ;; However if TERM continues to another tail, we have to generate the
+ ;; appropriate call for the continuation tail kind.
+ (define (rewrite-term cps term)
+ (match term
+ (($ $continue k src exp)
+ (let ((exp (rename-exp exp)))
+ (cond
+ ((eqv? k original-ktail)
+ (match exp
+ (($ $values args)
+ ;; The original term is a $values in tail position.
+ ;; Transform to pop the continuation stack and tail call
+ ;; it.
+ (with-cps cps
+ (letv ret)
+ (letk kcall ($kargs ('ret) (ret)
+ ($continue local-ktail src
+ ($calli args ret))))
+ (build-term ($continue kcall src
+ ($primcall 'restore1 'ptr ())))))
+ ((or ($ $call) ($ $callk) ($ $calli))
+ ;; Otherwise the original term was a tail call.
+ (with-cps cps
+ (build-term ($continue local-ktail src ,exp))))))
+ ((intset-ref body k)
+ ;; Continuation within current tail.
+ (with-cps cps
+ (build-term ($continue k src ,exp))))
+ (else
+ (match (intmap-ref cps k)
+ (($ $kreceive)
+ ;; A non-tail-call: push the pending continuation and tail
+ ;; call instead.
+ (match exp
+ ((or ($ $call) ($ $callk) ($ $calli))
+ (call-with-values (lambda ()
+ (compute-saved-vars fresh-names k))
+ (lambda (reprs vars)
+ (with-cps cps
+ (letk kexp ($kargs () ()
+ ($continue local-ktail src ,exp)))
+ (letv cont)
+ (letk kcont ($kargs ('cont) (cont)
+ ($continue kexp src
+ ($primcall 'save
+ (append reprs (list 'ptr))
+ ,(append vars (list
cont))))))
+ (build-term ($continue kcont src
+ ($code (intmap-ref entries k))))))))))
+ (($ $kargs names vars)
+ ;; Calling a join continuation. This is one of those
+ ;; cases where it might be nice in CPS to have names for
+ ;; phi predecessor values. Ah well.
+ (let ((vars' (map (lambda (_) (fresh-var)) vars)))
+ (with-cps cps
+ (letk kvals
+ ($kargs names vars'
+ ($continue local-ktail src
+ ,(compute-join-call (map cons vars vars') k))))
+ (build-term
+ ($continue kvals src ,exp))))))))))
+ (($ $branch kf kt src op param args)
+ (with-cps cps
+ (let$ kf (rewrite-branch-target src kf))
+ (let$ kt (rewrite-branch-target src kf))
+ (build-term
+ ($branch kf kt src op param ,(rename-vars args)))))
+ (($ $switch kf kt* src arg)
+ (with-cps cps
+ (let$ kf (rewrite-branch-target src kf))
+ (let$ kt* (rewrite-branch-targets src kt*))
+ (build-term ($switch kf kt* src (rename-var arg)))))
+ (($ $prompt k kh src escape? tag)
+ (call-with-values (lambda () (compute-saved-vars fresh-names kh))
+ (lambda (reprs vars)
+ (with-cps cps
+ (letv handler)
+ (let$ k (rewrite-branch-target src k))
+ (letk kpush ($kargs ('handler) (handler)
+ ($continue k src
+ ($primcall 'push-prompt escape?
+ ((rename-var tag) handler)))))
+ (letk kcode ($kargs () ()
+ ($continue kpush src ($code (intmap-ref entries
kh)))))
+ (build-term ($continue kpush src
+ ($primcall 'save reprs vars)))))))
+ (($ $throw src op param args)
+ (with-cps cps
+ (build-term ($throw src op param ,(rename-vars args)))))))
+
+ ;; A prompt body begins with a $prompt, may contain nested prompt
+ ;; bodies, and continues until a corresponding 'unwind primcall.
+ ;; Leaving a prompt body may or may not correspond to leaving the
+ ;; current tail. Leaving the prompt body must remove the handler from
+ ;; the stack. Removing the handler must happen before leaving the
+ ;; tail, and notably must happen before pushing saved state for a
+ ;; non-tail-call continuation.
+ (define (maybe-unwind-prompt cps label term)
+ (define (not-a-prompt-unwind) (with-cps cps term))
+ (define (pop-prompt kh)
+ (call-with-values (lambda () (compute-saved-vars fresh-names kh))
+ (lambda (reprs vars)
+ (with-cps cps
+ (letk kterm ($kargs () () ,term))
+ (build-term ($continue kterm #f
+ ($primcall 'drop reprs ())))))))
+ (cond
+ ((intmap-ref winds label (lambda (_) #f))
+ => (lambda (wind)
+ (match (intmap-ref cps wind)
+ (($ $prompt k kh) (pop-prompt kh))
+ (_ (not-a-prompt-unwind)))))
+ (else (not-a-prompt-unwind))))
+
+ ;; The entry for the current tail has already been rewritten, so here
+ ;; we just rewrite all the body conts.
+ (intset-fold
+ (lambda (label cps)
+ (match (intmap-ref cps label)
+ ((or ($ $kfun) ($ $kclause) ($ $ktail)) cps) ;; Unchanged.
+ (($ $kreceive) cps) ;; Dead.
+ (($ $kargs names vals term)
+ (with-cps cps
+ (let$ term (rewrite-term term))
+ (let$ term (maybe-unwind-prompt label term))
+ (setk label ($kargs names vals ,term))))))
+ body cps))
+
+(define (tailify-tails cps winds live-in constants reprs tails)
+ "Given that the conts in a function were partitioned into tails in the
+intmap TAILS, mapping tail entries to tail bodies, return a new CPS
+program in which the tails have been split to separate functions in
+which all calls are tail calls.
+
+WINDS associates 'unwind primcalls with the corresponding conts that
+pushes on the dynamic stack.
+
+LIVE-IN indicates the variables that are live at tail heads.
+
+CONSTANTS is an intmap associating vars known to be constant with their
+values.
+
+REPRS holds the representation of each var."
+
+ (define (cont-source label)
+ (match (intmap-ref cps label)
+ (($ $kargs _ _ term)
+ (match term
+ (($ $continue k src) src)
+ (($ $branch k kt src) src)
+ (($ $switch k kt* src) src)
+ (($ $prompt k kh src) src)
+ (($ $throw src) src)))))
+
+ ;; For live values that flow into a tail, each tail will need to give
+ ;; them unique names.
+ (define fresh-names-per-tail
+ (intmap-map (lambda (head body)
+ (intset-fold (lambda (var fresh)
+ (intmap-add fresh var (fresh-var)))
+ (intmap-ref live-in head)
+ empty-intmap))
+ tails))
+
+ (define (compute-saved-vars head)
+ (compute-saved-vars* (intmap-ref fresh-names-per-tail head)
+ live-in constants reprs head))
+
+ ;; For a tail whose head in the source program is HEAD, rewrite to be
+ ;; a $kfun. For the "main" tail, no change needed. For join tails,
+ ;; we make an unchecked $kfun-to-$kargs function to which live
+ ;; variables are received directly as arguments. For return tails,
+ ;; the live vars are restored from the stack. In all cases, adjoin a
+ ;; HEAD->ENTRY mapping to ENTRIES, where ENTRY is the $kfun label for
+ ;; the tail.
+ (define (add-entry head cps entries)
+ (define fresh-names (intmap-ref fresh-names-per-tail head))
+ ;; Constants don't need to be passed from tail to tail; rather they
+ ;; are rebound locally.
+ (define (restore-constants cps term)
+ (intmap-fold (lambda (var exp cps term)
+ (define var' (intmap-ref fresh-names var))
+ (with-cps cps
+ (letk k ($kargs ('const) (var') ,term))
+ (build-term ($continue k #f ,exp))))
+ (live-constants live-in constants head)
+ cps term))
+ (define (restore-saved cps term)
+ (call-with-values (lambda () (compute-saved-vars head))
+ (lambda (reprs vars)
+ (define names (map (lambda (_) 'restored) vars))
+ (if (null? names)
+ (with-cps cps term)
+ (with-cps cps
+ (letk krestore ($kargs names vars ,term))
+ (build-term ($continue krestore #f
+ ($primcall 'restore reprs ()))))))))
+ (match (intmap-ref cps head)
+ (($ $kfun)
+ ;; The main entry.
+ (values cps (intmap-add entries head head)))
+ (($ $kreceive ($ $arity req () rest () #f) kargs)
+ ;; The continuation of a non-tail call, or a prompt handler.
+ (match (intmap-ref cps kargs)
+ (($ $kargs names vars)
+ (let ((vars' (map (lambda (_) (fresh-var)) vars))
+ (src (cont-source kargs)))
+ (with-cps cps
+ (letk ktail ($ktail))
+ (let$ term (restore-constants
+ (build-term
+ ($continue kargs src ($values vars')))))
+ (let$ term (restore-saved term))
+ (letk krestore ($kargs names vars' ,term))
+ (letk kclause ($kclause (req '() rest '() #f) krestore #f))
+ (letk kfun ($kfun src #f #f ktail kclause))
+ (intmap-add entries head kfun))))))
+ (($ $kargs names vars term)
+ ;; A join point.
+ (call-with-values (lambda () (compute-saved-vars head))
+ (lambda (reprs vars')
+ (define names'
+ (let ((names (map acons vars names)))
+ (map (lambda (var) (assq-ref names var))
+ vars')))
+ (define meta `((arg-representations . ,reprs)))
+ (with-cps cps
+ (letk ktail ($ktail))
+ (let$ term (restore-constants term))
+ (letk kargs ($kargs names' vars' ,term))
+ (letk kfun ($kfun (cont-source head) meta #f ktail kargs))
+ (intmap-add entries head kfun)))))))
+
+ (define original-ktail
+ (match (intmap-ref cps (intmap-next tails))
+ (($ $kfun src meta self ktail kentry)
+ ktail)))
+ (call-with-values (lambda ()
+ (intmap-fold (lambda (head body cps entries)
+ (add-entry head cps entries))
+ tails cps empty-intmap))
+ (lambda (cps entries)
+ (intmap-fold
+ (lambda (head body cps)
+ (define fresh-names (intmap-ref fresh-names-per-tail head))
+ (tailify-tail cps head body fresh-names winds live-in constants
+ reprs entries original-ktail))
+ tails cps))))
+
+(define (compute-tails kfun body preds cps)
+ "Compute the set of tails in the function with entry KFUN and body
+BODY. Return as an intset mapping the head label for each tail to its
+body, as an intset."
+ ;; Initially, we start with the requirement that kfun and kreceive
+ ;; labels are split heads.
+ (define (initial-split label splits)
+ (match (intmap-ref cps label)
+ ((or ($ $kfun) ($ $kreceive))
+ (intmap-add splits label label))
+ (_
+ splits)))
+ ;; Then we build tails by propagating splits forward in the CFG,
+ ;; possibly creating new split heads at the dominance frontier.
+ (define (compute-split label splits)
+ (define (split-head? label)
+ (eqv? label (intmap-ref splits label (lambda (_) #f))))
+ (define (ktail? label)
+ (match (intmap-ref cps label)
+ (($ $ktail) #t)
+ (_ #f)))
+ (cond
+ ((split-head? label)
+ ;; Once a label is a split head, it stays a split head.
+ splits)
+ ((ktail? label)
+ ;; ktail always part of root tail.
+ (intmap-add splits label kfun))
+ (else
+ (match (intset-fold
+ (lambda (pred pred-splits)
+ (define pred-split
+ (intmap-ref splits pred (lambda (_) #f)))
+ (match pred-split
+ (#f pred-splits)
+ (split (cons split pred-splits))))
+ (intmap-ref preds label) '())
+ ((split)
+ ;; If all predecessors in same split, label is too.
+ (intmap-add splits label split (lambda (old new) new)))
+ ((_ _ . _)
+ ;; Otherwise this is a new split.
+ (intmap-add splits label label (lambda (old new) new)))))))
+ ;; label -> split head
+ (define initial-splits
+ (intset-fold initial-split body empty-intmap))
+ (cond
+ ((trivial-intmap initial-splits)
+ ;; There's only one split head, so only one tail.
+ (intmap-add empty-intmap kfun body))
+ (else
+ ;; Otherwise, assign each label to a tail, identified by the split
+ ;; head, then collect the tails by split head.
+ (let ((splits (fixpoint
+ (lambda (splits)
+ (intset-fold compute-split body splits))
+ initial-splits)))
+ (intmap-fold
+ (lambda (label head split-bodies)
+ (intmap-add split-bodies head (intset label) intset-union))
+ splits
+ empty-intmap)))))
+
+(define (intset-pop set)
+ "Return two values: all values in intset SET except the first one, and
+first value in SET, or #f if SET was empty."
+ (match (intset-next set)
+ (#f (values set #f))
+ (i (values (intset-remove set i) i))))
+
+(define (identify-winds cps kfun body succs)
+ "For each unwind primcall in BODY, adjoin an entry mapping it to the
+corresponding wind expression."
+ (define (visit-label label exits bodies)
+ (define wind (intmap-ref bodies label))
+ (match (intmap-ref cps label)
+ (($ $kargs _ _ ($ $prompt k kh))
+ (let* ((bodies (intmap-add bodies k label))
+ (bodies (intmap-add bodies kh wind)))
+ (values exits bodies)))
+ (($ $kargs _ _ ($ $continue k _ ($ $primcall 'wind)))
+ (let ((bodies (intmap-add bodies k wind)))
+ (values exits bodies)))
+ (($ $kargs _ _ ($ $continue k _ ($ $primcall 'unwind)))
+ (let* ((exits (intmap-add exits label wind))
+ (bodies (intmap-add bodies k (intmap-ref bodies wind))))
+ (values exits bodies)))
+ (else
+ (let ((bodies (intset-fold (lambda (succ bodies)
+ (intmap-add bodies succ wind))
+ (intmap-ref succs label)
+ bodies)))
+ (values exits bodies)))))
+ (values
+ (worklist-fold
+ (lambda (to-visit exits bodies)
+ (call-with-values (lambda () (intset-pop to-visit))
+ (lambda (to-visit label)
+ (call-with-values (lambda () (visit-label label exits bodies))
+ (lambda (exits* bodies*)
+ (if (and (eq? exits exits*) (eq? bodies bodies*))
+ (values to-visit exits bodies)
+ (values (intset-union to-visit (intmap-ref succs label))
+ exits* bodies*)))))))
+ (intset kfun)
+ empty-intmap
+ (intmap-add empty-intmap kfun #f))))
+
+(define (compute-live-in cps body preds)
+ "Return an intmap associating each label in BODY with an intset of
+live variables flowing into the label."
+ (let ((function (intmap-select cps body)))
+ (call-with-values
+ (lambda ()
+ (call-with-values (lambda () (compute-defs-and-uses function))
+ (lambda (defs uses)
+ ;; Unlike the use of compute-live-variables in
+ ;; slot-allocation.scm, we don't need to add prompt
+ ;; control-flow edges, as the prompt handler is in its own
+ ;; tail and therefore $prompt will push the handler
+ ;; continuation (including its needed live vars) before
+ ;; entering the prompt body.
+ (compute-live-variables preds defs uses))))
+ (lambda (live-in live-out)
+ live-in))))
+
+(define (compute-constants cps preds)
+ "Return an intmap associating each variables BODY to their defining
+expression, for all variables binding constant expressions."
+ (define (constant? exp)
+ (match exp
+ ((or ($ $const) ($ $prim) ($ $const-fun) ($ $code)) #t)
+ (_ #f)))
+ (intmap-fold
+ (lambda (label preds constants)
+ (cond
+ ((trivial-intset preds)
+ => (lambda (pred)
+ (match (intmap-ref cps pred)
+ (($ $continue _ _ (? constant? exp))
+ (match (intmap-ref cps label)
+ (($ $kargs (_) (var) _)
+ (intmap-add constants var exp))))
+ (_
+ constants))))
+ (else constants)))
+ preds empty-intmap))
+
+(define (tailify-trivial-tail body cps)
+ "For the function with body BODY and only one tail, rewrite any return
+to tail-call the saved continuation."
+ (define (ktail? k)
+ (match (intmap-ref cps k)
+ (($ $ktail) #t)
+ (_ #f)))
+ (define (rewrite-return-to-pop-and-calli label cps)
+ (match (intmap-ref cps label)
+ (($ $kargs names vars
+ ($ $continue (? ktail? k) src ($ $values args)))
+ ;; The original term is a $values in tail position.
+ ;; Transform to pop the continuation stack and tail
+ ;; call it.
+ (with-cps cps
+ (letv ret)
+ (letk kcall ($kargs ('ret) (ret)
+ ($continue k src ($calli args ret))))
+ (setk label ($kargs names vars
+ ($continue kcall src
+ ($primcall 'restore1 'ptr ()))))))))
+ (intset-fold rewrite-return-to-pop-and-calli body cps))
+
+(define (tailify-function kfun body cps)
+ "Partition the function with entry of KFUN into tails. Rewrite all
+tails in such a way that they enter via a $kfun and leave only via tail
+calls."
+ (define succs (compute-successors cps kfun))
+ (define preds (invert-graph succs))
+ (define tails (compute-tails kfun body preds cps))
+ (cond
+ ((trivial-intmap tails)
+ (tailify-trivial-tail body cps))
+ (else
+ ;; Otherwise we apply tailification.
+ (let ((winds (identify-winds cps kfun body succs))
+ (live-in (compute-live-in cps body preds))
+ (constants (compute-constants cps preds))
+ (reprs (compute-var-representations cps)))
+ (tailify-tails cps winds live-in constants reprs tails)))))
+
+(define (tailify cps)
+ ;; Renumber so that label order is topological order.
+ (let ((cps (renumber cps)))
+ (with-fresh-name-state cps
+ (intmap-fold tailify-function
+ (compute-reachable-functions cps)
+ cps))))