(load "cpsC.ss")

;;; Nothing below this line changes for the remaining two algorithms.
;;;-----------------------------------------------------------------------------
;;; run changes to include the pre-pass full-let->single-lets.

(define d? #t)

;; This turns let into let* (represented as a nested let).  Often this
;; has no effect.  It changes things slightly by introducing a
;; generated variable whenever there is an unintended variable
;; capture.  The goal is to maintain the old variable name if it is
;; possible.  When not possible, a carat is added at the end of the
;; variable name.  We reserve the carat as an unused character for
;; purposes of explaining this algorithm and maintaining the
;; readability of the output of this algorithm.  full-let->single-lets
;; could be part of chapter 2 (and then ask for it to be revised as we
;; introduce more forms in chapter 3). We know that when we cps the
;; let, that we are going to be linearizing the let expression anyway,
;; so we are doing it just a bit earlier than might be expected.  This
;; has profound positive consequences on the readabilty of the final
;; product. Also we treat "((lambda" as let.

;; An alternate approach would be to put this file "D" before "A".
;; The reason I hesitate to do that is I like knowing that "A"
;; without the little improvements in possible-letk and CLAM is
;; all that you need to know.  I also like tying let to cps-of-rands.
;; If we put "D" first, then we should feel obligated, to treat let's
;; with only one declaration, since that is what the language looks
;; like from then on.  The cps-of-expression still works, but it never
;; sees more than a single declaration.  That does not bother me.  So,
;; I have left this order, but you may feel that full-let->single-lets
;; should just be inside their heads by this point.

(define full-let->single-lets
  (lambda (pgm)
    (linearize-let pgm)))

(define linearize-let
  (lambda (exp)
    (match exp
      [,constant
        (guard (or (number? constant) (boolean? constant)))
        constant]
      [,variable (guard (symbol? variable)) variable]
      [(APk ,cont ,rand)
       `(APk ,(linearize-let cont) ,(linearize-let rand))]
      [(CLAM (,id) ,body)
       `(CLAM (,id)
              ,(linearize-let body))]
      [((lambda (,id* ...) ,body) ,rand* ...)
       (linearize-let
         `(let ,(map list id* rand*) ,body))]
      [(lambda (,id* ...) ,body)
       `(lambda ,id*
          ,(linearize-let body))]
      [(if ,test-exp ,true-exp ,false-exp)
       `(if ,(linearize-let test-exp)
            ,(linearize-let true-exp)
            ,(linearize-let false-exp))]
      [(,prim ,rand* ...)
       (guard (primitive? prim))
       `(,prim ,@(map linearize-let rand*))]
      [(let () ,body) (linearize-let body)]
      [(let ((,id ,rand)) body)
       `(let ((,id ,(linearize-let rand)))
          ,(linearize-let body))]
      [(let ((,id1 ,rand1) (,id2* ,rand2*) ...) ,body)
       (guard (ormap (lambda (rand) (free? id1 rand)) rand2*))
       (let ((g (fresh-var id1)))
         (linearize-let
           `(let ((,g ,rand1))
              (let ,(map list id2* rand2*)
                ,(beta-subst g id1 body)))))]
      [(let ((,id1 ,rand1) (,id2* ,rand2*) ...) ,body)
       `(let ((,id1 ,(linearize-let rand1)))
          ,(linearize-let
             `(let ,(map list id2* rand2*)
                ,body)))]
      [(letrec ((,name* (lambda (,ids* ...) ,body*)) ...) ,body)
       `(letrec ,(map list name*
                      (map (lambda (ids body)
                             `(lambda ,ids
                                ,(linearize-let
                                   body)))
                           ids* body*))
          ,(linearize-let body))]
      [(letcc ,id ,body)
       `(letcc ,id ,(linearize-let body))]
      [(throw ,value-exp ,cont-exp)
       `(throw
          ,(linearize-let value-exp)
          ,(linearize-let cont-exp))]
      [(,rator ,rand* ...)
       `(,(linearize-let rator)
          ,@(map linearize-let rand*))]
      [,other (error 'linearize-let "Invalid expression  ~s" other)])))

(define free?
  (lambda (var exp)
    (memq var (free-vars exp))))

(define free-vars
  (lambda (exp)
    (remove-duplicates (free-var-list exp))))

(define remove-duplicates
  (lambda (ls)
    (cond
      ((null? ls) '())
      ((memq (car ls) (cdr ls)) (remove-duplicates (cdr ls)))
      (else (cons (car ls) (remove-duplicates (cdr ls)))))))

(define list-diff
  (lambda (ls set)
    (cond
      ((null? ls) '())
      ((memq (car ls) set) (list-diff (cdr ls) set))
      (else (cons (car ls) (list-diff (cdr ls) set))))))

(define free-var-list
  (lambda (exp)
    (match exp
      [,constant
        (guard (or (number? constant) (boolean? constant)))
        '()]
      [,variable (guard (symbol? variable)) `(,variable)]
      [(APk ,cont ,rand)
       `(,@(free-var-list cont)
           ,@(free-var-list rand))]
      [(CLAM (,id) ,body)
       (list-diff (free-var-list body) (list id))]
      [(lambda (,id* ...) ,body)
       (list-diff (free-var-list body) id*)]
      [(if ,test-exp ,true-exp ,false-exp)
       `(,@(free-var-list test-exp)
           ,@(free-var-list true-exp)
           ,@(free-var-list false-exp))]
      [(,prim ,rand* ...)
       (guard (primitive? prim)) (map-append free-var-list rand*)]
      [(let ((,id* ,rand*) ...) ,body)
       `(,@(map-append free-var-list rand*)
           ,@(list-diff (free-var-list body) id*))]
      [(letrec ((,name* (lambda (,ids* ...) ,body*)) ...) ,body)
       (list-diff
         `(,@(map-append2 (lambda (ids body)
                            (list-diff (free-var-list body) ids))
                          ids* body*)
             ,@(free-var-list body))
         name*)]
      [(letcc ,id ,body)
       (list-diff (free-var-list body) (list id))]
      [(throw ,value-exp ,cont-exp)
       `(,@(free-var-list value-exp) ,@(free-var-list cont-exp))]
      [(,rator ,rand* ...)
       `(,@(free-var-list rator)
           ,@(map-append free-var-list rand*))]
      [,other (error 'free? "Invalid expression ~s%" other)])))

(define map-append
  (lambda (f args)
    (cond
      ((null? args) '())
      (else (append (f (car args)) (map-append f (cdr args)))))))

(define map-append2
  (lambda (f args1 args2)
    (cond
      ((null? args1) '())
      (else (append (f (car args1) (car args2))
                    (map-append2 f (cdr args1) (cdr args2)))))))

(define beta-subst
  (lambda (e v-id exp)
    (match exp
      [,constant
        (guard (or (number? constant) (boolean? constant)))
        constant]
      [,var (guard (symbol? var)) (if (eq? var v-id) e var)]
      [(APk ,cont ,rand)
       `(APk ,(beta-subst e v-id cont)
	  ,(beta-subst e v-id rand))]
      [(CLAM (,id) ,body)
       `(CLAM (,id) ,(beta-subst e v-id body))]
      [(lambda (,id* ...) ,body)
       (guard (memq v-id id*))
       `(lambda (,id* ...) ,body)]
      [(lambda (,id* ...) ,body)
       (guard (disjoint? id* (free-vars e)))
       `(lambda (,id* ...) ,(beta-subst e v-id body))]
      [(lambda (,id* ...) ,body)
       (beta-subst e v-id
	 (let ((x (car (intersection id* (free-vars e))))
	       (g (fresh-var (fresh-var x))))
	   `(lambda ,(subst g x id*)
	      ,(beta-subst g x body))))]
      [(if ,test-exp ,true-exp ,false-exp)
       `(if ,(beta-subst e v-id test-exp)
	  ,(beta-subst e v-id true-exp)
	  ,(beta-subst e v-id false-exp))]
      [(,prim ,rand* ...)
       (guard (primitive? prim))
       `(,prim ,@(map (lambda (rand) (beta-subst e v-id rand)) rand*))]
      [(let ((,id* ,rand*) ...) ,body)
       (guard (memq v-id id*))
       `(let ,(map list id*
		(map (lambda (rand)
		       (beta-subst e v-id rand))
		  rand*))
          ,body)]
      [(let ((,id* ,rand*) ...) ,body)
       (guard (disjoint? id* (free-vars e)))
       `(let ,(map list id*
		(map (lambda (rand)
		       (beta-subst e v-id rand))
		  rand*))
          ,(beta-subst e v-id body))]
      [(let ((,id* ,rand*) ...) ,body)
       (beta-subst e v-id
	 (let ((x (car (intersection id* (free-vars e)))))
	   (let ((g (fresh-var (fresh-var x))))
	     `(let ,(map list (subst g x id*)
		      (map (lambda (rand)
			     (beta-subst e v-id rand))
			rand*))
		,(beta-subst g x body)))))]
      [(letrec ((,name* (lambda (,ids* ...) ,body*)) ...)
         ,body)
       (guard (memq v-id name*))
       `(letrec ,(map list name*
		   (map (lambda (id* body)
			  `(lambda (,id* ...) ,body))
		     ids* body*))
          ,body)]
      [(letrec ((,name* (lambda (,ids* ...) ,body*)) ...)
         ,body)
       (guard (disjoint? name* (free-vars e)))
       (beta-subst e v-id
	 (let ((x (car (intersection name* (free-vars e)))))
	   (let ((g-id (fresh-var (fresh-var x))))
	     `(letrec (map list (subst g x name*)
			(map (lambda (ids body)
			       (if (memq x ids)
				 `(lambda ,ids ,body)
				 `(lambda ,(subst g-id x ids)
				    ,(beta-subst g-id x body))))
			  ids* body*))
		,(beta-subst g-id x body)))))]
      [(letrec ((,name* (lambda (,ids* ...) ,body*)) ...)
         ,body)
       `(letrec ,(map list name*
		   (map (lambda (ids body)
			  (beta-subst e v-id `(lambda (,ids ...) ,body)))
		     ids* body*))
          ,(beta-subst e v-id body))]
      [(letcc ,id ,body)
       (guard (eq? v-id id))
       `(letcc ,id ,body)]
      [(letcc ,id ,body)
       (guard (memq id (free-vars e)))
       (beta-step
         (let ((g (fresh-var (fresh-var id))))
           `(letcc ,g ,(beta-subst g id body))))]
      [(letcc ,id ,body)
       `(letcc ,id ,(beta-subst e v-id body))]
      [(throw ,value-exp ,cont-exp)
       `(throw
          ,(beta-subst e v-id value-exp)
          ,(beta-subst e v-id cont-exp))]
      [(,rator ,rand* ...)
       `(,(beta-subst e v-id rator)
	 ,@(map (lambda (x) (beta-subst e v-id x)) rand*))]
      [,other (error 'beta-subst "Invalid expression ~s" other)])))

(define disjoint?
  (lambda (set1 set2)
    (cond
      ((null? set1) #t)
      ((memq (car set1) set2) #f)
      (else (disjoint? (cdr set1) set2)))))

(define intersection
  (lambda (set1 set2)
    (cond
      ((null? set1) '())
      ((memq (car set1) set2)
       (cons (car set1) (intersection (cdr set1) set2)))
      (else (intersection (cdr set1) set2)))))

(load "cps-tester.ss")
(define header "******************************  D  *********************")
(newline)
(test-all header)

(define resultsD test-results)