Remember the location of return statements

Author: vbgl

Diff for: compiler/src/compile.ml

@@ -196,8 +196,8 @@ let compile (type reg regx xreg rflag cond asm_op extra_op)
         let (_, va) = Hv.find harrs (L.unloc x) in
         List.init (Array.length va) (fun _ -> [])
       with Not_found -> [a] in
-    let f_outannot = List.flatten (List.map2 do_outannot fd.f_ret fd.f_outannot) in
-    let finfo = fd.f_loc, fd.f_annot, f_cc, f_outannot in
+    let ret_annot = List.flatten (List.map2 do_outannot fd.f_ret fd.f_ret_info.ret_annot) in
+    let finfo = fd.f_loc, fd.f_annot, f_cc, { fd.f_ret_info with ret_annot } in
     { Array_expansion.vars; arrs = !arrs; finfo }
   in
 

Diff for: compiler/src/conv.ml

@@ -228,7 +228,7 @@ and stmt_of_cstmt c =
 (* ------------------------------------------------------------------------ *)
 let cufdef_of_fdef fd =
   let fn = fd.f_name in
-  let f_info = fd.f_loc, fd.f_annot, fd.f_cc, fd.f_outannot in
+  let f_info = fd.f_loc, fd.f_annot, fd.f_cc, fd.f_ret_info in
   let f_params =
     List.map (fun x -> cvari_of_vari (L.mk_loc L._dummy x)) fd.f_args in
   let f_body = cstmt_of_stmt fd.f_body in
@@ -244,7 +244,7 @@ let cufdef_of_fdef fd =
 
 
 let fdef_of_cufdef (fn, fd) =
-  let f_loc, f_annot, f_cc, f_outannot = fd.C.f_info in
+  let f_loc, f_annot, f_cc, f_ret_info = fd.C.f_info in
   { f_loc;
     f_annot;
     f_cc;
@@ -254,7 +254,7 @@ let fdef_of_cufdef (fn, fd) =
     f_args = List.map (fun v -> L.unloc (vari_of_cvari v)) fd.C.f_params;
     f_body = stmt_of_cstmt fd.C.f_body;
     f_tyout = List.map ty_of_cty fd.C.f_tyout;
-    f_outannot; 
+    f_ret_info;
     f_ret  = List.map (vari_of_cvari) fd.C.f_res;
   }
 

Diff for: compiler/src/ct_checker_forward.ml

@@ -415,7 +415,7 @@ let get_annot ensure_annot f =
   in
   let ain  = List.mapi process_argument f.f_args in
   let ainlevels = List.map (fun (_, x) -> x) ain in
-  let aout = List.mapi process_result f.f_outannot in
+  let aout = List.mapi process_result f.f_ret_info.ret_annot in
 
   let check_defined msg l =
     if List.exists (fun a -> a = None) l then

Diff for: compiler/src/fInfo.ml

@@ -42,5 +42,16 @@ let is_export = function
   | Export _ -> true
   | _ -> false
 
+(* ------------------------------------------------------------------------ *)
+type return_info = {
+    ret_annot : Annotations.annotations list;
+    (* annotation attached to return type *)
+    ret_loc : Location.t; (* location of the return statement *)
+  }
+
+(* ------------------------------------------------------------------------ *)
 type t =
-  Location.t * f_annot * call_conv * Annotations.annotations list
+  Location.t * f_annot * call_conv * return_info
+
+let ret_info (fi: t) : IInfo.t =
+  let (_, _, _, ri) = fi in (Location.i_loc0 ri.ret_loc, [])

Diff for: compiler/src/latex_printer.ml

@@ -387,7 +387,7 @@ let pp_funbody fmt { pdb_instr ; pdb_ret } =
         indent 1
         kw "return"
         (pp_list ", " pp_var) ret;
-  ) fmt pdb_ret
+  ) fmt (L.unloc pdb_ret)
 
 let pp_fundef fmt { pdf_cc ; pdf_name ; pdf_args ; pdf_rty ; pdf_body ; pdf_annot } =
   F.fprintf

Diff for: compiler/src/parser.mly

@@ -445,7 +445,7 @@ annot_pparamdecl:
 pfunbody :
 | LBRACE
     is = pinstr*
-    rt = option(RETURN vs=tuple(var) SEMICOLON { vs })
+    rt = loc(option(RETURN vs=tuple(var) SEMICOLON { vs }))
   RBRACE
     { { pdb_instr = is;
         pdb_ret   = rt; } }

Diff for: compiler/src/pretyping.ml

@@ -1983,10 +1983,11 @@ and tt_cmd arch_info env c =
 (* -------------------------------------------------------------------- *)
 let tt_funbody arch_info env (pb : S.pfunbody) =
   let env, bdy = tt_cmd arch_info env pb.S.pdb_instr in
