#lang racket

(define translate (hash `> `gt
                        `>= `ge
                        `< `lt
                        `<= `le
                        `= `equal
                        `and `land
                        `or `lor
                        `not `lnot
                        `+ `add
                        `- `sub
                        `* `mul
                        `div `div
                        `mod `mod))

(define (prefix-symbol prefix symbol)
  (string->symbol (string-append prefix (symbol->string symbol))))

(define (compile-simp prog)

  (define a-primp empty)
  (define functions (make-hash))

  ; GENERAL HELPERS
  (define (add-line line)
    (set! a-primp (cons line a-primp)))

  (define (add-label sym)
    (add-line `(label ,sym)))

  (define (val? exp)
    (or (number? exp) (boolean? exp)))

  ; GET FUNCTION NAMES AND STACK WIDTHS
  (for ([f prog])
    (match f
      [`(fun (,name ,args ...) (vars [,bindings ...] ,stmts ...))

       ; CHECK FOR DUPLICATE NAMES
       (if (hash-has-key? functions name)
           (error "duplicate")
           (void))

       ; CREATE FUNCTION BINDING
       (hash-set! functions name (list (length args) (+ 3 (length args) (length bindings))))

       ; CHECK FOR RETURNS
       (match stmts
             [`(,s ... (return ,aexp)) (void)]
             [_ (error "return " name)])

       ; CHECK FOR DUPLICATE ARGUMENTS
       (if (< (length (remove-duplicates (append args (map first bindings))))
              (length (append args (map first bindings))))
           (error "duplicates")
           (void))

       ]
      [_ (void)]))
  ; BUILD FUNCTIONS
  (for ([f prog])
    (match f
      [`(fun (,name ,args ...) (vars [,bindings ...] ,stmts ...))

       ; CREATE SUBROUTINE LABEL
       (define label (prefix-symbol "FUNCTION_" name))
       (add-label label)

       (define vars (make-hash))
       (define c 3)

       (define (get-var var)
         `(,(hash-ref vars var (lambda () (error "undefined"))) BSP))

       ; 0: RETURN VALUE
       ; 1: RETURN PC
       ; 2: RETURN BSP

       ; CREATE ARGUMENT TABLE
       (for ([arg args])
         (hash-set! vars arg c)
         (set! c (add1 c)))

       ; CREATE VAR BINDINGS
       (for ([binding bindings])
         (hash-set! vars (first binding) c)
         (add-line `(move (,c BSP) ,(second binding)))
         (set! c (add1 c)))

       ; DEBUG FUNCTION POINTER
       ;(add-line `(print-string "FUNCTION BEING CALLED:\n"))
       ;(for ([i c])
       ;  (add-line `(print-val (,i BSP)))
       ;  (add-line `(print-string "\n")))

       ; DO COMPUTATION

       ; COMPILE FUNCTION
       (define (compile-function exp)
         (match exp
           [`(,id ,ids ...)
            ; LOOKUP FUNCTION
            (define info (hash-ref functions id (lambda () (error "undefined" id))))
            (if (= (length ids) (first info))
                (void)
                (error "arguments"))
            ; SET IMPORTANT VALUES
            ; 0: RETURN VALUE
            ; 1: RETURN PC
            (add-line `(move (2 SP) BSP)) ; 2: RETURN BSP
            ; SET FAKE BSP (AT RETURN VALUE)
            (add-line `(move (0 SP) SP))
            ; ADVANCE SP TO PARAMS
            (add-line `(add SP SP 3))
            (define bk 3)
            (for ([i ids])
              (compile-exp i) ; 2+n: nth argument
              (set! bk (add1 bk))
              )
            ; SET NEW SP
            (add-line `(add SP SP ,(- (second info) (+ 3 (first info)))))
            ; SET REAL BSP
            (add-line `(move BSP (,(- 0 (second info)) SP)))
            ; JSR MESS
            (add-line `(jsr (1 BSP) ,(prefix-symbol "FUNCTION_" id)))
            ]
           [x x]))

       ; COMPILE EXPRESSION
       (define (compile-exp exp)
         ;(printf "  compiling expression ~a\n" exp)
         (match exp
           [(? val? exp)
            (add-line `(move (0 SP) ,exp))
            (add-line `(add SP SP 1))]
           [(? symbol? exp)
            (add-line `(move (0 SP) ,(get-var exp)))
            (add-line `(add SP SP 1))]
           [`(,bin ,exp1 ,exp2)
            (if (hash-has-key? translate bin)
                (begin (compile-exp exp1)
                       (compile-exp exp2)
                       (add-line `(,(hash-ref translate bin) (-2 SP) (-2 SP) (-1 SP)))
                       (add-line `(sub SP SP 1)))
                (compile-function exp))]
           [`(,un ,ex)
            (if (hash-has-key? translate un)
                (begin (compile-exp ex)
                       (add-line `(,(hash-ref translate un) (-1 SP) (-1 SP))))
                (compile-function exp))]
           [x (compile-function exp)]))

       ; COMPILE STATEMENT
       (define (compile-stmt stmt)
         ;(printf "compiling statement ~a\n" stmt)
         (match stmt
           ; SEQ
           [`(seq ,stmts ...)
            (for [(s stmts)]
              (compile-stmt s))]
           ; SET
           [`(set ,id ,exp)
            ;(printf "  compiling set expression ~a ~a\n" id exp)
            (compile-exp exp)
            (add-line `(sub SP SP 1))
            (add-line `(move ,(get-var id) (0 SP)))]
           ; IIF
           [`(iif ,exp ,s1 ,s2)
            ;   create label names
            (define label (gensym "IF_"))
            (define true-label (prefix-symbol "TRUE_" label))
            (define end-label (prefix-symbol "END_" label))
            ;   handle check
            (compile-exp exp)
            (add-line `(sub SP SP 1))
            (add-line `(branch (0 SP) ,true-label))
            ;   false case
            (compile-stmt s2)
            (add-line `(jump ,end-label))
            ;   true case
            (add-label true-label)
            (compile-stmt s1)
            (add-label end-label)
            ]
           ; WHILE
           [`(while ,exp ,stmts ...)
            ;   create label names
            (define label (gensym "WHILE_"))
            (define top-label (prefix-symbol "TOP_" label))
            (define cont-label (prefix-symbol "CONT_" label))
            (define end-label (prefix-symbol "END_" label))
            ;   handle check
            (add-label top-label)
            (compile-exp exp)
            (add-line `(sub SP SP 1))
            (add-line `(branch (0 SP) ,cont-label))
            (add-line `(jump ,end-label))
            ;   continue
            (add-label cont-label)
            ;   compile all statements
            (for [(s stmts)]
              (compile-stmt s))

            (add-line `(jump ,top-label))
            (add-label end-label)]
           ; PRINT
           [`(print ,exp)
            (if (string? exp)
                (add-line `(print-string ,exp))
                (begin (compile-exp exp)
                       (add-line `(sub SP SP 1))
                       (add-line `(print-val (0 SP)))))]
           ; RETURN
           [`(return ,exp)
            (compile-exp exp)
            (add-line `(move (0 BSP) (-1 SP)))  ; SET RETURN VAL
            (add-line `(move JMP (1 BSP)))      ; SET JUMP POINT
            (add-line `(add SP BSP 1))          ; RESET SP
            (add-line `(move BSP (1 SP)))       ; RESET BSP
            (add-line `(jump JMP))              ; END FUNCTION
            ]
           ; SKIP (do nothing)
           [`(skip) (void)]
           ; OOPS
           [_ ;(add-line stmt)
              (error "test")
            ]))

       ; PARSE STATEMENTS
       (for ([stmt stmts])
         (compile-stmt stmt))
       ]
      [_ (void)]))

  ; APPLY MAIN IF IT EXISTS

  (append (if (hash-has-key? functions `main)
              `((add SP SP ,(second (hash-ref functions `main)))
                (jsr (1 BSP) FUNCTION_main)
                (halt)
                )
              `((halt)))
          (reverse a-primp)
          `((data JMP 0)
            (data BSP X)
            (data SP X)
            (data X 0)))

  )


