#lang typed/racket
(require "../semantics.rkt")
(require "../ast.rkt")
(require "../printer.rkt")
(require typed/rackunit)

(: exp->program (-> CoreExpr CoreProgram))
(define (exp->program i)
  (list `(define (main) ,i)))

(define-syntax ap-exp
  (syntax-rules ()
    [(ap-exp a) a]
    [(ap-exp a b c ...)
     (ap-exp (list a b) c ...)]))

(: program-final-node (-> CoreProgram Node))
(define (program-final-node prog)
  (let* ([st (last (eval-state (compile-core prog)))]
         [stack (State-stack st)]
         [heap (State-heap st)])
    (if (null? stack)
        (error "impossible: not a final node")
        (lookup-node heap (first stack)))))

(: exp-final-node (-> CoreExpr Node))
(define (exp-final-node e)
  (program-final-node (exp->program e)))

(define rec-program
  '((define (pair x y f) ((f x) y))
    (define (fst p) (p K))
    (define (snd p) (p K1))
    (define (f x y)
      (letrec ((a ((pair x) b)) (b ((pair y) a))) (fst (snd (snd (snd a))))))
    (define (main) ((f 3) 4))))

(check-equal? (exp-final-node (ap-exp 'S 'K 'K 3)) 3)
(check-equal? (exp-final-node (ap-exp 'K 100 99)) 100)
(check-equal? (exp-final-node `(let ((x ,(ap-exp 'S 'K 'K 3)) (y 5)) ,(ap-exp 'K 'x 'x))) 3)
(check-equal? (exp-final-node `(let ((x ,(ap-exp 'S 'K 'K 3)) (y 5)) ,(ap-exp 'K 'y 'x))) 5)
(check-equal? (exp-final-node `(letrec ((y x) (x ,(ap-exp 'S 'K 'K 3))) ,(ap-exp 'K 'y 'x))) 3)
(check-equal? (program-final-node rec-program) 4)