+  let ret_loc = L.loc pb.pdb_ret in
   let ret =
     let for1 x = L.mk_loc (L.loc x) (tt_var `AllVar env x) in
-    List.map for1 (Option.default [] pb.pdb_ret) in
-  (bdy, ret)
+    List.map for1 (Option.default [] (L.unloc pb.pdb_ret)) in
+  (bdy, ret_loc, ret)
 
 
 (* -------------------------------------------------------------------- *)
@@ -2155,16 +2156,16 @@ let tt_fundef arch_info (env0 : 'asm Env.env) loc (pf : S.pfundef) : 'asm Env.en
   let env = Env.Vars.clear_locals env0 in
   if is_combine_flags pf.pdf_name then
     rs_tyerror ~loc:(L.loc pf.pdf_name) (string_error "invalid function name");
-  let inret = Option.map_default (List.map L.unloc) [] pf.pdf_body.pdb_ret in
+  let inret = Option.map_default (List.map L.unloc) [] (L.unloc pf.pdf_body.pdb_ret) in
   let dfl_mut x = List.mem x inret in
 
   let envb, args =
     let env, args = List.map_fold (tt_annot_paramdecls dfl_mut arch_info.pd) env pf.pdf_args in
     let env = add_known_implicits arch_info env pf.pdf_body.pdb_instr in
     env, List.flatten args in
   let rty  = Option.map_default (List.map (tt_type arch_info.pd env |- snd |- snd)) [] pf.pdf_rty in
-  let oannot = Option.map_default (List.map fst) [] pf.pdf_rty in
-  let body, xret = tt_funbody arch_info envb pf.pdf_body in
+  let ret_annot = Option.map_default (List.map fst) [] pf.pdf_rty in
+  let body, ret_loc, xret = tt_funbody arch_info envb pf.pdf_body in
   let f_cc = tt_call_conv loc args xret pf.pdf_cc in
   let args = List.map L.unloc args in
   let name = L.unloc pf.pdf_name in
@@ -2178,7 +2179,7 @@ let tt_fundef arch_info (env0 : 'asm Env.env) loc (pf : S.pfundef) : 'asm Env.en
       P.f_args  = args;
       P.f_body  = body;
       P.f_tyout = rty;
-      P.f_outannot = oannot;
+      P.f_ret_info = { ret_annot; ret_loc };
       P.f_ret   = xret; } in
 
   check_return_statement ~loc fdef.P.f_name rty

Diff for: compiler/src/prog.ml

@@ -111,6 +111,12 @@ and ('len,'info,'asm) ginstr = {
 
 and ('len,'info,'asm) gstmt = ('len,'info,'asm) ginstr list
 
+(* ------------------------------------------------------------------------ *)
+type return_info = {
+    ret_annot : Annotations.annotations list; (* annotation attached to return type *)
+    ret_loc : Location.t; (* location of the return statement *)
+  }
+
 (* ------------------------------------------------------------------------ *)
 type ('len,'info,'asm) gfunc = {
     f_loc  : L.t;
@@ -122,7 +128,7 @@ type ('len,'info,'asm) gfunc = {
     f_args : 'len gvar list;
     f_body : ('len,'info,'asm) gstmt;
     f_tyout : 'len gty list;
-    f_outannot : Annotations.annotations list; (* annotation attach to return type *)
+    f_ret_info : FInfo.return_info;
     f_ret  : 'len gvar_i list
   }
 

Diff for: compiler/src/prog.mli

@@ -86,7 +86,7 @@ type ('len,'info,'asm) gfunc = {
     f_args : 'len gvar list;
     f_body : ('len,'info,'asm) gstmt;
     f_tyout : 'len gty list;
-    f_outannot : Annotations.annotations list; (* annotation attach to return type *)
+    f_ret_info : FInfo.return_info;
     f_ret  : 'len gvar_i list
   }
 

Diff for: compiler/src/sct_checker_forward.ml

@@ -1309,7 +1309,7 @@ let init_constraint fenv f =
     mk_vty ~is_local:false loc ~msf:(not export) x ls an in
 
   (* process function outputs *)
-  let tyout = List.map2i process_return f.f_ret f.f_outannot in
+  let tyout = List.map2i process_return f.f_ret f.f_ret_info.ret_annot in
 
   (* infer msf_oracle info *)
   let msfs =
@@ -1511,7 +1511,7 @@ let compile_infer_msf (prog:('info, 'asm) prog) =
     in
 
     (* process function outputs *)
-    let tyout = List.mapi process_return f.f_outannot in
+    let tyout = List.mapi process_return f.f_ret_info.ret_annot in
 
     (* infer the set of input variables that need to be msf *)
     let msfin =

Diff for: compiler/src/syntax.ml

@@ -271,7 +271,7 @@ type pparam = {
 (* -------------------------------------------------------------------- *)
 type pfunbody = {
   pdb_instr : pinstr list;
-  pdb_ret   : pident list option;
+  pdb_ret   : pident list option L.located;
 }
 
 (* -------------------------------------------------------------------- *)