;  JAMES YU
;     +
; FRANK LIU
;
;      :)

#lang racket

(struct psymbol (type val) #:mutable)

(define (primp-assemble prog)
  ; LINE COUNTER
  (define lc 0)
  ; BINDINGS
  (define heap (make-hash))
  (define (add-to-heap k v)
    (cond [(hash-has-key? heap k) (error "duplicate")]
          [else (hash-set! heap k v)]))
  (define (lookup key)
    (hash-ref heap key (lambda () (error "undefined"))))
  ; FIRST PASS, GET VARIABLE NAMES
  (define (get-vars p)
    (if (empty? p)
        empty
        (match (first p)
          ; HALT
          [`(halt) (set! lc (add1 lc))
                   (cons 0 (get-vars (rest p)))]
          ; CONST
          [`(const ,psym ,val) (add-to-heap psym (psymbol `imm val))
                               (get-vars (rest p))]
          ; DATA
          [`(data ,psym (,n ,val)) (add-to-heap psym (psymbol `ind lc))
                                   (set! lc (+ n lc))
                                   (append (for/list ([i n])
                                             `(lit ,val))
                                           (get-vars (rest p)))]
          [`(data ,psym ,val ...) (add-to-heap psym (psymbol `ind lc))
                                  (append (for/list ([i val])
                                            (set! lc (add1 lc))
                                            `(lit ,i))
                                          (get-vars (rest p)))]
          ; LABEL
          [`(label ,psym) (add-to-heap psym (psymbol `tgt lc))
                          (get-vars (rest p))]
          ; REGULAR PRIMP
          [x (set! lc (add1 lc))
             (cons x (get-vars (rest p)))])))

  ; EXECUTE PASS 1
  (set! prog (get-vars prog))
  ; RESOLVE NAMES
  (define (dfs-resolve k)
    (define l (lookup k))
    (cond [(symbol? (psymbol-val l)) (define v (psymbol-val l))
                                     (set-psymbol-val! l empty) ; placeholder value
                                     (set-psymbol-val! l (dfs-resolve v))
                                     (if (empty? (psymbol-val l))
                                         (error "circular")
                                         (void))]
          [else (psymbol-val l)]))
  (for ([(key value) heap])
    (dfs-resolve key))
  ; TEST PRINT
  ;(for ([(key value) heap])
  ;å  (printf "~a:(~a,~a)\n" key (psymbol-type value) (psymbol-val value)))

  ; TYPE CHECKERS
  (define (get-lit lit)
    (cond [(number? lit) lit]
          [(boolean? lit) lit]
          [else (psymbol-val (lookup lit))]))
  (define (get-tgt tgt)
    (cond [(number? tgt) tgt]
          [(symbol? tgt) (define l (lookup tgt))
                         (if (symbol=? (psymbol-type l) `tgt)
                             (psymbol-val l)
                             `(,(psymbol-val l)))]
          [else (get-opd tgt)]
          ))
  (define (get-opd opd)
    (match opd
      [(? number? x) x]
      [(? boolean? x) x]
      [`(,imm ,ind) `(,(cond [(number? imm) imm]
                             [else (define l (lookup imm))
                                   (if (symbol=? (psymbol-type l) `tgt)
                                       (error "incorrect")
                                       (psymbol-val l))])
                      ,(match ind
                         [`(,x) (if (number? x)
                                    ind
                                    (error "incorrect"))]
                         [x (define l (lookup x))
                            (if (symbol=? (psymbol-type l) `ind)
                                `(,(psymbol-val l))
                                (error "incorrect"))]))]
      [`(,imm) (if (number? imm)
                 opd
                 (error "incorrect"))]
      [x (define l (lookup x))
         (cond [(symbol=? (psymbol-type l) `imm) (psymbol-val l)]
               [(symbol=? (psymbol-type l) `ind) `(,(psymbol-val l))]
               [else (error "incorrect")])]))
  (define (get-dest dest)
    (cond [(list? dest) (get-opd dest)]
          [else (define l (lookup dest))
                (if (or (symbol=? (psymbol-type l) `imm)
                        (symbol=? (psymbol-type l) `tgt))
                    (error "incorrect")
                    `(,(psymbol-val l)))]))
  ; SUBSTITUTION
  (define (subst p)
    (cond [(empty? p) empty]
          [else (match (first p)
                  [`(lit ,sym) (cons (get-lit sym)
                                     (subst (rest p)))]
                  [`(jump ,tgt) (cons `(jump ,(get-tgt tgt))
                                      (subst (rest p)))]
                  [`(print-val ,opd) (cons `(print-val ,(get-opd opd))
                                           (subst (rest p)))]
                  [`(branch ,opd ,tgt) (cons `(branch ,(get-opd opd) ,(get-tgt tgt))
                                             (subst (rest p)))]
                  [`(,bin ,dest ,opd) (cons `(,bin ,(get-dest dest) ,(get-opd opd))
                                            (subst (rest p)))]
                  [`(,tern ,dest ,opd1 ,opd2) (cons `(,tern ,(get-dest dest) ,(get-opd opd1) ,(get-opd opd2))
                                            (subst (rest p)))]
                  [(? symbol? x) (error "incorrect")]
                  [x (cons x (subst (rest p)))])]))
  (set! prog (subst prog))
  prog)