;  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 (resolve k v key n)
    (cond [(and (symbol? (psymbol-val v))
                (symbol=? key (psymbol-val v)))
           (if (and (symbol? n) (symbol=? n k))
               (error "circular")
               (set-psymbol-val! v n))]
          [else (void)]))
  (for ([(key value) heap])
    (cond [(symbol? (psymbol-val value))
           (define n (psymbol-val (lookup (psymbol-val value))))
           (set-psymbol-val! value n)
           (if (and (symbol? n)
                    (symbol=? key n))
               (error "circular")
               (for ([(k v) heap])
                 (resolve k v key n)))
           ]
          [else (void)]))
  ; 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-imm imm)
    (cond [(number? imm) imm]
          [(boolean? imm) imm]
          [else (define l (lookup imm))
                (if (symbol=? (psymbol-type l) `tgt)
                    (error "incorrect")
                    (psymbol-val l))]))
  (define (get-ind ind)
    (match ind
      [`(,x) (if (number? x)
                 ind
                 (error "incorrect"))]
      [x (define l (lookup ind))
         (if (symbol=? (psymbol-type l) `ind)
             `(,(psymbol-val l))
             (error "incorrect")
             )]))
  (define (get-opd opd)
    (match opd
      [(? number? x) x]
      [(? boolean? x) x]
      [`(,imm) (if (number? imm)
                 opd
                 (error "incorrect"))]
      [`(,imm ,ind) `(,(get-imm imm) ,(get-ind ind))]
      [x (cond [(number? x) x]
               [(boolean? x) x]
               [else (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)))]
                  [x (cons x (subst (rest p)))])]))
  (set! prog (subst prog))
  prog)