open Types
open Lam

exception Could_not_infer

let rec typecheck (g : gam) (e : lam) (expected_t : ty) : bool =
  match e with
  Var x ->
    begin try (List.assoc x g) = expected_t
    with Not_found -> false
    end
  | 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, t) when t = expected_t ->
    begin match expected_t with
      Or (t_l, _) -> typecheck g e t_l
      | _ -> false
    end
  | Right (e, t) when t = expected_t ->
    begin match expected_t with
      Or (_, t_r) -> typecheck g e t_r
      | _ -> false
    end
  | Left _ | Right _ -> false


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 raise Could_not_infer
    | _ -> raise Could_not_infer
    end
  | Exf (e, t) ->
    if typecheck g e t
    then Bot
    else raise Could_not_infer
  | Pair (e1, e2) ->
    And (
      typeinfer g e1,
      typeinfer g e2
    )
  | Left (e, t) ->
    begin match t with
    Or (t_l, _) ->
      if typecheck g e t_l
      then t
      else raise Could_not_infer
    | _ -> raise Could_not_infer
    end
  | Right (e, t) ->
    begin match t with
    Or (_, t_r) ->
      if typecheck g e t_r
      then t
      else raise Could_not_infer
    | _ -> raise Could_not_infer
    end