(module static-semantics mzscheme
  (require (planet "environment.ss" ("cobbe" "environment.plt" 3 0))
           (lib "match.ss")
           (lib "etc.ss")
           "core-language.ss")

  (define (compatible? t1 t2)
    (match (list t1 t2)
      ['(bool bool) #t]
      ['(int int) #t]
      [`(dynamic ,_) #t]
      [`(,_ dynamic) #t]
      [`((,s1 -> ,s2) (,t1 -> ,t2))
        (and (compatible? t1 s1)
             (compatible? s2 t2))]
      [else #f]))

  ;; type-of : intermediate-language * (env symbol type) -> type
  (define (type-of term env)
    (match term
      [(? symbol?) (lookup env term (lambda (id)
                                      (error 'lookup "type-check: ~a is not bound" id)))]
      [(? number?) 'int]
      [(? boolean?) 'bool]
      [`(lambda (,x : ,t) ,body)
       `(,t -> ,(type-of body (extend-env (list x) (list t) env)))]
      [`(let (,x : ,t = ,e1) ,e2)
       (type-of e2 (extend-env (list x) (list t) env))]
      [`(letrec (,bindings ...) ,e)
       (type-of e (extend-env (map binding-variable bindings)
                              (map binding-type bindings)
                              env))]
      [`(= ,e1 ,e2)
       'bool]
      [`(,(or '+ '-) ,e1 ,e2)
       'int]
      [`(if ,e1 ,e2 ,e3)
       (type-of e2 env)]
      [`(cast (,t <= ,c) ,e)
       t]
      [`(cast ,t ,e)
       t]
      [`(,e1 ,e2)
       (let ([type-of-e1 (type-of e1 env)])
         (if (eq? type-of-e1 'dynamic)
             'dynamic
             (function-type-range type-of-e1)))]
      ))

  ;; insert-casts : (type * type -> coercion) -> source-language-term * (env symbol type) -> intermediate-language-term
  (define (insert-casts coerce)
    (rec (insert-casts term env)
      ;; insert-casts/binding : source-language-binding * (env symbol type) -> intermediate-language-binding
      (define (insert-casts/binding binding env)
        (match binding
          [`(,x : ,t = ,e)
           (let ([e-with-casts (insert-casts e env)])
             (let ([type-of-e (type-of e-with-casts env)])
               (unless (compatible? type-of-e t)
                 (error 'type-check "letrec binding does not match declared type: ~a" e-with-casts))
               `(,x : ,t = ,(if (equal? type-of-e t)
                                e-with-casts
                                `(cast ,(coerce type-of-e t) ,e-with-casts)))))]
          [`(,x = ,e)
           (let ([e-with-casts (insert-casts e env)])
             (let ([type-of-e (type-of e-with-casts env)])
               `(,x : dynamic = ,(if (eq? type-of-e 'dynamic)
                                     e-with-casts
                                     `(cast ,(coerce type-of-e 'dynamic) ,e-with-casts)))))]))
      (match term
        [(? symbol?) term]
        [(? number?) term]
        [(? boolean?) term]
        [`(lambda (,x) ,body)
         `(lambda (,x : dynamic) ,(insert-casts body (extend-env (list x) (list 'dynamic) env)))]
        [`(lambda (,x : ,t) ,body)
         `(lambda (,x : ,t) ,(insert-casts body (extend-env (list x) (list t) env)))]
        [`(let (,x = ,e1) ,e2)
         (let ([e1-with-casts (insert-casts e1 env)])
           (let ([type-of-e1 (type-of e1-with-casts env)])
             `(let (,x : dynamic = ,(if (eq? type-of-e1 'dynamic)
                                        e1-with-casts
                                        `(cast ,(coerce type-of-e1 'dynamic) ,e1-with-casts)))
                ,(insert-casts e2 (extend-env (list x) (list 'dynamic) env)))))]
        [`(let (,x : ,t = ,e1) ,e2)
         (let ([e1-with-casts (insert-casts e1 env)])
           (let ([type-of-e1 (type-of e1 env)])
             (unless (compatible? type-of-e1 t)
               (error 'type-check "let binding does not match declared type: ~a" e1-with-casts))
             `(let (,x : ,t = ,(if (equal? type-of-e1 t)
                                   e1-with-casts
                                   `(cast ,(coerce type-of-e1 t) ,e1-with-casts)))
                ,(insert-casts e2 (extend-env (list x) (list t) env)))))]
        [`(letrec (,bindings ...) ,e)
         (let ([extended-env (extend-env (map binding-variable bindings)
                                         (map binding-type bindings)
                                         env)])
           (let ([bindings-with-casts (map (lambda (binding)
                                             (insert-casts/binding binding extended-env))
                                           bindings)])
             `(letrec ,bindings-with-casts ,(insert-casts e extended-env))))]
        [`(,(and op (? operator?)) ,e1 ,e2)
         (let ([e1-with-casts (insert-casts e1 env)])
           (let ([type-of-e1 (type-of e1-with-casts env)])
             (unless (compatible? type-of-e1 'int)
               (error 'type-check "not an integer: ~a" type-of-e1))
             (let ([e2-with-casts (insert-casts e2 env)])
               (let ([type-of-e2 (type-of e2-with-casts env)])
                 (unless (compatible? type-of-e2 'int)
                   (error 'type-check "not an integer: ~a" type-of-e2))
                 (let ([e1 (if (eq? type-of-e1 'int)
                               e1-with-casts
                               `(cast ,(coerce type-of-e1 'int) ,e1-with-casts))]
                       [e2 (if (eq? type-of-e2 'int)
                               e2-with-casts
                               `(cast ,(coerce type-of-e2 'int) ,e2-with-casts))])
                   `(,op ,e1 ,e2))))))]
        [`(if ,e1 ,e2 ,e3)
         (let ([e1-with-casts (insert-casts e1 env)])
           (let ([type-of-e1 (type-of e1 env)])
             (unless (compatible? type-of-e1 'bool)
               (error 'type-check "not a boolean: ~a" e1-with-casts))
             (let ([e2-with-casts (insert-casts e2 env)]
                   [e3-with-casts (insert-casts e3 env)])
               (let ([type-of-e2 (type-of e2-with-casts env)]
                     [type-of-e3 (type-of e3-with-casts env)])
                 (unless (and (compatible? type-of-e2 type-of-e3)
                              (compatible? type-of-e3 type-of-e2))
                   (error 'type-check "incompatible if clauses: ~a ~a" e2-with-casts e3-with-casts))
                 `(if ,(if (eq? type-of-e1 'bool)
                           e1-with-casts
                           `(cast ,(coerce type-of-e1 'bool) ,e1-with-casts))
                      ,e2-with-casts
                      ,e3-with-casts)))))]
        [`(,e1 ,e2)
         (let ([e1-with-casts (insert-casts e1 env)]
               [e2-with-casts (insert-casts e2 env)])
           (let ([type-of-e1 (type-of e1-with-casts env)]
                 [type-of-e2 (type-of e2-with-casts env)])
             (cond
               [(eq? type-of-e1 'dynamic)
                `((cast ,(coerce 'dynamic `(,type-of-e2 -> dynamic)) ,e1-with-casts) ,e2-with-casts)]
               [(and (function-type? type-of-e1)
                     (equal? (function-type-domain type-of-e1) type-of-e2))
                `(,e1-with-casts ,e2-with-casts)]
               [(and (function-type? type-of-e1)
                     (not (equal? (function-type-domain type-of-e1) type-of-e2))
                     (compatible? type-of-e2 (function-type-domain type-of-e1)))
                `(,e1-with-casts (cast ,(coerce type-of-e2 (function-type-domain type-of-e1)) ,e2-with-casts))]
               [else
                (error 'type-check "bad application: ~a" term)])))])))

  (provide (all-defined)))