#|
(define prog '((fun (test id) (vars [(i 10)] (return (+ 5 4))))
               (fun (main) (vars [(i 1) (j 0) (acc 0)]
                    (while (<= i 1000)
                           (set j 1)
                           (set acc 0)
                           (while (< j i)
                                  (iif (= (mod i j) 0)
                                       (set acc (+ acc j))
                                       (skip))
                                  (set j (+ j 1)))
                           (iif (= acc i)
                                (seq
                                 (print i)
                                 (print "\n"))
                                (skip))
                           (set i (+ i 1))))))
  )
|#
#|
(define prog '((fun (add a1 a2) (vars [] (return (+ a1 a2))))
               (fun (main) (vars [(i 1) (j 2) (k 3)]
                                 (print i)
                                 (print " ")
                                 (print j)
                                 (print " ")
                                 (print k)
                                 (print "\n")
                                 (print (add i j))
                                 (print "\n")
                                 (print i)
                                 (print " ")
                                 (print j)
                                 (print " ")
                                 (print k)
                                 (print "\n")
                                 (print (add j k))
                                 (print "\n")
                                 (print (add (add i j) k)))))


  )|#

#|
(define prog `((fun (fact n) (vars [(r 1)]
                                   (iif (> n 1)
                                        (set r (* n (fact (- n 1))))
                                        (skip))
                                   (return r)))
               (fun (main) (vars [(n 5)]
                                 (print (fact n)))))
  )|#
#|
(define prog `(
               (fun (main) (vars [] (print (fib (fib 8)))))
               (fun (fib n) (vars [(r 0)]
                                  (iif (= n 0)
                                       (set r 0)
                                       (iif (= n 1)
                                            (set r 1)
                                            (set r (+ (fib (- n 1)) (fib (- n 2))))))
                                  (return r)))

               )
  )
