(define-library (csc codegen)
  (export
    ir2->ir3)
  (import (scheme base)
          (only (csc format)
            sprintf)
          (only (csc 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?
            fix-body
            fix-functions
            globals?
            label-gensym
            label?
            library-ref-library
            library-ref-name
            library-ref?
            make-apply
            make-constant
            make-primitive
            variable-gensym
            variable?)
          (only (csc loop)
            loop)
          (only (csc match)
            match))
  (begin


    (define *next-label-id* 0)


    (define (new-label)
      (define id *next-label-id*)
      (set! *next-label-id* (+ 1 *next-label-id*))
      id)


    (define *label-map* (make-map (make-comparer
                                    (lambda (x) (gensym->int (label-gensym x)))
                                    (lambda (x y) (- (gensym->int (label-gensym y)) (gensym->int (label-gensym x)))))))


    (define (translate-label x)
      (define new-id (new-label))
      (set! *label-map* (insert *label-map* x new-id))
      new-id)


    (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"))
            ((integer? x)
              (list 'const x))
            (else (error "Only integer 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-label 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 (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
            for easy-result = (get-satisfying results (lambda (x) (not (lookup in->out x #f))))
            until (empty? results)
            if easy-result
              collect (list 'mov (list 'local easy-result) (list 'local (lookup out->in easy-result)))
              and do (set! results (delete out->in easy-result))
            else if currently-in-temp
              collect (list 'mov (list 'local (lookup in->out currently-in-temp)) (list 'local currently-in-temp))
              and do (set! currently-in-temp #f)
            else
              append (let ((any-result (get-satisfying results (lambda (x) #t))))
                        (set! currently-in-temp any-result)
                        (list
                          (list 'mov (list 'local *temp-reg*) (list 'local any-result))
                          (list 'mov (list 'local any-result) (list 'local (lookup out->in any-result)))))
              and do (set! results (delete out->in any-result))))


    (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 temp1 (new-label))
          (define temp2 (new-label))
          (append
            (list
              (list 'jmpif (a->b atom) temp1))
            (ir2->bytecode false translate-local)
            (list
              (list 'jmp (list 'label temp2))
              (list 'label temp1))
            (ir2->bytecode true translate-local)
            (list
              (list 'label temp2))))
        ((% %apply proc args)
          (define in->out (make-map compare-numbers))
          (define constants
            (loop for arg in args
                  for i from 1
                  if (variable? arg)
                    do (set! in->out (insert in->out (translate-local arg) i))
                  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 arg))
                  else if (label? arg)
                    collect (list 'mov (list 'local i) (list 'label (translate-label 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)))
          (append (chains in->out)
                  constants
                  (list
                    (if (label? proc)
                      (list 'jmp (list 'label (translate-label proc)))
                      (list 'jmp (a->b proc))))))
        (_ (error "Unexpected form in ir2->bytecode expr"))))


    (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 (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)
          (define res (lookup locals-map x #f))
          (if res
            res
            (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))))))))