12-08.rkt

#lang racket

;; Thursday, Dec 8, 2022
;; Class 2 -- Type Inference
;; 
;; In this exercise, we will study type inference, and
;; a variety of type system features.

(provide (all-defined-out))


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

;;
;; Constraint-Based Type Inference
;;

;; symbolic types may have free variables
(define (sym-type? t)
  (match t
    ;; May just be a concrete type
    [(? type? t) #t]
    ;; Arrow types: t0 -> t1 (also may be symbolic)
    [`(,(? sym-type? t0) -> ,(? sym-type? t1)) #t]
    ;; or a type variable
    ;; a type variable looks like '(type-var N) where N is a number
    [`(type-var ,(? nonnegative-integer? number)) #t]
    [_ #f]))

;; mutable
(define ty-var-ct 0)
(define/contract (fresh-tyvar)
  (-> sym-type?)
  (set! ty-var-ct (add1 ty-var-ct))
  `(type-var ,ty-var-ct))

;; Constraints assert equality of symbolic types
(define (constraint? t)
  (match t
    [`(= ,(? sym-type? st0) ,(? sym-type? st1)) #t]
    [_ #f]))

;; Full expresssions also include unannotated lambdas,
;; which must be filled in later. As before, any subterm
;; may be annotated
(define (unannotated-expr? e)
  (match e
    [(? expr?) #t]
    [`(lambda ,(? symbol? x) ,(? unannotated-expr? e)) #t]))

;; (build-constraints env e) consumes a term e and produces a
;; *fully-annotated* term, where each subterm has been explictly
;; typed with a type variable. However, the type variable
;; will be "fresh" (using an impure function). We also return
;; a set of *constraints*, which constrain equalities on these
;; type variables
;;
;; In sum, the return value is a cons cell of (a) the fully-annotated
;; term and (b) a set of constraints

;; Starter code, no need to modify
(define (build-constraints env e)
  (match e
    ;; Literals
    [(? integer? i) (cons `(,i : int) (set))]
    [(? boolean? b) (cons `(,b : bool) (set))]
    ;; Look up a type variable in an environment
    [(? symbol? x) (cons `(,x : ,(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 (fresh-tyvar)) ;; (type-var 0)
     (match (build-constraints (hash-set env x T1) e)
       [(cons `(,e+ : ,T2) S)
        (cons `((lambda (,x : ,T1) ,e+) : (,T1 -> ,T2)) S)])]
    ;; Application: constrain input matches, return output
    [`(,e1 ,e2)
     (match (build-constraints env e1)
       [(cons `(,e1+ : ,T1) C1)
        (match (build-constraints env e2)
          [(cons `(,e2+ : ,T2) C2)
           (define X (fresh-tyvar))
           (cons `(((,e1+ : ,T1) (,e2+ : ,T2)) : ,X)
                 (set-union C1 C2 (set `(= ,T1 (,T2 -> ,X)))))])])]
    ;; Type stipulation against t--constrain
    [`(,e : ,t)
     (match (build-constraints env e)
       [(cons `(,e+ : ,T) C)
        (define X (fresh-tyvar))
        (cons `((,e+ : ,T) : ,X) (set-add (set-add C `(= ,X ,T)) `(= ,X ,t)))])]
    ;; If: the guard must evaluate to bool, branches must be
    ;; of equal type.
    [`(if ,e1 ,e2 ,e3)
     (match-define (cons `(,e1+ : ,T1) C1) (build-constraints env e1))
     (match-define (cons `(,e2+ : ,T2) C2) (build-constraints env e2))
     (match-define (cons `(,e3+ : ,T3) C3) (build-constraints env e3))
     (cons `((if (,e1+ : ,T1) (,e2+ : ,T2) (,e3+ : ,T3)) : ,T2)
           (set-union C1 C2 C3 (set `(= ,T1 bool) `(= ,T2 ,T3))))]))


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

;; Return the free type variables of a given symbolic type

(define (free-type-vars T)
  (match T
    [`(type-var ,N) (set T)]
    [`(,T1 -> ,T2) (set-union (free-type-vars T1) (free-type-vars T2))]
    [_ (set)]))

(free-type-vars '(int -> (type-var 0)))
(free-type-vars '((type-var 0) -> ((type-var 2) -> (type-var 1))))
(free-type-vars '(bool -> (bool -> (bool -> ((type-var 0) -> (type-var 1))))))

(define (ty-var? tv)
    (match tv
      [`(type-var ,n) #t]
      [_ #f]))

;; ((type-var 0) -> (type-var 0)) if I say to replace '(type-var 0) with int
;; you should return '(int -> int)
;; ((type-var 0) -> (type-var 1)) if I say to replace '(type-var 1) with bool
;; you should return '((type-var 0) -> bool)

;; within the constraint constr, substitute S for T
(define (ty-subst ty X T)
  (match ty
    [(? ty-var? Y) #:when (equal? X Y) T]
    [(? ty-var? Y) ty]
    ['bool 'bool]
    ['int 'int]
    ;; '((type-var 0) -> int)
    [`(,T0 -> ,T1) `(,(ty-subst T0 X T) -> ,(ty-subst T1 X T))]))

;; Unification takes a set of constraints and
;; walks over it to produce a hash from variable names
;; to their "true" types

(define (unify constraints)
  ;; Substitute into a constraint
  (define (constr-subst constr S T)
    (match constr
      [`(= ,C0 ,C1) `(= ,(ty-subst C0 S T) ,(ty-subst C1 S T))]))
  ;; Is t an arrow type?
  (define (arrow? t)
    (match t [`(,_ -> ,_) #t] [_ #f]))
  ;; Walk over constraints one at a time
  (define (for-each constraints)
    (match constraints
      ;; when no more constraints left, just return empty hash
      ;; (no more variables to assign)
      ['() (hash)]
      ;; otherwise, the first constraint is *some* symbolic
      ;; equality between two symbolic types S and T
      [`((= ,S ,T) . ,rest)
       ;; So, consider a few possibilities
       (cond [(equal? S T)
              ;; If S and T *are syntactically identical*, then, just skip
              ;; (= int int) (= ((type-var 0) -> bool) ((type-var 0) -> bool))
              (for-each rest)]
             ;; (= S T) and S is a type variable that is not present in T
             ;; In this case, we map all elements of rest (which are type constraints)
             ;; so that all occurrences of S become remapped to T
             ;; and then, we report back that S |-> T
             [(and (ty-var? S) (not (set-member? (free-type-vars T) S)))
              ;; build a new set of constraints where all occurences of S have been replaced by T
              ;; (this is fine, because S is a type-variable, that doesn't occur in T)
              (define new-constraints (map (lambda (constr) (constr-subst constr S T)) rest))
              (hash-set (for-each new-constraints) S T)]
             ;; If (= S T), S is *not* a type variable,
             ;; maybe it's something like int, or (... -> ...),
             ;; but if T *is* a type variable, we can do the same thing
             [(and (ty-var? T) (not (set-member? (free-type-vars S) T)))
              (hash-set (for-each (map (lambda (constr) (constr-subst constr T S)) rest)) T S)]
             ;; If we want to unify (= (T1 -> T2) (S1 -> S2))
             ;; then, we need to check that T1 = S1 and T2 = S2
             ;; so to do that, we call for each, with some *new* constraints
             ;; stipulating equality of T1/S2 and T2/S2
             [(and (arrow? S) (arrow? T))
              (match-define `(,S1 -> ,S2) S)
              (match-define `(,T1 -> ,T2) T)
              (for-each (cons `(= ,S1 ,T1) (cons `(= ,S2 ,T2) rest)))]
             ;; Otherwise, there is no solution and type error should be thrown
             [else (error "type failure")])]))

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

;; Returns 'yes when the expression e is typeable
;; Returns 'type-error otherwise
(define (typeable? e)
  (with-handlers ([exn:fail? (lambda (exn)
                             'type-error)])
    (unify (cdr (build-constraints (hash '+ '(int -> (int -> int))
                                         'is-zero? '(int -> bool)) e)))
    'yes))

;; (pretty-print (typeable? '(lambda (f) (lambda (y) (lambda (x) (if (is-zero? 0) ((+ x) ((+ (f x)) (f y))) (f y)))))))
;; (pretty-print (typeable? '(lambda (f) (lambda (y) (lambda (x) (if (is-zero? 0) ((+ x) ((+ (f x)) (f y))) ((f y) : bool)))))))