Add the type an expression was checked against to tannots, and use for Coq

Tweak extra Coq files to match.
Tweak early return rewrite to use declared return type, which can always
be put into an E_cast.

4 files changed, 79 insertions, 55 deletions

diff --git a/src/pretty_print_coq.ml b/src/pretty_print_coq.ml
index fb53db96..39a3e18a 100644
--- a/src/pretty_print_coq.ml
+++ b/src/pretty_print_coq.ml
@@ -825,6 +825,15 @@ let is_range_from_atom env (Typ_aux (argty,_)) (Typ_aux (fnty,_)) =
      Type_check.prove env (nc_and (nc_eq nexp low) (nc_eq nexp high))
   | _ -> false
 
+(* Get a more general type for an annotation/expression - i.e.,
+   like typ_of but using the expected type if there was one *)
+let general_typ_of_annot annot =
+  match expected_typ_of annot with
+  | None -> typ_of_annot annot
+  | Some typ -> typ
+
+let general_typ_of (E_aux (_,annot)) = general_typ_of_annot annot
+
 let prefix_recordtype = true
 let report = Reporting_basic.err_unreachable
 let doc_exp_lem, doc_let_lem =
@@ -842,7 +851,11 @@ let doc_exp_lem, doc_let_lem =
     let wrap_parens doc = if aexp_needed then parens (doc) else doc in
     let maybe_add_exist epp =
       let env = env_of full_exp in
-      let typ = Env.expand_synonyms env (typ_of full_exp) in
+      let typ = Env.expand_synonyms env (general_typ_of full_exp) in
+      let () =
+        debug ctxt (lazy ("Considering build_ex for " ^ string_of_exp full_exp));
+        debug ctxt (lazy (" at type " ^ string_of_typ typ))
+      in
       let typ = expand_range_type typ in
       match destruct_exist env typ with
       | None -> epp
@@ -1078,7 +1091,7 @@ let doc_exp_lem, doc_let_lem =
             (* Insert existential unpacking of arguments where necessary *)
             let doc_arg arg typ_from_fn =
               let arg_pp = expY arg in
-              let arg_ty_plain = Env.expand_synonyms (env_of arg) (typ_of arg) in
+              let arg_ty_plain = Env.expand_synonyms (env_of arg) (general_typ_of arg) in
               let arg_ty = expand_range_type arg_ty_plain in
               let typ_from_fn_plain = subst_unifiers inst typ_from_fn in
               let typ_from_fn_plain = Env.expand_synonyms (env_of arg) typ_from_fn_plain in
@@ -1106,7 +1119,7 @@ let doc_exp_lem, doc_let_lem =
                  subst_unifiers inst ret_typ
             in
             let unpack,build_ex,autocast =
-              let ann_typ = Env.expand_synonyms env (typ_of_annot (l,annot)) in
+              let ann_typ = Env.expand_synonyms env (general_typ_of_annot (l,annot)) in
               let ann_typ = expand_range_type ann_typ in
               let ret_typ_inst = expand_range_type (Env.expand_synonyms env ret_typ_inst) in
               let ret_typ_inst =
@@ -1185,7 +1198,7 @@ let doc_exp_lem, doc_let_lem =
        let epp = epp ^/^ doc_tannot_lem ctxt (env_of e) (effectful (effect_of e)) typ in
        let env = env_of_annot (l,annot) in
        let unpack,build_ex =
-         let ann_typ = Env.expand_synonyms env (typ_of_annot (l,annot)) in
+         let ann_typ = Env.expand_synonyms env (general_typ_of_annot (l,annot)) in
          let ann_typ = expand_range_type ann_typ in
          let cast_typ = expand_range_type (Env.expand_synonyms env typ) in
          match cast_typ, ann_typ with
@@ -1939,7 +1952,7 @@ let doc_val pat exp =
                     "Top-level value definition with complex pattern not supported for Coq yet")
   in
   let env = env_of exp in
