(*
 * File: mp5-sol.ml
 *)


(*now to create the proof + gather the constraints
  gather_ty_constraints : judgment -> (proof * (expType * expType) list) option
*)

open Mp5common

let rec gather_ty_constraints judgment =
    let {gamma = gamma; exp = exp; expType = expType} = judgment in
    match exp
    with ConstExp c ->
      let c_ty = const_signature c in 
         Some ({antecedents = []; conclusion = judgment}, [(expType, c_ty)])
    | VarExp x ->
      (match lookup_env gamma x
       with Some x_ty ->  
            Some ({antecedents = []; conclusion = judgment}, [(expType, x_ty)])
       | None -> None)
    | BinOpExp bin_op ->
       let bin_op_ty = bin_op_signature bin_op in
           Some ({antecedents = []; conclusion = judgment}, [(expType, bin_op_ty)])
    | IfExp (bool_exp, then_exp, else_exp) ->
      (match
        (gather_ty_constraints {gamma = gamma; exp = bool_exp; expType = bool_ty},
         gather_ty_constraints {gamma = gamma; exp = then_exp; expType = expType},
         gather_ty_constraints {gamma = gamma; exp = else_exp; expType = expType})
       with (Some(bool_pf, bool_const), Some(then_pf, then_const), Some(else_pf, else_const))
           -> Some({antecedents = [bool_pf; then_pf; else_pf]; conclusion = judgment},
                   (bool_const @ (then_const @ else_const)))
       | _ -> None)
    | AppExp (e1, e2) ->
      let (tau1, tau2) = (fresh(),fresh()) in
      (match
        (gather_ty_constraints {gamma = gamma; exp = e1; expType = tau1},
         gather_ty_constraints {gamma = gamma; exp = e2; expType = tau2})
       with (Some(e1_pf, e1_const), Some(e2_pf, e2_const))
           -> Some({antecedents = [e1_pf; e2_pf]; conclusion = judgment},
                   (tau1, mk_fun_ty tau2 expType) :: (e1_const @ e2_const))
       | _ -> None)
    | FunExp (x,e) ->
      let (tau1, tau2) = (fresh(), fresh()) in
      (match gather_ty_constraints {gamma = (ins_env gamma x tau1); exp = e; expType = tau2}
       with Some (pf, const)
           -> Some ({antecedents = [pf]; conclusion = judgment},
		    (expType, mk_fun_ty tau1 tau2) :: const)
        | None -> None)
    | LetInExp (x, e1, e2) ->
      let tau1 = fresh() in
      (match (gather_ty_constraints {gamma = gamma; exp = e1; expType = tau1},
              gather_ty_constraints{gamma =(ins_env gamma x tau1); exp = e2; expType = expType})
       with (Some(e1_pf, e1_const), Some(e2_pf, e2_const))
           -> Some({antecedents = [e1_pf; e2_pf]; conclusion = judgment},
                   (e1_const @ e2_const))
       | _ -> None)
    | LetRecInExp (x, e1, e2) ->
      let tau1 = fresh() in
      (match (gather_ty_constraints{gamma = (ins_env gamma x tau1); exp = e1; expType = tau1},
              gather_ty_constraints{gamma = (ins_env gamma x tau1); exp = e2; expType = expType})
       with (Some(e1_pf, e1_const), Some(e2_pf, e2_const))
           -> Some({antecedents = [e1_pf; e2_pf]; conclusion = judgment},
                  (e1_const @ e2_const))
       | _ -> None)
    | RaiseExp e ->
      (match gather_ty_constraints {gamma = gamma; exp = e; expType = int_ty}
       with Some (pf, const)
           -> Some ({antecedents = [pf]; conclusion = judgment}, const)
       | None -> None)
    | TryWithExp (e, match1, more_matches) ->
      let get_pf_const (int_opt, exp) =
          (match gather_ty_constraints {gamma = gamma; exp = exp; expType = expType}
           with Some (pf, const) -> Some([pf], const)
           | None -> None) in
      let rec gather_pfs_consts matches =
          (match matches with [] -> raise (Failure "Shouldn't be able to happen")
           | [m] -> get_pf_const m
           | (m::ms) -> 
             (match get_pf_const m
              with None -> None
              | Some(m_pf,m_const) ->
                (match gather_pfs_consts ms with None -> None
                 | Some (ms_pf, ms_const) -> Some(m_pf@ms_pf, m_const @ ms_const))))   
      in
        (match (gather_ty_constraints {gamma = gamma; exp = e; expType = expType},
		gather_pfs_consts (match1::more_matches))
         with (Some (e_pf, e_const), Some (pfs, match_consts)) 
            -> Some({antecedents=e_pf::pfs; conclusion = judgment},
                    (e_const @ match_consts))
         | _ -> None);;