#lang plait
;; Question 1
(define-type TIFLANG
  [numE  (val : Number)]
  [plusE  (l : TIFLANG) (r : TIFLANG)]
  [divE  (l : TIFLANG) (r : TIFLANG)]
  [varE   (name : Symbol)]
  [let1E (id : Symbol) (named-expr : TIFLANG) (bound-body : TIFLANG)]
  [lamE  (param : Symbol) (param-type : TE) (body : TIFLANG)]
  [iLamE  (param : Symbol) (zero-body : TIFLANG) (non-zero-body : TIFLANG)]
  [appE (fun : TIFLANG) (val : TIFLANG)]) ; first type!

(define-type TE
  [numTE]
  [arrowTE (arg : TE)
           (result : TE)])

(define-type Type
  [numT]
  [arrowT (arg : Type)
          (result : Type)])


(define (expect2num l r env)
  (let ([left-t (typecheck l env)]
        [right-t (typecheck r env)])
    (if (and (equal? (numT) left-t) (equal? (numT) right-t))
        (numT)
        (error 'typecheck "expected 2 num"))))

(define-type-alias TypeEnv (Hashof Symbol Type))

(define mt-type-env (hash empty)) ;; "empty type environment"
(define (type-lookup (s : Symbol) (n : TypeEnv))
  (type-case (Optionof Type) (hash-ref n s)
    [(none) (error s "not bound")]
    [(some b) b]))

(define (type-extend (env : TypeEnv) (s : Symbol) (t : Type))
  (hash-set env s t))

(define (type-error iflang msg)
  (error 'typecheck (string-append
                     "no type: "
                     (string-append
                      (to-string iflang)
                      (string-append " not "
                                     msg)))))

(define (type-to-string [type : Type])
  (type-case Type type
    [(numT) "num"]
    [else (to-string type)]))

(define (parse-type te)
  (type-case TE te
    [(numTE) (numT)]
    [(arrowTE a b) (arrowT (parse-type a)
                           (parse-type b))]))
(define (typecheck [iflang : TIFLANG] [env : TypeEnv]) : Type
  (type-case TIFLANG iflang
    [(numE n) (numT)]
    [(plusE l r) (expect2num  l r env)]
    [(divE l r) (expect2num  l r env)]
    [(varE s) (type-lookup s env)]
    [(let1E name expr body)
     (let ([new-env (type-extend env name (typecheck expr env))])
       (typecheck body new-env))]
    [(iLamE name zero-body non-zero-body) ....]
    [(lamE name te body)
     (let* ([arg-type (parse-type te)]
            [body-type (typecheck body (type-extend env name arg-type))])
       (arrowT arg-type body-type))]
    [(appE fn arg)
     (type-case Type (typecheck fn env)
       [(arrowT arg-type result-type)
        (let ([given-type (typecheck arg env)])
          (if (equal? given-type arg-type) result-type
              (type-error arg "type mismatch")))]
       [else (type-error fn "function")])]
    ))

(module+ test
  (define (check iflang)
    (typecheck iflang mt-type-env))

  (test (check
         (appE
          (appE
           (lamE 'x (arrowTE (numTE) (arrowTE (numTE) (numTE))) (appE (varE 'x) (numE 1)))
           (lamE 'x (numTE) (iLamE 'y (numE -1) (plusE (varE 'x) (varE 'y)))))
          (numE 123)))
        (numT))

  (test (check
         (let1E 'f (iLamE 'x (varE 'x) (divE (numE 1) (varE 'x)))
                (plusE (appE (varE 'f) (numE 3)) (appE (varE 'f) (numE 0)))))
        (numT))

  (test (check
         (let1E 'x (numE 3)
                (let1E 'f (iLamE 'y (varE 'x) (divE (varE 'x) (varE 'y)))
                       (let1E 'x (numE 5)
                              (plusE (appE (varE 'f)
                                           (numE 3)) (appE (varE 'f) (numE 0)))))))
        (numT))

  (test/exn
   (check (iLamE 'x (numE 0) (lamE 'x (numTE) (varE 'x)))) "no type")

  (test/exn
   (check
    (appE (iLamE 'x (numE 0) (lamE 'x (numTE) (varE 'x))) (numE 0)))
   "no type")

  (test/exn
   (check
    (appE (iLamE 'x (numE 0) (numE 1)) (lamE 'x (numTE) (varE 'x))))
   "no type")
  )

;; Question 2
(define-type-alias BoxList (Listof (Boxof (Optionof Number))))

(define (unify-lists! l1 l2)
  (cond
    [(and (empty? l1) (empty? l2)) (void)]
    [(or (empty? l1) (empty? l2)) (error 'unify-lists! "length mismatch")]
    [else ....]))

(module+ test
  (let* ([l1 (list (box (some 0)) (box (some 1)) (box (none)))]
         [l2 (list (box (some 0)) (box (none)) (box (some 2)))]
         [result (list (box (some 0)) (box (some 1)) (box (some 2)))])
    (begin
      (unify-lists! l1 l2)
      (test l1 result)
      (test l2 result)))
  (test/exn (unify-lists! (list (box (some 1)) (box (none)))
                          (list (box (none)))) "length mismatch")
  (test/exn (unify-lists! (list (box (some 1)))
                          (list (box (some 2)))) "value mismatch"))