(library (Compiler verify-scheme)
  (export verify-scheme)
  (import (chezscheme)
    (Framework helpers)
    (Framework match))

  ;; Verify that the program is valid. We try to catch as many errors
  ;; as possible in this pass, so that we do not encounter errors from
  ;; the middle of the compiler.

  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;; Assignment 3 updates                                       ;;
  ;;                                                            ;;
  ;; implementing env-passing-logic from a2sol for Body         ;;
  ;;                                                            ;;
  ;; Still need :: Bools and other values (and debugging)       ;;
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  (define-who verify-scheme
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    ;; Environment checkers
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    (define verify-label-list
      (lambda (label*)
        (let loop ([label* label*] [idx* '()])
          (unless (null? label*)
            (let ([label (car label*)] [label* (cdr label*)])
              (unless (label? label)
                (error who "invalid label ~s found" label))
              (let ([idx (extract-suffix label)])
                (when (member idx idx*)
                  (error who "non-unique label suffix ~s found" idx))
                (loop label* (cons idx idx*))))))))
    (define verify-uvar-list
      (lambda (uvar*)
        (let loop ([uvar* uvar*] [idx* '()])
          (unless (null? uvar*)
            (let ([uvar (car uvar*)] [uvar* (cdr uvar*)])
              (unless (uvar? uvar)
                (error who "invalid uvar ~s found" uvar))
              (let ([idx (extract-suffix uvar)])
                (when (member idx idx*)
                  (error who "non-unique uvar suffix ~s found" idx))
                (loop uvar* (cons idx idx*))))))))
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    ;; SIDE-effects and Predicates
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    (define Var
      (lambda (var)
        (unless (or (register? var) (frame-var? var) (uvar? var))
          (error who "invalid variable ~s" var))
        var)) ;; this returns var because of Effect logic

    (define Loc
      (lambda (l)
        (unless (or (register? l) (frame-var? l))
          (error who "invalid location variable ~s" l))))

    (define verify-binds
      (lambda (binds)
        (match binds
          (((,uvar* ,[Loc -> loc*]) ...)
           (verify-uvar-list uvar*))
          (,x (error who "invalid uvar binding ~s" x)))))

    (define Bool?
      (lambda (b)
        (or (eqv? b 'true)
            (eqv? b 'false))))

    (define Relop
      (lambda (r)
        (memq r (list '< '<= '= '>= '> '=))))
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    ;; BNF constructs
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    (define Body
      (lambda (label*)
        (lambda (body)
          (match body
            ((locate (,ubinds* ...) ,tail)
             (verify-binds ubinds*)
             ((Tail label* ubinds*) tail))
            (,x (error who "invalid body ~s" x))))))

    (define Tail
      (lambda (label* ubinds*)
        (lambda (tail)
          (match tail
            ((begin ,[(Effect label* ubinds*) -> ef*] ... ,tail) ((Tail label* ubinds*) tail))
            ((if ,pred ,[tail1] ,[tail2])
             ((Pred label* ubinds*) pred))
            ((,[(Triv label* ubinds*) -> t])
             (when (integer? t)
               (error who "~s violates machine constraints" `(,t))))
            (,tail (error who "invalid Tail ~s" tail))))))

    (define Pred
      (lambda (label* ubinds*)
        (lambda (pred)
          (match pred
            ((,bool) (unless (Bool? bool) (error who "invalid boolean variable ~s" bool)))
            ((,r ,[(Triv label* ubinds*) -> triv1] ,[(Triv label* ubinds*) -> triv2])
             (unless (and (Relop r)
                          (or (and (register? triv1)
                                 (or (register? triv2)
                                     (frame-var? triv2)
                                     (int32? triv2)
                                     (label? triv2)))
                            (and (frame-var? triv1)
                                 (or (register? triv2)
                                     (int32? triv2)))))
               (error who "invalid relational expression ~s" r)))
            ((if ,[p1] ,[p2] ,p3)
             ((Pred label* ubinds*) p3))
            ((begin ,ef* ... ,[p])
             (for-each (Effect label* ubinds*) ef*))
            (,pred (error who "invalid Pred ~s" pred))))))

    (define Effect
      (lambda (label* ubinds*)
        (lambda (ef)
          (match ef
            ((set! ,[(Triv label* ubinds*) -> x]
                   (,binop ,[(Triv label* ubinds*) -> y] ,[(Triv label* ubinds*) -> z]))
             (unless (and (eq? y x)
                          (case binop
                                [(+ - logand logor)
                                 (or (and (register? x)
                                          (or (register? z)
                                              (frame-var? z)
                                              (int32? z)))
                                     (and (frame-var? x)
                                          (or (register? z)
                                              (int32? z))))]
                                [(*)
                                 (and (register? x)
                                      (or (register? z)
                                          (frame-var? z)
                                          (int32? z)))]
                                [(sra)
                                 (and (or (register? x) (frame-var? x))
                                      (uint6? z))]
                                [else (error who "invalid binary operator ~s" binop)]))
               (error who "~s violates machine constraints" `(set! ,x (,binop ,y ,z)))))
            ((set! ,[(Triv label* ubinds*) -> x] ,[(Triv label* ubinds*) -> y])
             (unless (or (and (register? x)
                              (or (register? y)
                                  (frame-var? y)
                                  (int64? y)
                                  (label? y)))
                         (and (frame-var? x)
                              (or (register? y)
                                  (int32? y))))
                     (error who "~s violates machine constraints" `(set! ,x ,y))))
            ((nop) #t)
            ((if ,pred ,[e1] ,[e2])
             ((Pred label* ubinds*) pred))
            ((begin ,[ef*] ... ,ef)
             ((Effect label* ubinds*) ef))
            (,ef (error who "invalid Effect ~s" ef))))))

    (define Triv
      (lambda (label* ubinds*)
        (lambda (t)
          (unless (or (register? t) (frame-var? t) (label? t) (uvar? t)
                      (and (integer? t) (exact? t)))
            (error who "invalid Triv ~s" t))
          (when (or (label? t) (uvar? t))
            (if (label? t)
                (unless (memq t label*)
                  (error who "unbound label ~s" t))
                (unless (assq t ubinds*)
                  (error who "unbound uvar ~s" t))))
          (if (uvar? t)
              (cadr (assv t ubinds*))
              t))))

    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    ;; primary
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
    (lambda (x)
      (match x
        ((letrec ((,label* (lambda () ,body*)) ...) ,body)
         (verify-label-list label*)
         (for-each (Body label*) body*)
         ((Body label*) body))
        [,x (error who "invalid Program. Expected letrec ~s" x)])
      x))
  )