#lang scheme/base

(require scheme/list)
(require scheme/pretty)
(require (file "primitives.scm"))

(provide trace
         subst
         find
         variable-value?
         literal-value?
         simple-value?
         value-form?
         statement-form?
         primitive-form?
         pure-form?
         bind-form?
         map-form
         search-form
         form-sets?
         form-uses?
         form-captures?
         form-captures-input?
         form-captures-output?
         narrow-binds
         subst-var
         flatten-binds
         free-variables
         free-input-variables
         free-output-variables
         value-used?)

(define (trace fn . args)
  (let ([x (apply fn args)])
    (pretty-print (list fn x))
    x))

(define (subst old new lst)
  (foldr (lambda (x rst)
           (cons (if (eq? x old)
                     new
                     x)
                 rst))
         '()
         lst))

(define (find x lst)
  (let/cc return
    (for ([i (in-naturals 0)]
          [y (in-list lst)])
      (when (eq? y x) (return i)))
    #f))

(define (variable-value? form)
  (and (symbol? form)
       (not (eq? form '%undef))))

(define (literal-value? form)
  (and (not (variable-value? form))
       (or (not (pair? form))
           (eq? (first form) 'quote)
           (eq? (first form) '%template))))

(define (simple-value? form)
  (or (variable-value? form)
      (literal-value? form)))

; A value-form is any simple form which can appear on the right-hand side of a (set! ...).
; If there are any side-effect they occur before the variable is updated.
(define (value-form? form)
  (define complex-values '(%bind %apply %call/cc %values))
  (or (simple-value? form)
      (memq (first form) complex-values)
      (memq (first form) (map first value-primitives))))

; A statement-form is any simple form which has, or may have, side-effects.
(define (statement-form? form)
  (define complex-statements '(%set! %apply %call/cc %tail-call))
  (and (not (simple-value? form))
       (or (memq (first form) complex-statements)
           (memq (first form) (map first statement-primitives)))))

(define (primitive-form? form)
  (and (pair? form) (memq (first form) (map first all-primitives))))

; A pure form is any form known to be free of side-effects.
(define (pure-form? form)
  (and (value-form? form)
       (not (statement-form? form))))

(define (bind-form? form)
  (and (pair? form) (eq? (first form) '%bind)))

(define (map-form form
                  #:bind       [bind-fn       (lambda (recurse op vars . subforms)
                                                `(,op ,vars ,@(map recurse subforms)))]
                  #:lambda     [lambda-fn     (lambda (recurse op g-vars i-vars bind)
                                                `(,op ,g-vars ,i-vars ,(recurse bind)))]
                  #:set        [set-fn        (lambda (recurse op var value)
                                                `(,op ,var ,(recurse value)))]

                  #:primitive  [primitive-fn  (lambda (recurse op . simple-values)
                                                `(,op ,@(map recurse simple-values)))]
                  #:values     [values-fn     primitive-fn]
                  #:call       [call-fn       primitive-fn]
                  #:apply      [apply-fn      call-fn]
                  #:call/cc    [call/cc-fn    call-fn]
                  #:tail-call  [tail-call-fn  call-fn]

                  #:simple     [simple-fn     (lambda (recurse kind form) form)]
                  #:variable   [variable-fn   simple-fn]
                  #:literal    [literal-fn    simple-fn]

                  #:other      [other-fn      (lambda (recurse . form)
                                                (error "Unsimplified form:" form))])
  (define (recurse subform)
    (map-form subform
              #:bind      bind-fn
              #:lambda    lambda-fn
              #:set       set-fn
              #:primitive primitive-fn
              #:values    values-fn
              #:call      call-fn
              #:apply     apply-fn
              #:call/cc   call/cc-fn
              #:tail-call tail-call-fn
              #:simple    simple-fn
              #:variable  variable-fn
              #:literal   literal-fn
              #:other     other-fn))

  (cond
    [(variable-value? form) (variable-fn recurse 'variable form)]
    [(literal-value?  form) (literal-fn  recurse 'literal  form)]
    [else
     (let ([handler (case (first form)
                      [(%bind)      bind-fn]
                      [(%lambda)    lambda-fn]
                      [(%set!)      set-fn]
                      [(%values)    values-fn]
                      [(%apply)     apply-fn]
                      [(%call/cc)   call/cc-fn]
                      [(%tail-call) tail-call-fn]
                      [else         (if (primitive-form? form)
                                        primitive-fn
                                        other-fn)])])
       (apply handler recurse form))]))

; Like map-form, but intended for boolean results. (Just different defaults.)
(define (search-form form
                     #:merge-with [merge-fn      ormap]
                     #:base-value [base-value    #f]
                     #:bind       [bind-fn       (lambda (recurse op vars . subforms)
                                                   (merge-fn recurse subforms))]
                     #:lambda     [lambda-fn     (lambda (recurse op g-vars i-vars bind)
                                                   (recurse bind))]
                     #:set        [set-fn        (lambda (recurse op var value)
                                                   (recurse value))]

                     #:primitive  [primitive-fn  (lambda (recurse op . simple-values)
                                                   (merge-fn recurse simple-values))]
                     #:values     [values-fn     primitive-fn]
                     #:call       [call-fn       primitive-fn]
                     #:apply      [apply-fn      call-fn]
                     #:call/cc    [call/cc-fn    call-fn]
                     #:tail-call  [tail-call-fn  call-fn]

                     #:simple     [simple-fn     (lambda (recurse kind form) base-value)]
                     #:variable   [variable-fn   simple-fn]
                     #:literal    [literal-fn    simple-fn]

                     #:other      [other-fn      (lambda (recurse . form)
                                                   (error "Unsimplified form:" form))])
  (map-form form
            #:bind      bind-fn
            #:lambda    lambda-fn
            #:set       set-fn
            #:primitive primitive-fn
            #:values    values-fn
            #:call      call-fn
            #:apply     apply-fn
            #:call/cc   call/cc-fn
            #:tail-call tail-call-fn
            #:simple    simple-fn
            #:variable  variable-fn
            #:literal   literal-fn
            #:other     other-fn))

(define (form-sets? form variable [call-may-set? #t])
  (search-form form
               #:bind      (lambda (recurse op vars . subforms)
                             (and (not (memq variable vars))
                                  (ormap recurse subforms)))
               #:lambda    (lambda _ #f)
               #:set       (lambda (recurse op var complex-value)
                             (eq? var variable))
               #:call      (lambda _ call-may-set?)))

(define (form-uses? form variable [call-may-use? #t] [descend? #t])
  (search-form form
               #:bind      (lambda (recurse op vars . subforms)
                             (and (not (memq variable vars))
                                  (ormap recurse subforms)))
               #:lambda    (lambda (recurse op g-vars i-vars bind)
                             (and descend? (recurse bind)))
               #:call      (lambda (recurse op . simple-values)
                             (or call-may-use? (ormap recurse simple-values)))
               #:variable  (lambda (recurse op var) (eq? var variable))))

(define (form-captures? form variable [input? #t] [output? #t])
  (search-form form
               #:bind      (lambda (recurse op vars . subforms)
                             (and (not (memq variable vars))
                                  (ormap recurse subforms)))
               #:lambda    (lambda (recurse op g-vars i-vars bind)
                             (and (memq variable (free-variables bind input? output?)) #t))))

(define (form-captures-input? form var)
  (form-captures? form var #t #f))

(define (form-captures-output? form var)
  (form-captures? form var #f #t))

(define (value-used? variable forms)
  (cond
    [(null? forms) #f]
    [(form-uses? (first forms) variable #f #t) #t]
    [(form-sets? (first forms) variable #f) #f]
    [else (value-used? variable (cdr forms))]))

(define (subst-var old-var new-var form)
  (map-form form
            #:bind     (lambda (recurse op vars . subforms)
                         `(%bind ,(subst old-var new-var vars) ,@(map recurse subforms)))
            #:set      (lambda (recurse op var value)
                         `(,op ,(if (eq? var old-var) new-var var) ,(recurse value)))
            #:variable (lambda (recurse op var)
                         (if (eq? var old-var) new-var var))))

(define (narrow-binds simple-lambda-form)
  (define bind (fourth simple-lambda-form))

  (define (at-top-level? var)
    (or (ormap (lambda (x) (form-sets? x var #f)) (cddr bind))
        (ormap (lambda (x) (form-uses? x var #f #f)) (cddr bind))))

  (define (captured-twice? var)
    (let/cc return
      (foldl (lambda (subform once?)
               (if (form-captures? subform var)
                   (if once? (return #t) #t)
                   once?))
             (at-top-level? var)
             (cddr bind))
      #f))

  (define extra-binds
    (filter-not captured-twice?
                (filter-not at-top-level?
                            (second bind))))

  `(%lambda ,(second simple-lambda-form) ,(third simple-lambda-form)
     (%bind ,(remove* extra-binds (second bind))
       ,@(map (lambda (subform)
                (if (and (pair? subform)
                         (eq? (first subform) '%set!)
                         (pair? (third subform))
                         (eq? (first (third subform)) '%lambda))
                    (let* ([dest (second subform)]
                           [value (third subform)]
                           [g-vars (second value)]
                           [i-vars (third value)])
                      `(%set! ,dest ,(foldl (lambda (var temp-value)
                                              (define temp-bind (fourth temp-value))
                                              (if (form-captures? temp-value var)
                                                  (narrow-binds
                                                    `(%lambda ,g-vars ,i-vars
                                                       (%bind (,@(second temp-bind) ,var)
                                                         ,@(cddr temp-bind))))
                                                  temp-value))
                                            value
                                            extra-binds)))
                    subform))
              (cddr bind)))))

(define (flatten-binds form)
  (define (make-bindings-unique bind rename-vars)
    (define (needs-rename? var) (memq var rename-vars))
    (define (make-binding-unique var bind)
      (let* ([prefix (string-append (symbol->string var) "->g")]
             [unique-var (gensym prefix)])
        (subst-var var unique-var bind)))
    (foldr make-binding-unique bind (filter needs-rename? (second bind))))

  (map-form form
            #:bind   (lambda (recurse op bound-vars . original-subforms)
                       (define rename-vars
                          (remove-duplicates
                            (append (free-variables `(,op ,bound-vars ,@original-subforms))
                                    bound-vars)))
                       (define (form->list subform)
                         (if (bind-form? subform)
                             (let ([unique-form (make-bindings-unique
                                                  (recurse subform)
                                                  rename-vars)])
                               (set! bound-vars (append (second unique-form) bound-vars))
                               (cddr unique-form))
                             (list subform)))
                       (let ([subforms (append-map form->list original-subforms)])
                         `(%bind ,bound-vars ,@subforms)))
            #:lambda (lambda (recurse . form) form)))

(define (free-variables form [input? #t] [output? #t])
  (map-form form
            #:bind      (lambda (recurse op vars . subforms)
                          (remove-duplicates (remove* vars (append-map recurse subforms))))
            #:lambda    (lambda (recurse op g-vars i-vars bind)
                          (recurse bind))
            #:set       (lambda (recurse op var value)
                          (let ([value-free (recurse value)])
                            (if output?
                                (cons var value-free)
                                value-free)))
            #:primitive (lambda (recurse op . simple-values)
                          (remove-duplicates (append-map recurse simple-values)))
            #:simple    (lambda (recurse kind form)
                          (if (and input?
                                   (variable-value? form)
                                   (not (memq form '(%nil %self %argv %ctx %k))))
                              (list form)
                              '()))))

(define (free-input-variables form)
  (free-variables form #t #f))

(define (free-output-variables form)
  (free-variables form #f #t))

; vim:set sw=2 expandtab: