(module mp6 (lib "eopl.ss" "eopl")
  (require "lang.scm")             ; for expression datatype

  (provide exp-to-mnum
           transform-diff-zero-simple
           transform-diff-zero*
           transform-diff-consts-simple ; XXXX
           transform-diff-consts*       ; XXXXX
           transform-rename             ; XXXXXX 
           transform-constprop*         ; XXXXXXX
           )

  ;; exp-to-mnum : Exp -> Maybe[Number]
  ;; usage: (exp-to-mnum e) returns n if e = (const-exp n); otherwise #f

  (define exp-to-mnum
    (lambda (e)
      (cases expression e
        (const-exp (n) n)
        (else #f))))


  ;; transform-diff-zero-simple : Exp -> Maybe[Exp]
  ;; usage: (transform-diff-zero-simple e) returns e' if 
  ;;    e = (diff-exp e' (const-exp 0)), otherwise returns #f

  (define transform-diff-zero-simple
    (lambda (e)
      (cases expression e

        ;; The core of the transformation
        (diff-exp (lft rgt)
          (cond ((and (exp-to-mnum rgt)
                      (= 0 (exp-to-mnum rgt)))
                 lft)
                (else #f)))

        (else #f))))

  ;; transform-diff-zero* : Exp -> Exp
  ;; usage: (transform-diff-zero* exp) returns exp' 
  ;; where exp' is the result of iteratively replacing -(e,0) with e
  ;; in exp as many times as possible.

  (define transform-diff-zero*
    (lambda (e)
      (cases expression e

        (const-exp (n) e) ;; no subexpressions, and transform inapplicable.

        (diff-exp (lft rgt)
          (let ((lft* (transform-diff-zero* lft))
                (rgt* (transform-diff-zero* rgt)))
            ;; Note that we check applicability to lft* and rgt*
            ;; rather than lft and rgt.  What would happen if we 
            ;; instead checked applicability to lft and rgt below?
            (cond ((and (exp-to-mnum rgt*)
                        (= 0 (exp-to-mnum rgt*)))
                   lft*)
                  (else (diff-exp lft* rgt*)))))

        (zero?-exp (e) (zero?-exp (transform-diff-zero* e)))

        (if-exp (test then else)
          (if-exp (transform-diff-zero* test)
                  (transform-diff-zero* then)
                  (transform-diff-zero* else)))

        (var-exp (id) e) ;; no subexpressions, and transform inapplicable.

        (let-exp (x binding-exp body)
          (let-exp x 
                   (transform-diff-zero* binding-exp)
                   (transform-diff-zero* body)))

        (proc-exp (x body)
          (proc-exp x (transform-diff-zero* body)))

        (call-exp (rator-exp rand-exp) 
          (call-exp (transform-diff-zero* rator-exp)
                    (transform-diff-zero* rand-exp)))

        (letrec-exp (fs xs es body)
          (letrec-exp fs xs (map transform-diff-zero* es)
            (transform-diff-zero* body)))

        (begin-exp (exp remaining-exps)
          (begin-exp (transform-diff-zero* exp) 
                     (map transform-diff-zero* remaining-exps)))

        (newref-exp (exp)
          (newref-exp (transform-diff-zero* exp)))

        (deref-exp (exp)
          (deref-exp (transform-diff-zero* exp)))

        (setref-exp (lhs-exp rhs-exp)
          (setref-exp (transform-diff-zero* lhs-exp)
                      (transform-diff-zero* rhs-exp)))
        )))

;;; START OF STUDENT PROVIDED CODE

  (define ... "unfinished")

  ;; transform-rename : Exp Id Id -> Exp
  ;; usage: (transform-rename e x y) returns e[x:=y] 
  ;;    if x not in FV(e), y not in Vars(e)
  ;;    otherwise signals an error.
  (define transform-rename
    (lambda (express old-id new-id)
      (let recur ((e express) 
                  (bound #f))

        (define recur-prop  ;; Exp -> Exp, propagating bound
          (lambda (e) (recur e bound)))

        (define recur-found ;; Exp -> Exp, upgrading bound to #t
          (lambda (e) (recur e #t)))

        (define replace-var ;; Sym -> Sym, dropping old-id for new-id
          (lambda (id) (if (eq? old-id id) new-id id)))

        (cases expression e
          (const-exp (n) e)
          (diff-exp (l r) (diff-exp (recur-prop l) (recur-prop r)))
          (zero?-exp (e) (zero?-exp (recur-prop e)))
          (if-exp (b t e) (if-exp (recur-prop b) 
                                  (recur-prop t)
                                  (recur-prop e)))
          (var-exp (id) 
            (cond ((eq? id new-id)
                   (eopl:error 'alpha-rename "~a occurs in ~a"
                               new-id express))
                  ((eq? id old-id)
                   (if (not bound)
                       (eopl:error 'alpha-rename "~a is free in ~a"
                                   id express)
                       (var-exp new-id)))
                  (else
                   e)))
          (let-exp (x e b) 
            (cond ((eq? x new-id)
                   (eopl:error 'alpha-rename "~a occurs in ~a"
                               new-id express))
                  ((eq? x old-id)
                   (let-exp new-id (recur-prop e) (recur-found b)))
                  (else
                   (let-exp x      (recur-prop e) (recur-prop b)))))
          (proc-exp (x b)
            (cond ((eq? x new-id)
                   (eopl:error 'alpha-rename "~a occurs in ~a"
                               new-id express))
                  ((eq? x old-id)
                   (proc-exp new-id (recur-found b)))
                  (else
                   (proc-exp x      (recur-prop b)))))
          (call-exp (r d) (call-exp (recur-prop r) 
                                    (recur-prop d)))
          (letrec-exp (fs xs es b) 
            (cond ((or (member new-id fs) 
                       (member new-id xs))
                   (eopl:error 'alpha-rename "~a occurs in ~a" new-id express))
                  ((member old-id fs)
                   (letrec-exp 
                    (map replace-var fs)
                    (map replace-var xs)
                    (map recur-found es) 
                    (recur-found b)))
                  ((member old-id xs)
                   (letrec-exp 
                    fs
                    (map replace-var xs)
                    (map (lambda (e x)
                           ;; ooh tricky case in mapping
                           (if (eq? x old-id)
                               (recur-found e)
                               (recur-prop  e)))
                         es xs)
                    (recur-prop b)))
                  (else 
                   ;; the easy case 
                   ;; (all three cases could be folded together, but Felix
                   ;;  finds it easier to check the test-suite with 
                   ;;  them separated)
                   (letrec-exp fs xs (map recur-prop es) (recur-prop b)))
                  ))

          (begin-exp (e es) (begin-exp (recur-prop e) (map recur-prop es)))
          (newref-exp (e)   (newref-exp (recur-prop e)))
          (deref-exp (e)    (deref-exp (recur-prop e)))
          (setref-exp (l r) (setref-exp (recur-prop l) (recur-prop r)))
          ))))


  (define traversal-template
    (lambda (express  ...)
      (let recur ((e express))
        (cases expression e
          (const-exp (n) ... )
          (diff-exp (l r) ... (recur l) ... (recur r) ... )
          (zero?-exp (e) ... (recur e) ... )
          (if-exp (b t e) ... (recur b) ... (recur t) ... (recur e) ... )
          (var-exp (id) ... )
          (let-exp (x e b) ... (recur e) ... (recur b) ... )
          (proc-exp (x b) ... (recur b) ... )
          (call-exp (r d) ... (recur r) ... (recur d) ... )
          (letrec-exp (fs xs es b) ... (map recur es) ... (recur b) ... )
          (begin-exp (e es) ... (recur e) ... (map recur es) ... )
          (newref-exp (e) ... (recur e) ... )
          (deref-exp (e) ... (recur e) ... )
          (setref-exp (l r) ... (recur l) ... (recur r) ...)
          ))))


  ;; transform-diff-consts-simple : Exp -> Maybe[Exp]
  ;; usage: (transform-diff-consts-simple e) returns (const-exp n_3) if 
  ;;    e = -((const-exp n_1), (const-exp n_2)), where n_3 = n_1 + n_2,
  ;;    otherwise returns #f

  (define transform-diff-consts-simple
    (lambda (e)
      (cases expression e
        (diff-exp (lft rgt)
          (let ((ln (exp-to-mnum lft))
                (rn (exp-to-mnum rgt)))
            (and ln rn (const-exp (- ln rn)))))
        (else #f))))

  ;; Exp -> Exp
  ;; Iteratively applies transform-diff-simple to e until it cannot 
  ;; be applied anymore
  (define transform-diff-consts*
    (lambda (e)
      (cases expression e
        (const-exp (n) (const-exp n))
        (diff-exp (l r) 
          (let ((l* (transform-diff-consts* l))
                (r* (transform-diff-consts* r)))
            (or (transform-diff-consts-simple (diff-exp l* r*))
                (diff-exp l* r*))))
        (zero?-exp (e)
          (zero?-exp (transform-diff-consts* e)))
        (if-exp (b t e)
          (if-exp (transform-diff-consts* b) 
                  (transform-diff-consts* t)
                  (transform-diff-consts* e)))
        (var-exp (v) (var-exp v))
        (let-exp (x e b)  (let-exp x 
                                   (transform-diff-consts* e)
                                   (transform-diff-consts* b)))
        (proc-exp (x b)   (proc-exp x (transform-diff-consts* b)))
        (call-exp (f e)   (call-exp (transform-diff-consts* f) 
                                    (transform-diff-consts* e)))
        (letrec-exp (fs xs es b) 
          (letrec-exp fs xs 
                      (map transform-diff-consts* es)
                      (transform-diff-consts* b)))
        (begin-exp (e es) 
          (begin-exp (transform-diff-consts* e) 
                     (map transform-diff-consts* es)))
        (newref-exp (e)   (newref-exp (transform-diff-consts* e)))
        (deref-exp (e)    (deref-exp (transform-diff-consts* e)))
        (setref-exp (a b) (setref-exp (transform-diff-consts* a) 
                                      (transform-diff-consts* b)))
        )))

  ;; remf : (X -> Bool) Listof[X] -> Listof[X]
  (define remf
    (lambda (p l)
      (let recur ((l l))
        (cond 
         ((null? l) '())
         (else (if (p (car l))
                   (recur (cdr l))
                   (cons (car l) (recur (cdr l)))
                   ))))))

  (define transform-constprop*
    (lambda (express)
      (let recur-gen ((e express) (mapping '()))

        (define recur ;; recur : Exp -> Exp
          (lambda (e) (recur-gen e mapping)))
        (define recur-bind ;; recur-bind : Exp Id Num -> Exp
          (lambda (e x n) 
            (recur-gen e
                       (cons (list x n) mapping))))
        (define recur-unbind ;; recur-unbind : Exp Id -> Exp
          (lambda (e x)
            (recur-gen e (remf (lambda (entry) (eq?  (car entry) x))
                               mapping))))
        (define recur-unbinds ;; recur-unbind : Exp Listof[Id] -> Exp
          (lambda (e xs)
            (recur-gen e (remf (lambda (entry) (memq (car entry) xs))
                               mapping))))


        (cases expression e
          (const-exp (n)  e)
          (diff-exp (l r) (diff-exp (recur l) (recur r)))
          (zero?-exp (e)  (zero?-exp (recur e)))
          (if-exp (b t e) (if-exp (recur b) (recur t) (recur e)))
          (var-exp (id)
            (cond ((assq id mapping) 
                   => (lambda (e) (const-exp (cadr e))))
                  (else e)))
          (let-exp (x e b) 
            (let ((e* (recur e)))
              (cond ((exp-to-mnum e*)
                     => (lambda (n)
                          (recur-bind b x n)))
                    (else
                     (let-exp x e*
                              (recur-unbind b x))))))
          (proc-exp (x b) 
            (proc-exp x (recur-unbind b x)))
          (call-exp (r d) (call-exp (recur r) (recur d)))
          (letrec-exp (fs xs es b) 
            (letrec-exp fs xs 
                        (map (lambda (f x e)
                               ;; Again, tricky case (worth at point at least)
                               (recur-unbinds e (cons x fs)))
                             fs xs es)
                        (recur-unbinds b fs)))
          (begin-exp (e es) (begin-exp (recur e) (map recur es)))
          (newref-exp (e)   (newref-exp (recur e)))
          (deref-exp (e)    (deref-exp  (recur e)))
          (setref-exp (l r) (setref-exp (recur l) (recur r)))
          ))))

;;; END OF STUDENT PROVIDED CODE

)