-  let typ = expand_range_type (Env.expand_synonyms env (typ_of exp)) in
+  let typ = expand_range_type (Env.expand_synonyms env (general_typ_of exp)) in
   let ctxt = { empty_ctxt with debug  = List.mem (string_of_id id) (!opt_debug_on) } in
   let id, opt_unpack =
     match destruct_exist env typ with
diff --git a/src/rewrites.ml b/src/rewrites.ml
index 47c7e923..689c7897 100644
--- a/src/rewrites.ml
+++ b/src/rewrites.ml
@@ -1963,7 +1963,9 @@ let rewrite_defs_early_return (Defs defs) =
        let env = env_of_annot (l, tannot) in
        let eff = effect_of_annot tannot in
        let tannot' = mk_tannot env typ (union_effects eff (mk_effect [BE_escape])) in
-       let exp' = add_e_cast (typ_of exp) exp in
+       let exp' = match Env.get_ret_typ env with
+         | Some typ -> add_e_cast typ exp
+         | None -> exp in
        E_aux (E_app (mk_id "early_return", [exp']), (l, tannot'))
     | _ -> full_exp in
 
diff --git a/src/type_check.ml b/src/type_check.ml
index 84c21861..be6950e9 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -1940,17 +1940,18 @@ let subtype_check env typ1 typ2 =
 (**************************************************************************)
 
 (* The type checker produces a fully annoted AST - tannot is the type
-   of these type annotations. *)
-type tannot = (Env.t * typ * effect) option
+   of these type annotations.  The extra typ option is the expected type,
+   that is, the type that the AST node was checked against, if there was one. *)
+type tannot = ((Env.t * typ * effect) * typ option) option
 
-let mk_tannot env typ effect = Some (env, typ, effect)
+let mk_tannot env typ effect : tannot = Some ((env, typ, effect), None)
 
 let empty_tannot = None
 let is_empty_tannot = function
   | None -> true
   | Some _ -> false
 
-let destruct_tannot tannot = tannot
+let destruct_tannot tannot = Util.option_map fst tannot
 
 let string_of_tannot tannot =
   match destruct_tannot tannot with
@@ -1959,7 +1960,7 @@ let string_of_tannot tannot =
   | None -> "None"
 
 let replace_typ typ = function
-  | Some (env, _, eff) -> Some (env, typ, eff)
+  | Some ((env, _, eff), _) -> Some ((env, typ, eff), None)
   | None -> None
 
 let infer_lit env (L_aux (lit_aux, l) as lit) =
@@ -2071,17 +2072,17 @@ let destruct_vec_typ l env typ =
 
 
 let env_of_annot (l, tannot) = match tannot with
-  | Some (env, _, _) -> env
+  | Some ((env, _, _),_) -> env
   | None -> raise (Reporting_basic.err_unreachable l "no type annotation")
 
 let env_of (E_aux (_, (l, tannot))) = env_of_annot (l, tannot)
 
 let typ_of_annot (l, tannot) = match tannot with
-  | Some (_, typ, _) -> typ
+  | Some ((_, typ, _), _) -> typ
   | None -> raise (Reporting_basic.err_unreachable l "no type annotation")
 
 let env_of_annot (l, tannot) = match tannot with
-  | Some (env, _, _) -> env
+  | Some ((env, _, _), _) -> env
   | None -> raise (Reporting_basic.err_unreachable l "no type annotation")
 
 let typ_of (E_aux (_, (l, tannot))) = typ_of_annot (l, tannot)
@@ -2108,6 +2109,10 @@ let lexp_typ_of (LEXP_aux (_, (l, tannot))) = typ_of_annot (l, tannot)
 
 let lexp_env_of (LEXP_aux (_, (l, tannot))) = env_of_annot (l, tannot)
 
+let expected_typ_of (l, tannot) = match tannot with
+  | Some ((_, _, _), exp_typ) -> exp_typ
+  | None -> raise (Reporting_basic.err_unreachable l "no type annotation")
+
 (* Flow typing *)
 
 let rec big_int_of_nexp (Nexp_aux (nexp, _)) = match nexp with
@@ -2260,9 +2265,9 @@ let fresh_var =
             mk_id ("v#" ^ string_of_int n)
 
 let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ_aux, _) as typ) : tannot exp =
-  let annot_exp_effect exp typ eff = E_aux (exp, (l, Some (env, Env.expand_synonyms env typ, eff))) in
+  let annot_exp_effect exp typ' eff = E_aux (exp, (l, Some ((env, Env.expand_synonyms env typ', eff),Some typ))) in
   let add_effect exp eff = match exp with
-    | (E_aux (exp, (l, Some (env, typ, _)))) -> E_aux (exp, (l, Some (env, typ, eff)))
+    | (E_aux (exp, (l, Some ((env, typ, _), otyp)))) -> E_aux (exp, (l, Some ((env, typ, eff),otyp)))
     | _ -> failwith "Tried to add effect to unannoted expression"
   in
   let annot_exp exp typ = annot_exp_effect exp typ no_effect in
@@ -2270,8 +2275,8 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
   | E_block exps, _ ->
      begin
        let rec check_block l env exps typ =
-         let annot_exp_effect exp typ eff = E_aux (exp, (l, Some (env, typ, eff))) in
-         let annot_exp exp typ = annot_exp_effect exp typ no_effect in
+         let annot_exp_effect exp typ eff exp_typ = E_aux (exp, (l, Some ((env, typ, eff), exp_typ))) in
+         let annot_exp exp typ exp_typ = annot_exp_effect exp typ no_effect exp_typ in
          match exps with
          | [] -> typ_equality l env typ unit_typ; []
          | [exp] -> [crule check_exp env exp typ]
@@ -2287,7 +2292,7 @@ let rec check_exp env (E_aux (exp_aux, (l, ())) as exp : unit exp) (Typ_aux (typ
                  Env.add_constraint nc env
               | None -> env
             in
-            let texp = annot_exp_effect (E_assert (constr_exp, checked_msg)) unit_typ (mk_effect [BE_escape]) in
+            let texp = annot_exp_effect (E_assert (constr_exp, checked_msg)) unit_typ (mk_effect [BE_escape]) (Some unit_typ) in
             texp :: check_block l env exps typ
          | (exp :: exps) ->
             let texp = crule check_exp env exp (mk_typ (Typ_id (mk_id "unit"))) in
@@ -2538,9 +2543,9 @@ and check_mpexp other_env env mpexp typ =
    or throws a type error if the coercion cannot be performed. *)
 and type_coercion env (E_aux (_, (l, _)) as annotated_exp) typ =
   let strip exp_aux = strip_exp (E_aux (exp_aux, (Parse_ast.Unknown, None))) in
-  let annot_exp exp typ = E_aux (exp, (l, Some (env, typ, no_effect))) in
-  let switch_typ exp typ = match exp with
-    | (E_aux (exp, (l, Some (env, _, eff)))) -> E_aux (exp, (l, Some (env, typ, eff)))
+  let annot_exp exp typ' = E_aux (exp, (l, Some ((env, typ', no_effect), Some typ))) in
+  let switch_exp_typ exp = match exp with
+    | (E_aux (exp, (l, Some ((env, typ', eff), _)))) -> E_aux (exp, (l, Some ((env, typ', eff), Some typ)))
     | _ -> failwith "Cannot switch type for unannotated function"
   in
   let rec try_casts trigger errs = function
@@ -2557,7 +2562,7 @@ and type_coercion env (E_aux (_, (l, _)) as annotated_exp) typ =
   begin
     try
       typ_debug (lazy ("PERFORMING TYPE COERCION: from " ^ string_of_typ (typ_of annotated_exp) ^ " to " ^ string_of_typ typ));
-      subtyp l env (typ_of annotated_exp) typ; annotated_exp
+      subtyp l env (typ_of annotated_exp) typ; switch_exp_typ annotated_exp
     with
     | Type_error (_, trigger) when Env.allow_casts env ->
        let casts = filter_casts env (typ_of annotated_exp) typ (Env.get_casts env) in
@@ -2572,7 +2577,7 @@ and type_coercion env (E_aux (_, (l, _)) as annotated_exp) typ =
    throws a unification error *)
 and type_coercion_unify env (E_aux (_, (l, _)) as annotated_exp) typ =
   let strip exp_aux = strip_exp (E_aux (exp_aux, (Parse_ast.Unknown, None))) in
-  let annot_exp exp typ = E_aux (exp, (l, Some (env, typ, no_effect))) in
+  let annot_exp exp typ' = E_aux (exp, (l, Some ((env, typ', no_effect), Some typ))) in
   let switch_typ exp typ = match exp with
     | (E_aux (exp, (l, Some (env, _, eff)))) -> E_aux (exp, (l, Some (env, typ, eff)))
     | _ -> failwith "Cannot switch type for unannotated expression"
@@ -2608,9 +2613,9 @@ and bind_pat_no_guard env (P_aux (_,(l,_)) as pat) typ =
 and bind_pat env (P_aux (pat_aux, (l, ())) as pat) (Typ_aux (typ_aux, _) as typ) =
   let (Typ_aux (typ_aux, _) as typ), env = bind_existential l typ env in
   typ_print (lazy ("Binding " ^ string_of_pat pat ^  " to " ^ string_of_typ typ));
-  let annot_pat pat typ = P_aux (pat, (l, Some (env, typ, no_effect))) in
+  let annot_pat pat typ' = P_aux (pat, (l, Some ((env, typ', no_effect), Some typ))) in
   let switch_typ pat typ = match pat with
-    | P_aux (pat_aux, (l, Some (env, _, eff))) -> P_aux (pat_aux, (l, Some (env, typ, eff)))
+    | P_aux (pat_aux, (l, Some ((env, _, eff), exp_typ))) -> P_aux (pat_aux, (l, Some ((env, typ, eff), exp_typ)))
     | _ -> typ_error l "Cannot switch type for unannotated pattern"
   in
   let bind_tuple_pat (tpats, env, guards) pat typ =
@@ -2810,7 +2815,7 @@ and bind_pat env (P_aux (pat_aux, (l, ())) as pat) (Typ_aux (typ_aux, _) as typ)
         | _ -> raise typ_exn
 
 and infer_pat env (P_aux (pat_aux, (l, ())) as pat) =
-  let annot_pat pat typ = P_aux (pat, (l, Some (env, typ, no_effect))) in
+  let annot_pat pat typ = P_aux (pat, (l, Some ((env, typ, no_effect), None))) in
   match pat_aux with
   | P_id v ->
      begin
