(import (scheme base)
        (only (csc gensym)
          gensym
          gensym?)
        (only (csc ir1)
          %constant
          %lexical-ref
          %library-ref
          constant?
          lexical-ref?
          library-ref?
          make-call
          make-constant
          make-define-syntax
          make-if
          make-lambda
          make-letrec
          make-lexical-ref
          make-lexical-set
          make-library-ref
          make-sequence)
        (only (csc ir2)
          %apply
          %branch
          %closure
          %fix
          %primitive
          %variable
          apply?
          branch?
          closure?
          fix?
          make-apply
          make-atom
          make-branch
          make-call-closure
          make-closure
          make-fix
          make-kargs
          make-klabel
          make-ktail
          make-primitive
          make-variable
          primitive?
          variable?)
        (only (csc testing)
          assert-equal
          test)
        (csc cps))


(define transform-ir2
  (list
    (cons constant? %constant)
    (cons lexical-ref? %lexical-ref)
    (cons library-ref? %library-ref)
    (cons variable? %variable)
    (cons primitive? %primitive)
    (cons branch? %branch)
    (cons apply? %apply)
    (cons closure? %closure)
    (cons fix? %fix)
    (cons gensym? (lambda (x) 'gensym))))


(define (test-ref name)
  (make-lexical-ref name (gensym)))


(define (tail x)
  (make-apply (test-ref 'tail) (list x)))


(test atom-const
  (assert-equal
    (make-apply (test-ref 'tail) (list (make-constant 5)))
    (ir1->ir2 (make-constant 5) tail)
    transform-ir2))


(test atom-lexical-ref
  (assert-equal
    (make-apply (test-ref 'tail) (list (test-ref 'var)))
    (ir1->ir2 (test-ref 'var) tail)
    transform-ir2))


(test atom-library-ref
  (assert-equal
    (make-primitive 'peek (list (make-library-ref 'var '(csc builtins)) (make-constant 0))
                          (list (test-ref 'generated-symbol))
      (make-apply (test-ref 'tail) (list (test-ref 'generated-symbol))))
    (ir1->ir2 (make-library-ref 'var '(csc builtins)) tail)
    transform-ir2))


(test lexical-set
  (assert-equal
    (make-primitive 'poke (list (make-constant 5) (test-ref 'var) (make-constant 0)) '()
      (make-apply (test-ref 'tail) (list (make-constant #f))))
    (ir1->ir2 (make-lexical-set (test-ref 'var) (make-constant 5))
              tail)
    transform-ir2))


(test no-op-define-syntax
  (assert-equal
    (make-apply (test-ref 'tail) (list (make-constant #f)))
    (ir1->ir2 (make-define-syntax 'name '(transformer))
              tail)
    transform-ir2))


(test branch
  (assert-equal
    (make-fix
      (list
        (make-closure (test-ref 'generated-symbol) (list (test-ref 'generated-symbol)) #f
          (make-apply (test-ref 'tail) (list (test-ref 'generated-symbol)))))
      (make-branch (make-constant #t)
        (make-apply (test-ref 'generated-symbol) (list (make-constant 1)))
        (make-apply (test-ref 'generated-symbol) (list (make-constant 2)))))
    (ir1->ir2 (make-if (make-constant #t)
                (make-constant 1)
                (make-constant 2))
              tail)
    transform-ir2))


(test call-closure
  (assert-equal
    (make-fix
      (list
        (make-closure (test-ref 'generated-symbol) (list (test-ref 'generated-symbol)) #f
          (make-apply (test-ref 'tail) (list (test-ref 'generated-symbol)))))
      (make-apply (test-ref 'f) (list (test-ref 'generated-symbol)
                                                 (make-constant 1)
                                                 (make-constant 2))))
    (ir1->ir2 (make-call (test-ref 'f) (list (make-constant 1) (make-constant 2)))
              tail)
    transform-ir2))


(test sequence
  (assert-equal
    (make-primitive 'poke (list (make-constant 5) (test-ref 'a) (make-constant 0)) '()
      (make-primitive 'poke (list (make-constant 6) (test-ref 'b) (make-constant 0)) '()
        (make-apply (test-ref 'tail) (list (make-constant #f)))))
    (ir1->ir2 (make-sequence (make-lexical-set (test-ref 'a) (make-constant 5))
                             (make-lexical-set (test-ref 'b) (make-constant 6)))
              tail)
    transform-ir2))


(test closure
  (assert-equal
    (make-fix
      (list
        (make-closure (test-ref 'generated-symbol)
                      (list (test-ref 'generated-symbol)
                            (test-ref 'a)
                            (test-ref 'b))
                      (test-ref 'c)
          (make-apply (test-ref 'generated-symbol) (list (make-constant 5)))))
      (make-apply (test-ref 'tail) (list (test-ref 'generated-symbol))))
    (ir1->ir2 (make-lambda
                (list (test-ref 'a) (test-ref 'b))
                (test-ref 'c)
                (make-constant 5))
              tail)
    transform-ir2))


(test letrec-functions
  (define x (test-ref 'x))
  (define f (gensym))
  (assert-equal
    (make-fix
      (list
        (make-closure (test-ref 'f) (list (test-ref 'generated-symbol) (test-ref 'x)) #f
          (make-apply (test-ref 'generated-symbol) (list (test-ref 'x)))))
      (make-fix
        (list
          (make-closure (test-ref 'generated-symbol) (list (test-ref 'generated-symbol)) #f
            (make-apply (test-ref 'tail) (list (test-ref 'generated-symbol)))))
        (make-apply (test-ref 'f) (list (test-ref 'generated-symbol) (make-constant 10)))))
    (ir1->ir2
      (make-letrec #f '(f) (list f)
                   (list (make-lambda (list x) #f x))
        (make-call (make-lexical-ref 'f f) (list (make-constant 10))))
      tail)
    transform-ir2))


(test letrec-in-order
  (assert-equal
    (make-primitive 'alloc (list (make-constant 1)) (list (test-ref 'b))
      (make-primitive 'alloc (list (make-constant 1)) (list (test-ref 'a))
        (make-fix
          (list
            (make-closure (test-ref 'f) (list (test-ref 'generated-symbol) (test-ref 'x)) #f
              (make-apply (test-ref 'generated-symbol) (list (make-constant 5)))))
          (make-primitive 'poke (list (make-constant 1) (test-ref 'a) (make-constant 0)) '()
            (make-primitive 'poke (list (make-constant 2) (test-ref 'b) (make-constant 0)) '()
              (make-apply (test-ref 'tail) (list (make-constant 10))))))))
    (ir1->ir2
      (make-letrec #t
        '(a f b)
        (list (gensym) (gensym) (gensym))
        (list (make-constant 1)
              (make-lambda (list (test-ref 'x)) #f (make-constant 5))
              (make-constant 2))
        (make-constant 10))
      tail)
    transform-ir2))


; What does the following letrec return?
; (letrec* ((f (lambda () x))
;           (x (f)))
;   x)
(test letrec-very-cool
  (define f (gensym))
  (define x (gensym))
  (assert-equal
    (make-primitive 'alloc (list (make-constant 1)) (list (test-ref 'x))
      (make-fix
        (list
          (make-closure (test-ref 'f) (list (test-ref 'generated-symbol)) #f
            (make-primitive 'peek (list (test-ref 'x) (make-constant 0)) (list (test-ref 'generated-symbol))
              (make-apply (test-ref 'generated-symbol) (list (test-ref 'generated-symbol))))))
        (make-fix
          (list
            (make-closure (test-ref 'generated-symbol) (list (test-ref 'generated-symbol)) #f
              (make-primitive 'poke (list (test-ref 'generated-symbol) (test-ref 'x) (make-constant 0)) '()
                (make-primitive 'peek (list (test-ref 'x) (make-constant 0)) (list (test-ref 'generated-symbol))
                  (make-apply (test-ref 'tail) (list (test-ref 'generated-symbol)))))))
          (make-apply (test-ref 'f) (list (test-ref 'generated-symbol))))))
    (ir1->ir2
      (make-letrec #t
        '(f x)
        (list f x)
        (list (make-lambda '() #f (make-lexical-ref 'x x))
              (make-call (make-lexical-ref 'f f) '()))
        (make-lexical-ref 'x x))
      tail)
    transform-ir2))


(test set-argument
  (define test-sym (gensym))
  (assert-equal
    (make-fix
      (list
        (make-closure (test-ref 'f) (list (test-ref 'generated-symbol)
                                          (test-ref 'generated-symbol)) #f
          (make-primitive 'alloc (list (make-constant 1))
                                 (list (test-ref 'x))
            (make-primitive 'poke (list (test-ref 'generated-symbol)
                                        (test-ref 'x)
                                        (make-constant 0)) '()
              (make-primitive 'poke (list (make-constant 10)
                                          (test-ref 'x)
                                          (make-constant 0)) '()
                (make-apply (test-ref 'generated-symbol) (list (make-constant #f))))))))
      (make-apply (test-ref 'tail) (list (make-constant 5))))
    (ir1->ir2 (make-letrec
                #f
                '(f)
                (list (gensym))
                (list (make-lambda (list (make-lexical-ref 'x test-sym)) #f
                        (make-lexical-set (make-lexical-ref 'x test-sym) (make-constant 10))))
                (make-constant 5))
              tail)
    transform-ir2))

(test set-function
  (define test-sym (gensym))
  (assert-equal
    (make-primitive 'alloc (list (make-constant 1)) (list (test-ref 'f))
      (make-fix
        (list
          (make-closure (test-ref 'generated-symbol) (list (test-ref 'generated-symbol)) #f
            (make-apply (test-ref 'generated-symbol) (list (make-constant 10)))))
        (make-primitive 'poke (list (test-ref 'generated-symbol) (test-ref 'f) (make-constant 0)) '()
          (make-primitive 'poke (list (make-constant 5) (test-ref 'f) (make-constant 0)) '()
            (make-apply (test-ref 'tail) (list (make-constant #f)))))))
    (ir1->ir2 (make-letrec
                #f
                '(f)
                (list test-sym)
                (list (make-lambda '() #f (make-constant 10)))
                (make-lexical-set (make-lexical-ref 'f test-sym) (make-constant 5)))
              tail)
    transform-ir2))


(define (test-var)
  (make-variable (gensym)))


(test closure-convert-primitive
  (define a-sym (gensym))
  (define f-sym (gensym))
  (define ret-sym (gensym))
  (define x-sym (gensym))
  (assert-equal
    (make-primitive 'alloc (list (make-constant 1)) (list (test-var))
      (make-fix
        (list
          (make-closure (test-var) (list (test-var) (test-var) (test-var)) #f
            (make-primitive 'peek (list (test-var) (make-constant 0)) (list (test-var))
              (make-primitive 'poke (list (test-var) (test-var) (make-constant 0)) '()
                (make-primitive 'peek (list (test-var) (make-constant 0)) (list (test-var))
                  (make-apply (test-var) (list (test-var) (make-constant #f))))))))
        (make-primitive 'alloc (list (make-constant 2)) (list (test-var))
          (make-primitive 'poke (list (test-var) (test-var) (make-constant 0)) '()
            (make-primitive 'poke (list (test-var) (test-var) (make-constant 1)) '()
              (make-primitive 'peek (list (test-var) (make-constant 0)) (list (test-var))
                (make-apply (test-var) (list (test-var) (make-library-ref 'tail '(csc builtins)) (make-constant 10)))))))))
    (closure-convert (make-primitive 'alloc (list (make-constant 1)) (list (make-lexical-ref 'a a-sym))
                       (make-fix
                         (list
                           (make-closure (make-lexical-ref 'f f-sym) (list (make-lexical-ref 'ret ret-sym) (make-lexical-ref 'x x-sym)) #f
                             (make-primitive 'poke (list (make-lexical-ref 'x x-sym) (make-lexical-ref 'a a-sym) (make-constant 0)) '()
                               (make-apply (make-lexical-ref 'ret ret-sym) (list (make-constant #f))))))
                         (make-apply (make-lexical-ref 'f f-sym) (list (make-library-ref 'tail '(csc builtins)) (make-constant 10))))))
    transform-ir2))