#lang scheme/base

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

(provide trace
         subst
         subst*
         find
         curry-keywords
         variable-value?
         literal-value?
         simple-value?
         value-form?
         statement-form?
         primitive-form?
         pure-form?
         bind-form?
         traverse-form
         map-form
         search-form
         form-sets?
         form-uses?
         form-captures?
         form-captures-input?
         form-captures-output?
         subst-var
         subst-var*
         flatten-binds
         free-variables
         free-input-variables
         free-output-variables
         free-variable?
         free-input-variable?
         free-output-variable?
         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 (subst* old->new lst)
  (foldr (lambda (x rst)
           (let ([item (assoc x old->new)])
             (cons (if item (second item) 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))

;; Combines two sorted keyword-value list pairs into a single
;; sorted list-pair, taking the values from the second list
;; where the same keyword exists in both lists.
(define (merge-keywords keywords-1 kw-values-1 keywords-2 kw-values-2)
  (let iter ([kw   '()]        [kv   '()]
             [kw-1 keywords-1] [kv-1 kw-values-1]
             [kw-2 keywords-2] [kv-2 kw-values-2])
    (cond
      [(null? kw-2) (values (append (reverse kw) kw-1) (append (reverse kv) kv-1))]
      [(null? kw-1) (values (append (reverse kw) kw-2) (append (reverse kv) kv-2))]
      [(eq? (car kw-1) (car kw-2))
       (iter (cons (car kw-2) kw) (cons (car kv-2) kv)
             (cdr kw-1)           (cdr kv-1)
             (cdr kw-2)           (cdr kv-2))]
      [(keyword<? (car kw-1) (car kw-2))
       (iter (cons (car kw-1) kw) (cons (car kv-1) kv)
             (cdr kw-1)           (cdr kv-1)
             kw-2                 kv-2)]
      [else
       (iter (cons (car kw-2) kw) (cons (car kv-2) kv)
             kw-1                 kv-1
             (cdr kw-2)           (cdr kv-2))])))

;; Like (curry ...) from scheme/function, but with much better support for
;; keyword arguments. Keyword arguments supplied to curry-keywords act as
;; defaults, which may be overridden by keywords passed to the resulting
;; function. It can also curry non-keyword arguments, naturally.
(define curry-keywords
  (make-keyword-procedure
    (lambda (curried-keywords curried-kw-values fn . curried-non-keywords)
      (make-keyword-procedure
        (lambda (keywords kw-values . non-keywords)
          (let-values ([(all-keywords all-kw-values)
                        (merge-keywords curried-keywords curried-kw-values keywords kw-values)])
            (keyword-apply fn all-keywords all-kw-values
                           (append curried-non-keywords non-keywords))))))))

(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 %lambda %apply %call/cc %values %value-list))
  (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 (traverse-form form
          #:bind       [bind-fn       (lambda (recurse op vars . subforms)
                                        (for ([subform (in-list subforms)])
                                          (recurse subform)))]
          #:lambda     [lambda-fn     (lambda (recurse op g-vars i-vars bind)
                                        (recurse bind))]
          #:set        [set-fn        (lambda (recurse op var value-form)
                                        (recurse value-form))]
          #:value-list [value-list-fn (lambda (recurse op values-form)
                                        (recurse values-form))]
          #:primitive  [primitive-fn  (lambda (recurse op . simple-values)
                                        (for ([val (in-list simple-values)])
                                          (recurse val)))]
          #:simple     [simple-fn     (lambda (recurse kind simple-value) (void))]
          #:other      [other-fn      (lambda (recurse . form)
                                        (error "Unsimplified form:" form))]
          #:values     [values-fn     primitive-fn]
          #:call       [call-fn       primitive-fn]
          #:variable   [variable-fn   simple-fn]
          #:literal    [literal-fn    simple-fn]
          #:apply      [apply-fn      call-fn]
          #:call/cc    [call/cc-fn    call-fn]
          #:tail-call  [tail-call-fn  call-fn])
  (define (recurse subform)
    ; Can skip #:primitive, #:call, and #:simple, which only set defaults.
    (traverse-form subform
      #:bind       bind-fn
      #:lambda     lambda-fn
      #:set        set-fn
      #:value-list value-list-fn
      #:primitive  primitive-fn
      #:other      other-fn
      #:values     values-fn
      #:variable   variable-fn
      #:literal    literal-fn
      #:apply      apply-fn
      #:call/cc    call/cc-fn
      #:tail-call  tail-call-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]
                      [(%value-list) value-list-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))]))