@@ -2917,8 +2922,8 @@ and bind_typ_pat_arg env (TP_aux (typ_pat_aux, l) as typ_pat) (Typ_arg_aux (typ_
   | _, _ -> typ_error l ("Couldn't bind type argument " ^ string_of_typ_arg typ_arg ^ " with " ^ string_of_typ_pat typ_pat)
 
 and bind_assignment env (LEXP_aux (lexp_aux, _) as lexp) (E_aux (_, (l, ())) as exp) =
-  let annot_assign lexp exp = E_aux (E_assign (lexp, exp), (l, Some (env, mk_typ (Typ_id (mk_id "unit")), no_effect))) in
-  let annot_lexp_effect lexp typ eff = LEXP_aux (lexp, (l, Some (env, typ, eff))) in
+  let annot_assign lexp exp = E_aux (E_assign (lexp, exp), (l, Some ((env, mk_typ (Typ_id (mk_id "unit")), no_effect), None))) in
+  let annot_lexp_effect lexp typ eff = LEXP_aux (lexp, (l, Some ((env, typ, eff), None))) in
   let annot_lexp lexp typ = annot_lexp_effect lexp typ no_effect in
   let has_typ v env =
     match Env.lookup_id v env with
@@ -2987,7 +2992,7 @@ and bind_assignment env (LEXP_aux (lexp_aux, _) as lexp) (E_aux (_, (l, ())) as
 
 and bind_lexp env (LEXP_aux (lexp_aux, (l, ())) as lexp) typ =
   typ_print (lazy ("Binding mutable " ^ string_of_lexp lexp ^  " to " ^ string_of_typ typ));
-  let annot_lexp_effect lexp typ eff = LEXP_aux (lexp, (l, Some (env, typ, eff))) in
+  let annot_lexp_effect lexp typ eff = LEXP_aux (lexp, (l, Some ((env, typ, eff),None))) in
   let annot_lexp lexp typ = annot_lexp_effect lexp typ no_effect in
   match lexp_aux with
   | LEXP_cast (typ_annot, v) ->
@@ -3044,7 +3049,7 @@ and bind_lexp env (LEXP_aux (lexp_aux, (l, ())) as lexp) typ =
      inferred_lexp, env
 
 and infer_lexp env (LEXP_aux (lexp_aux, (l, ())) as lexp) =
-  let annot_lexp_effect lexp typ eff = LEXP_aux (lexp, (l, Some (env, typ, eff))) in
+  let annot_lexp_effect lexp typ eff = LEXP_aux (lexp, (l, Some ((env, typ, eff), None))) in
   let annot_lexp lexp typ = annot_lexp_effect lexp typ no_effect in
   match lexp_aux with
   | LEXP_id v ->
@@ -3128,7 +3133,7 @@ and infer_lexp env (LEXP_aux (lexp_aux, (l, ())) as lexp) =
   | _ -> typ_error l ("Unhandled l-expression " ^ string_of_lexp lexp)
 
 and infer_exp env (E_aux (exp_aux, (l, ())) as exp) =
-  let annot_exp_effect exp typ eff = E_aux (exp, (l, Some (env, typ, eff))) in
+  let annot_exp_effect exp typ eff = E_aux (exp, (l, Some ((env, typ, eff),None))) in
   let annot_exp exp typ = annot_exp_effect exp typ no_effect in
   match exp_aux with
   | E_nondet exps ->
@@ -3325,7 +3330,7 @@ and instantiation_of_without_type (E_aux (exp_aux, (l, _)) as exp) =
   | _ -> invalid_arg ("instantiation_of expected application,  got " ^ string_of_exp exp)
 
 and infer_funapp' l env f (typq, f_typ) xs ret_ctx_typ =
-  let annot_exp exp typ eff = E_aux (exp, (l, Some (env, typ, eff))) in
+  let annot_exp exp typ eff = E_aux (exp, (l, Some ((env, typ, eff), ret_ctx_typ))) in
   let switch_annot env typ = function
     | (E_aux (exp, (l, Some (_, _, eff)))) -> E_aux (exp, (l, Some (env, typ, eff)))
     | _ -> failwith "Cannot switch annot for unannotated function"
@@ -3486,9 +3491,9 @@ and infer_funapp' l env f (typq, f_typ) xs ret_ctx_typ =
 and bind_mpat allow_unknown other_env env (MP_aux (mpat_aux, (l, ())) as mpat) (Typ_aux (typ_aux, _) as typ) =
   let (Typ_aux (typ_aux, _) as typ), env = bind_existential l typ env in
   typ_print (lazy ("Binding " ^ string_of_mpat mpat ^  " to " ^ string_of_typ typ));
-  let annot_mpat mpat typ = MP_aux (mpat, (l, Some (env, typ, no_effect))) in
+  let annot_mpat mpat typ' = MP_aux (mpat, (l, Some ((env, typ', no_effect), Some typ))) in
   let switch_typ mpat typ = match mpat with
