(module b-add-transform (lib "eopl.ss" "eopl")
  (require "top.scm")
  
  (provide b_add-matches?
           b_add-transform)
  
  (provide simple-non-add-src
           simple-add-src
           slightly-complex-add-src)
  
  ;; > (run "let add = proc(n) proc(m) -(n,-(0,m)) in ((add 3) 4)")
  ;; > (scan&parse "let x = 3 in x")
  ;; > (scan&parse "let add = proc(n) proc(m) -(n,-(0,m)) in ((add 3) 4)")
  
  (define simple-non-add-src
    "let x = 3 in x")
  (define simple-add-src
    "let add = proc(n) proc(m) -(n,-(0,m)) in ((add 3) 4)")
  (define slightly-complex-add-src
  "
let add = proc(n) proc(m) -(n,-(0,m)) 
in let x = 3
   in if zero?(x)
      then 4
      else -(11,((add x) x))")
  
  (define       simple-non-add-pgm
    (scan&parse simple-non-add-src))
  (define       simple-add-pgm
    (scan&parse simple-add-src))
  (define       slightly-complex-add-pgm
    (scan&parse slightly-complex-add-src))
  
  ;; Program -> Exp
  ;; usage: (program->exp p) produces the expression for p.
  (define (program->exp pgm)
    (cases program pgm
      (a-program (e) e)))
  
  (define         simple-non-add-exp 
    (program->exp simple-non-add-pgm))
  (define         simple-add-exp     
    (program->exp simple-add-pgm))
  (define         slightly-complex-add-exp
    (program->exp slightly-complex-add-pgm))
  
    
  ;; > (value-of-program slightly-complex-add-pgm)
  
  
  ;; expression-is-add-invocation : Exp -> Boolean or (list Exp Exp)
  ;; usage: (expression-is-add-invocation? <<((add E1) E2)>>) 
  ;;        returns (list E1 E2), otherwise returns #f
  (define expression-is-add-invocation 
    (lambda (exp)
      (cases expression exp
        (call-exp (rator-1 rand-1) 
          (cases expression rator-1
            (call-exp (rator-2 rand-2)
              (cases expression rator-2
                (var-exp (id) (if (equal? 'add id) (list rand-1 rand-2) #f))
                (else #f)))
            (else #f)))
        (else #f))))
  
  ;; expression-is-let-of-add : Exp -> Boolean or (list Exp)
  ;; usage: (expression-is-let-of-add? 
  ;;         <<let add = proc(n)) proc(m) -(n,-(0,m)) in E>>)
  ;;        returns (list E), otherwise returns #f
  (define expression-is-let-of-add
    (let ((bound-exp-for-add 
           (cases program (scan&parse 
                           "let add = proc(n) proc(m) -(n,-(0,m)) in 0")
             (a-program (e)
               (cases expression e
                 (let-exp (id bound-exp body-exp)
                    bound-exp)
                 (else (eopl:error 'expression-is-let-of-add 
                                   "this should never happen.")))))))
      (lambda (exp)
        (cases expression exp
          (let-exp (id bound-exp body)
            (if (and (equal? id 'add)
                     (equal? bound-exp bound-exp-for-add))
                (list body)
                #f))
          (else #f)))))
  
  ;; expression-count-add-invocations : Exp -> Nat
  ;; usage: (expression-count-add-invocations e) 
  ;;        returns the number of times e contains a subexpression
  ;;        of the form ((add E1) E2)
  (define expression-count-add-invocations
    (lambda (e)
      (cond
        ((expression-is-add-invocation e)
         => (lambda (subexps)
              (+ 1 
                 (expression-count-add-invocations (car subexps))
                 (expression-count-add-invocations (cadr subexps)))))
        (else 
         (cases expression e
           (const-exp (num) 0)
           (var-exp (var) 0)
           (diff-exp (exp1 exp2)
             (+ (expression-count-add-invocations exp1)
                (expression-count-add-invocations exp2)))
           (zero?-exp (exp1)
             (expression-count-add-invocations exp1))
           (if-exp (exp1 exp2 exp3)
             (+ (expression-count-add-invocations exp1)
                (expression-count-add-invocations exp2)
                (expression-count-add-invocations exp3)))
           (let-exp (var exp1 exp2)
             (+ (expression-count-add-invocations exp1)
                (expression-count-add-invocations exp2)))
           (proc-exp (var exp1)
             (expression-count-add-invocations exp1))
           (call-exp (exp1 exp2)
             (+ (expression-count-add-invocations exp1)
                (expression-count-add-invocations exp2)))
           (letrec-exp (p-name b-var exp1 exp2)
             (+ (expression-count-add-invocations exp1)
                (expression-count-add-invocations exp2)))
           (begin-exp (exp1 exps)
              (letrec ((count-adds 
                        (lambda (e1 es count)
                          (let ((new-count 
                                 (+ count
                                    (expression-count-add-invocations e1))))
                            (if (null? es)
                                new-count
                                (count-adds (car es)
                                            (cdr es)
                                            new-count))))))
                (count-adds exp1 exps)))
           (newref-exp (exp1)
              (expression-count-add-invocations exp1))
           (deref-exp (exp1)
              (expression-count-add-invocations exp1))
           (setref-exp (exp1 exp2)
             (+ (expression-count-add-invocations exp1)
                (expression-count-add-invocations exp2))))))))
      
  ;; expression-matches-b_add-pattern? : Exp -> Boolean
  ;; usage: (expression-matches-b_add-pattern? e) 
  ;;   returns #t iff e defines a curried add procedure and 
  ;;   has exactly one expression invoking add on two arguments.
  (define expression-matches-b_add-pattern?
    (lambda (e)
      (let ((match-1 (expression-is-let-of-add e)))
        (if (not match-1)
            #f
            (let ((body (car match-1)))
              (= 1 (expression-count-add-invocations body)))))))
  
  ;; program-matches-b_add-pattern? : Program -> Boolean
  (define program-matches-b_add-pattern?
    (lambda (p)
      (cases program p
        (a-program (e)
          (expression-matches-b_add-pattern? e)))))
  
  
  (define expression-body-transform-by-b_add
    (lambda (e)
      (let ((match (expression-is-add-invocation e)))
        (if match
            (let ((exp1 (car match))
                  (exp2 (cadr match)))
              (begin-exp (call-exp (var-exp 'b_add)
                                   exp1)
                         (list (call-exp (var-exp 'b_add)
                                         exp2))))
            (cases expression e
              (const-exp (num) e)
              (var-exp (var) e)
              (diff-exp (exp1 exp2)
                (diff-exp (expression-body-transform-by-b_add exp1)
                          (expression-body-transform-by-b_add exp2)))
              (zero?-exp (exp1)
                (zero?-exp (expression-body-transform-by-b_add exp1)))
              (if-exp (exp1 exp2 exp3)
                (if-exp (expression-body-transform-by-b_add exp1)
                        (expression-body-transform-by-b_add exp2)
                        (expression-body-transform-by-b_add exp3)))
              (let-exp (var exp1 body)
                 (let-exp var
                          (expression-body-transform-by-b_add exp1)
                          (expression-body-transform-by-b_add body)))
              (proc-exp (var body)
                 (proc-exp var
                           (expression-body-transform-by-b_add body)))
              (call-exp (rator rand)
                 (call-exp (expression-body-transform-by-b_add rator)
                           (expression-body-transform-by-b_add rand)))
              (letrec-exp (p-name b-var p-body letrec-body)
                 (letrec-exp 
                  p-name
                  b-var
                  (expression-body-transform-by-b_add p-body)
                  (expression-body-transform-by-b_add letrec-body)))
              (begin-exp (exp1 exps)
                 (begin-exp (expression-body-transform-by-b_add exp1)
                            (map expression-body-transform-by-b_add
                                 exps)))
              (newref-exp (exp1)
                 (newref-exp (expression-body-transform-by-b_add exp1)))
              (deref-exp (exp1)
                 (deref-exp (expression-body-transform-by-b_add exp1)))
              (setref-exp (exp1 exp2)
                 (setref-exp (expression-body-transform-by-b_add exp1)
                             (expression-body-transform-by-b_add exp2)))
              )))))
  
  (define expression-unparse
    (lambda (e)
      (cases expression e
        (const-exp (num) (number->string num))
        (var-exp (var) (symbol->string var))
        (diff-exp (exp1 exp2)
          (string-append "-(" (expression-unparse exp1)
                         ","  (expression-unparse exp2)
                         ")"))
        (zero?-exp (exp1)
          (string-append "zero?(" (expression-unparse exp1)
                         ")"))
        (if-exp (exp1 exp2 exp3)
          (string-append "if "    (expression-unparse exp1)
                         " then " (expression-unparse exp2)
                         " else " (expression-unparse exp3)))
        (let-exp (var exp1 body)
          (string-append "let " (symbol->string var)
                         " = "  (expression-unparse exp1)
                         "\n in " (expression-unparse body)))
        (proc-exp (var body)
          (string-append "proc (" (symbol->string var)
                         ") "     (expression-unparse body)))
        (call-exp (rator rand)
          (string-append "("  (expression-unparse rator)
                         " "  (expression-unparse rand)
                         ")"))
        (letrec-exp (p-name b-var p-body letrec-body)
          (string-append "letrec " (symbol->string b-var)
                         " (" (symbol->string b-var)
                         ") = " (expression-unparse p-body)
                         "\n in " (expression-unparse letrec-body)))
        (begin-exp (exp1 exps)
           (letrec ((prepend-semicolons
                     (lambda (exps)
                       (if (null? exps)
                           ""
                           (string-append "; "
                                          (expression-unparse (car exps))
                                          (prepend-semicolons (cdr exps)))))))
             (string-append "begin "
                            (expression-unparse exp1)
                            (prepend-semicolons exps)
                            " end")))
        (newref-exp (exp1)
           (string-append "newref( " (expression-unparse exp1) ")"))
        (deref-exp (exp1)
           (string-append "deref( " (expression-unparse exp1) ")"))
        (setref-exp (exp1 exp2)
           (string-append "setref( " (expression-unparse exp1)
                          ", " (expression-unparse exp2) ")"))
        )))
  
  (define program-unparse
    (lambda (p)
      (cases program p
        (a-program (e)
          (expression-unparse e)))))
  
  ;; b_add-matches? : String -> Boolean
  (define (b_add-matches? s)
    (program-matches-b_add-pattern? (scan&parse s)))

  ;; > (b_add-matches? "let add = x in x")
  ;; > (b_add-matches? "let add = proc(m) proc(n) -(m,-(0,n)) in x")
  ;; > (b_add-matches? "let add = proc(n) proc(m) -(m,-(0,n)) in x")
  ;; > (b_add-matches? "let add = proc(n) proc(m) -(n,-(0,m)) in x")
  ;; > (b_add-matches? "let add = proc(n) proc(m) -(n,-(0,m)) in ((add 1) 2)")
  ;; > (b_add-matches? "let add = proc(m) proc(n) -(n,-(0,m)) in ((add 1) 2)")
  
  ;; b_add-transform : String -> String
  (define (b_add-transform s)
    (let ((parsed (scan&parse s)))
      (if (program-matches-b_add-pattern? parsed)
          (string-append 
           "let b_add = let sum = newref(0)
                        in proc(n) begin setref(sum, -(deref(sum),-(0,n))); 
                                         deref(sum) 
                                   end
            in "
           (expression-unparse (expression-body-transform-by-b_add 
                                (program->exp parsed))))
          (eopl:error 'b_add-transform 
                      " input ~s does not match b_add input." s))))
     
  )