Type inference in Typed Racket

inference.rkt

#lang typed/racket

;;; Typed Racket version of Martin Grabmü̈ller's "Algorithm W Step by Step"
;;; http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.65.7733

;;; This is my first use of Typed Racket. I am looking to change this from
;;; following Haskell idiom to following Racket idiom.


;; An expression is a variable, a literal, an application,
;; an abstraction or a let form
(define-type Exp (U EVar ELit EApp EAbs ELet))
(struct: EVar ([name : Symbol]) #:transparent)
(struct: ELit ([val : Lit]) #:transparent)
(struct: EApp ([rator : Exp] [rand : Exp]) #:transparent)
(struct: EAbs ([arg : Symbol] [body : Exp]) #:transparent)
(struct: ELet ([var : Symbol] [intermed : Exp] [body : Exp]) #:transparent)

;; A literal is an integer or a boolean
(define-type Lit (U LInt LBool))
(struct: LInt ([num : Integer]))
(struct: LBool ([bool : Boolean]))

;; A type is a variable, Integer, Boolean, or a function
(define-type Type (U TVar TInt TBool TFun))
(struct: TVar ([name : Symbol]) #:transparent)
(struct: TInt () #:transparent)
(struct: TBool () #:transparent)
(struct: TFun ([arg : Type] [ret : Type]) #:transparent)

;; A type scheme is a list of bound type variables and a (possibly open) type
(struct: Scheme ([vars : (Listof Symbol)] [type : Type]) #:transparent)

;; A substitution is a mapping from names to types
(define-type Subst (HashTable Symbol Type))
(define: null-subst : Subst (make-immutable-hash '()))

;; A type environment is a mapping from names to type schemes
(define-type TypeEnv (HashTable Symbol Scheme))
(define: empty-type-env : TypeEnv (make-immutable-hash '()))
(define: example-type-env : TypeEnv
  `#hash([q . ,(Scheme '(a) (TFun (TVar 'a) (TInt)))]
         [x . ,(Scheme '() (TBool))]
         [y . ,(Scheme '(a) (TFun (TVar 'a) (TVar 'a)))]
         [z . ,(Scheme '() (TFun (TVar 'a) (TVar 'a)))]))

;; Identify the free type variables
(: free-type-vars ((U Type Scheme TypeEnv) -> (Setof Symbol)))
(define (free-type-vars t)
  (match t
    [(TVar n) (set n)]
    [(TInt) (set)]
    [(TBool) (set)]
    [(TFun a r) (set-union (free-type-vars a) (free-type-vars r))]
    [(Scheme vs t) (set-subtract (free-type-vars t) (list->set vs))]
    [(? hash? (app hash->list
                   (list (cons vars #{schemes : (Listof Scheme)}) ...)))
     (for/fold: ([ftvs : (Setof Symbol) (set)])
       ([s schemes])
       (set-union ftvs (free-type-vars s)))]))

;; Apply a substitution
(: apply-subst (case-> [Subst Type -> Type]
                       [Subst Scheme -> Scheme]
                       [Subst TypeEnv -> TypeEnv]))
(define (apply-subst s t)
  (match t
    [(TVar n) (cond [(hash-has-key? s n) (hash-ref s n)]
                    [else t])]
    [(TFun a r) (TFun (apply-subst s a) (apply-subst s r))]
    [(Scheme vs t2) (Scheme vs (apply-subst (hash-remove-all s vs) t2))]
    [(? hash? (app hash->list
                   (list (cons #{vars : (Listof Symbol)}
                               #{schemes : (Listof Scheme)}) ...)))
     ; apply the substitution to each value in the hash
     (for/hash: : TypeEnv [(v vars) (scm schemes)]
       (values v (apply-subst s scm)))]
    [anything anything]))

;; Compose two substitutions to produce a new substitution
(: compose-subst (Subst Subst -> Subst))
(define (compose-subst s1 s2)
  ; apply the first substitution to the result side of the second
  (define mod-s2
    (for/hash: : Subst [(pr (in-hash-pairs s2))]
      (values (car pr) (apply-subst s1 (cdr pr)))))
  ; add each modified mapping into the first substitution to get the result
  (for/fold: : Subst [(accum s1)]
    [(pr (in-hash-pairs mod-s2))]
    (hash-set accum (car pr) (cdr pr))))

;; Remove a binding from a type environment
;; If the environment does not have that binding, keep it as is
(: env-remove (TypeEnv Symbol -> TypeEnv))
(define (env-remove env var)
  (hash-remove env var))

;; Remove multiple keys (and associated values) from a hash
(: hash-remove-all (All (X Y) ((HashTable X Y) (Listof X) -> (HashTable X Y))))
(define (hash-remove-all h ks)
  (match ks
    ['() h]
    [(cons key keys) (hash-remove-all (hash-remove h key) keys)]))

;; Convert a type into a type scheme by having the scheme bind its free
;; type variables (i.e. those not bound by the environment)
(: generalize-type (TypeEnv Type -> Scheme))
(define (generalize-type env t)
  (Scheme (set->list (set-subtract (free-type-vars t) (free-type-vars env))) t))

;; Convert a type scheme into a concrete type by replacing the scheme-bound
;; variables with fresh variables
(: instantiate-scheme (Scheme -> Type))
(define (instantiate-scheme scm)
  (match scm
    [(Scheme vars t)
     (define: new-vars : (Listof Symbol) (for/list [(i vars)] (gensym)))
     (define: inst-subst : Subst
       (for/hash: : Subst [(ov vars) (nv new-vars)]
         (values ov (TVar nv))))
     (apply-subst inst-subst t)]))

;; Construct a substitution that will convert both input types to the same type
;; Note: This is based on the code given in AWSbS, which does not actually
;; have the behavior described in all cases.
(: unify-types (Type Type -> Subst))
(define (unify-types t1 t2)
  (match* (t1 t2)
    [((TFun l1 r1) (TFun l2 r2))
     (let*: [(left-subst : Subst (unify-types l1 l2))
             (right-subst : Subst (unify-types (apply-subst left-subst r1)
                                               (apply-subst left-subst r2)))]
       (compose-subst left-subst right-subst))]
    [((TVar n) t) (var-bind n t)]
    [(t (TVar n)) (var-bind n t)]
    [((TInt) (TInt)) null-subst]
    [((TBool) (TBool)) null-subst]
    [(_ _) (error (format "Types do not unify: ~v vs. ~v" t1 t2))]))

;; Construct a substitution that maps the given name to the given type,
;; unless that name already appears free in that type
(: var-bind (Symbol Type -> Subst))
(define (var-bind s t)
  (cond [(equal? t (TVar s)) null-subst]
        [(set-member? (free-type-vars t) s)
         (error (format "Occurs check fails: ~v vs. ~v" s t))]
        [else (make-immutable-hash (list (cons s t)))]))


;; Intermediate result for type inference
(: infer-type/lit (TypeEnv Lit -> (Pairof Subst Type)))
(define (infer-type/lit env lit)
  (match lit
    [(LInt _) (cons null-subst (TInt))]
    [(LBool _) (cons null-subst (TBool))]))

;; Intermediate result for type inference
(: infer-type/h (TypeEnv Exp -> (Pairof Subst Type)))
(define (infer-type/h env expr)
  (match expr
    [(EVar n) (cons null-subst (instantiate-scheme (hash-ref env n)))]
    [(ELit l) (infer-type/lit env l)]
    [(EAbs n e)
     (define tv (TVar (gensym)))
     (define new-env (hash-set env n (Scheme '() tv)))
     (match-define (cons subst type) (infer-type/h new-env e))
     (cons subst (TFun (apply-subst subst tv) type))]
    [(EApp e1 e2)
     (define tv (TVar (gensym)))
     (match-define (cons sub1 type1) (infer-type/h env e1))
     (match-define (cons sub2 type2) (infer-type/h (apply-subst sub1 env) e2))
     (define sub3 (unify-types (apply-subst sub2 type1) (TFun type2 tv)))
     (cons (compose-subst sub3 (compose-subst sub2 sub1))
           (apply-subst sub3 tv))]
    [(ELet x e1 e2)
     (match-let* ([(cons sub1 type1) (infer-type/h env e1)]
                  [new-env (hash-set
                            env x
                            (generalize-type (apply-subst sub1 env) type1))]
                  [(cons sub2 type2) (infer-type/h (apply-subst sub1 new-env)
                                                 e2)])
       (cons (compose-subst sub1 sub2) type2))]))

;; Infer a type for an expression
(: infer-type (TypeEnv Exp -> Type))
(define (infer-type env exp)
  (match-let ([(cons subst type) (infer-type/h env exp)])
    (apply-subst subst type)))


;; Example cases:
(infer-type empty-type-env
            (ELet 'id (EAbs 'x (EVar 'x))
                  (EVar 'id)))
(infer-type empty-type-env
            (ELet 'id (EAbs 'x (EVar 'x))
                  (EApp (EVar 'id) (EVar 'id))))
(infer-type empty-type-env
            (ELet 'id (EAbs 'x (ELet 'y (EVar 'x) (EVar 'y)))
                  (EApp (EVar 'id) (EVar 'id))))
(infer-type empty-type-env
            (ELet 'id (EAbs 'x (ELet 'y (EVar 'x) (EVar 'y)))
                  (EApp (EApp (EVar 'id) (EVar 'id)) (ELit (LInt 2)))))
(infer-type empty-type-env
            (EAbs 'm (ELet 'y (EVar 'm)
                           (ELet 'x (EApp (EVar 'y) (ELit (LBool #t)))
                                 (EVar 'x)))))
;; In this case, no type can be inferred for 
;;   (EAbs 'x (EApp (EVar 'x) (EVar 'x)))
;; so an error would be raised.
#; (infer-type empty-type-env
               (ELet 'id (EAbs 'x (EApp (EVar 'x) (EVar 'x)))
                     (EVar 'id)))