#lang plait

;; =============================================================================
;; Type Inference: support.rkt
;; =============================================================================

;; DO NOT EDIT THIS FILE =======================================================

(require (typed-in racket/base [gensym : (-> Symbol)]))

(define-type-alias Label Symbol)

(define-type-alias LEnv (Hashof Symbol Label))

(define-type LExpr
  (e-num [label : Label] [x : Number])
  (e-bool [label : Label] [x : Boolean])
  (e-str [label : Label] [x : String])
  (e-op [label : Label]
        [op : Operator]
        [left : LExpr]
        [right : LExpr])
  (e-un-op [label : Label]
           [op : UnaryOperator]
           [LExpr : LExpr])
  (e-if [label : Label]
        [cond : LExpr]
        [consq : LExpr]
        [altern : LExpr])
  (e-lam [label : Label]
         [param-label : Label]
         [param : Symbol]
         [body : LExpr])
  (e-app [label : Label]
         [func : LExpr]
         [arg : LExpr])
  (e-id [label : Label] [name : Symbol])
  (sugar-and [label : Label]
             [left : LExpr]
             [right : LExpr])
  (sugar-or [label : Label]
            [left : LExpr]
            [right : LExpr])
  (sugar-let [label : Label]
             [id : Symbol]
             [value : LExpr]
             [body : LExpr])
  (e-empty [label : Label]))

(define-type Operator
  (op-plus)
  (op-num-eq)
  (op-append)
  (op-str-eq)
  (op-link))

(define-type UnaryOperator
  (op-first)
  (op-rest)
  (op-is-empty))

(define-type Term
  (t-var [label : Label])
  (t-con [head : Symbol] [args : (Listof Term)]))

(define-type Constraint
  (eq [l : Term] [r : Term]))

(define-type-alias Substitution (Hashof Label Term))

(define-type TypeInferenceError
  (ti-err-fails-occurs-check [replace : Term] [with : Term])
  (ti-err-constructor-mismatch [t1 : Term] [t2 : Term])
  (ti-err-unbound-id [name : Symbol]))

(define-syntax-rule (~a arg ...)
  (foldl (lambda (val string)
                 (string-append string val))
         ""
         (list (to-string arg) ...)))

(define (type-inference-error->string type-inference-error)
  (type-case TypeInferenceError type-inference-error
    [(ti-err-fails-occurs-check replace with) (~a replace " occurs in " with)]
    [(ti-err-constructor-mismatch t1 t2) (~a t1 " does not unify with " t2)]
    [(ti-err-unbound-id name) (~a name " is unbound")]))

