;; -*- mode: scheme -*-
#lang racket
(require redex)
(require rackunit)

;; Straightforwardish translation from this tutorial:
;; Bidirectional Typing Rules: A Tutorial, David Raymond Christiansen

;; Simply typed lambda-calculus
(define-language L
  (t ::=
     x                                  ; Variables
     v                                  ; Values
     (lambda (x) t)                          ; Lambda abstractions
     (t t)                              ; Application
     (t : τ)                            ; Type annotation
     (if t t t))
  (v ::= (lambda (x) t) true false)
  (x ::= variable-not-otherwise-mentioned)
  (τ ::=
     bool
     (τ -> τ))
  ;; Type environment
  (Γ ::= ((x τ) ...)))

;;; Helpers for type environments
(define-language REDEX)
(define-judgment-form REDEX
  #:mode (lookup I I O)
  #:contract (lookup ((any any) ...) any any)
  [(lookup (_ ... (any any_0) _ ...) any any_0)])
;; Assign or set a type of variable
(define-metafunction REDEX
  ext1 : ((any any) ...) (any any) -> ((any any) ...)
  ;; any_k exists -> assign
  [(ext1 (any_0 ... (any_k any_v0) any_1 ...) (any_k any_v1))
   (any_0 ... (any_k any_v1) any_1 ...)]
  ;; Add new
  [(ext1 (any_0 ...) (any_k any_v1))
   ((any_k any_v1) any_0 ...)])

;;; The typing rules
;; Need two judgment relations because of different modes:
;; - Checking (:-c ... <= ...) takes everything as input
;; - Inferring (:- ... => ...) outputs the type

(define-judgment-form L
  ;; (:-c Γ t <= τ) "t can be checked to have the given type τ in context Γ"
  #:mode (:-c I I I I) #:contract (:-c Γ t any τ)
  [(:-c Γ t_1 <= bool) (:-c Γ t_2 <= τ) (:-c Γ t_3 <= τ)
   --- bt-if
   (:-c Γ (if t_1 t_2 t_3) <= τ)]
  [(:-c (ext1 Γ (x τ_1)) t <= τ_2)
   --- bt-abs
   (:-c Γ (lambda (x) t) <= (τ_1 -> τ_2))]
  [(:- Γ t => τ)
   --- bt-checkinfer
   (:-c Γ t <= τ)]
  )

(define-judgment-form L
  ;; (:- Γ t => τ) "t's type can be inferred to be τ in context Γ"
  #:mode (:- I I I O) #:contract (:- Γ t any τ)
  [(lookup Γ x τ)
   ---- bt-var
   (:- Γ x => τ)]
  [----
   bt-true
   (:- Γ true => bool)]
  [----
   bt-false
   (:- Γ false => bool)]
  [(:-c Γ t <= τ)
   --- bt-ann
   (:- Γ (t : τ) => τ)]
  )

;;; Tests
(check-equal? (judgment-holds (:- ((x bool)) x => any_1) any_1) '(bool))

(check-true (judgment-holds (:-c () (if true false true) <= bool)))
(check-equal? (judgment-holds (:- () ((if true false true) : bool) => any_1) any_1) '(bool))
(check-true (judgment-holds (:-c ((x bool)) (if x false true) <= bool)))

(check-true (judgment-holds (:-c ((x bool)) x <= bool)))
(check-true (judgment-holds (:-c () (lambda (x) x) <= (bool -> bool))))
(check-true (judgment-holds (:-c () (lambda (x) (if x false true)) <= (bool -> bool))))

(check-equal? (judgment-holds
               (:- ()
                   ((lambda (x) (if x false true))
                    : (bool -> bool))
                   => any_1)
               any_1)
              '((bool -> bool)))

;; Show me the derivation tree of the proof:
;; (show-derivations
;;  (build-derivations
;;   (:- ()
;;       ((lambda (x) (if x false true))
;;        : (bool -> bool))
;;       => any_1)))