#lang plait
(define-type TypeExp
  [numTE]
  [boolTE]
  [arrowTE (arg : TypeExp)
           (result : TypeExp)])

(define-type Exp
  [numE (n : Number)]
  [boolE (b : Boolean)]
  [plusE (left : Exp) (right : Exp)]
  [timesE (left : Exp) (right : Exp)]
  [minusE (left : Exp) (right : Exp)]
  [leqE (left : Exp) (right : Exp)]
  [lamE (var : Symbol) (te : TypeExp) (body : Exp)]
  [appE (fun : Exp) (arg : Exp)]
  [varE (name : Symbol)]
  [ifE (check : Exp) (zero : Exp) (non-zero : Exp)]
  [let1E (var : Symbol) (te : TypeExp) (value : Exp) (body : Exp)]
  [recE (var : Symbol) (te : TypeExp) (value : Exp) (body : Exp)])

(define (sx-ref sx n) (list-ref (s-exp->list sx) n))

(define (parse-te sx)
  (cond
    [(s-exp-symbol? sx)
     (case (s-exp->symbol sx)
       [(num) (numTE)]
       [(bool) (boolTE)])]
    [(s-exp-match? `(ANY -> ANY) sx)
     (arrowTE (parse-te (sx-ref sx 0)) (parse-te (sx-ref sx 2)))]))

(define (parse s)
  (local
      [(define (sx n) (list-ref (s-exp->list s) n))
       (define (px n) (parse (sx n)))
       (define (? pat) (s-exp-match? pat s))
       (define (parse-let)
         (let* ([def (sx 1)]
                [val (parse (sx-ref def 1))]
                [var-type (sx-ref def 0)]
                [var (s-exp->symbol (sx-ref var-type 0))]
                [te  (parse-te (sx-ref var-type 2))]
                [body (px 2)])
           (values var te val body)))]
    (cond
      [(? `true) (boolE #t)]
      [(? `false) (boolE #f)]
      [(? `SYMBOL) (varE (s-exp->symbol s))]
      [(? `NUMBER) (numE (s-exp->number s))]
      [(? `(+ ANY ANY)) (plusE (px 1) (px 2))]
      [(? `(- ANY ANY)) (minusE (px 1) (px 2))]
      [(? `(* ANY ANY)) (timesE (px 1) (px 2))]
      [(? `(<= ANY ANY)) (leqE (px 1) (px 2))]
      [(? `(if ANY ANY ANY))
       (ifE (px 1) (px 2) (px 3))]
      [(? `(rec ([SYMBOL : ANY] ANY) ANY))
       (local [(define-values (var te val body) (parse-let))]
         (recE var te val body))]
      [(? `(let1 ([SYMBOL : ANY] ANY) ANY))
       (local [(define-values (var te val body) (parse-let))]
         (let1E var te val body))]
      [(? `(lam (SYMBOL : ANY) ANY))
       (let* ([def (sx 1)]
              [parts (s-exp->list def)]
              [var (s-exp->symbol (list-ref parts 0))]
              [te (parse-te (list-ref parts 2))]
              [body (px 2)])
         (lamE var te body))]
      [(? `(ANY ANY)) (appE (px 0) (px 1))]
      [else (error 'parse (to-string s))])))


;; Coverage for parser
(test/exn (parse `"strings are not in our language") "parse")
(test (parse `false) (boolE #f))
(test (parse-te `bool) (boolTE))

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

(define (interp-te te)
  (type-case TypeExp te
    [(numTE) (numT)]
    [(boolTE) (boolT)]
    [(arrowTE a b) (arrowT (interp-te a)
                           (interp-te b))]))

(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]))

(test/exn (type-lookup 'x mt-type-env) "not bound")

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

(define (typecheck [exp : Exp] [env : TypeEnv]) : Type
  (local
      [(define (num2 l r type)
         (let ([left-t (typecheck l env)]
               [right-t (typecheck r env)])
           (if (and (equal? (numT) left-t) (equal? (numT) right-t))
               type
               (error 'typecheck "expected 2 num"))))]
    (type-case Exp exp
      [(numE n) (numT)]
      [(boolE b) (boolT)]
      [(plusE l r) (num2 l r (numT))]
      [(minusE l r) (num2 l r (numT))]
      [(timesE l r) (num2 l r (numT))]
      [(leqE l r) (num2 l r (boolT))]
      [(varE s) (type-lookup s env)]
      [(lamE name te body) ....]
      [(appE fn arg) ....]
      [(ifE c t f) ....]
      [(let1E var te val body) ....]
      [(recE var te val body) ....]
)))

(tc : (S-Exp -> Type))
(define (tc s)
  (typecheck (parse s) mt-type-env))

(test (tc `{rec {[fact : (num -> num)]
                {lam {n : num} {if {<= n 0} 1 {* n {fact {- n 1}}}}}}
                {fact 10}})
      (numT))

(test (tc `{lam {n : num} {if {<= n 0} 1 2}}) (arrowT (numT) (numT)))

(test (tc `{let1 {[x : num] 3} {lam {y : num} {+ x y}}}) (arrowT (numT) (numT)))
(test/exn (tc `{lam {n : bool} {if {<= n 0} 1 2}}) "typecheck")

(test/exn (tc `{lam {n : num} {if n 1 2}}) "boolean")

(test/exn (tc `{1 1}) "function")