-    | MP_aux (pat_aux, (l, Some (env, _, eff))) -> MP_aux (pat_aux, (l, Some (env, typ, eff)))
+    | MP_aux (pat_aux, (l, Some ((env, _, eff), exp_typ))) -> MP_aux (pat_aux, (l, Some ((env, typ, eff), exp_typ)))
     | _ -> typ_error l "Cannot switch type for unannotated mapping-pattern"
   in
   let bind_tuple_mpat (tpats, env, guards) mpat typ =
@@ -3672,7 +3677,7 @@ and bind_mpat allow_unknown other_env env (MP_aux (mpat_aux, (l, ())) as mpat) (
            typed_mpat, env, guard::guards
         | _ -> raise typ_exn
 and infer_mpat allow_unknown other_env env (MP_aux (mpat_aux, (l, ())) as mpat) =
-  let annot_mpat mpat typ = MP_aux (mpat, (l, Some (env, typ, no_effect))) in
+  let annot_mpat mpat typ = MP_aux (mpat, (l, Some ((env, typ, no_effect), None))) in
   match mpat_aux with
   | MP_id v ->
      begin
@@ -3769,13 +3774,13 @@ and infer_mpat allow_unknown other_env env (MP_aux (mpat_aux, (l, ())) as mpat)
 (**************************************************************************)
 
 let effect_of_annot = function
-| Some (_, _, eff) -> eff
+| Some ((_, _, eff), _) -> eff
 | None -> no_effect
 
 let effect_of (E_aux (exp, (l, annot))) = effect_of_annot annot
 
 let add_effect_annot annot eff = match annot with
-  | Some (env, typ, eff') -> Some (env, typ, union_effects eff eff')
+  | Some ((env, typ, eff'), exp_typ) -> Some ((env, typ, union_effects eff eff'), exp_typ)
   | None -> None
 
 let add_effect (E_aux (exp, (l, annot))) eff =
@@ -3950,8 +3955,8 @@ and propagate_pexp_effect = function
        let p_exp = propagate_exp_effect exp in
        let p_eff = union_effects (effect_of_pat p_pat) (effect_of p_exp) in
        match annot with
-       | Some (typq, typ, eff) ->
-          Pat_aux (Pat_exp (p_pat, p_exp), (l, Some (typq, typ, union_effects eff p_eff))),
+       | Some ((typq, typ, eff), exp_typ) ->
+          Pat_aux (Pat_exp (p_pat, p_exp), (l, Some ((typq, typ, union_effects eff p_eff), exp_typ))),
          union_effects eff p_eff
        | None -> Pat_aux (Pat_exp (p_pat, p_exp), (l, None)), p_eff
      end
@@ -3964,8 +3969,8 @@ and propagate_pexp_effect = function
                                           (union_effects (effect_of p_guard) (effect_of p_exp))
        in
        match annot with
-       | Some (typq, typ, eff) ->
-          Pat_aux (Pat_when (p_pat, p_guard, p_exp), (l, Some (typq, typ, union_effects eff p_eff))),
+       | Some ((typq, typ, eff), exp_typ) ->
+          Pat_aux (Pat_when (p_pat, p_guard, p_exp), (l, Some ((typq, typ, union_effects eff p_eff), exp_typ))),
           union_effects eff p_eff
        | None -> Pat_aux (Pat_when (p_pat, p_guard, p_exp), (l, None)), p_eff
      end
@@ -3976,8 +3981,8 @@ and propagate_mpexp_effect = function
        let p_mpat = propagate_mpat_effect mpat in
        let p_eff = effect_of_mpat p_mpat in
        match annot with
-       | Some (typq, typ, eff) ->
-          MPat_aux (MPat_pat p_mpat, (l, Some (typq, typ, union_effects eff p_eff))),
+       | Some ((typq, typ, eff), exp_typ) ->
+          MPat_aux (MPat_pat p_mpat, (l, Some ((typq, typ, union_effects eff p_eff), exp_typ))),
          union_effects eff p_eff
        | None -> MPat_aux (MPat_pat p_mpat, (l, None)), p_eff
      end
@@ -3988,8 +3993,8 @@ and propagate_mpexp_effect = function
        let p_eff = union_effects (effect_of_mpat p_mpat) (effect_of p_guard)
        in
        match annot with
-       | Some (typq, typ, eff) ->
-          MPat_aux (MPat_when (p_mpat, p_guard), (l, Some (typq, typ, union_effects eff p_eff))),
+       | Some ((typq, typ, eff), exp_typ) ->
+          MPat_aux (MPat_when (p_mpat, p_guard), (l, Some ((typq, typ, union_effects eff p_eff), exp_typ))),
           union_effects eff p_eff
        | None -> MPat_aux (MPat_when (p_mpat, p_guard), (l, None)), p_eff
      end
@@ -4081,7 +4086,7 @@ and propagate_mpat_effect_aux = function
 and propagate_letbind_effect (LB_aux (lb, (l, annot))) =
   let p_lb, eff = propagate_letbind_effect_aux lb in
   match annot with
-  | Some (typq, typ, eff) -> LB_aux (p_lb, (l, Some (typq, typ, eff))), eff
+  | Some ((typq, typ, eff), exp_typ) -> LB_aux (p_lb, (l, Some ((typq, typ, eff), exp_typ))), eff
   | None -> LB_aux (p_lb, (l, None)), eff
 and propagate_letbind_effect_aux = function
   | LB_val (pat, exp) ->
@@ -4163,7 +4168,7 @@ let check_funcl env (FCL_aux (FCL_Funcl (id, pexp), (l, _))) typ =
          else env
        in
        let typed_pexp, prop_eff = propagate_pexp_effect (check_case env typ_arg (strip_pexp pexp) typ_ret) in
-       FCL_aux (FCL_Funcl (id, typed_pexp), (l, Some (env, typ, prop_eff)))
+       FCL_aux (FCL_Funcl (id, typed_pexp), (l, Some ((env, typ, prop_eff), Some typ)))
      end
   | _ -> typ_error l ("Function clause must have function type: " ^ string_of_typ typ ^ " is not a function type")
 
@@ -4181,20 +4186,20 @@ let check_mapcl : 'a. Env.t -> 'a mapcl -> typ -> tannot mapcl =
 
              let typed_mpexp1, prop_eff1 = propagate_mpexp_effect (check_mpexp right_id_env env (strip_mpexp mpexp1) typ1) in
              let typed_mpexp2, prop_eff2 = propagate_mpexp_effect (check_mpexp left_id_env env (strip_mpexp mpexp2) typ2) in
-             MCL_aux (MCL_mapcl (typed_mpexp1, typed_mpexp2), (l, Some (env, typ, union_effects prop_eff1 prop_eff2)))
+             MCL_aux (MCL_mapcl (typed_mpexp1, typed_mpexp2), (l, Some ((env, typ, union_effects prop_eff1 prop_eff2), Some typ)))
            end
         | _ -> typ_error l ("Function clause must have function type: " ^ string_of_typ typ ^ " is not a function type")
 
 
 let funcl_effect (FCL_aux (FCL_Funcl (id, typed_pexp), (l, annot))) =
   match annot with
-  | Some (_, _, eff) -> eff
+  | Some ((_, _, eff), _) -> eff
   | None -> no_effect (* Maybe could be assert false. This should never happen *)
 
 
 let mapcl_effect (MCL_aux (MCL_mapcl _, (l, annot))) =
   match annot with
-  | Some (_, _, eff) -> eff
+  | Some ((_, _, eff), _) -> eff
   | None -> no_effect (* Maybe could be assert false. This should never happen *)
 
 
@@ -4225,7 +4230,7 @@ let mk_val_spec env typq typ id =
     | Typ_aux (Typ_fn (_,_,eff),_) -> eff
     | _ -> no_effect
   in
-  DEF_spec (VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (typq, typ), Parse_ast.Unknown), id, (fun _ -> None), false), (Parse_ast.Unknown, Some (env,typ,eff))))
+  DEF_spec (VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (typq, typ), Parse_ast.Unknown), id, (fun _ -> None), false), (Parse_ast.Unknown, Some ((env,typ,eff),None))))
 
 let check_tannotopt env typq ret_typ = function
   | Typ_annot_opt_aux (Typ_annot_opt_none, _) -> ()
@@ -4319,7 +4324,7 @@ let check_mapdef env (MD_aux (MD_mapping (id, tannot_opt, mapcls), (l, _)) as md
    context. We have to destructure the various kinds of val specs, but
    the difference is irrelevant for the typechecker. *)
 let check_val_spec env (VS_aux (vs, (l, _))) =
-  let annotate vs typ eff = DEF_spec (VS_aux (vs, (l, Some (env,typ,eff)))) in
+  let annotate vs typ eff = DEF_spec (VS_aux (vs, (l, Some ((env,typ,eff), None)))) in
   let (id, quants, typ, env) = match vs with
     | VS_val_spec (TypSchm_aux (TypSchm_ts (quants, typ), _) as typschm, id, ext_opt, is_cast) ->
        typ_debug (lazy ("VS typschm: " ^ string_of_id id ^ ", " ^ string_of_typschm typschm));
@@ -4471,11 +4476,11 @@ and check_def : 'a. Env.t -> 'a def -> (tannot def) list * Env.t =
   | DEF_overload (id, ids) -> [DEF_overload (id, ids)], Env.add_overloads id ids env
   | DEF_reg_dec (DEC_aux (DEC_reg (typ, id), (l, _))) ->
      let env = Env.add_register id (mk_effect [BE_rreg]) (mk_effect [BE_wreg]) typ env in
