#lang plaitypus
(print-only-errors #t)

(define-type TIFAE
  [num (n : number)]
  [add (l : TIFAE)
       (r : TIFAE)]
  [sub (l : TIFAE)
       (r : TIFAE)]
  [id (name : symbol)]
  [fun (param-name : symbol)
       (param-te : TE) ; type expression! will need to be turned into a type
       (body : TIFAE)]
  [app (fun-expr : TIFAE)
       (arg-expr : TIFAE)])

(define-type TE
  [numberTE]
  [arrowTE (dom : TE)
           (rng : TE)]
  [guessTE])

(define-type Type
  [numberT]
  [arrowT (dom : Type)
          (rng : Type)]
  [varT (is : (boxof MaybeType))])
(define-type MaybeType
  [none]
  [some (t : Type)])

(define-type TypeEnv
  [mtEnv]
  [aBind (name : symbol)
         (type : Type)
         (rest : TypeEnv)])

;; ----------------------------------------------------------------------

;; Look away... Go back... Do not see! Do not see! The vengeance of the
;; infinite abysses...
;;   - H. P. Lovecraft, The Other Gods
(define parse : (s-expression -> TIFAE)
  (lambda (s-exp)
    (cond
      [(s-exp-number? s-exp)
       (num (s-exp->number s-exp))]
      [(s-exp-symbol? s-exp)
       (id (s-exp->symbol s-exp))]
      [(s-exp-list? s-exp)
       (define as-list (s-exp->list s-exp))
       (cond [(empty? as-list)
              (error 'parse "the empty list is not a valid TIFAE")]
             [(s-exp-symbol? (first as-list))
              (case (s-exp->symbol (first as-list))
                [(+)
                 (check-pieces as-list "add" 3)
                 (add (parse (second as-list))
                      (parse (third as-list)))]
                [(-)
                 (check-pieces as-list "sub" 3)
                 (sub (parse (second as-list))
                      (parse (third as-list)))]
                [(fun)
                 (check-pieces as-list "fun" 3)
                 (unless (s-exp-list? (second as-list))
                   (error 'parse "expected parameter list"))
                 (define param-list (s-exp->list (second as-list)))
                 (check-pieces param-list "parameter list" 3)
                 (unless (s-exp-symbol? (first param-list))
                   (error 'parse "expected symbol as parameter name"))
                 (unless (and (s-exp-symbol? (second param-list))
                              (equal? ': (s-exp->symbol (second param-list))))
                   (error 'parse "expected `:`"))
                 (fun (s-exp->symbol (first param-list))
                      (parse-type (third param-list))
                      (parse (third as-list)))]
                [else (parse-app as-list)])]
             [else (parse-app as-list)])]
      [else
       (error 'parse "wat")])))

(define parse-app : ((listof s-expression) -> TIFAE)
  (lambda (s-exps)
    (check-pieces s-exps "app" 2)
    (app (parse (first  s-exps))
         (parse (second s-exps)))))

(define parse-type : (s-expression -> TE)
  (lambda (s-exp)
    (cond [(and (s-exp-symbol? s-exp)
                (equal? 'number (s-exp->symbol s-exp)))
           (numberTE)]
          [(and (s-exp-symbol? s-exp)
                (equal? '? (s-exp->symbol s-exp)))
           (guessTE)]
          [(s-exp-list? s-exp)
           (define as-list (s-exp->list s-exp))
           (check-pieces as-list "function type" 3)
           (if (and (s-exp-symbol? (second as-list))
                    (equal? '-> (s-exp->symbol (second as-list))))
               (arrowTE (parse-type (first as-list))
                        (parse-type (third as-list)))
               (error 'parse "expected `->`"))])))

(define check-pieces : ((listof s-expression) string number -> void)
  (lambda (s-exps expected n-pieces)
    (unless (= n-pieces (length s-exps))
      ;; With a typechecker, niceties like `error` taking 2 *or more* arguments
      ;; are gone. Hard for a typechecker to prove it's ok!
      (error 'parse
             (string-append (string-append "expected " expected)
                            (string-append " got " (to-string s-exps)))))))

;; ----------------------------------------------------------------------

(define typecheck : (TIFAE TypeEnv -> Type)
  (lambda (a-tifae gamma)
    (type-case TIFAE a-tifae
      [num (n) (numberT)]
      [id (name) (get-type name gamma)]
      [add (l r)
           (define l-type (typecheck l gamma))
           (define r-type (typecheck r gamma))
           (unify! l-type (numberT))
           (unify! r-type (numberT))
           (numberT)]
      [sub (l r)
           (define l-type (typecheck l gamma))
           (define r-type (typecheck r gamma))
           (unify! l-type (numberT))
           (unify! r-type (numberT))
           (numberT)]
      [fun (param-name param-te body)
           (define param-type (instantiate-type param-te))
           (arrowT param-type
                   (typecheck body (aBind param-name
                                          param-type
                                          gamma)))]
      [app (fun arg)
           (define return-type (varT (box (none))))
           (define fun-type (typecheck fun gamma))
           (define arg-type (typecheck arg gamma))
           (unify! fun-type (arrowT arg-type return-type))
           return-type])))

(define unify! : (Type Type -> void)
  (lambda (t1 t2)
    (type-case Type t1
      [numberT ()
               (type-case Type t2
                 [numberT () (void)]
                 [arrowT (d2 r2) (error 'typecheck "type mismatch")]
                 [varT (b2) (unify! t2 t1)])]
      [arrowT (d1 r1)
              (type-case Type t2
                [numberT () (error 'typecheck "type mismatch")]
                [arrowT (d2 r2)
                        (unify! d1 d2)
                        (unify! r1 r2)]
                [varT (b2) (unify! t2 t1)])]
      [varT (b1)
            (type-case MaybeType (unbox b1)
              [some (t1-2)
                    (unify! t1-2 t2)]
              [none ()
                    (type-case Type t2
                      [varT (b2)
                            (type-case MaybeType (unbox b2)
                              [none ()
                                    (if (eq? b1 b2)
                                        (void)
                                        (set-box! b2 (some t1)))]
                              [some (t2-2)
                                    (unify! t1 t2-2)])]
                      [else
                       (if (occurs? b1 t2)
                           (error 'typecheck "occurs check failed")
                           (set-box! b1 (some t2)))])])])))

(define occurs? : ((boxof MaybeType) Type -> boolean)
  (lambda (b1 t2)
    (type-case Type t2
      [numberT () #f]
      [arrowT (d2 r2) (or (occurs? b1 d2)
                          (occurs? b1 r2))]
      [varT (b2) (if (eq? b1 b2)
                     #t
                     (type-case MaybeType (unbox b2)
                       [none () #f]
                       [some (t2-2) (occurs? b1 t2-2)]))])))

(define instantiate-type : (TE -> Type)
  (lambda (te)
    (type-case TE te
      [numberTE () (numberT)]
      [arrowTE (d r) (arrowT (instantiate-type d)
                             (instantiate-type r))]
      [guessTE () (varT (box (none)))])))

(define get-type : (symbol TypeEnv -> Type)
  (lambda (name gamma)
    (type-case TypeEnv gamma
      [mtEnv ()
             (error 'typecheck "free identifier")]
      [aBind (name2 type rest)
             (if (equal? name name2)
                 type
                 (get-type name rest))])))

;; ----------------------------------------------------------------------

;; helper for tests, now that we have type variables and boxes in the way

(define unwrap-type : (Type -> Type)
  (lambda (t)
    (type-case Type t
      [numberT () t]
      [arrowT (d r) (arrowT (unwrap-type d) (unwrap-type r))]
      [varT (b) (type-case MaybeType (unbox b)
                  [none () (error 'unwrap-type "wat")]
                  [some (t) (unwrap-type t)])])))

(test (unwrap-type (typecheck (parse `{fun {x : ?} {+ x 5}})
                              (mtEnv)))
      (instantiate-type (parse-type `(number -> number))))
(test (unwrap-type (typecheck (parse `{{fun {x : ?} {+ x 5}} 5})
                              (mtEnv)))
      (numberT))
(test (unwrap-type (typecheck (parse `{{fun {f : ?} {f 5}}
                                       {fun {x : ?} x}})
                              (mtEnv)))
      (numberT))
(test/exn (unwrap-type (typecheck (parse `{{fun {f : ?} {f 5}}
                                           4})
                                  (mtEnv)))
          "type mismatch")


;; ----------------------------------------------------------------------

(test (unwrap-type
       (typecheck (parse `5)
                  (mtEnv)))
      (numberT))

(test (unwrap-type
       (typecheck (parse `{+ 2 3})
                  (mtEnv)))
      (numberT))

(test (unwrap-type
       (typecheck (parse `{fun {x : number} {+ x 5}})
                  (mtEnv)))
      (arrowT (numberT) (numberT)))

(test (unwrap-type
       (typecheck (parse `{fun {f : (number -> number)}
                               {fun {x : number} {f x}}})
                  (mtEnv)))
      ;(arrowT (arrowT (numberT) (numberT))
      ;        (arrowT (numberT) (numberT)))
      (instantiate-type
       (parse-type `((number -> number)
                     -> (number -> number)))))
(test/exn (unwrap-type
           (typecheck (parse `{{fun {f : (number -> number)}
                                    {fun {x : number} {f x}}}
                               3})
                      (mtEnv)))
          "type mismatch")
(test (unwrap-type
       (typecheck (parse `{{fun {f : (number -> number)}
                                {fun {x : number} {f x}}}
                           {fun {x : number} {+ x 5}}})
                  (mtEnv)))
      (instantiate-type
       (parse-type `(number -> number))))
(test (unwrap-type
       (typecheck (parse `{{{fun {f : (number -> number)}
                                 {fun {x : number} {f x}}}
                            {fun {x : number} {+ x 5}}}
                           5})
                  (mtEnv)))
      (numberT))

(test (unwrap-type
       (typecheck (parse `5)
                  (mtEnv)))
      (numberT))
(test (unwrap-type
       (typecheck (parse `{+ 4 5})
                  (mtEnv)))
      (numberT))
(test (unwrap-type
       (typecheck (parse `{- 4 5})
                  (mtEnv)))
      (numberT))
(test/exn (unwrap-type
           (typecheck (parse `{+ 4 {fun {x : number} x}})
                      (mtEnv)))
          "type mismatch")

(test (unwrap-type
       (typecheck (parse `{fun {x : number} x})
                  (mtEnv)))
      (instantiate-type
       (parse-type `(number -> number))))
(test (unwrap-type
       (typecheck (parse `{fun {x : number} {+ x 5}})
                  (mtEnv)))
      (instantiate-type
       (parse-type `(number -> number))))
(test (unwrap-type
       (typecheck (parse `{fun {f : (number -> number)}
                               {f 5}})
                  (mtEnv)))
      (instantiate-type
       (parse-type `((number -> number) -> number))))

(test (unwrap-type
       (typecheck (parse `{{fun {x : number} x}
                           5})
                  (mtEnv)))
      (numberT))
(test (unwrap-type
       (typecheck (parse `{{fun {f : (number -> number)}
                                {f 5}}
                           {fun {x : number} x}})
                  (mtEnv)))
      (numberT))
(test/exn (unwrap-type
           (typecheck (parse `{4 5})
                      (mtEnv)))
          "type mismatch")
(test/exn (unwrap-type
           (typecheck (parse `{{fun {f : (number -> number)}
                                    {f 5}}
                               4})
                      (mtEnv)))
          "type mismatch")

(test (unwrap-type
       (typecheck (parse `{fun {x : ?} {+ x 3}})
                  (mtEnv)))
      (instantiate-type
       (parse-type `(number -> number))))

(test (unwrap-type
       (typecheck (parse `{{fun {x : ?} x}
                           5})
                  (mtEnv)))
      (numberT))
(test (unwrap-type
       (typecheck (parse `{{fun {f : ?} {f 5}}
                           {fun {x : ?} x}})
                  (mtEnv)))
      (numberT))
(test/exn (unwrap-type
           (typecheck (parse `{{fun {f : ?} {f 5}}
                               4})
                      (mtEnv)))
          "type mismatch")


;; Sadly, we expect this one to fail: the type variables come from different
;; places. But really, we just "renamed" them consistently... Binding structure
;; is the same!
;; What we really want here is not equality between types, but rather
;; α-equivalence. Take our advanced PL classes to learn what that is!
(test (unwrap-type
       (typecheck (parse `{fun {f : ?} {f 3}})
                  (mtEnv)))
      (instantiate-type
       (parse-type `(number -> ?))))