#lang racket

(require "uniquify.scm")

(provide remove-complex-opera*)

; remove complex SUB-expressions:
; of course complex expression are possible
; but their argument can only be either
; - number
; - symbol
; - `(read)
; temporary variables are introduced to make this possible
; with the caveat that a binding should never associate
; an existing symbol to a temporary
; (Invariant of rco-arg)
(define (remove-complex-opera* p)
  (let ([uniq-res (uniquify p)])
       (match (cadr uniq-res)
         [`(program ,data ,exp) 
          (begin
            (define initial-tmpcount (hash-ref (car uniq-res) `tmp 0))
            (define-values (new-exp bla) (rco-exp exp initial-tmpcount))
            `(program ,data ,new-exp))])))

(define (let-binding var val exp)
  `(let ([,var ,val]) ,exp))

(define (wrap-associations assoc-list exp)
  (if (empty? assoc-list)
    exp
    (let ([binding (car assoc-list)])
      (let-binding (car binding) (cadr binding)
                   (wrap-associations (cdr assoc-list) exp)))))
  
(define (substitute-var old-var var expr)
  (match expr
    [(? fixnum?) expr]
    [`(read) expr]
    [(? symbol?) 
     (if (eq? old-var expr) var expr)]
    [`(- ,e) `(- ,(substitute-var old-var var e))]
    [`(+ ,e1 ,e2)
     (begin
      (define se1 (substitute-var old-var var e1))
      (define se2 (substitute-var old-var var e2))
      `(+ ,se1 ,se2))]
    [`(let ([,v ,e]) ,body)
     (begin
      (define se (substitute-var old-var var e))
      (define sbody (substitute-var old-var var body))
      `(let ([,v ,se]) ,sbody))]))

(define (change-assoc-key assoc-list old-key new-key)
  (if (empty? assoc-list)
    `()
    (let ([key (caar assoc-list)])
      (let ([new-pair (if (eq? key old-key)
                        `(,new-key . ,(cdar assoc-list))
                        (car assoc-list))])
       (cons new-pair (change-assoc-key (cdr assoc-list) old-key new-key))))))

(define (is-tmp exp tmpcount)
  (if (and (symbol? exp) (string-prefix? (symbol->string exp) "tmp."))
    (let ([tmpvalue (string->number (substring (symbol->string exp) 4))])
      (if (and tmpvalue (> tmpvalue tmpcount))
        #t
        #f))
    #f))


(define (rco-exp exp tmpcount)
  (match exp
    [(? fixnum?) (values exp tmpcount)]
    [(? symbol?) (values exp tmpcount)]
    [`(read) (values exp tmpcount)]
    [`(- ,e) 
     (begin
      (define-values (new-exp assoc-list new-tmpcount) (rco-arg e tmpcount))
      (values (wrap-associations assoc-list `(- ,new-exp)) new-tmpcount))]
    [`(,op ,e1 ,e2)
     #:when (or (eq? op `+) (eq? op `-))
     (begin
      (define-values (new-exp1 assoc-list1 new-tmpcount1) (rco-arg e1 tmpcount))
      (define-values (new-exp2 assoc-list2 new-tmpcount) (rco-arg e2 new-tmpcount1))
      (values (wrap-associations (append assoc-list1 assoc-list2)
                                 `(,op ,new-exp1 ,new-exp2)) new-tmpcount))]
    [`(let ([,var ,e]) ,body)
     (begin
      (define-values (exp-tmp assoc-list exp-tmpcount) (rco-arg e tmpcount))      
      (define-values (new-body new-tmpcount) (rco-exp body exp-tmpcount))
      (define elide-tmp (or (is-tmp exp-tmp tmpcount)))
      (set! assoc-list
        (if (or elide-tmp (fixnum? exp-tmp))
          (change-assoc-key assoc-list exp-tmp var)
          (append assoc-list (list `(,var ,exp-tmp)))))
      (set! new-body
        (if (symbol? exp-tmp)
          new-body ;(substitute-var exp-tmp var new-body) new-body)
          `(let ([,var ,exp-tmp]) ,new-body)))
      (values (wrap-associations assoc-list new-body)
              new-tmpcount))]))

(define (get-unique-symbol tmpcount)
  (string->symbol (format "tmp.~a" tmpcount)))

(define (rco-arg exp tmpcount)
  (match exp
   [(? fixnum?) (values exp '() tmpcount)]
   [(? symbol?) (values exp '() tmpcount)]
   [`(read) (values exp '() tmpcount)]
   [`(- ,e)
    (begin
     (define-values (new-exp assoc-list exp-tmpcount) (rco-arg e tmpcount))
     (define new-tmpcount (+ exp-tmpcount 1))
     (define tmpname (get-unique-symbol new-tmpcount))
     (set! assoc-list (append assoc-list (list `(,tmpname (- ,new-exp)))))
     (values tmpname
             assoc-list
             new-tmpcount))]
   [`(,op ,e1 ,e2)
    #:when (or (eq? op `+) (eq? op `-))
    (begin
     (define-values (new-exp1 assoc-1 exp1-tmpcount) (rco-arg e1 tmpcount))
     (define-values (new-exp2 assoc-2 exp2-tmpcount) (rco-arg e2 exp1-tmpcount))
     (define assoc-list (append assoc-1 assoc-2))
     (define new-tmpcount (+ exp2-tmpcount 1))
     (define tmpname (get-unique-symbol new-tmpcount))
     (set! assoc-list (append assoc-list (list `(,tmpname (,op ,new-exp1 ,new-exp2)))))
     (values tmpname
             assoc-list
             new-tmpcount))]
   [`(let ([,var ,rexp]) ,body)
    (begin
     (define-values (new-exp assoc-list-1 exp-tmpcount) (rco-arg rexp tmpcount))
     (define-values (body-tmp assoc-list new-tmpcount) (rco-arg body exp-tmpcount))
     (define new-assoc-list
        (if (not (is-tmp new-exp tmpcount))
          (cons `(,var ,new-exp) (append assoc-list-1 assoc-list))
          (change-assoc-key (append assoc-list-1 assoc-list) new-exp var)))
     (values body-tmp
             new-assoc-list
             new-tmpcount))]))

;(require racket/trace)
;(trace rco-exp)
;(trace rco-arg)