Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
Type inference in Typed Racket
#lang typed/racket
;;; Typed Racket version of Martin Grabmü̈ller's "Algorithm W Step by Step"
;;; 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)
(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)))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment