open Types
open Lam

let rec typecheck (g : gam) (e : lam) (expected_t : ty) : bool =
  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)
  | Pair (e1, e2) ->
    begin match expected_t with
      And (t1, t2) ->
        (typecheck g e1 t1) && (typecheck g e2 t2)
      | _ -> false
    end
  | Left e ->
    begin match expected_t with
      Or (t, _) -> typecheck g e t
      | _ -> false
    end
  | Right e ->
    begin match expected_t with
      Or (_, t) -> typecheck g e t
      | _ -> false
    end


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"
  | Pair (e1, e2) ->
    And (
      typeinfer g e1,
      typeinfer g e2
    )
  | Left e ->
    Or (typeinfer g e, Unknown)
  | Right e ->
    Or (Unknown, typeinfer g e)