Inferring monotonicity with modal types for a simple lambda-calculus

tone-inference.rkt

#lang racket
(require racket/hash rackunit)

;;; ---------- TONES and TONE CONTEXTS ---------- ;;;
(define tone? (or/c 'id 'op 'iso 'path))

(define (tone<=? a b)
  (match* (a b)
    [('iso _) #t] [(_ 'path) #t]
    [(x x) #t]
    [(_ _) #f]))

(define (tone-meet . tones)
  (match (set->list (set-remove (apply set tones) 'path))
    ['() 'path]
    [(list x) x]
    [_ 'iso]))

(define (tone-compose a b)
  (match* (a b)
    [(_ (or 'iso 'path)) b]
    [((or 'iso 'path) _) a]
    [('id b) b] [(a 'id) a]
    [('op 'op) 'id]))

(define (tones-apply T tones)
  (for/hash ([(x U) tones]) (values x (tone-compose T U))))

(define (tones-meet . tones)
  (apply hash-union #hash() tones #:combine tone-meet))

;;; ---------- BASE TYPES and LITERALS ---------- ;;;
(define base-type? (or/c 'bool 'real 'str))
(define (lit? x) (if (lit-type x) #t #f))
(define (lit-type l)
  (cond
    [(boolean? l) 'bool]
    [(real? l) 'real]
    [(string? l) 'str]
    [(null? l) '(*)]
    [#t #f]))

;;; ---------- POINTED TYPES ---------- ;;;
(struct exn:bottomless exn:fail () #:transparent)
(define (raise-bottomless msg) (raise (exn:bottomless msg (current-continuation-marks))))
(define (bottom tone tp)
  (match* (tone tp)
    [('id 'bool) #f] [('op 'bool) #t]
    [('id 'str) ""]
    [('id `(set ,A)) (set)]
    [(T `(,(? tone? U) ,A)) (bottom (tone-compose T U) A)]
    [(T `(* ,@As)) (map (curry bottom T) As)]
    [(T `(-> ,A ,B)) (let ([bot (bottom T B)]) (lambda (_) bot))]
    ;; singleton sums. I'm not sure this is a good idea. by adding branches,
    ;; subtyping could turn a "bottom" value into a non-bottom! i.e. now bottom
    ;; isn't covariant!
    [(T `(+ ,(hash-table (name A)))) (hash name (bottom tone A))]
    [(T A) (raise-bottomless (format "type ~s has no bottom at tone ~s" A T))]))

(define (type-pointed? tp)
  (with-handlers ([exn:bottomless? (lambda () #f)])
    (bottom 'id tp)
    #t))

;;; ---------- SUBTYPING ---------- ;;;
(define (left-adjoint tone)
  (match tone ['id 'id] ['op 'op] ['iso 'path]))

(define (strip-tone T type)
  (match type
    [`(,(? tone? U) ,A) (strip-tone (tone-compose T U) A)]
    [A (values T A)]))

(define (detune type)
  (define-values (T A) (strip-tone 'id type))
  (values (left-adjoint T) A))

;; Yields a tone T such that TA <: B, or errors.
(define (subtype A B)
  (define (fail) (error 'subtype "~s is not a subtype of ~s" A B))
  (match* (A B)
    ;; base cases.
    [(A A) #:when (base-type? A) 'id]
    ;; tone introduction/elimination
    [(A `(,(? tone? U) ,B)) (tone-compose U (subtype A B))]
    [(`(,(? tone? V) ,A) B) (tone-compose (subtype A B) (left-adjoint V))]
    ;; sets
    [(`(set ,A) `(set ,B)) (subtype A B) 'id]
    ;; distribution over products & sums.
    [(`(* ,@As) `(* ,@Bs)) (apply tone-meet (map subtype As Bs))]
    [(`(+ ,@As) `(+ ,@Bs))
     #:when (subset? (hash-keys As) (hash-keys Bs))
     (apply tone-meet
            (for/list ([(name A) As])
              (subtype A (hash-ref B name))))]
    ;; function cases
    [(`(-> ,A1 ,B1) `(-> ,A2 ,B2))
     (define dom (subtype A2 A1))
     (define cod (subtype B1 B2))
     (unless (match cod
               ;; TODO: can I simplify this to a single case? THEORY!
               ['path (eq? 'path (tone-compose 'path dom))]
               [_ (tone<=? (left-adjoint cod) dom)])
       (fail))
     cod]
    ;; failure case.
    [(A B) (fail)]))

;;; ---------- TYPE CHECKING ---------- ;;;
(define env (make-parameter #hash()))

(define (typerr msg . vals) (error (apply format msg vals)))
(define (unbound x) (typerr "unbound variable: ~a" x))

(define-syntax-rule (with-var x-expr A body ...)
  (let ([x x-expr])
    (define tones (parameterize ([env (hash-set (env) x A)]) body ...))
    (values (hash-ref tones x 'path) (hash-remove tones x))))

(define-syntax-rule (with-vars vars body ...)
  (with-vars-do vars (lambda () body ...)))

(define (with-vars-do vars thunk)
  (define tones
    (parameterize ([env (hash-union (env) vars #:combine (lambda (x y) y))])
      (thunk)))
  (values
   (for/hash ([(x T) tones] #:when (hash-has-key? vars x)) (values x T))
   (for/hash ([(x T) tones] #:unless (hash-has-key? vars x)) (values x T))))

;; Checks that `pat` can match values of `type`, returning the variables it
;; binds and their types.
(define (check-pat type pat #:tone [tone 'id])
  (define-values (T A) (strip-tone tone type))
  (set! tone T)
  (match* (A pat)
    [(A (? symbol? x)) (hash x `(,tone ,A))]
    [(A (app lit-type B)) #:when B (subtype B A) #hash()]
    [(`(* ,@As) `(tuple ,@Ps))
     (apply hash-union #hash() (map (curry check-pat #:tone tone) As Ps))]
    [(`(+ ,@arms) `(tag ,name ,P)) (=> fail)
     (check-pat (hash-ref arms name (lambda () (fail))) P #:tone tone)]
    [(A P) (typerr "pattern ~s cannot match type ~s" P A)]))

;; Infers a type for expr if it is a synthesizing expression
;; (variables, literals, annotated expressions, and applications).
(define (infer expr)
  (match expr
    [(? symbol? x)
     (values (hash-ref (env) x (lambda () (unbound x)))
             (hash x 'id))]
    [(app lit-type A) #:when A (values A #hash())]
    [`(the ,A ,M) (values A (check A M))]
    [`(,M ,N)
     (define-values (M-type M-tones) (infer M))
     (match-define-values (M-tone `(-> ,A ,B)) (detune M-type))
     (define N-tones (check A N))
     (values B (tones-meet (tones-apply M-tone M-tones) N-tones))]))

;; Checks that an expression has a type.
(define (check expected expr)
  (define-syntax-rule (mustbe pat)
    (match-define pat
     (match expected
       [pat expected]
       [_ (typerr "type has wrong form")])))

  (match expr
    ;; if the type is toned, immediately apply the corresponding intro rule.
    [(app (const expected) `(,(? tone? tone) ,A))
     (tones-apply tone (check A expr))]

    [`(let ([,x ,M]) ,N)
     (define-values (M-tp M-tones) (infer M))
     (define-values (x-tone N-tones) (with-var x M-tp (check expected N)))
     (tones-meet (tones-apply x-tone M-tones) N-tones)]

    [`(if ,M ,N ,O)
     (tones-meet (check '(iso bool) M) (check expected N) (check expected O))]

    [`(when ,M ,N)
     (unless (type-pointed? expected)
       (typerr "cannot use `when` at a type without a bottom"))
     (tones-meet (check 'bool M) (check expected N))]

    [`(lambda (,x) ,M)
     (mustbe `(-> ,A ,B))
     (define-values (x-tone M-tones) (with-var x A (check B M)))
     (unless (tone<=? 'id x-tone)
       (typerr "not used monotonically: ~a" x))
     M-tones]

    [`(tag ,name ,M)
     (mustbe `(+ ,arms))
     (define (uhoh) (typerr "invalid constructor for type"))
     (check (hash-ref arms name uhoh) M)]

    [`(tuple ,@Ms)
     (mustbe `(* ,@As))
     (apply tones-meet (map check As Ms))]

    [`(case ,M (,pats ,Ns) ...)
     (define-values (M-type M-tones) (infer M))
     (apply tones-meet
       (for/list ([pat pats] [N Ns])
         (define var-types (check-pat M-type pat))
         (define-values (var-tones N-tones)
           (with-vars var-types (check expected N)))
         (define tone (apply tone-meet (hash-values var-tones)))
         (tones-meet (tones-apply tone M-tones) N-tones)))]

    ;; otherwise, defer to type inference.
    [expr
     (define-values (inferred tones) (infer expr))
     (tones-apply (subtype inferred expected) tones)]))

;;; ---------- TESTS ---------- ;;;
(module+ test
  (define (checks! tp expr) (check-not-exn (lambda () (check tp expr))))
  (define (fails! tp expr) (check-exn any/c (lambda () (check tp expr))))

  ;; Type checking
  (checks! 'bool #t)
  (checks! 'real 2)
  (checks! '(-> bool real) '(lambda (x) 2))
  (checks! '(-> bool bool) '(lambda (x) x))

  (checks! '(iso (-> bool bool))      '(lambda (x) x))
  (checks! '(-> (op bool) (op bool))  '(lambda (x) x))
  (checks! '(-> (iso bool) (op bool)) '(lambda (x) x))

  (fails! '(-> bool (op bool))  '(lambda (x) x))
  (fails! '(-> bool (iso bool)) '(lambda (x) x))

  ;; let
  (checks! 'real '(let ([x 2]) x))
  (checks! '(-> bool bool) '(lambda (x) (let ([y x]) y)))
  (checks! '(-> (op bool) (iso real)) '(let ([x 2]) (lambda (y) x)))

  ;; if & when
  (checks! 'real '(if #t 0 1))
  (checks! 'str '(when #f "foo"))
  (fails!  '(-> bool real) '(lambda (x) (if x 1 0)))
  (checks! '(-> bool str) '(lambda (x) (when x "foo")))

  ;; case. TODO: more tests
  (checks! 'real '(case 2 [x x]))
  (checks! 'bool '(case #t [#t #f] [#f #t]))
  (checks! 'bool '(case (the (* bool real) (tuple #f 7))
                    [(tuple x y) x]))

  ;; TODO: tests for tag, tuple, the, function application.
  )