#lang scheme/base

(require scheme/list)
(require scheme/match)
(require scheme/pretty)
(require (file "utilities.scm"))

(provide reduce-variables
         reduce-set!
         propogate-set!
         optimize-function)

(define (optimize-function simple-lambda-form)
  ((compose reduce-variables
            reduce-set!
            propogate-set!)
   simple-lambda-form))

; Don't bind variables which aren't referenced.
(define (reduce-variables form)
  (define (bind-fn recurse op vars . subforms)
    (let* ([reduced-forms (map recurse subforms)]
           [ref-vars (remove-duplicates (append-map free-variables reduced-forms))])
      (define (referenced? var) (and (memq var ref-vars) #t))
      `(#%bind ,(filter referenced? vars)
         ,@reduced-forms)))
  (map-form form #:bind bind-fn))

; Don't set variables which won't be accessed later.
(define (reduce-set! form)
  (define (bind-fn recurse op vars . subforms)
    (define (prepend-if-used subform after)
      (if (and (pair? subform)
               (eq? (first subform) '#%set!)
               (or (memq (second subform) vars)
                   (error "Setting unbound (constant) variable:" subform))
               (not (value-used? (second subform) after)))
          after
          (cons subform after)))
    `(#%bind ,vars
       ,@(foldr prepend-if-used '() (map recurse subforms))))
  (map-form form #:bind bind-fn))

(define (propogate-value variable value invalidates? forms)
  (if (invalidates? `(#%set! ,variable ,value))
      forms
      (%propogate-value% variable value invalidates? forms)))

(define (%propogate-value% variable value invalidates? forms)
  (if (null? forms)
      forms
      (let* ([form (car forms)]
             [after (cdr forms)]
             [new-form (case (first form)
                         [(#%set!) (if (eq? (third form) variable)
                                      `(#%set! ,(second form) ,value)
                                      form)]
                         [else form])])
        (if (or (and (eq? (first (car forms)) '#%set!)
                     (eq? (second (car forms)) variable))
                (invalidates? new-form))
            (cons new-form after)
            (cons new-form (%propogate-value% variable value invalidates? after))))))

; Simple values (literals, variables) can replace arguments as well as #%set! values.
(define (propogate-simple-value variable value invalidates? forms)
  (if (null? forms)
      forms
      (let* ([form (car forms)]
             [after (cdr forms)]
             [new-form (case (first form)
                         [(#%set!)
                          (let ([set-value (if (eq? (third form) variable) value (third form))])
                            (if (simple-value? set-value)
                                `(#%set! ,(second form) ,set-value)
                                `(#%set! ,(second form)
                                        (,(first set-value)
                                         ,@(subst variable value (cdr set-value))))))]
                         [else `(,(first form) ,@(subst variable value (cdr form)))])])
        (if (or (and (eq? (first (car forms)) '#%set!)
                     (eq? (second (car forms)) variable))
                (invalidates? new-form))
            (cons new-form after)
            (cons new-form (propogate-simple-value variable value invalidates? after))))))

; When value of var2 is known, change (#%set! var1 var2) to (#%set! var1 value).
; Known values are:
;   literals, always
;   var, until (#%set! var ...)
;   (#%unbox var), until (#%set-box! ...) or (#%set! var)
;   (#%car var), until (#%set-car! ...) or (#%set! var)
;   (#%cdr var), until (#%set-cdr! ...) or (#%set! var)
(define (propogate-set! form)
  (define (bind-fn recurse op vars . subforms)
    (define (prepend subform after)
      (cons
       subform
       (match subform
         [`(#%set! ,var ,(? simple-value? value))
          (propogate-simple-value var value
                                  (lambda (form)
                                    (and (eq? (first form) '#%set!)
                                         (eq? (second form) value)))
                                  after)]
         [`(#%set! ,var ,(and value `(#%unbox ,box-var)))
          (propogate-value var value
                           (lambda (form)
                             (or (and (eq? (first form) '#%set!)
                                      (eq? (second form) box-var))
                                 (eq? (first form) '#%set-box!)))
                           after)]
         [`(#%set! ,var ,(and value `(#%car ,pair-var)))
          (propogate-value var value
                           (lambda (form)
                             (or (and (eq? (first form) '#%set!)
                                      (eq? (second form) pair-var))
                                 (eq? (first form) '#%set-car!)))
                           after)]
         [`(#%set! ,var ,(and value `(#%cdr ,pair-var)))
          (propogate-value var value
                           (lambda (form)
                             (or (and (eq? (first form) '#%set!)
                                      (eq? (second form) pair-var))
                                 (eq? (first form) '#%set-cdr!)))
                           after)]
         [_ after])))
    `(#%bind ,vars
       ,@(foldr prepend '() (map recurse subforms))))
  (map-form form #:bind bind-fn))

; vim:set sw=2 expandtab: