Update val specs after rewriting functions

3 files changed, 108 insertions, 76 deletions

diff --git a/src/rewriter.ml b/src/rewriter.ml
index bcfb731a..4dcc3404 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -1049,7 +1049,7 @@ let rewrite_trivial_sizeof, rewrite_trivial_sizeof_exp =
                        |> List.concat
             in
             match exps with
-            | (exp :: _) -> exp
+            | (exp :: _) -> check_exp env (strip_exp exp) (typ_of exp)
             | [] when split_sizeof ->
                fold_exp (rewrite_e_sizeof false) (check_exp env (split_nexp nexp) (typ_of orig_exp))
             | [] -> orig_exp
@@ -1318,8 +1318,8 @@ let rewrite_sizeof (Defs defs) =
                  | _ -> Typ_aux (Typ_tup (kid_typs @ [vtyp_arg]), vl)
                end in
              Typ_aux (Typ_fn (vtyp_arg', vtyp_ret, declared_eff), vl)
-          | _ -> raise (Reporting_basic.err_typ l
-                                                "val spec with non-function type") in
+          | _ ->
+            raise (Reporting_basic.err_typ l "val spec with non-function type") in
         TypSchm_aux (TypSchm_ts (tq, typ'), l)
       else ts in
     match def with
@@ -2401,40 +2401,58 @@ let rewrite_defs_early_return =
 
 (* Propagate effects of functions, if effect checking and propagation
    have not been performed already by the type checker. *)
-let rewrite_fix_fun_effs (Defs defs) =
-  let e_aux fun_effs (exp, (l, annot)) =
+let rewrite_fix_val_specs (Defs defs) =
+  let find_vs env val_specs id =
+    try Bindings.find id val_specs with
+    | Not_found ->
+       begin
+         try Env.get_val_spec id env with
+         | _ ->
+            raise (Reporting_basic.err_unreachable (Parse_ast.Unknown)
+              ("No val spec found for " ^ string_of_id id))
+       end
+  in
+
+  let add_eff_to_vs eff = function
+    | (tq, Typ_aux (Typ_fn (args_t, ret_t, eff'), a)) ->
+      (tq, Typ_aux (Typ_fn (args_t, ret_t, union_effects eff eff'), a))
+    | vs -> vs
+  in
+
+  let eff_of_vs = function
+    | (tq, Typ_aux (Typ_fn (args_t, ret_t, eff), a)) -> eff
+    | _ -> no_effect
+  in
+
+  let e_aux val_specs (exp, (l, annot)) =
     match fix_eff_exp (E_aux (exp, (l, annot))) with
     | E_aux (E_app_infix (_, f, _) as exp, (l, Some (env, typ, eff)))
-    | E_aux (E_app (f, _) as exp, (l, Some (env, typ, eff)))
-      when Bindings.mem f fun_effs ->
-      let eff' = Bindings.find f fun_effs in
-      let env =
-        try
-          match Env.get_val_spec f env with
-          | (tq, Typ_aux (Typ_fn (args_t, ret_t, eff), a)) ->
-            Env.update_val_spec f (tq, Typ_aux (Typ_fn (args_t, ret_t, union_effects eff eff'), a)) env
-          | _ -> env
-        with
-        | _ -> env in
-      E_aux (exp, (l, Some (env, typ, union_effects eff eff')))
-    | e_aux -> e_aux in
-
-  let rewrite_exp fun_effs = fold_exp { id_exp_alg with e_aux = e_aux fun_effs } in
-
-  let rewrite_funcl (fun_effs, funcls) (FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))) =
-    let exp = propagate_exp_effect (rewrite_exp fun_effs exp) in
-    let fun_eff =
-      try union_effects (effect_of exp) (Bindings.find id fun_effs)
-      with Not_found -> (effect_of exp) in
-    let annot =
-      match annot with
-      | Some (env, typ, eff) -> Some (env, typ, union_effects eff fun_eff)
-      | None -> None in
-    (Bindings.add id fun_eff fun_effs,
-     funcls @ [FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))]) in
-
-  let rewrite_fundef (fun_effs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) =
-    let (fun_effs, funcls) = List.fold_left rewrite_funcl (fun_effs, []) funcls in
+    | E_aux (E_app (f, _) as exp, (l, Some (env, typ, eff))) ->
+       let vs = find_vs env val_specs f in
+       let env = Env.update_val_spec f vs env in
+       E_aux (exp, (l, Some (env, typ, union_effects eff (eff_of_vs vs))))
+    | e_aux -> e_aux
+  in
+
+  let rewrite_exp val_specs = fold_exp { id_exp_alg with e_aux = e_aux val_specs } in
+
+  let rewrite_funcl (val_specs, funcls) (FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))) =
+    let exp = propagate_exp_effect (rewrite_exp val_specs exp) in
+    let vs, eff = match find_vs (env_of_annot (l, annot)) val_specs id with
+      | (tq, Typ_aux (Typ_fn (args_t, ret_t, eff), a)) ->
+         let eff' = union_effects eff (effect_of exp) in
+         let args_t' = pat_typ_of pat in
+         (tq, Typ_aux (Typ_fn (args_t', ret_t, eff'), a)), eff'
+    in
+    let annot = add_effect_annot annot eff in
+    (Bindings.add id vs val_specs,
+     funcls @ [FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))])
+  in
+
+  let rewrite_fundef (val_specs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) =
+    let (val_specs, funcls) = List.fold_left rewrite_funcl (val_specs, []) funcls in
+    (* Repeat once to cross-propagate effects between clauses *)
+    let (val_specs, funcls) = List.fold_left rewrite_funcl (val_specs, []) funcls in
     let is_funcl_rec (FCL_aux (FCL_Funcl (id, _, exp), _)) =
       fst (fold_exp
         { (compute_exp_alg false (||) ) with
@@ -2444,39 +2462,54 @@ let rewrite_fix_fun_effs (Defs defs) =
              E_app (f, es)));
           e_app_infix = (fun ((r1,e1), f, (r2,e2)) ->
             (r1 || r2 || (string_of_id f = string_of_id id),
-             E_app_infix (e1, f, e2))) } exp) in
-    let is_rec = List.exists is_funcl_rec funcls in
-    (* Repeat once for recursive functions:
-       propagates union of effects to all clauses *)
-    let recopt, (fun_effs, funcls) =
-      if is_rec then
-        Rec_aux (Rec_rec, Parse_ast.Unknown),
-        List.fold_left rewrite_funcl (fun_effs, []) funcls
-      else recopt, (fun_effs, funcls) in
-    (fun_effs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) in
-
-  let rec rewrite_fundefs (fun_effs, fundefs) =
+             E_app_infix (e1, f, e2))) }
+        exp)
+    in
+    let recopt =
+      if List.exists is_funcl_rec funcls then
+        Rec_aux (Rec_rec, Parse_ast.Unknown)
+      else recopt
+    in
+    (val_specs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) in
+
+  let rec rewrite_fundefs (val_specs, fundefs) =
     match fundefs with
     | fundef :: fundefs ->
-      let (fun_effs, fundef) = rewrite_fundef (fun_effs, fundef) in
-      let (fun_effs, fundefs) = rewrite_fundefs (fun_effs, fundefs) in
-      (fun_effs, fundef :: fundefs)
-    | [] -> (fun_effs, []) in
-
-  let rewrite_def (fun_effs, defs) = function
-  | DEF_fundef fundef ->
-    let (fun_effs, fundef) = rewrite_fundef (fun_effs, fundef) in
-    (fun_effs, defs @ [DEF_fundef fundef])
-  | DEF_internal_mutrec fundefs ->
-    let (fun_effs, fundefs) = rewrite_fundefs (fun_effs, fundefs) in
-    (fun_effs, defs @ [DEF_internal_mutrec fundefs])
-  | DEF_val (LB_aux (LB_val (pat, exp), a)) ->
-    (fun_effs, defs @ [DEF_val (LB_aux (LB_val (pat, rewrite_exp fun_effs exp), a))])
-  | def -> (fun_effs, defs @ [def]) in
+      let (val_specs, fundef) = rewrite_fundef (val_specs, fundef) in
+      let (val_specs, fundefs) = rewrite_fundefs (val_specs, fundefs) in
+      (val_specs, fundef :: fundefs)
+    | [] -> (val_specs, []) in
+
+  let rewrite_def (val_specs, defs) = function
+    | DEF_fundef fundef ->
+      let (val_specs, fundef) = rewrite_fundef (val_specs, fundef) in
+      (val_specs, defs @ [DEF_fundef fundef])
+    | DEF_internal_mutrec fundefs ->
+      let (val_specs, fundefs) = rewrite_fundefs (val_specs, fundefs) in
+      (val_specs, defs @ [DEF_internal_mutrec fundefs])
+    | DEF_val (LB_aux (LB_val (pat, exp), a)) ->
+      (val_specs, defs @ [DEF_val (LB_aux (LB_val (pat, rewrite_exp val_specs exp), a))])
+    | def -> (val_specs, defs @ [def])
+  in
+
+  let rewrite_val_specs val_specs = function
+    | DEF_spec (VS_aux (VS_val_spec (typschm, id, ext_opt, is_cast), a))
+      when Bindings.mem id val_specs ->
+       let typschm = match typschm with
+         | TypSchm_aux (TypSchm_ts (tq, typ), l) ->
+            let (tq, typ) = Bindings.find id val_specs in
+            TypSchm_aux (TypSchm_ts (tq, typ), l)
+       in
+       DEF_spec (VS_aux (VS_val_spec (typschm, id, ext_opt, is_cast), a))
+    | def -> def
+  in
+
+  let (val_specs, defs) = List.fold_left rewrite_def (Bindings.empty, []) defs in
+  let defs = List.map (rewrite_val_specs val_specs) defs in
 
   (* if !Type_check.opt_no_effects
   then *)
-  Defs (snd (List.fold_left rewrite_def (Bindings.empty, []) defs))
+  Defs defs
   (* else Defs defs *)
 
 (* Turn constraints into numeric expressions with sizeof *)
@@ -3540,7 +3573,7 @@ let rewrite_defs_lem = [
   ("guarded_pats", rewrite_defs_guarded_pats);
   (* ("recheck_defs", recheck_defs); *)
   ("early_return", rewrite_defs_early_return);
-  ("fix_fun_effs", rewrite_fix_fun_effs);
+  ("fix_val_specs", rewrite_fix_val_specs);
   ("exp_lift_assign", rewrite_defs_exp_lift_assign);
   ("remove_blocks", rewrite_defs_remove_blocks);
   ("letbind_effects", rewrite_defs_letbind_effects);
diff --git a/src/type_check.ml b/src/type_check.ml
index 270b2cf4..751e70e5 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -2858,24 +2858,22 @@ let effect_of_annot = function
 
 let effect_of (E_aux (exp, (l, annot))) = effect_of_annot annot
 
-let add_effect (E_aux (exp, (l, annot))) eff1 =
-  match annot with
-  | Some (env, typ, eff2) -> E_aux (exp, (l, Some (env, typ, union_effects eff1 eff2)))
-  | None -> assert false
+let add_effect_annot annot eff = match annot with
+  | Some (env, typ, eff') -> Some (env, typ, union_effects eff eff')
+  | None -> None
+
+let add_effect (E_aux (exp, (l, annot))) eff =
+  E_aux (exp, (l, add_effect_annot annot eff))
 
 let effect_of_lexp (LEXP_aux (exp, (l, annot))) = effect_of_annot annot
 
-let add_effect_lexp (LEXP_aux (lexp, (l, annot))) eff1 =
-  match annot with
-  | Some (env, typ, eff2) -> LEXP_aux (lexp, (l, Some (env, typ, union_effects eff1 eff2)))
-  | None -> assert false
+let add_effect_lexp (LEXP_aux (lexp, (l, annot))) eff =
+  LEXP_aux (lexp, (l, add_effect_annot annot eff))
 
 let effect_of_pat (P_aux (exp, (l, annot))) = effect_of_annot annot
 
-let add_effect_pat (P_aux (pat, (l, annot))) eff1 =
-  match annot with
-  | Some (env, typ, eff2) -> P_aux (pat, (l, Some (env, typ, union_effects eff1 eff2)))
-  | None -> assert false
+let add_effect_pat (P_aux (pat, (l, annot))) eff =
+  P_aux (pat, (l, add_effect_annot annot eff))
 
 let collect_effects xs = List.fold_left union_effects no_effect (List.map effect_of xs)
 
diff --git a/src/type_check.mli b/src/type_check.mli
index ff9eb74e..f3e5e861 100644
--- a/src/type_check.mli
+++ b/src/type_check.mli
@@ -217,6 +217,7 @@ val pat_env_of : tannot pat -> Env.t
 
 val effect_of : tannot exp -> effect
 val effect_of_annot : tannot -> effect
+val add_effect_annot : tannot -> effect -> tannot
 
 val destruct_atom_nexp : Env.t -> typ -> nexp option