(define-library (csc codegen)
  (export
    ir2->ir3)
  (import (scheme base)
          (only (csc format)
            sprintf)
          (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
            %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))
              (list 'const x))
            (else (error "Only small ints and bool constants are supported for now"))))
        ((% %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* 127)


    (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))))))
        (_ (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 0)
      (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
        (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 'init))
        (ir2->bytecode (fix-body expr) (make-locals-map '()))))))