#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?
         side-effect-form?
         primitive-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 (call-with-values (lambda () (time (apply fn args))) list)])
    (pretty-print (list fn x))
    (apply values 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)
  (symbol? form))

(define (literal-value? form)
  (and (not (variable-value? form))
       (or (not (pair? form))
           (eq? (first form) 'quote)
           (memq (first form) '(#%builtin #%include #%immutable #%struct #%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-effects they occur before the variable is updated.
(define (value-form? form)
  (define complex-values '(#%bind #%lambda #%apply #%call/cc #%values #%values->list))
  (or (simple-value? form)
      (memq (first form) complex-values)
      (memq (first form) (map first all-primitives))))

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

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

(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))]
          #:values->list [values->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)
    (cond
      [(variable-value? subform) (variable-fn recurse 'variable subform)]
      [(literal-value?  subform) (literal-fn  recurse 'literal  subform)]
      [else
       (let ([handler (case (first subform)
                        [(#%bind)         bind-fn]
                        [(#%lambda)       lambda-fn]
                        [(#%set!)         set-fn]
                        [(#%values->list) values->list-fn]
                        [(#%values)       values-fn]
                        [(#%apply)        apply-fn]
                        [(#%call/cc)      call/cc-fn]
                        [(#%tail-call)    tail-call-fn]
                        [else             (if (primitive-form? subform)
                                              primitive-fn
                                              other-fn)])])
         (apply handler recurse subform))]))
  (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)))
    #:values->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))
    #:values->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)))

(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)
                 (if (memq old-var vars)
                     `(,op ,vars ,@subforms)
                     `(#%bind ,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)
                 (let ([unbound-map (filter (lambda (item) (not (memq (car item) vars)))
                                            var-map)])
                   (if (null? unbound-map)
                       `(,op ,vars ,@subforms)
                       `(#%bind ,vars ,@(map (lambda (sf) (subst-var* unbound-map sf))
                                            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 (unique-var var)
      (let ([prefix (string-append (symbol->string var) "->g")])
        (list var (gensym prefix))))
    (let ([var-map (map unique-var (filter needs-rename? (second bind)))])
      `(#%bind ,(subst* var-map (second bind))
         ,@(map (lambda (sf) (subst-var* var-map sf)) (cddr 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: