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
You can’t perform that action at this time.