|#
#|
(define prog `((fun (main) (vars [(x 10) (y 20)]
                                 (set x (oh x y 10))
                                 (print x)))
               (fun (oh my g k) (vars [(l 1)]
                                    (set my (+ my (+ g k)))
                                    (return my)))
               ))
|#

(define prog '(
               (fun (perfect? n) (vars [(i 1) (acc 0)]
                                       (while (< i n)
                                              (iif (= (mod n i) 0)
                                                   (set acc (+ acc i))
                                                   (skip))
                                              (set i (+ i 1)))
                                       (iif (= n acc)
                                            (set acc n)
                                            (set acc 0))
                                       (return acc))
                                       )
               (fun (main) (vars [(i 1) (max 500)]
                                 (while (< i max)
                                        (iif (> (perfect? i) 0)
                                             (seq (print i)
                                                  (print " is perfect !!\n"))
                                             (skip))
                                        (set i (+ i 1)))
                                 (return 0)))
               )
  )
#|
(define prog '(
               (fun (tadd a1 a2 a3) (vars [(j 0)]
                                          (set a1 (+ a1 j))
                                          (set a2 (+ a2 a3))
                                          (set a1 (+ a1 a2))
                                          (return a1)))
               (fun (add a1 a2) (vars [(j 1)]
                                      (set a1 (+ a1 a2))
                                      (return a1)))
               (fun (main) (vars [(a 1) (b 2) (c 3)]
                                (set a (add a b))
                                (set b (add b c))
                                (set c (add a b))
                                (print c)
                                (print "\n")
                                (print (add (add a b) (tadd a (add a b) c)))
                                (print "\n")
                                (print a)
                                (print b)
                                (print c)
                                (return c)
                                ))
               ))|#

(compile-simp prog)