#lang racket
(require "lib.rkt" "base.rkt")
(require (only-in "stlc-rec.rkt" replace))
(require (only-in "stlc-ext.rkt" type->whnf))
(provide (all-defined-out))

;; The Simply-Typed Lambda Calculus with higher-order *impredicative* references,
;; plus sums products booleans ascryption etc, to implement doubly-linked lists

;; Checks an expression for syntactic well-formedness.
(define (stlc-dll/expr? expr)
  (match expr
    [x #:when (symbol? x) #t]
    [n #:when (natural? n) #t]
    [b #:when (boolean? b) #t]
    [`(,e : ,t)
      (and (stlc-dll/expr? e) (stlc-dll/type? t))]
    [`(type ,t1 ,t2 ,e)
      (and (stlc-dll/type? t1) (stlc-dll/type? t2) (stlc-dll/expr? e))]
    [(or
      `(inc ,e)
      `(car ,e) `(cdr ,e)
      `(inl ,e) `(inr ,e)
      `(new ,e)  `(! ,e)
      `(fold ,e) `(unfold ,e))
      (stlc-dll/expr? e)]
    [(or
      `(pair ,e1 ,e2)
      `(set ,e1 ,e2)
      `(,e1 ,e2))
      (and (stlc-dll/expr? e1) (stlc-dll/expr? e2))]
    [`(if ,c ,e1 ,e2)
      (and (stlc-dll/expr? c) (stlc-dll/expr? e1) (stlc-dll/expr? e2))]
    [`(case ,c (,x1 ⇒ ,e1) (,x2 ⇒ ,e2))
      (and (symbol? x1) (symbol? x2)
        (stlc-dll/expr? c) (stlc-dll/expr? e1) (stlc-dll/expr? e2))]
    [`(λ (,x : ,t) ,e)
      (and (symbol? x) (stlc-dll/type? t) (stlc-dll/expr? e))]
    [_ #f]))

;; Checks a type for syntactic well-formedness.
(define (stlc-dll/type? type)
  (match type
    ; Symbols are only valid if previously bound (by `type` or `μ`).
    ; We can't check that here, however.
    [x #:when (symbol? x) #t]
    [`(Ref ,t) (stlc-dll/type? t)]
    [(or `(,t1 × ,t2) `(,t1 ⊕ ,t2))
      (and (stlc-dll/type? t1) (stlc-dll/type? t2))]
    [`(,t1 → ,k ,t2)
      (and (stlc-dll/type? t1) (natural? k) (stlc-dll/type? t2))]
    [`(μ ,x ,t)
      (and (symbol? x) (stlc-dll/type? t))]
    [_ #f]))

;; Checks a value for syntactic well-formedness.
(define (stlc-dll/value? expr)
  (match expr
    [x #:when (symbol? x) #t]
    [n #:when (natural? n) #t]
    [b #:when (boolean? b) #t]
    [(or `(inl ,v) `(inr ,v))
      (stlc-dll/value? v)]
    [(or `(pair ,v1 ,v2) `(,v1 ,v2))
      (and (stlc-dll/value? v1) (stlc-dll/value? v2))]
    [`(λ (,x : ,t) ,e ,env)
      (and (symbol? x) (stlc-dll/type? t) (stlc-dll/expr? e) (dict? env))]
    [_ #f]))


(define (interpret expr)
  (interpret/core (desugar expr) #hash() (make-hash)))
;; Γ: a Table[Symbol, Expr] representing the context:
;;   the current bindings in scope introduced by λx.[]
;; Σ: a Table[Symbol, Expr] representing the heap:
;;   the current references on the heap generated by (gensym). mutable
;; Interprets a *desugared* expression *stripped* of type annotations.
(define/contract (interpret/core expr Γ Σ)
  (-> stlc-dll/expr? dict? dict? stlc-dll/value?)
  (match expr
    ['sole 'sole]
    [n #:when (natural? n) n]
    [b #:when (boolean? b) b]
    [r #:when (dict-has-key? Σ r) r]
    [x #:when (dict-has-key? Γ x) (dict-ref Γ x)]
    [f #:when (symbol? f) f]

    [`(inc ,e)
      (match (interpret/core e Γ Σ)
        [n #:when (natural? n) (+ n 1)]
        [e (format "incrementing an unknown value ~a" e)])]
    [`(if ,c ,e1 ,e2)
      (match (interpret/core c Γ Σ)
        ['#t (interpret/core e1 Γ Σ)]
        ['#f (interpret/core e2 Γ Σ)]
        [e (err (format "calling if on unknown expression ~a" e))])]

    [`(pair ,e1 ,e2)
      `(pair ,(interpret/core e1 Γ Σ) ,(interpret/core e2 Γ Σ))]
    [`(car ,e)
      (match (interpret/core e Γ Σ)
        [`(pair ,e1 ,e2) e1]
        [e (err (format "calling car on unknown expression ~a" e))])]
    [`(cdr ,e)
      (match (interpret/core e Γ Σ)
        [`(pair ,e1 ,e2) e2]
        [e (err (format "calling cdr on unknown expression ~a" e))])]

    [`(inl ,e) `(inl ,(interpret/core e Γ Σ))]
    [`(inr ,e) `(inr ,(interpret/core e Γ Σ))]
    [`(case ,e (,x1 ⇒ ,e1) (,x2 ⇒ ,e2))
      (match (interpret/core e Γ Σ)
        [`(inl ,e) (interpret/core e1 (dict-set Γ x1 e) Σ)]
        [`(inr ,e) (interpret/core e2 (dict-set Γ x2 e) Σ)]
        [e (err (format "calling case on unknown expression ~a" e))])]

    [`(new ,e)
      (let ([r (gensym)])
      (dict-set! Σ r e) r)]
    [`(! ,e)
      (let ([r (interpret/core e Γ Σ)])
      (if (dict-has-key? Σ r)
        (interpret/core (dict-ref Σ r) Γ Σ)
        (err (format "attempting to deref unknown reference ~a" r))))]
    [`(set ,e1 ,e2)
      (let ([r (interpret/core e1 Γ Σ)])
      (if (dict-has-key? Σ r) (dict-set! Σ r (interpret/core e2 Γ Σ))
        (err (format "attempting to update unknown reference ~a" r))))
      'sole]

    [`(fold ,e) `(fold ,(interpret/core e Γ Σ))]
    [`(unfold ,e)
      (match (interpret/core e Γ Σ)
        [`(fold ,e) e]
        [e (err (format "attempting to unfold unknown expression ~a" e))])]

    [`(λ (,x : ,t) ,e) `(λ ,x ,e ,Γ)]
    [`(,e1 ,e2)
      (match (interpret/core e1 Γ Σ)
        [`(λ ,x ,e1 ,env)
          (interpret/core e1 (dict-set env x (interpret/core e2 Γ Σ)) Σ)]
        [e1 (err (format "attempting to interpret arg ~a applied to unknown expression ~a" e2 e1))])]))

;; Checks that an expression is of a type, and returns #t or #f (or a bubbled-up error)
;; with: a type in weak-head normal form (!!)
;; Γ: a Table[Symbol, Expr ⊕ Type] representing the context:
;;   the current bindings in scope introduced by λx.[] and μx.[] and τx.[]
(define (check expr with)
  (check/core (desugar expr) with #hash()))
(define/contract (check/core expr with Γ)
  (-> stlc-dll/expr? stlc-dll/type? dict? boolean?)
  (match expr
    [`(type ,t1 ,t2 ,in)
      (check/core in with (dict-set Γ `(type ,t1) t2))]

    [`(if ,c ,e1 ,e2)
      (and (check/core c 'Bool Γ)
        (check/core e1 with Γ) (check/core e2 with Γ))]

    [`(pair ,e1 ,e2)
      (match with
        [`(,t1 × ,t2) (and (check/core e1 t1 Γ) (check/core e2 t2 Γ))]
        [_ #f])]

    [`(inl ,e)
      (match with
        [`(,t1 ⊕ ,t2) (check/core e t1 Γ)]
        [_ #f])]
    [`(inr ,e)
      (match with
        [`(,t1 ⊕ ,t2) (check/core e t2 Γ)]
        [_ #f])]
    [`(case ,e (,x1 ⇒ ,e1) (,x2 ⇒ ,e2))
      (match (infer/core e Γ) ; avoid needing type annotation on e
        [`(,a1 ⊕ ,a2)
          (and (check/core e1 with (dict-set Γ x1 a1))
            (check/core e2 with (dict-set Γ x2 a2)))]
        [_ #f])]

    [`(new ,e)
      (match with
        [`(Ref ,t) (check/core e t Γ)]
        [_ #f])]
    [`(! ,e)
      (check/core e `(Ref ,with) Γ)]

    [`(fold ,e)
      (match with
        [`(μ ,x ,t) (check/core e t (dict-set Γ `(type ,x) `(μ ,x ,t)))]
        [_ #f])]

    [`(λ (,x : ,t) ,e)
      (match with
        [`(,t1 → ,k ,t2)
          (and (equiv-type t t1 Γ) (check/core e t2 (dict-set Γ x t))
            (> k (level-body e (dict-set Γ x t1))))] ; KNOB
        [`(,t1 → ,t2) (err (format "missing level annotation on function type"))]
        [_ #f])]

    [_ (equiv-type (infer/core expr Γ) with Γ)]))

;; Checks if two types are equivalent up to α-conversion in context
(define (equiv-type e1 e2 Γ)
  (equiv-type/core e1 e2 Γ Γ))
(define/contract (equiv-type/core e1 e2 Γ1 Γ2)
  (-> stlc-dll/type? stlc-dll/type? dict? dict? boolean?)
  (match* (e1 e2)
    ; bound identifiers: if a key exists in the context, look it up
    [(x1 x2) #:when (dict-has-key? Γ1 x1)
      (equiv-type/core (dict-ref Γ1 x1) x2 Γ1 Γ2)]
    [(x1 x2) #:when (dict-has-key? Γ2 x2)
      (equiv-type/core x1 (dict-ref Γ2 x2) Γ1 Γ2)]

    ; recursive types: self-referential names can be arbitrary
    [(`(μ ,x1 ,t1) `(μ ,x2 ,t2))
      (let ([name gensym])
      (equiv-type/core t1 t2 (dict-set Γ1 x1 name) (dict-set Γ2 x2 name)))]

    ; check for syntactic equivalence on remaining forms
    [(`(,l1 ...) `(,l2 ...))
      (foldl (λ (x1 x2 acc) (if (equiv-type/core x1 x2 Γ1 Γ2) acc #f)) #t l1 l2)]
    [(v1 v2) (equal? v1 v2)]))

;; Infers a type from a given expression, if possible, or errors out.
;; Returns a type in weak-head normal form for structural matching.
(define (infer expr)
  (infer/core (desugar expr) #hash()))
(define/contract (infer/core expr Γ)
  (-> stlc-dll/expr? dict? stlc-dll/type?)
  (match expr
    ['sole 'Unit]
    [n #:when (natural? n) 'Nat]
    [b #:when (boolean? b) 'Bool]
    [x #:when (dict-has-key? Γ x)
      (type->whnf (dict-ref Γ x) Γ)]
    [f #:when (symbol? f)
      (err (format "attempting to infer type of free variable ~a" f))]

    [`(type ,t1 ,t2 ,in)
      (infer/core in (dict-set Γ `(type ,t1) t2))]
    [`(,e : ,t) ; we have a manual type annotation, so we must expand to weak-head normal form
      (if (check/core e (type->whnf t Γ) Γ) (type->whnf t Γ)
        (err (format "annotated expression ~a is not of annotated type ~a" e t)))]

    [`(inc ,e)
      (if (check/core e 'Nat Γ) 'Nat
        (err (format "calling inc on incorrect type ~a" (infer/core e Γ))))]
    [`(if ,c ,e1 ,e2)
      (if (check/core c 'Bool Γ)
        (let ([t (infer/core e1 Γ)])
        (if (check/core e2 t Γ) t
          (err (format "condition has branches of differing types ~a and ~a"
            t (infer/core e2 Γ)))))
        (err (format "condition ~a has incorrect type ~a" c (infer/core c Γ))))]

    [`(pair ,e1 ,e2)
      `(,(infer/core e1 Γ) × ,(infer/core e2 Γ))]
    [`(car ,e)
      (match (infer/core e Γ)
        [`(,t1 × ,t2) t1]
        [t (err (format "calling car on incorrect type ~a" t))])]
    [`(cdr ,e)
      (match (infer/core e Γ)
        [`(,t1 × ,t2) t2]
        [t (err (format "calling cdr on incorrect type ~a" t))])]

    [`(inl ,e)
      (err (format "unable to infer the type of a raw inl"))]
    [`(inr ,e)
      (err (format "unable to infer the type of a raw inr"))]
    [`(case ,e (,x1 ⇒ ,e1) (,x2 ⇒ ,e2))
      (match (infer/core e Γ)
        [`(,a1 ⊕ ,a2)
          (let ([b1 (infer/core e1 (dict-set Γ x1 a1))]
                [b2 (infer/core e2 (dict-set Γ x2 a2))])
            (if (equiv-type b1 b2 Γ) b1
              (err (format "case ~a is not of consistent type!"
                `(case (,a1 ⊕ ,a2) (,x1 ⇒ ,b1) (,x2 ⇒ ,b2))))))]
        [t (err (format "calling case on incorrect type ~a" t))])]

    [`(new ,e)
      `(Ref ,(infer/core e Γ))]
    [`(! ,e)
      (match (infer/core e Γ)
        [`(Ref ,t) t]
        [t (err (format "attempting to deref term ~a of type ~a" e t))])]
    [`(set ,e1 ,e2)
      (match (infer/core e1 Γ)
        [`(Ref ,t)
          (if (check/core e2 t Γ) 'Unit
            (err (format "attempting to update ~a: ~a with term ~a: ~a of differing type"
              e1 t e2 (infer/core e2 Γ))))]
        [t (err (format "attempting to update non-reference ~a: ~a" e1 t))])]

    [`(unfold ,e)
      (match (infer/core e Γ)
        [`(μ ,x ,t) (replace t x `(μ ,x ,t))]
        [t (err (format "expected ~a to be recursive, got ~a" e t))])]

    [`(λ (,x : ,t1) ,e)
      (let* ([t2 (infer/core e (dict-set Γ x t1))]
        [t1 (type->whnf t1 Γ)] ; type annotation, must expand
        [k (+ 1 (level-body e (dict-set Γ x t1)))]) ; KNOB
        `(,t1 → ,k ,t2))]
    [`(,e1 ,e2)
      (match (infer/core e1 Γ)
        [`(,t1 → ,k ,t2)
          (if (check/core e2 t1 Γ) t2
            (err (format "inferred argument type ~a does not match arg ~a of type ~a" t1 e2 (infer/core e2 Γ))))]
        [`(,t1 → ,t2) (err (format "missing level annotation on function type"))]
        [t (err (format "expected → type on application body, got ~a" t))])]))

;; Checks if a type is well-formed in the current context.
;; BIG ASSUMPTION: types in the current context are well-formed
(define/contract (well-formed t Γ)
  (-> stlc-dll/type? dict? boolean?)
  (match t
    [x #:when (dict-has-key? Γ x) #t]
    [(or 'Unit 'Nat 'Bool) #t]
    [`(Ref ,t) (well-formed t Γ)]
    [`(μ ,x ,t) (well-formed t (dict-set Γ x `(μ ,x ,t)))]
    [`(type ,x ,t) (well-formed t (dict-set Γ x `(μ ,x ,t)))]
    [(or `(,t1 → ,_ ,t2) `(,t1 × ,t2) `(,t1 ⊕ ,t2))
      (and (well-formed t1 Γ) (well-formed t2 Γ))]
    [_ #f]))

;; Checks if a type is well-kinded with respect to a level in the current context
;; BIG ASSUMPTION: types in the current context are well-formed
(define/contract (well-kinded t l Γ)
  (-> stlc-dll/type? natural? dict? boolean?)
  (match t
    [x #:when (dict-has-key? Γ x) #t]
    [(or 'Unit 'Nat 'Bool) (>= l 0)]
    [`(Ref ,t)
      (if (zero? l)
        (well-kinded t l Γ)
        (well-kinded t (- l 1) Γ))]
    [`(μ ,x ,t)
      (well-kinded t l (dict-set Γ x `(μ ,x ,t)))]
    [(or `(,t1 × ,t2) `(,t1 ⊕ ,t2))
      (and (well-kinded t1 l Γ) (well-kinded t2 l Γ))]
    [`(,t1 → ,k ,t2)
      (and (>= l k) (well-kinded t1 k Γ) (well-kinded t2 k Γ))]
    [_ #f]))

;; Infers the level of a (well-formed) type.
(define/contract (level-type t Γ)
  (-> stlc-dll/type? dict? natural?)
  (match t
    [x #:when (dict-has-key? Γ x)
      (level-type (dict-ref Γ x) Γ)]
    [(or 'Unit 'Nat) 0]
    [(or `(,t1 × ,t2) `(,t1 ⊕ ,t2))
      (max (level-type t1 Γ) (level-type t2 Γ))]
    [`(μ ,x ,t) ; note: correct but VERY WEIRD
      (level-type t Γ)]
    [`(,t1 → ,k ,t2)
      (if (and (>= k (level-type t1 Γ)) (>= k (level-type t2 Γ))) k ; KNOB
        (err (format "annotated level ~a is less than inferred levels of ~a and ~a!" k t1 t2)))]
    [`(Ref ,t)
      (let ([k (level-type t Γ)])
      (if (zero? k) 0 (+ 1 k)))] ; KNOB
    [t #:when (symbol? t) 0])) ; μ-type variables, not in Γ

;; Infers the level of a (well-formed) expression.
(define/contract (level-body e Γ)
  (-> stlc-dll/expr? dict? natural?)
  (match e
    ['sole 0]
    [n #:when (natural? n) 0]
    [x #:when (dict-has-key? Γ x) ; free variables, get their level
      (level-type (type->whnf (dict-ref Γ x) Γ) Γ)]
    [(or `(,e : ,_) `(λ (,_ : ,_) ,e)
      `(inc ,e) `(new ,e) `(! ,e) `(car ,e) `(cdr ,e) `(inl ,e) `(inr ,e)
      `(fold ,e) `(unfold ,e) `(fold (μ ,_ ,_) ,e) `(unfold (μ ,_ ,_) ,e))
      (level-body e Γ)]
    [(or `(set ,e1 ,e2) `(pair ,e1 ,e2) `(,e1 ,e2))
      (max (level-body e1 Γ) (level-body e2 Γ))]
    [(or `(if ,c ,e1 ,e2) `(case ,c (,_ ⇒ ,e1) (,_ ⇒ ,e2)))
      (max (level-body c Γ) (level-body e1 Γ) (level-body e2 Γ))]
    [x #:when (symbol? x) 0])) ; local variables, not in Γ