racket/collects/frtime/opt/frtime-opt.rkt
2010-04-27 16:50:15 -06:00

773 lines
37 KiB
Racket

;; This module serves as a language module for an optimized version of FrTime.
;; The only thing it exports is a #%module-begin macro that knows how to
;; perform optimization of FrTime functions. The rest of the language
;; (i.e. all the functions needed to actually write FrTime programs) is
;; provided by the frtime-opt-lang module, which is automatically imported.
(module frtime-opt mzscheme
(provide (rename my-module-begin #%module-begin)
#%app #%top #%datum optimize-expr optimize-module dont-optimize)
(require-for-syntax frtime/opt/lowered-equivs)
(require-for-syntax (only srfi/1 lset-union lset-difference every))
(require-for-syntax mzlib/list)
(require (only frtime/core/frp super-lift undefined undefined?))
(require (rename frtime/lang-ext frtime:lift lift)
(rename frtime/lang-core frtime:if if)
(only frtime/lang-core frp:copy-list))
; (require mzlib/unit mzlib/unitsig)
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Helper functions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; prevent a subexpression from being optimized. Calls to this function
;; are specifically recognized, and handled specially.
(define (dont-optimize x) x)
;; like module-identifier=?, but returns #f for non-identifier arguments
;; instead of throwing an exception
(define-for-syntax (safe-module-identifier=? id1 id2)
(and (identifier? id1)
(identifier? id2)
(module-identifier=? id1 id2)))
;; Convert a syntax-object to a datum and back again. This replaces all
;; the context information. It's necessary to get generated require
;; statements to work, for some reason.
;; See http://list.cs.brown.edu/pipermail/plt-scheme/2006-July/014163.html
(define-for-syntax (so->d->so ref-stx stx)
(datum->syntax-object ref-stx (syntax-object->datum stx)))
;; Convert a syntactic module reference to a module-path-index
(define-for-syntax (module-stx-to-path-index mod-stx)
(let ([mod (syntax-object->datum mod-stx)])
(if (symbol? mod)
mod
(module-path-index-join mod #f))))
;; Convert a syntactic module reference to a module-path (this is
;; subtly different from a module-path-index)
(define-for-syntax (module-stx-to-path mod-stx)
(syntax-object->datum mod-stx))
;; Does the given module export the given id?
(define-for-syntax (module-exports-id? mod-stx id-stx)
(not (module-provide-protected? (module-stx-to-path-index mod-stx)
(syntax-e id-stx))))
;; Returns all the identifiers exported by a module
(define-for-syntax (all-provided-ids mod-stx ref-stx)
;; we get exn:fail:filesystem if the module doesn't exist.
;; we get exn:fail:contract during compilation if compiling a module that
;; imports mred.
(with-handlers ([exn:fail:filesystem? (lambda (exn) null)]
[exn:fail:contract? (lambda (exn) null)])
(let ([mod-path (module-stx-to-path mod-stx)]
[mod-path-index (module-stx-to-path-index mod-stx)])
;; instantiate the module so we can call module->namespace on it
(dynamic-require mod-path #f)
;; get the list of provided symbols
(let* ([all-symbols (namespace-mapped-symbols (module->namespace mod-path))]
[exported-symbols (filter
(lambda (id) (not (module-provide-protected? mod-path-index id)))
all-symbols)]
[exported-ids (map
(lambda (symbol) (datum->syntax-object ref-stx symbol))
exported-symbols)])
exported-ids))))
;; This macro takes a list of variables and an expression.
;; The variables are projected before evaluating the expression,
;; and the result is then injected into the dataflow graph as a
;; single node.
(require (only frtime/core/frp proc->signal value-now))
(define-syntax (dip stx)
(syntax-case stx (begin)
;; special case: don't dip lone identifiers
[(_ (VAR ...) VAR2)
(identifier? #'VAR2)
#'VAR2]
;; special case: strip off unnecessary begins
[(_ (VAR ...) (begin E))
#'(dip (VAR ...) E)]
;; special case: don't bother dipping if there are no dependencies
[(_ () EXPR)
#'EXPR]
;; general case: wrap the subexpression in a lambda, and lift it
[(_ (VAR ...) EXPR)
#'(frtime:lift
#t
(lambda (VAR ...) EXPR)
VAR ...)]))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Module-level optimization
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Wrap the entire module body in an optimize-module macro.
(define-syntax (my-module-begin stx)
(syntax-case stx ()
[(_ FORMS ...)
(let* (;; get a list of all the symbols provided by the frtime-opt-lang module
[lang-symbols (all-provided-ids #'frtime/opt/frtime-opt-lang stx)]
;; convert those symbols into an equiv-map by pairing up functions
;; with their lowered equivalents
[lang-equiv-map (symbol-list-to-equiv-map lang-symbols)]
;; convert the equiv-map to syntax so we can embed it in the call
;; to optimize-module
[equiv-map-stx (equiv-map-to-stx lang-equiv-map)])
#`(#%plain-module-begin
(require-for-syntax #,(so->d->so stx #`mzscheme))
(require #,(so->d->so stx #`frtime/opt/frtime-opt-lang))
(optimize-module #,equiv-map-stx FORMS ...)))]))
;; Expand a module body until it's just top-level definitions, and then
;; separately optimize each kind of top-level definition.
(define-syntax (optimize-module stx)
(syntax-case stx ()
[(_ EQUIV-MAP)
#`(begin)]
[(_ EQUIV-MAP FORM FORMS ...)
(let ([expanded-form
(local-expand #'FORM 'module
(list #'begin #'begin0 #'#%provide #'#%require
#'define-syntaxes #'define-values-for-syntax
#'define-values #'#%app #'unit #'unit/sig))])
(syntax-case expanded-form (begin begin0 #%provide #%require
define-syntaxes define-values-for-syntax
define-values #%app)
;; explode top-level begin statements
[(begin MORE-FORMS ...)
#`(optimize-module EQUIV-MAP MORE-FORMS ... FORMS ...)]
;; require
[(#%require . __)
#`(optimize-require EQUIV-MAP #,expanded-form FORMS ...)]
;; provide
[(#%provide . __)
;; TBD: provide lowered equivs as well.
;; TBD: support frtime-specific provide specs (lifted, etc)
#`(begin #,expanded-form
(optimize-module EQUIV-MAP FORMS ...))]
;; syntax definitions
[(define-syntaxes . __)
#`(begin #,expanded-form
(optimize-module EQUIV-MAP FORMS ...))]
[(define-values-for-syntax . __)
#`(begin #,expanded-form
(optimize-module EQUIV-MAP FORMS ...))]
;; top-level variable definitions
[(define-values (ID ...) VAL)
#`(begin
(optimize-definition EQUIV-MAP #,expanded-form FORMS ...))]
;; expressions
[else
#`(begin
(optimize-expr EQUIV-MAP #,expanded-form)
(optimize-module EQUIV-MAP FORMS ...))]))]))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Require/Provide handling
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Rewrite require forms so that they import not only regular identifiers,
;; but also lowered equivalents, and add those lowered equivs to the equiv map.
(define-syntax (optimize-require stx)
;; Take a module and a list of identifiers, and return the list of
;; associated lowered equivalents (for those identifiers that have one).
(define (module-lowered-equivs mod-stx id-stx-list)
(map make-lowered-equiv-id
(filter (lambda (id-stx)
(module-exports-id? mod-stx (make-lowered-equiv-id id-stx)))
id-stx-list)))
(syntax-case stx (#%require)
[(_ EQUIV-MAP (#%require SPEC) FORMS ...)
;; In the code below, we must convert syntax objects to datums, and then
;; back again. Otherwise the generated require statement doesn't
;; seem to work properly (it doesn't import the correct identifiers).
;; See http://list.cs.brown.edu/pipermail/plt-scheme/2006-July/014163.html
(syntax-case #'SPEC (only prefix all-except prefix-all-except rename
lifted lifted:nonstrict as-is as-is:unchecked)
[(only MOD ID ...)
;; Add lowered equivs to the list of ids to import
(let ([lowered-equivs
(module-lowered-equivs #'MOD (syntax->list #'(ID ...)))]
[new-equiv-map #'EQUIV-MAP])
#`(begin
(require #,(so->d->so #'MOD #`(only MOD ID ... . #,lowered-equivs)))
(optimize-module #,new-equiv-map FORMS ...)))]
[(prefix PFX MOD)
;; Requiring the entire module with a prefix will automatically import
;; the lowered-equiv bindings with the same prefix.
(let ([new-equiv-map #'EQUIV-MAP])
#`(begin
(require #,(so->d->so #'MOD #`(prefix PFX MOD)))
(optimize-module #,new-equiv-map FORMS ...)))]
[(all-except MOD ID ...)
;; Add the lowered-equiv ids to the exclude list. If they aren't actually
;; provided by the module then they will be silently ignored.
(let ([lowered-equivs
(module-lowered-equivs #'MOD (syntax->list #'(ID ...)))]
[new-equiv-map #'EQUIV-MAP])
#`(begin
(require #,(so->d->so #'MOD #`(all-except MOD ID ... . #,lowered-equivs)))
(optimize-module #,new-equiv-map FORMS ...)))]
[(prefix-all-except PFX MOD ID ...)
;; Add the lowered-equiv bindings to the exclude list.
(let ([lowered-equivs
(module-lowered-equivs #'MOD (syntax->list #'(ID ...)))]
[new-equiv-map #'EQUIV-MAP])
#`(begin
(require #,(so->d->so #'MOD #`(prefix-all-except PFX MOD ID ... . #,lowered-equivs)))
(optimize-module #,new-equiv-map FORMS ...)))]
[(rename MOD LOCAL-ID EXPORTED-ID)
;; Rename the lowered-equiv binding as well.
(let* ([exported-lowered-equiv-id (make-lowered-equiv-id #'EXPORTED-ID)]
[local-lowered-equiv-id (make-lowered-equiv-id #'LOCAL-ID)]
[has-lowered-equiv (module-exports-id? #'MOD
exported-lowered-equiv-id)])
(if has-lowered-equiv
#`(begin
(require #,(so->d->so #'MOD #`(rename MOD LOCAL-ID EXPORTED-ID)))
(require #,(so->d->so #'MOD #`(rename MOD
#,local-lowered-equiv-id
#,exported-lowered-equiv-id)))
(optimize-module ((LOCAL-ID #,local-lowered-equiv-id) . EQUIV-MAP) FORMS ...))
#`(begin
(require #,(so->d->so #'MOD #`(rename MOD LOCAL-ID EXPORTED-ID)))
(optimize-module EQUIV-MAP FORMS ...))))]
[(LIFTED MOD ID)
;; import the identifier itself as the lowered equiv, and define a new, lifted, version.
(and (identifier? #'LIFTED)
(or (free-identifier=? #'LIFTED #'lifted)
(free-identifier=? #'LIFTED #'lifted:nonstrict)))
(let* ([lowered-equiv-id (make-lowered-equiv-id #'ID)]
[strict? (datum->syntax-object #'MOD (free-identifier=? #'LIFTED #'lifted))])
#`(begin
(require #,(so->d->so #'MOD #`(rename MOD #,lowered-equiv-id ID)))
(define (ID . args) (apply frtime:lift #,strict? #,lowered-equiv-id args))
(optimize-require ((ID #,lowered-equiv-id) . EQUIV-MAP) FORMS ...)))]
[(LIFTED MOD ID IDS ...)
;; only import one lifted identifier at a time
(and (identifier? #'LIFTED)
(or (free-identifier=? #'LIFTED #'lifted)
(free-identifier=? #'LIFTED #'lifted:nonstrict)))
#`(optimize-require EQUIV-MAP
(require (LIFTED MOD ID))
(require (LIFTED MOD IDS ...))
FORMS ...)]
[(AS-IS MOD IDS ...)
;; as-is and as-is:unchecked are treated just like (require (only ...))
(and (identifier? #'AS-IS)
(or (free-identifier=? #'AS-IS #'as-is)
(free-identifier=? #'AS-IS #'as-is:unchecked)))
#`(optimize-require EQUIV-MAP
(require #,(so->d->so #'MOD #`(only MOD IDS ...)))
FORMS ...)]
[MOD
;; Requiring an entire module will automatically import the lowered-equiv
;; bindings, so we don't need to change the require directive itself, just
;; the equiv map.
;; Vote: it's very important to use #'SPEC as the ref-stx here. If you
;; use #'MOD or stx instead, then the identifiers returned by all-provided-ids
;; will not be module-identifier=? to the identifiers in the actual code.
;; TBD: make the initial import of frtime-opt-lang use this same mechanism.
(let* ([additional-equiv-map (symbol-list-to-equiv-map
(all-provided-ids #'MOD #'SPEC))]
[new-equiv-map-stx (equiv-map-to-stx
(union-equiv-maps additional-equiv-map
(stx-to-equiv-map #'EQUIV-MAP)))])
#`(begin (require MOD)
(optimize-module #,new-equiv-map-stx FORMS ...)))]
)]
[(_ EQUIV-MAP (#%require SPECS ...) FORMS ...)
;; Process each require spec individually.
#`(optimize-require EQUIV-MAP (require SPECS) ... FORMS ...)]
))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Expression optimization
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Helper functions for working with dipped expressions
(define-for-syntax (dipped? stx)
(syntax-case stx (dip)
[(dip . _) #t]
[ELSE #f]))
(define-for-syntax (get-dipped-expr stx)
(syntax-case stx (dip)
[(dip DEPS EXPR)
#'EXPR]))
(define-for-syntax (get-dipped-deps stx)
(syntax-case stx (dip)
[(dip (DEP ...) EXPR)
(syntax->list #'(DEP ...))]))
(define-for-syntax (union-id-lists . id-lists)
(foldl (lambda (l1 l2)
(lset-union bound-identifier=? l1 l2))
()
id-lists))
(define-for-syntax (diff-id-lists id-list1 id-list2)
(lset-difference bound-identifier=? id-list1 id-list2))
;; return the list of identifiers in an argument list.
;; #'(a b c) => (list #'a #'b #'c)
;; #'(x y . rest) => (list #'x #'y #'rest)
;; #'args => (list #'args)
(define-for-syntax (extract-args stx)
(syntax-case stx ()
[ID
(identifier? #'ID)
(list #'ID)]
[(ID . REST)
(append (extract-args #'ID)
(extract-args #'REST))]
[()
()]
[ELSE
(raise-syntax-error #f "doesn't look like an arg list" stx)]))
;; return the rest arg from an argument list.
;; #'(a b . rest) => #'rest
;; #'args => #'args
(define-for-syntax (extract-rest-arg stx)
(syntax-case stx ()
[ID
(identifier? #'ID)
#'ID]
[(ID . REST)
(identifier? #'REST)
#'REST]
[(ID . REST)
(extract-rest-arg #'REST)]
[ELSE
(raise-syntax-error #f "cannot extract rest arg" stx)]))
;; search and replace an identifier in a syntax object
(define-for-syntax (replace-id stx old-id new-id)
(syntax-case stx ()
[ID
(and (identifier? #'ID)
(module-identifier=? old-id #'ID))
new-id]
[(X . Y)
(datum->syntax-object
stx
(cons (replace-id #'X old-id new-id)
(replace-id #'Y old-id new-id)))]
[()
()]
[ELSE
stx]))
;; return #t if a syntax object refers to a given id
(define-for-syntax (refers-to-id? stx id)
(syntax-case stx ()
[ID
(and (identifier? #'ID)
(module-identifier=? id #'ID))
#t]
[(X . Y)
(or (refers-to-id? #'X id)
(refers-to-id? #'Y id))]
[ELSE #f]))
;; Optimize an expression by dipping subexpressions wherever possible.
;;
;; Note: this function does not intercept syntax errors due to misused
;; certificates. It's up to the caller to handle that.
(define-for-syntax (recursively-optimize-expr stx equiv-map lower-lambda)
(syntax-case stx (#%top #%app quote begin begin0 lambda case-lambda
let-values letrec-values letrec-syntaxes+values
unit unit/sig if super-lift undefined? undefined rename
frp:copy-list frp:->boolean dont-optimize)
[(#%top . X)
#`(dip (X) X)]
[X
(identifier? #'X)
#`(dip (X) X)]
[(quote EXPR)
#`(dip () #,stx)]
[(BEGIN EXPR ...)
(and (identifier? #'BEGIN)
(or (module-identifier=? #'BEGIN #'begin)
(module-identifier=? #'BEGIN #'begin0)))
(let* ([optimized-exprs (map (lambda (expr)
(recursively-optimize-expr expr equiv-map #f))
(syntax->list #'(EXPR ...)))]
[all-exprs-were-dipped (every dipped? optimized-exprs)])
(if all-exprs-were-dipped
(let ([deps (apply union-id-lists (map get-dipped-deps optimized-exprs))]
[lowered-exprs (map get-dipped-expr optimized-exprs)])
#`(dip #,deps (begin . #,lowered-exprs)))
;; we can't dip the entire subexpression, but we may have been able
;; to optimize some of the exprs.
#`(begin . #,optimized-exprs)))]
[(lambda ARGS (let-values (((REST) (#%app frp:copy-list REST_))) EXPRS ...))
;; If frtime sees a lambda with a rest arg, it inserts a call to
;; frp:copy-list. Recognize this pattern and ignore it if we're lowering
;; instead of dipping.
(let* ([rest-arg (extract-rest-arg #`ARGS)]
;; have to replace REST to avoid "identifier used out of context" errors
[body (replace-id #`(begin EXPRS ...) #'REST rest-arg)]
[optimized-body (recursively-optimize-expr body equiv-map #f)])
(if (and lower-lambda (dipped? optimized-body))
(let* ([lowered-body (get-dipped-expr optimized-body)]
[body-deps (get-dipped-deps optimized-body)]
[new-deps (diff-id-lists body-deps (extract-args #`ARGS))])
#`(dip #,new-deps (lambda ARGS #,lowered-body)))
(if (refers-to-id? optimized-body rest-arg)
#`(lambda ARGS
(let-values (((#,rest-arg) (frp:copy-list #,rest-arg)))
#,optimized-body))
#`(lambda ARGS #,optimized-body))))]
[(lambda ARGS EXPRS ...)
;; If "lower-lambda" is true, then we know that it's safe to lower the
;; body of the lambda. "lower-lambda" should only be true if we know
;; that the result will be immediately bound to a variable that is
;; guaranteed never to be called except in a lowered context. In contrast,
;; if we always lowered the body of a lambda, the closure might escape
;; to where it could be called with time-varying arguments.
(let* ([body #`(begin EXPRS ...)]
[optimized-body (recursively-optimize-expr body equiv-map #f)])
(if (and lower-lambda (dipped? optimized-body))
(let* ([lowered-body (get-dipped-expr optimized-body)]
[body-deps (get-dipped-deps optimized-body)]
[new-deps (diff-id-lists body-deps (extract-args #'ARGS))])
#`(dip #,new-deps (lambda ARGS #,lowered-body)))
#`(lambda ARGS #,optimized-body)))]
[(case-lambda (FORMALS EXPR ...) ...)
;; dip each clause separately. TODO: pay attention to lower-lambda
(let* ([bodies (syntax->list #`((begin EXPR ...) ...))]
[optimized-bodies (map (lambda (expr)
(recursively-optimize-expr expr equiv-map #f))
bodies)]
[args (syntax->list #`(FORMALS ...))]
[clauses (map list args optimized-bodies)])
#`(case-lambda . #,clauses))]
;; special case: recognize the expanded version of "let loop" and allow
;; it to be fully lowered. Otherwise it ends up calling itself, and since
;; we don't define a fully lowered equivalent of the loop itself, then
;; whole thing ends up using signal:switching, which is slow.
;; TBD: "let loop" embeds its letrec-values in an #%app -- we could
;; recognize this and avoid defining the upper version at all.
[(#%app (letrec-values (((LOOP) BODY))
LOOP_) ARG ...)
(and (identifier? #'LOOP)
(identifier? #'LOOP_)
(module-identifier=? #'LOOP #'LOOP_))
(let* ([optimized-args (map (lambda (e)
(recursively-optimize-expr e equiv-map #f))
(syntax->list #'(ARG ...)))]
[loop-lowered-id (make-lowered-equiv-id #'LOOP)]
[extended-equiv-map (add-equiv-map equiv-map #'LOOP loop-lowered-id)]
;; allow lambdas to be lowered, because "let loop" generates a lambda
[lowered-body (recursively-optimize-expr #'BODY extended-equiv-map #t)])
;; if the body can be fully lowered, and the lower version doesn't refer to
;; the upper version, then we can define both an upper and a lower definition,
;; and return the upper one.
(if (dipped? lowered-body)
(if (every dipped? optimized-args)
;; all the args were dippable -- we can lower the *entire* thing, and
;; do away with the upper version of the loop.
(let* ([lowered-expr (get-dipped-expr lowered-body)]
[lowered-deps (get-dipped-deps lowered-body)]
[deps (apply union-id-lists lowered-deps
(map get-dipped-deps optimized-args))])
#`(dip #,deps
(#%app (letrec-values (((#,loop-lowered-id) #,lowered-expr))
#,loop-lowered-id)
#,@(map get-dipped-expr optimized-args))))
;; at least one arg wasn't dippable, so we have to keep the upper version,
;; and we can't dip the entire expression.
(let* ([optimized-body (recursively-optimize-expr #'BODY extended-equiv-map #f)]
[lowered-expr (get-dipped-expr lowered-body)]
[lowered-deps (get-dipped-deps lowered-body)]
[deps (diff-id-lists lowered-deps (list #'LOOP))])
#`(#%app (dip #,deps
(letrec-values (((#,loop-lowered-id) #,lowered-expr)
((LOOP) #,optimized-body))
LOOP))
#,@optimized-args)))
(let* ([optimized-body (recursively-optimize-expr #'BODY equiv-map #f)])
#`(#%app (letrec-values (((LOOP) #,optimized-body))
LOOP)
#,@optimized-args))))]
[(LET-VALUES ((VARS VALS) ...) EXPR ...)
(and (identifier? #'LET-VALUES)
(or (module-identifier=? #'LET-VALUES #'let-values)
(module-identifier=? #'LET-VALUES #'letrec-values)))
(let* ([bindings (syntax->list #'(VARS ...))]
[flattened-bindings (apply append (map syntax->list bindings))]
[body #`(begin EXPR ...)]
[optimized-body (recursively-optimize-expr body equiv-map lower-lambda)]
;; TBD: consider defining lowered equivs for local bindings
[optimized-vals (map (lambda (arg)
(recursively-optimize-expr arg equiv-map #f))
(syntax->list #'(VALS ...)))]
[all-exprs-were-dipped (and (dipped? optimized-body)
(every dipped? optimized-vals))])
(if all-exprs-were-dipped
(let* ([val-deps (apply append (map get-dipped-deps optimized-vals))]
[new-deps (union-id-lists
(diff-id-lists (get-dipped-deps optimized-body)
flattened-bindings)
val-deps)]
[lowered-body (get-dipped-expr optimized-body)]
[lowered-vals (map get-dipped-expr optimized-vals)])
#`(dip #,new-deps
(LET-VALUES #,(map list bindings lowered-vals)
#,lowered-body)))
#`(LET-VALUES #,(map list bindings optimized-vals)
#,optimized-body)))]
[(letrec-syntaxes+values SYNTAX-STUFF ((IDS VALS) ...) EXPR ...)
(let* ([optimized-vals
(map (lambda (e)
(recursively-optimize-expr e equiv-map #f))
(syntax->list #'(VALS ...)))]
[optimized-bindings
(map list (syntax->list #'(IDS ...)) optimized-vals)]
[body #`(begin EXPR ...)]
[optimized-body (recursively-optimize-expr body equiv-map #f)])
#`(letrec-syntaxes+values SYNTAX-STUFF #,optimized-bindings #,optimized-body))]
[(if . ARGS)
(let* ([optimized-args (map (lambda (expr)
(recursively-optimize-expr expr equiv-map #f))
(syntax->list #'ARGS))]
[all-args-were-dipped (every dipped? optimized-args)])
(if all-args-were-dipped
;; we can dip the entire subexpression
(let ([deps (apply union-id-lists (map get-dipped-deps optimized-args))]
[lowered-args (map get-dipped-expr optimized-args)])
#`(dip #,deps (if . #,lowered-args)))
;; we can't dip the entire subexpression, but we may have been able
;; to optimize some of the args.
#`(if . #,optimized-args)))]
;; frtime's if expands into a complicated expression involving super-lift.
;; recognize this pattern, and treat it like an if statement.
[(#%app super-lift
(lambda (_B) (if (#%app undefined? __B)
(begin undefined)
(if ___B
TRUE-CASE
FALSE-CASE)))
(#%app _TO_BOOLEAN CONDITIONAL))
(let* ([optimized-condition (recursively-optimize-expr #'CONDITIONAL equiv-map #f)]
[optimized-true-case (recursively-optimize-expr #'TRUE-CASE equiv-map #f)]
[optimized-false-case (recursively-optimize-expr #'FALSE-CASE equiv-map #f)])
(if (and (dipped? optimized-condition)
(dipped? optimized-true-case)
(dipped? optimized-false-case))
(let ([lowered-condition (get-dipped-expr optimized-condition)]
[lowered-true-case (get-dipped-expr optimized-true-case)]
[lowered-false-case (get-dipped-expr optimized-false-case)]
[deps (union-id-lists (get-dipped-deps optimized-condition)
(get-dipped-deps optimized-true-case)
(get-dipped-deps optimized-false-case))])
#`(dip #,deps (if #,lowered-condition #,lowered-true-case #,lowered-false-case)))
#`(frtime:if #,optimized-condition
#,optimized-true-case
#,optimized-false-case)))]
[(#%app dont-optimize EXPR)
#'EXPR]
[(#%app FUNC . ARGS)
(identifier? #'FUNC)
(let* ([lowered-equiv (lookup-lowered-equiv equiv-map #'FUNC)]
[optimized-args (map (lambda (expr)
(recursively-optimize-expr expr equiv-map #f))
(syntax->list #'ARGS))]
[all-args-were-dipped (every dipped? optimized-args)])
(if (and lowered-equiv all-args-were-dipped)
;; we can dip the entire subexpression
(let ([deps (apply union-id-lists (map get-dipped-deps optimized-args))]
[lowered-args (map get-dipped-expr optimized-args)])
#`(dip #,deps (#%app #,lowered-equiv #,@lowered-args)))
;; we can't dip the entire subexpression, but we may have been able
;; to optimize some of the args.
#`(#%app FUNC #,@optimized-args)))]
;; first-class functions -- we have no idea what function is being called,
;; so we can't optimize it. But try to optimize the subexpressions.
[(#%app FUNC ARG ...)
#`(#%app #,(recursively-optimize-expr #'FUNC equiv-map #f)
#,@(map (lambda (expr)
(recursively-optimize-expr expr equiv-map #f))
(syntax->list #'(ARG ...))))]
;; units expand into complicated code that uses set! to define identifiers.
;; Since we don't support set!, we optimize units based on their *unexpanded*
;; syntax, which is much easier to work with.
[(unit IMPORTS EXPORTS EXPRS ...)
#`(unit IMPORTS EXPORTS
(optimize-module #,(equiv-map-to-stx equiv-map) EXPRS ...))]
[(unit/sig SIG IMPORTS (rename RENAMES ...) EXPRS ...)
#`(unit/sig SIG IMPORTS (rename RENAMES ...)
(optimize-module #,(equiv-map-to-stx equiv-map) EXPRS ...))]
[(unit/sig SIG IMPORTS EXPRS ...)
#`(unit/sig SIG IMPORTS
(optimize-module #,(equiv-map-to-stx equiv-map) EXPRS ...))]
[ELSE
(raise-syntax-error #f
(format "recursively-lower-expr: unrecognized syntax: ~a" (syntax-object->datum stx))
stx)]))
;; Optimize a single expression. Raises exn:fail:syntax if the optimized version
;; would use a protected identifier in an uncertified context.
(define-for-syntax (optimize-expr-helper stx equiv-map allow-lambda)
;; Expand everything but references to identifiers, and units.
;; References to identifiers can't be expanded because they might not be
;; defined yet, in which case we'll get a syntax error. Units can't be
;; expanded because they expand into highly indirect code that uses set!
;; to define structures, and we don't support set! -- so instead we handle
;; units explicitly.
(let* ([expanded-stx (local-expand stx 'top-level (list #'#%top #'unit #'unit/sig))]
[optimized-stx (recursively-optimize-expr expanded-stx equiv-map allow-lambda)])
;; expand the result so that if we happen to have a reference to a
;; protected identifier in an uncertified context, then we can
;; trigger an exception, thus giving us the chance to rollback our
;; changes and return the original code unmodified.
(local-expand optimized-stx 'top-level (list #'#%top)) ;; expand units, too
;; return the unexpanded result, so that callers can figure out whether
;; the expression was completely dipped or not.
optimized-stx))
;; Macro to optimize a single expression. The expression will be dipped
;; wherever possible, and its observable semantics will remain unchanged.
(define-syntax (optimize-expr stx)
(syntax-case stx ()
[(_ EQUIV-MAP EXPR)
(with-handlers ([exn:fail:syntax?
(lambda (exn)
#`(begin #,(exn-message exn) EXPR))])
(let* ([equiv-map (stx-to-equiv-map #'EQUIV-MAP)])
(optimize-expr-helper #'EXPR equiv-map #f)))]))
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Definitions
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; returns multiple values:
;; * new equiv map
;; * optimized value
;; * lowered equiv id (or #f if unlowerable)
;; * lowered equiv value (or #f if unlowerable)
(define-for-syntax (optimize-definition-helper equiv-map id val)
(with-handlers
([exn:fail:syntax? (lambda (exn)
(values equiv-map
#`(begin #,(exn-message exn) #,val)
#f
#f))])
(let* ([lowered-equiv-id (make-lowered-equiv-id id)]
[new-equiv-map (add-equiv-map equiv-map id lowered-equiv-id)]
[lowered-val (optimize-expr-helper val new-equiv-map #t)])
(if (and (dipped? lowered-val)
(null? (get-dipped-deps lowered-val)))
;; use new-equiv-map in order to allow the dipped version to call
;; the lowered version for recursive calls.
(let ([optimized-val (optimize-expr-helper val new-equiv-map #f)])
(values new-equiv-map
optimized-val
lowered-equiv-id
(get-dipped-expr lowered-val)))
;; use the old equiv-map, since there is no lowered version so the
;; dipped version can't call it
(let ([optimized-val (optimize-expr-helper val equiv-map #f)])
(values equiv-map
optimized-val
#f
#f))))))
;; Optimize a top-level variable binding, by lowering its value.
(define-syntax (optimize-definition stx)
(syntax-case stx (define-values values)
[(_ EQUIV-MAP (define-values (ID) VAL) FORMS ...)
;; if this is a lowered equiv definition, add it to our equiv map
(lowered-equiv-id? #'ID)
(let ([lifted-id (lowered-equiv-id->lifted-id #'ID)])
#`(begin
(define-values (ID) VAL)
(optimize-module ((#,lifted-id ID) . EQUIV-MAP) FORMS ...)))]
[(_ EQUIV-MAP (define-values (ID) VAL) FORMS ...)
;; We're defining a single identifier
(let-values ([(new-equiv-map optimized-val lowered-equiv-id lowered-val)
(optimize-definition-helper (stx-to-equiv-map #'EQUIV-MAP)
#'ID
#'VAL)])
(if (and lowered-equiv-id lowered-val)
#`(begin
(define ID #,optimized-val)
(define #,lowered-equiv-id #,lowered-val)
(optimize-module #,(equiv-map-to-stx new-equiv-map) FORMS ...))
#`(begin
(define ID #,optimized-val)
(optimize-module #,(equiv-map-to-stx new-equiv-map) FORMS ...))))]
[(_ EQUIV-MAP (define-values (ID ...) (values VAL ...)) FORMS ...)
(= (length (syntax->list #'(ID ...)))
(length (syntax->list #'(VAL ...))))
;; Very common special case: the multiple identifiers are immediately combined
;; with multiple values. Common enough that it's worth handling specially.
(let ([new-equiv-map (stx-to-equiv-map #'EQUIV-MAP)]
[ids null]
[vals null])
(for-each (lambda (id val)
(let-values ([(em optimized-val lowered-equiv-id lowered-val)
(optimize-definition-helper new-equiv-map id val)])
(set! new-equiv-map em)
(when (and lowered-equiv-id lowered-val)
(set! ids (cons lowered-equiv-id ids))
(set! vals (cons lowered-val vals)))
(set! ids (cons id ids))
(set! vals (cons optimized-val vals))))
(syntax->list #'(ID ...))
(syntax->list #'(VAL ...)))
#`(begin
(define-values #,ids (values . #,vals))
(optimize-module #,(equiv-map-to-stx new-equiv-map) FORMS ...)))]
[(_ EQUIV-MAP (define-values (ID ...) EXPR) FORMS ...)
;; We're binding multiple identifiers with some arbitrary set of values,
;; so we can't define lowered-equivs because we can't tease apart the
;; various values. So just optimize the value expression as a whole.
#`(begin
(define-values (ID ...) (optimize-expr EQUIV-MAP EXPR))
(optimize-module EQUIV-MAP FORMS ...))]))
)