open Types
open Lam

let rec typecheck (g : gam) (e : lam) (expected_t : ty) =
  match e with
  Var x -> (List.assoc x g) = expected_t
  | Fun ((x, actual_t1), e) ->
    begin match expected_t with
    Arr (expected_t1, expected_t2) ->
      let g' = (x, expected_t1)::g in
      (actual_t1 = expected_t1) && (typecheck g' e expected_t2)
    | _ -> false
    end
  | App (e1, e2) ->
    begin match typeinfer g e1 with
    Arr (expected_t2, actual_t) ->
        (actual_t = expected_t) && (typecheck g e2 expected_t2)
    | _ -> false
    end
  | Exf (e, t) ->
    (expected_t = Bot) && (typecheck g e t)

and typeinfer (g : gam) (e : lam) : ty =
  match e with
  Var x -> List.assoc x g
  | Fun ((x, t1), e) ->
    let g' = (x, t1)::g in
    let t2 = typeinfer g' e in
    Arr (t1, t2)
  | App (e1, e2) ->
    begin match typeinfer g e1 with
    Arr (t1, t2) ->
      if typecheck g e2 t1
      then t2
      else failwith "couldn't infer"
    | _ -> failwith "couldn't infer"
    end
  | Exf (e, t) ->
    if typecheck g e t
    then Bot
    else failwith "couldn't infer"