(define map-form
  (curry-keywords traverse-form
    #:bind       (lambda (recurse op vars . subforms)
                   `(,op ,vars ,@(map recurse subforms)))
    #:lambda     (lambda (recurse op g-vars i-vars bind)
                   `(,op ,g-vars ,i-vars ,(recurse bind)))
    #:set        (lambda (recurse op var value-form)
                   `(,op ,var ,(recurse value-form)))
    #:value-list (lambda (recurse op values-form)
                   `(,op ,(recurse values-form)))
    #:primitive  (lambda (recurse op . simple-values)
                   `(,op ,@(map recurse simple-values)))
    #:simple     (lambda (recurse kind form) form)))

; Like map-form, but intended for boolean results.
(define search-form
  (curry-keywords traverse-form
    #:bind       (lambda (recurse op vars . subforms)
                   (ormap recurse subforms))
    #:lambda     (lambda (recurse op g-vars i-vars bind)
                   (recurse bind))
    #:set        (lambda (recurse op var value)
                   (recurse value))
    #:value-list (lambda (recurse op var values-form)
                   (recurse values-form))
    #:primitive  (lambda (recurse op . simple-values)
                   (ormap recurse simple-values))
    #:simple     (lambda (recurse kind form) #f)))

(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 value-form)
                  (or (eq? var variable)
                      (recurse value-form)))
    #:primitive (lambda _ #f)
    #:call      (lambda _ call-may-set?)))

(define (form-uses? form variable [call-may-use? #t])
  (search-form form
    #:bind      (lambda (recurse op vars . subforms)
                  (and (not (memq variable vars))
                       (ormap recurse subforms)))
    #:lambda    (lambda (recurse . form) #f)
    #:primitive (lambda (recurse op . simple-values)
                  (and (memq variable simple-values) #t))
    #:call      (lambda (recurse op . simple-values)
                  (or call-may-use? (and (memq variable simple-values) #t)))
    #: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)
                  (free-variable? variable bind input? output?))
    #:primitive (lambda _ #f)
    #:simple    (lambda _ #f)))

(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)
  (and (not (null? forms))
       (let ([form (first forms)])
         (or (form-captures-input? form variable)
             (form-uses? form variable #f)
             (and (not (form-sets? form variable #f))
                  (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 (subst-var* var-map form)
  (map-form form
    #:bind     (lambda (recurse op vars . subforms)
                 `(%bind ,(subst* var-map vars) ,@(map recurse subforms)))
    #:set      (lambda (recurse op var value)
                 (let ([item (assoc var var-map)])
                   `(,op ,(if item (second item) var) ,(recurse value))))
    #:variable (lambda (recurse op var)
                 (let ([item (assoc var var-map)])
                   (if item (second item) var)))))

(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])
  (define (append-map/unique fn . lsts)
    (remove-duplicates (apply append-map fn lsts)))
  (traverse-form form
    #:bind       (lambda (recurse op vars . subforms)
                   (remove* vars (append-map/unique recurse subforms)))
    #:set        (lambda (recurse op var value)
                   (let ([value-free (recurse value)])
                     (if (and output? (not (memq var value-free)))
                         (cons var value-free)
                         value-free)))
    #:primitive  (lambda (recurse op . simple-values)
                   (append-map/unique recurse simple-values))
    #:simple     (lambda (recurse kind form) '())
    #:variable   (lambda (recurse kind var)
                   (if (and input? (not (special-variable? var)))
                       (list var)
                       '()))))

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

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

(define (free-variable? variable form [input? #t] [output? #t])
  (search-form form
    #:bind       (lambda (recurse op vars . subforms)
                   (and (not (memq variable vars))
                        (ormap recurse subforms)))
    #:set        (lambda (recurse op var value)
                   (or (and output? (eq? var variable))
                       (recurse value)))
    #:variable   (lambda (recurse kind var)
                   (and input? (eq? var variable)))))

(define (free-input-variable? var form)
  (free-variable? var form #t #f))

(define (free-output-variable? var form)
  (free-variable? var form #f #t))

; vim:set sw=2 expandtab: