(define-library (csc codegen)
  (export
    ir2->ir3)
  (import (scheme base)
          (only (csc gensym)
            gensym
            gensym->int)
          (only (csc hash-map)
            compare-numbers
            delete
            hash-bytevector
            insert
            lookup
            make-comparer
            make-map
            map-for-each)
          (only (csc ir2)
            %apply
            %branch
            %constant
            %globals
            %label
            %library-ref
            %primitive
            %tail
            %variable
            closure-arguments
            closure-body
            closure-name
            constant-expression
            constant?
            fix-body
            fix-functions
            globals?
            label-gensym
            label?
            library-ref-library
            library-ref-name
            library-ref?
            make-apply
            make-constant
            make-label
            make-primitive
            variable-gensym
            variable?)
          (only (csc loop)
            loop)
          (only (csc match)
            match))
  (begin


    (define (atom->bytecode atom translate-local)
      (match atom
        ((% %constant x)
          (cond
            ((and (integer? x)
                  (> x (- (expt 2 30) 1)))  ; out of range for a small int
              (error "I don't support big ints yet"))
            ((or (integer? x)
                 (boolean? x)
                 (null? x))
              (list 'const x))
            (else (error "Only small ints and bool constants are supported for now" x))))
        ((% %library-ref x lib)
          (list 'global x lib))
        ((% %variable sym)
          (list 'local (translate-local atom)))
        ((% %globals)
          ; The globals array is stored in register 0.
          (list 'local 0))
        ((% %label sym)
          (list 'label (translate-local atom)))
        (_ (error "Unexpected form in atom->bytecode" atom))))


    (define-record-type <not-empty>
      (make-not-empty)
      not-empty?)


    (define *not-empty* (make-not-empty))


    (define (empty? m)
      (guard (e ((not-empty? e) #f))
        (map-for-each (lambda (k v)
                        (raise *not-empty*))
          m)
        #t))


    (define *temp-reg* 255)


    (define (get-satisfying m pred)
      (define elem #f)
      (guard (e ((not-empty? e) elem))
        (map-for-each (lambda (k v)
                        (when (pred k)
                          (set! elem k)
                          (raise *not-empty*)))
          m)
        #f))


    (define (get-least m)
      (define least #f)
      (map-for-each (lambda (k v)
                      (when (or (not least)
                              (< k least))
                        (set! least k)))
        m)
      least)


    (define (chains in->out)
      (define out->in (make-map compare-numbers))
      (map-for-each (lambda (k v)
                      (set! out->in (insert out->in v k)))
        in->out)
      (define currently-in-temp #f)
      (loop with results = out->in
            with save-regs = in->out
            for easy-result = (get-satisfying results (lambda (x) (not (lookup save-regs x #f))))
            until (empty? results)
            if easy-result
              collect (let ((in (lookup out->in easy-result)))
                        (set! results (delete results easy-result))
                        (set! save-regs (delete save-regs in))
                        (list 'mov (list 'local easy-result) (list 'local (lookup out->in easy-result))))
            else if currently-in-temp
              collect (let ((target (lookup in->out currently-in-temp)))
                        (set! results (delete results target))
                        (list 'mov (list 'local target) (list 'local *temp-reg*)))
              and do (set! currently-in-temp #f)
            else
              append (let* ((any-result (get-least results))
                            (in (lookup out->in any-result)))
                        (set! currently-in-temp any-result)
                        (set! results (delete results any-result))
                        (set! save-regs (delete save-regs any-result))
                        (list
                          (list 'mov (list 'local *temp-reg*) (list 'local any-result))
                          (list 'mov (list 'local any-result) (list 'local in))))))


    (define (hash-symbol s)
      (hash-bytevector (string->utf8 (symbol->string s))))


    (define (cmp-symbols s1 s2)
      (cond
        ((symbol=? s1 s2) 0)
        ((string<? (symbol->string s1) (symbol->string s2)) -1)
        (else 1)))


    (define (ir2->bytecode expr translate-local)
      (define (a->b atom)
        (atom->bytecode atom translate-local))
      (match expr
        ((% %primitive op args res cont)
          (cons
            (append (list op) (map a->b res) (map a->b args))
            (ir2->bytecode cont translate-local)))
        ((% %branch atom true false)
          (define temp (translate-local (make-label (gensym))))
          (append
            (list
              (list 'jmpif (a->b atom) (list 'label temp)))
            (ir2->bytecode false translate-local)
            (list
              (list 'label temp))
            (ir2->bytecode true translate-local)))
        ((% %apply proc args)
          (define in->out (make-map compare-numbers))
          (define constants
            (loop for arg in args
                  for i from 1
                  if (variable? arg)
                    unless (= (translate-local arg) i)
                      do (set! in->out (insert in->out (translate-local arg) i))
                    end
                  else if (globals? arg)
                    do (set! in->out (insert in->out 0 i))
                  else if (constant? arg)
                    collect (list 'mov (list 'local i) (list 'const (constant-expression arg)))
                  else if (label? arg)
                    collect (list 'mov (list 'local i) (list 'label (translate-local arg)))
                  else if (library-ref? arg)
                    collect (list 'mov (list 'local i) (list 'global
                                                             (library-ref-name arg)
                                                             (library-ref-library arg)))
                  else
                    do (error "Unexpected form in arguments list" arg)))
          (define proc-temp (translate-local proc))
          (when (and (variable? proc)
                     (<= proc-temp (length args)))
            (let ((available-reg (+ 1 (length args))))
              (set! in->out (insert in->out proc-temp available-reg))
              (set! proc-temp available-reg)))
          (append
            (chains in->out)
            constants
            (list
              (if (label? proc)
                (list 'jmp (list 'label proc-temp))
                (list 'jmp (list 'local proc-temp))))))
        ((% %tail)
          (list
            (list 'jmp (list 'label 0))))
        (_ (error "Unexpected form in ir2->bytecode" expr))))


    (define compare-labels
      (make-comparer
        (lambda (x) (gensym->int (label-gensym x)))
        (lambda (x y) (- (gensym->int (label-gensym y)) (gensym->int (label-gensym x))))))


    (define compare-variables
      (make-comparer
        (lambda (x) (gensym->int (variable-gensym x)))
        (lambda (x y) (- (gensym->int (variable-gensym y)) (gensym->int (variable-gensym x))))))


    ; Converts an IR2 program into bytecode.
    (define (ir2->ir3 expr)
      (define label-map (make-map compare-labels))
      (define next-label-id 1)  ; start at 1 because label 0 is used for tail.
      (define (translate-label x)
        (or (lookup label-map x #f)
            (let ((id next-label-id))
              (set! next-label-id (+ 1 next-label-id))
              (set! label-map (insert label-map x id))
              id)))
      (define (make-locals-map args)
        (define locals-map (make-map compare-variables))
        (loop for arg in args
              for i from 1
              do (set! locals-map (insert locals-map arg i)))
        (define local-count (length args))
        (lambda (x)
          (if (label? x)
            (translate-label x)
            (or (lookup locals-map x #f)
                (begin
                  (set! local-count (+ 1 local-count))
                  ; start at 1, because register 0 holds the globals array
                  (set! locals-map (insert locals-map x local-count))
                  local-count)))))
      (append
        (ir2->bytecode (fix-body expr) (make-locals-map '()))
        (loop for func in (fix-functions expr)
              collect (list 'label (translate-label (closure-name func)))
              append (ir2->bytecode (closure-body func) (make-locals-map (closure-arguments func))))
        (list
          (list 'label 0))))))