tychk.rkt

#lang racket

;; Class 1 -- Type Checking
;;
;; In this exercise we will check fully-annotated terms
;; In the next project we will look at synthesis

;; We will build a type system based on the
;; Simply-Typed λ-calculus (STLC)

(provide (all-defined-out))

;;
;; Type Checking
;;

;; Primitive literals
(define bool-lit? boolean?)
(define int-lit? integer?)

;; Expressions are ifarith, with several special builtins
(define (expr? e)
  (match e
    ;; Variables
    [(? symbol? x) #t]
    ;; Literals
    [(? bool-lit? b) #t]
    [(? int-lit? i) #t]
    ;; Applications
    [`(,(? expr? e0) ,(? expr? e1)) #t]
    ;; Annotated expressions
    [`(,(? expr? e) : ,(? type? t)) #t]
    ;; Anotated lambdas
    [`(lambda (,(? symbol? x) : ,(? type? t)) ,(? expr? e)) #t]))

;; Types for STLC
(define (type? t)
  (match t
    ;; Base types: int and bool
    ['int #t]
    ['bool #t]
    ;; Arrow types: t0 -> t1
    [`(,(? type? t0) -> ,(? type? t1)) #t]
    [_ #f]))

;; Synthesize a type for e in the environment env
;; Returns a type
;;
;; Note: synthesis in this setting is *syntax directed*, because
;; the type of the lambda's argument is forced to be annotated.
;; Annotations on arguments are the primary source of nondeterminism
;; when we think about reconstructing a type's proof
(define (synthesize-type env e)
  (match e
    ;; Literals
    [(? integer? i) 'int]
    [(? boolean? b) 'bool]
    ;; Look up a type variable in an environment
    [(? symbol? x) (hash-ref env x)]
    ;; Lambda w/ annotation
    [`(lambda (,x : ,A) ,e)
     `(,A -> ,(synthesize-type (hash-set env x A) e))]
    ;; Arbitrary expression
    [`(,e : ,t) (let ([e-t (synthesize-type env e)])
                  (if (equal? e-t t)
                    t
                    (error (format "types ~a and ~a are different" e-t t))))]
    ;; Application
    [`(,e1 ,e2)
     (match (synthesize-type env e1)
       [`(,A -> ,B)
        (let ([t-2 (synthesize-type env e2)])
          (if (equal? t-2 A)
            B
            (error (format "types ~a and ~a are different" A t-2))))])]))

;; To synthesize an STLC type, synthesize a type with a starting environment
(define (synthesize-stlc-type e)
  (synthesize-type (hash '+ '(int -> (int -> int))
                         'is-zero? '(int -> bool))
                   e))

;; Tests 1: public-synthesize
#;(synthesize-stlc-type '((lambda (x : int) x) 23))
#;(synthesize-stlc-type '(lambda (x : int) (lambda (y : bool) (lambda (z : int) x))))
#;(synthesize-stlc-type '(lambda (x : (int -> int)) (lambda (y : bool) (lambda (z : int) (x (13 : int))))))

;; To type check: synthesize then match

;; Check that, in an environment env, e's type is t
;; Returns a bool
(define (check-stlc-type e t)
  (equal? (synthesize-stlc-type e) t))

;; Tests 2: public-typecheck
#;(check-stlc-type '((lambda (x : int) x) 23) 'int)
#;(check-stlc-type '(lambda (x : int) (lambda (y : bool) (lambda (z : int) x)))
                 '(int -> (bool -> (int -> int))))
#;(check-stlc-type
 '(lambda (x : (int -> int)) (lambda (y : bool) (lambda (z : int) (x (13 : int)))))
 '((int -> int) -> (bool -> (int -> int))))

;; Erase types: transform to Racket code
(define (erase-types e)
  (match e
    [`((+ ,e0) ,e1) `(+ ,(erase-types e0) ,(erase-types e1))]
    [`(lambda (,x : ,t) ,e) `(lambda (,x) ,(erase-types e))]
    [(? symbol? x) x]
    [(? bool-lit? b) b]
    [(? int-lit? i) i]
    [`(,e : ,t) (erase-types e)]
    [`(,e0 ,e1) `(,(erase-types e0) ,(erase-types e1))]))

;; returns either 'type-error or the expression
;; Converting STLC to Racket means first (a) run the typechecker and
;; (b) then erase the types
(define (stlc->racket e)
  (with-handlers ([exn:fail? (lambda (exn) 'type-error)])
    (synthesize-stlc-type e)
    (erase-types e)))

;; Tests 3: type erasure

#;(stlc->racket '(lambda (x : (int -> int)) (lambda (y : bool) (lambda (z : int) (x y)))))


;;
;; Constraint-Based Type Inference
;;

;; Full expresssions also include unannotated lambdas
(define (full-expr? e)
  (match e
    [(? expr?) #t]
    [`(lambda ,(? symbol? x) ,(? full-expr? e)) #t]))


;; Generate a pair of the *resulting type* and a *set of constraints*
(define (build-constraints env e)
  (match e
    ;; Literals
    [(? integer? i) (cons 'int (set))]
    [(? boolean? b) (cons 'bool (set))]
    ;; Look up a type variable in an environment
    [(? symbol? x) (cons (hash-ref env x) (set))]
    ;; Lambda w/o annotation
    [`(lambda (,x) ,e)
     ;; Generate a new type variable using gensym
     ;; gensym creates a unique symbol
     (define T1 (gensym 'ty-var))
     (match (build-constraints (hash-set env x T1) e)
       [(cons T2 s) (cons `(,T1 -> ,T2) s)])]
    [`(,e1 ,e2)
     (match (build-constraints env e1)
       [(cons T1 C1)
        (match (build-constraints env e2)
          [(cons T2 C2)
           (define X (gensym 'ty-var))
           (cons X (set-union C1 C2 (set `(= ,T1 (,T2 -> ,X)))))])])]
    [`(,e : ,t)
     (match (build-constraints env e)
       [(cons T C)
        (define X (gensym 'ty-var))
        (cons X (set-add C `(= ,X ,T)))])]
    ;; Application
    [`(if ,e1 ,e2 ,e3)
     (match-define (cons T1 C1) (build-constraints env e1))
     (match-define (cons T2 C2) (build-constraints env e2))
     (match-define (cons T3 C3) (build-constraints env e3))
     (cons T2 (set-union C1 C2 C3 (set `(= ,T1 'bool) `(= ,T2 ,T3))))]))

;; Test 4
#;(build-constraints (hash))
#;
'((((lambda (x) (lambda (y) (lambda (z) (x 13))))
    (lambda (x) 5))
   (lambda (y) (#t : bool)))
  (lambda (z) 3)))

(define (unify constraints)
  ;; Free (type) variables of a type
  (define (free-vars T)
    (match T
      [(? symbol? X) (set X)]
      [`(,T1 -> ,T2) (set-union (free-vars T1) (free-vars T2))]))
  ;; within the constraint constr, substitute S for T
  (define (ty-subst ty X T)
    (match ty
      [(? symbol? Y) #:when (equal? X Y) T]
      [(? symbol? Y) Y]
      ['bool 'bool]
      ['int 'int]
      [`(,T0 -> ,T1) `(,(ty-subst T0 X T) -> ,(ty-subst T1 X T))]))
  (define (constr-subst constr S T)
    (match constr
      [`(= ,C0 ,C1) `(= ,(ty-subst C0 S T) ,(ty-subst C1 S T))]))
  ;; Walk over constraints one at a time
  (define (for-each constraints)
    (match constraints
      ['() (hash)]
      [`((= ,S ,T) . ,rest)
       (cond [(equal? S T)
              (for-each constraints)]
             [(and (symbol? S) (not (set-member? (free-vars T) S)))
              (hash-set (unify (map (lambda (constr) (constr-subst constr S T)) rest)) S T)]
             [(and (symbol? T) (not (set-member? (free-vars S) T)))
              (hash-set (unify (map (lambda (constr) (constr-subst constr S T)) rest)) T S)]
             [else
              (match-define `(,S1 -> ,S2) S)
              (match-define `(,T1 -> ,T2) T)
              (unify (cons `(= ,S1 ,T1) (cons `(= ,S2 ,T2) rest)))])]))

  ;; Turn it into a list of constraints to walk over
  (define constrs (set->list constraints))
  (for-each constrs))


(define (elaborate-types expr)
  (match (build-constraints (hash) expr)
    [(cons ty constraints)
     (define assignment (unify constraints))
     (displayln ty)
     assignment]))