#lang scheme/base

(require scheme/list)
(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)
                   ; Top-level (free) variables are presumed to be
                   ; constant. The alternative was to assume them
                   ; to be boxes, which has its own complications.
                   (error "Setting unbound var:" subform))
               (not (value-used? (second subform) after)))
          after
          (cons subform after)))
    `(%bind ,vars
       ,@(foldr prepend-if-used '() (map recurse subforms))))
  (narrow-binds
    (map-form form #:bind bind-fn)))

(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! var ...) or (%set! var)
;   (%car var), until (%set-car! var) or (%set! var)
;   (%cdr var), until (%set-cdr! var) or (%set! var)
(define (propogate-set! form)
  (define (bind-fn recurse op vars . subforms)
    (define (prepend subform after)
      (if (eq? (first subform) '%set!)
          (let ([var (second subform)]
                [value (third subform)])
            (cons
             subform
             (cond
               [(simple-value? value)
                (propogate-simple-value var value
                                        (lambda (form)
                                          (and (eq? (first form) '%set!)
                                               (eq? (second form) value)))
                                        after)]
               [(eq? (first value) '%unbox)
                (let ([box-var (second value)])
                  (propogate-value var value
                                   (lambda (form)
                                     (or (and (eq? (first form) '%set!)
                                              (eq? (second form) box-var))
                                         (and (eq? (first form) '%set-box!)
                                              (eq? (second form) box-var))))
                                   after))]
               [(eq? (first value) '%car)
                (let ([pair-var (second value)])
                  (propogate-value var value
                                   (lambda (form)
                                     (or (and (eq? (first form) '%set!)
                                              (eq? (second form) pair-var))
                                         (and (eq? (first form) '%set-car!)
                                              (eq? (second form) pair-var))))
                                   after))]
               [(eq? (first value) '%cdr)
                (let ([pair-var (second value)])
                  (propogate-value var value
                                   (lambda (form)
                                     (or (and (eq? (first form) '%set!)
                                              (eq? (second form) pair-var))
                                         (and (eq? (first form) '%set-cdr!)
                                              (eq? (second form) pair-var))))
                                   after))]
              [else after])))
          (cons subform after)))
    `(%bind ,vars
       ,@(foldr prepend '() (map recurse subforms))))
  (map-form form #:bind bind-fn))