(define (member? a b)
  (cond ((null? b) #f)
        ((equal? a (car b)) #t)
        (else (member? a (cdr b)))))

(define (adjoin-set x set)
  (if (member? x set)
      set
      (cons x set)))

(define (intersection-set set1 set2)
  (cond ((or (null? set1) (null? set2)) '())
	((member? (car set1) set2)
	 (cons (car set1) (intersection-set (cdr set1) set2)))
	(else (intersection-set (cdr set1) set2))))

(define (union-set set1 set2)
  (cond ((null? set1) set2)
	(else (union (cdr set1)
		     (adjoin-set (car set1) set2)))))

(define compile-port
  (make-parameter
   (current-output-port)
   (lambda (p)
     (unless (output-port? p)
             (error 'compile-port (format #t "Not an output port ~s." p)))
     p)))

(define (emit . args)
  (apply format (compile-port) args)
  (format (compile-port) ";")
  (newline (compile-port)))

(define (emit-no-colon . args)
  (apply format (compile-port) args)
  (format (compile-port) "")
  (newline (compile-port)))


(define (emit-no-newline . args)
  (apply format (compile-port) args)
  (format (compile-port) ""))

(define (compile-label exp)
  (emit-no-colon "~s:" (label-name exp)))

(define (compile-save exp)
  (emit "push(~s)" (register-name exp)))

(define (extract-registers prog)
  (define (extract-acc prog acc)
    (cond
     ((null? prog) acc)
     ((eq? (caar prog) 'assign)
      (extract-acc
       (cdr prog)
       (adjoin-set (cadar prog) acc)))
     (else
      (extract-acc
       (cdr prog)
       acc))))
  (extract-acc prog '()))

(define (tagged-list? x sym)
  (eq? (car x) sym))
(define (register? x)
  (tagged-list? x 'reg))
(define (label? x)
  (tagged-list? x 'label))
(define (register-name x)
  (and (register? x) (cadr x)))
(define (label-name x)
  (and (label? x) (cadr x)))
(define (assign? x)
  (tagged-list? x 'assign))
(define (branch? x)
  (tagged-list? x 'branch))
(define (goto? x)
  (tagged-list? x 'goto))
(define (test? x)
  (tagged-list? x 'test))
(define (test-op x)
  (cadadr x))
(define (test-args x)
  (cddr x))
(define (op? x)
  (tagged-list? x 'op))
(define (op-arg x)
  (and (op? x) (cadr x)))
(define (const? x)
  (tagged-list? x 'const))
(define (const-name x)
  (and (const? x) (cadr x)))
(define (save? x)
  (tagged-list? x 'save))
(define (save-name x)
  (and (save? x) (cadr x)))
(define (restore? x)
  (tagged-list? x 'restore))
(define (restore-name x)
  (and (restore? x) (cadr x)))
(define (assign? x)
  (tagged-list? x 'assign))
(define (assign-reg-name x)
  (and (assign? x) (cadr x)))
(define (assign-args x)
  (and (assign? x) (cddr x)))
(define (compile-goto exp)
  (let ((arg (cadr exp)))
    (cond ((label? arg)
           (emit "goto ~s" (label-name arg)))
          ((and (register? arg)
                (eq? (register-name arg) 'continue))
           (emit "goto *cont->g")))))

(define (display? x)
  (tagged-list? x 'display))
(define (display-reg x)
  (and (display? x) (cadr x)))

(define (continue? x)
  (eq? x 'continue))
;; Assume (test (op <op>) <reg or const> <reg or const>)

(define (compile-reg-or-const exp)
  (cond ((const? exp)
         (let ((var (const-name exp)))
           (cond ((number? var)
                  (emit-no-newline "make_fixnum(~a)" var))
                 ((symbol? var)
                  (emit-no-newline "make_symbol(~a)" var))
                 ((boolean? var)
                  (emit-no-newline "make_boolean(~a)" (if var 1 0)))
                 ((string? var)
                  (emit-no-newline "make_string(\"~a\")"
                                   (string->symbol var))))))
        ((register? exp)
         (emit-no-newline "(reg*)~a" (register-name exp)))))

(define (compile-test exp)
  (begin
    (let ((compare-op (op-arg (cadr exp))))
      (cond ((eq? '< compare-op)
             (emit "test = &less_than_primop"))
            ((eq? '> compare-op)
             (emit "test = &greater_than_primop"))
            ((eq? '= compare-op)
             (emit "test = &equal_to_primop"))))
    (emit-no-newline "(*test)")
    (compile-op-args (test-args exp) 1)))

(define (compile-branch exp)
  (cond ((label? (cadr exp))
         (emit-no-colon "if (al) {goto ~a;}" (label-name (cadr exp))))))

(define (compile-save exp)
  (let ((register (save-name exp)))
    (emit "save(~a)" (if (continue? register)
                         'cont
                         register))))

(define (compile-restore exp)
  (let ((register (restore-name exp)))
    (emit "restore(~a)" (if (eq? register 'continue)
                            'cont
                            register))))

(define (compile-assign exp)
  (let ((reg (assign-reg-name exp))
        (args (assign-args exp)))
    (cond ((eq? reg 'continue)
           (emit "cont->g = &&~a" (label-name (car args))))
          ((= 3 (length args))
           (begin
             (emit-no-newline "~a = " reg)
             (let ((op (op-arg (car args)))
                   (op-args (cdr args)))
               (cond ((eq? '+ op)
                      (emit-no-newline "(**add_primop)"))
                     ((eq? '- op)
                      (emit-no-newline "(**sub_primop)"))
                     ((eq? '* op)
                      (emit-no-newline "(**mult_primop)")))
               (compile-op-args op-args 1))))
          ((= 1 (length args))
           (begin (emit-no-newline "~a = " reg)
                  (compile-reg-or-const (car args))
                  (emit ""))))))

(define (compile-declare-reg x)
  (if (continue? x)
      (begin (emit "reg *cont")
             (emit "cont = alloc_reg()")
             (emit "cont->t = GOTO"))
      (emit "reg *~a = alloc_reg()" x)))

(define fib-prog
  '((assign n (const 10))
    (assign continue (label fib_done))
    (assign val (const 0))
    (label fib_loop)
    (test (op <) (reg n) (const 2))
    (branch (label immediate_answer))
    (save continue)
    (assign continue (label afterfib_n_1))
    (save n) 
    (assign n (op -) (reg n) (const 1))
    (goto (label fib_loop))
    
    (label afterfib_n_1)
    (restore n)
    (restore continue)
    (assign n (op -) (reg n) (const 2))
    (save continue)
    (assign continue (label afterfib_n_2))
    (save val)
    (goto (label fib_loop))
    (label afterfib_n_2)
    (assign n (reg val))
    (restore val)
    (restore continue)
    (assign val (op +) (reg val) (reg n))
    (goto (reg continue))
    (label immediate_answer)
    (assign val (reg n))
    (goto (reg continue))
    (label fib_done)
    (display val)))


(define (compile-prog prog)
  (emit-no-colon "int main(void) {")
  ;; These are the registers that are common to all register machine
  ;; programs.
  (emit "stack = calloc(1, sizeof(stack_s))")
  (emit "void (*test)()")
  (for-each (lambda (x) (compile-declare-reg x))
            (extract-registers prog))
  (for-each (lambda (x)
              (cond ((test? x)
                     (compile-test x))
                    ((restore? x)
                     (compile-restore x))
                    ((branch? x)
                     (compile-branch x))
                    ((label? x)
                     (compile-label x))
                    ((goto? x)
                     (compile-goto x))
                    ((assign? x)
                     (compile-assign x))
                    ((save? x)
                     (compile-save x))
                    ((display? x)
                     (compile-display x))))
            prog)
  (emit-no-colon ";\n}"))

;; Need to compile operation expressions...
;; ((op ⟨operation-name⟩) ⟨input1⟩ . . . ⟨inputn⟩)
;; 


(define (compile-op-args exp headp)
  (cond ((null? exp)
         (emit ")"))
        ((eq? headp 1)
         (begin
           (emit-no-newline "(")
           (compile-reg-or-const (car exp))
           (compile-op-args (cdr exp) 0)))
        ((eq? headp 0)
         (begin
           (emit-no-newline ",")
           (compile-reg-or-const (car exp))
           (compile-op-args (cdr exp) 0)))))

(define simple-prog
  '((assign n (const 5))
    (display n)
    (save n)
    (assign n (op +) (reg n) (const 3))
    (display n)
    (restore n)
    (save n)))

;; The following program pushes the numbers 1 to 9 onto the stack and
;; restores them again.

(define counter-prog
  '((assign n (const 1))
    (label start_save)
    (display n)
    (test (op =) (reg n) (const 10))
    (branch (label done_save))
    (save n)
    (assign n (op +) (reg n) (const 1))
    (goto (label start_save))
    (label done_save)
    
    (label print_n)
    (display n)
    (restore n)
    (test (op =) (reg n) (const 1))
    (branch (label print_done))
    
    (goto (label print_n))
    (label print_done)))

;; Will go into an infinite loop printing 5 if the goto statement
;; wasn't saved successfully onto the stack.

(define goto-prog
  '((assign foo (const 5))
    (assign bar (const 1000))
    (label start)
    (display foo)
    (assign continue (label doge))
    (save continue)
    (assign continue (label start))
    (restore continue)
    (goto (reg continue))
    (label doge)
    (display bar)))

(define (compile-display x)
  (emit "print(*~a)" (display-reg x))
  (emit "puts(\"\")"))

(define stack-prog
  '((assign n (const 1))
    (save n)
    (display n)

    (assign n (op +) (reg n) (const 1))
    (save n)
    (display n)
    
    (assign n (op +) (reg n) (const 1))
    (save n)
    (display n)

    (restore n)
    (display n)
    (restore n)
    (display n)
    (restore n)
    (display n)))


(define fact-prog
  '((assign continue (label fact_done))
    (assign n (const 5))
    (label fact_loop)
    (test (op =) (reg n) (const 1))
    (branch (label base_case))
    (save continue)
    (save n)
    (assign n (op -) (reg n) (const 1))
    (assign continue (label after_fact))
    (goto (label fact_loop))
    (label after_fact)
    (restore n)
    (restore continue)
    (assign val (op *) (reg n) (reg val))
    (goto (reg continue))
    (label base_case)
    (assign val (const 1))
    (goto (reg continue))
    (label fact_done)
    (assign msg (const "Value: "))
    (display msg)
    (display val)))