(define (raise-error [err : TypeInferenceError])
  (error 'TypeInferenceError (type-inference-error->string err)))

(define (parse [input : S-Exp]): LExpr
  (cond
    [(s-exp-match? `true input) (e-bool (gensym) #t)]
    [(s-exp-match? `false input) (e-bool (gensym) #f)]
    [(s-exp-number? input) (e-num (gensym) (s-exp->number input))]
    [(s-exp-string? input) (e-str (gensym) (s-exp->string input))]
    [(s-exp-match? `{if ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 4)
           (e-if (gensym) (parse (second inlst)) (parse (third inlst)) (parse (fourth inlst)))
           (error '+ "incorrect number of args to +")))]
    [(s-exp-match? `{+ ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-plus) (parse (second inlst)) (parse (third inlst)))
           (error '+ "incorrect number of args to +")))]
    [(s-exp-match? `{num= ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-num-eq) (parse (second inlst)) (parse (third inlst)))
           (error 'num= "incorrect number of args to num=")))]
    [(s-exp-match? `{++ ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-append) (parse (second inlst)) (parse (third inlst)))
           (error '++ "incorrect number of args to ++")))]
    [(s-exp-match? `{str= ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-str-eq) (parse (second inlst)) (parse (third inlst)))
           (error 'str= "incorrect number of args to str=")))]
    [(s-exp-match? `{lam SYMBOL ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-lam (gensym) (gensym) (s-exp->symbol (second inlst)) (parse (third inlst)))
           (error 'lam "lambdas should only have one body")))]
    [(s-exp-match? `{lam ANY ...} input)
     (error 'lam "lambda parameters must be symbols")]
    [(s-exp-match? `{and ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (sugar-and (gensym) (parse (second inlst)) (parse (third inlst)))
           (error 'and "incorrect number of args to and")))]
    [(s-exp-match? `{or ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (sugar-or (gensym) (parse (second inlst)) (parse (third inlst)))
           (error 'or "incorrect number of args to or")))]
    [(s-exp-match? `{let {SYMBOL ANY} ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (sugar-let
            (gensym)
            (s-exp->symbol (first (s-exp->list (second inlst))))
            (parse (second (s-exp->list (second inlst))))
            (parse (third inlst)))
           (error 'let "incorrect number of args to let")))]
    [(s-exp-match? `empty input) (e-empty (gensym))]
    [(s-exp-match? `{link ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 3)
           (e-op (gensym) (op-link) (parse (second inlst)) (parse (third inlst)))
           (error 'link "incorrect number of args to link")))]
    [(s-exp-match? `{is-empty ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-un-op (gensym) (op-is-empty) (parse (second inlst)))
           (error 'is-empty "incorrect number of args to is-empty")))]
    [(s-exp-match? `{first ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-un-op (gensym) (op-first) (parse (second inlst)))
           (error 'first "incorrect number of args to first")))]
    [(s-exp-match? `{rest ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-un-op (gensym) (op-rest) (parse (second inlst)))
           (error 'rest "incorrect number of args to rest")))]
    [(s-exp-match? `{ANY ...} input)
     (let ([inlst (s-exp->list input)])
       (if (equal? (length inlst) 2)
           (e-app (gensym) (parse (first inlst)) (parse (second inlst)))
           (error 'app "incorrect number of args to app")))]
    [(s-exp-symbol? input) (e-id (gensym) (s-exp->symbol input))]))

; Helper functions that simulate the Type language seen in previous assignments
; in the language of Term.
(define (t-bool) (t-con 'Bool empty))
(define (t-num) (t-con 'Num empty))
(define (t-str) (t-con 'Str empty))
(define (t-list [elem : Term]) (t-con 'List (list elem)))
(define (t-fun [arg : Term] [ret : Term]) (t-con '-> (list arg ret)))

; Consumes an LExpr and returns its Label.
(define (get-label [e : LExpr]) : Label
  (type-case LExpr e
    [(e-bool lbl v) lbl]
    [(e-num lbl v) lbl]
    [(e-str lbl v) lbl]
    [(e-op lbl op l r) lbl]
    [(e-un-op lbl op e) lbl]
    [(e-if lbl cond consq altern) lbl]
    [(e-lam lbl param-lbl param body) lbl]
    [(e-app lbl func arg) lbl]
    [(e-id lbl name) lbl]
    [(sugar-and lbl l r) lbl]
    [(sugar-or lbl l r) lbl]
    [(sugar-let lbl id value body) lbl]
    [(e-empty lbl) lbl]))

(define (get-t-var [e : LExpr]) (t-var (get-label e)))

; Consumes a Term `t` and produces a Term which is structurally identical to `t`
; but renames the `t-var` labels such that:
;
;     1) The first t-var label is renamed throughout with 'a.
;     2) If β is renamed to γ,
;         - And (t-var 'δ) is the next unseen t-var after (t-var 'β)...
;         - And ϵ directly follows γ in the English alphabet...
;        then δ will be renamed to ϵ.
;
; You should use this function to help you write deterministic test cases in
; your testing file.
(define (normalize [t : Term]) : Term
  (letrec ([char->symbol (λ (c) (string->symbol (list->string (cons c empty))))]
           [αs (box (map char->symbol
                         (string->list "abcdefghijklmnopqrstuvwxyz")))]
           [pop
            (λ ()
              (let ([α (first (unbox αs))])
                (begin (set-box! αs (rest (unbox αs))) α)))]
           [θ (make-hash empty)]
           [normalize-impl
            (λ (t)
              (type-case Term t
                [(t-var lbl)
                 (type-case (Optionof Symbol) (hash-ref θ lbl)
                   [(none)
                    (let ([α (pop)])
                      (begin (hash-set! θ lbl α) (t-var α)))]
                   [(some α) (t-var α)])]
                [(t-con head args) (t-con head (map normalize-impl args))]))])
    (normalize-impl t)))

; Set operations.

; Type for immutable sets.
(define-type-alias (Setof 'a) (Hashof 'a Boolean))

; Convenience definition to create an empty set.
(define empty-set (hash (list)))

; Returns a (Setof 'a) containing all of the elements in the input (Listof 'a).
(define (list->set [items : (Listof 'a)]): (Setof 'a)
  (hash
    (map (lambda ([item : 'a]): ('a * Boolean)
           (pair item #t)) items)))

; Returns a (Listof 'a) containing all of the elements in the input (Setof 'a).
(define (set->list [set : (Setof 'a)]): (Listof 'a)
  (hash-keys set))

; Convenience macro equivalent to the built-in (list ...) macro, but for
; creating (Setof 'a)'s.
(define-syntax set
  (syntax-rules ()
    [(set) (list->set empty)]
    [(set args ...) (list->set (list args ...))]))

; Returns true if `set` contains no elements; otherwise, returns false.
(define (set-empty? [set : (Setof 'a)]): Boolean
  (empty? (set->list set)))

; Returns the number of elements in the input (Setof 'a).
(define (set-count [set : (Setof 'a)]): Number
  (length (set->list set)))

; Returns true if `set` contains `thing`; otherwise, returns false.
(define (set-member? [set : (Setof 'a)] [thing : 'a]): Boolean
  (type-case (Optionof Boolean) (hash-ref set thing)
    [(some v) #t]
    [(none) #f]))

; Returns a new, immutable (Setof 'a) containing all of the elements in `set`
; plus the element `thing`.
(define (set-add [set : (Setof 'a)] [thing : 'a]): (Setof 'a)
  (hash-set set thing #t))

; Returns a new, immutable (Setof 'a) containing all of the elements in `set`
; except the element `thing`.
(define (set-remove [set : (Setof 'a)] [thing : 'a]): (Setof 'a)
  (hash-remove set thing))

; The return type of the `set-pick` function.
(define-type (Pick 'a)
  (pick-none)
  (pick-some [element : 'a] [rest : (Setof 'a)]))

; Picks an arbitrary element out of the input (Setof 'a) and returns a
; (Pick 'a) structure.
(define (set-pick [set : (Setof 'a)]): (Pick 'a)
  (type-case (Listof 'a) (set->list set)
    [empty (pick-none)]
    [(cons item rst-items) (pick-some item (set-remove set item))]))

; Returns a new, immutable (Setof 'a) containing all of the elements in
; the given sets.
(define (set-union [sets : (Listof (Setof 'a))]): (Setof 'a)
  (foldl (lambda ([new-set : (Setof 'a)] [partial-union : (Setof 'a)])
           (foldl (lambda ([item : 'a] [set : (Setof 'a)]): (Setof 'a)
                    (set-add set item))
                  partial-union
                  (set->list new-set)))
         empty-set
         sets))

; "maps" the contents of a set:
; (equal?
;    (map-set (λ (e) (* e 2)) (set 1 2))
;    (set 2 4))
;
(define (map-set [f : ('a -> 'b)] [set : (Setof 'a)]) : (Setof 'b)
  (list->set (map f (set->list set))))

; maps the values of a table
; (equal?
;  (map-table-values (λ (e) (* e 2))
;                    (hash (list (pair 'a 1) (pair 'b 2))))
;  (hash (list (pair 'a 2) (pair 'b 4))))
;
(define (map-table-values [f : ('a -> 'b)] [table : (Hashof 'k 'a)]) : (Hashof 'k 'b)
  (foldl
   (λ (key table) (hash-set table key (f (some-v (hash-ref table key)))))
   table
   (hash-keys table)))