-     [DEF_reg_dec (DEC_aux (DEC_reg (typ, id), (l, Some (env, typ, no_effect))))], env
+     [DEF_reg_dec (DEC_aux (DEC_reg (typ, id), (l, Some ((env, typ, no_effect), Some typ))))], env
   | DEF_reg_dec (DEC_aux (DEC_config (id, typ, exp), (l, _))) ->
      let checked_exp = crule check_exp env (strip_exp exp) typ in
      let env = Env.add_register id no_effect (mk_effect [BE_config]) typ env in
-     [DEF_reg_dec (DEC_aux (DEC_config (id, typ, checked_exp), (l, Some (env, typ, no_effect))))], env
+     [DEF_reg_dec (DEC_aux (DEC_config (id, typ, checked_exp), (l, Some ((env, typ, no_effect), Some typ))))], env
   | DEF_reg_dec (DEC_aux (DEC_alias (id, aspec), (l, annot))) -> cd_err ()
   | DEF_reg_dec (DEC_aux (DEC_typ_alias (typ, id, aspec), (l, tannot))) -> cd_err ()
   | DEF_scattered _ -> raise (Reporting_basic.err_unreachable Parse_ast.Unknown "Scattered given to type checker")
diff --git a/src/type_check.mli b/src/type_check.mli
index f1886f50..507a85b0 100644
--- a/src/type_check.mli
+++ b/src/type_check.mli
@@ -332,6 +332,10 @@ val effect_of_pat : tannot pat -> effect
 val effect_of_annot : tannot -> effect
 val add_effect_annot : tannot -> effect -> tannot
 
+(* Returns the type that an expression was checked against, if any.
+   Note that these may be removed if an expression is rewritten. *)
+val expected_typ_of : Ast.l * tannot -> typ option
+
 (** {2 Utilities } *)
 
 val destruct_atom_nexp : Env.t -> typ -> nexp option