diff --git a/analysis/src/CompletionFrontEnd.ml b/analysis/src/CompletionFrontEnd.ml
index 65a26a1467..6622b1aa62 100644
--- a/analysis/src/CompletionFrontEnd.ml
+++ b/analysis/src/CompletionFrontEnd.ml
@@ -277,6 +277,7 @@ let rec exprToContextPathInner ~(inJsxContext : bool) (e : Parsetree.expression)
             {funct = d; args = (Nolabel, lhs) :: args; partial; transformed_jsx};
         pexp_loc;
         pexp_attributes;
+        pexp_is_return = false;
       }
   | Pexp_apply
       ({
@@ -294,11 +295,18 @@ let rec exprToContextPathInner ~(inJsxContext : bool) (e : Parsetree.expression)
           Pexp_apply
             {
               app with
-              funct = {pexp_desc = Pexp_ident id; pexp_loc; pexp_attributes};
+              funct =
+                {
+                  pexp_desc = Pexp_ident id;
+                  pexp_loc;
+                  pexp_attributes;
+                  pexp_is_return = false;
+                };
               args = [(Nolabel, lhs)];
             };
         pexp_loc;
         pexp_attributes;
+        pexp_is_return = false;
       }
   | Pexp_apply {funct = e1; args} -> (
     match exprToContextPath ~inJsxContext e1 with
diff --git a/compiler/frontend/ast_compatible.ml b/compiler/frontend/ast_compatible.ml
index 8d6ac2d427..c092c570fb 100644
--- a/compiler/frontend/ast_compatible.ml
+++ b/compiler/frontend/ast_compatible.ml
@@ -35,6 +35,7 @@ let apply_simple ?(loc = default_loc) ?(attrs = []) (fn : expression)
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc =
       Pexp_apply
         {
@@ -49,6 +50,7 @@ let app1 ?(loc = default_loc) ?(attrs = []) fn arg1 : expression =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc =
       Pexp_apply
         {
@@ -63,6 +65,7 @@ let app2 ?(loc = default_loc) ?(attrs = []) fn arg1 arg2 : expression =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc =
       Pexp_apply
         {
@@ -77,6 +80,7 @@ let app3 ?(loc = default_loc) ?(attrs = []) fn arg1 arg2 arg3 : expression =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc =
       Pexp_apply
         {
@@ -91,6 +95,7 @@ let fun_ ?(loc = default_loc) ?(attrs = []) ?(async = false) ~arity pat exp =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc =
       Pexp_fun
         {
@@ -108,6 +113,7 @@ let const_exp_string ?(loc = default_loc) ?(attrs = []) ?delimiter (s : string)
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc = Pexp_constant (Pconst_string (s, delimiter));
   }
 
@@ -115,6 +121,7 @@ let const_exp_int ?(loc = default_loc) ?(attrs = []) (s : int) : expression =
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc = Pexp_constant (Pconst_integer (string_of_int s, None));
   }
 
@@ -123,6 +130,7 @@ let apply_labels ?(loc = default_loc) ?(attrs = []) fn
   {
     pexp_loc = loc;
     pexp_attributes = attrs;
+    pexp_is_return = false;
     pexp_desc =
       Pexp_apply
         {
diff --git a/compiler/frontend/ast_exp_apply.ml b/compiler/frontend/ast_exp_apply.ml
index 04034fad51..d39f634de7 100644
--- a/compiler/frontend/ast_exp_apply.ml
+++ b/compiler/frontend/ast_exp_apply.ml
@@ -96,6 +96,7 @@ let app_exp_mapper (e : exp) (self : Bs_ast_mapper.mapper) : exp =
             {funct = fn1; args = (Nolabel, a) :: args; partial; transformed_jsx};
         pexp_loc = e.pexp_loc;
         pexp_attributes = e.pexp_attributes @ f.pexp_attributes;
+        pexp_is_return = false;
       }
     | Pexp_tuple xs ->
       bound a (fun bounded_obj_arg ->
@@ -122,11 +123,13 @@ let app_exp_mapper (e : exp) (self : Bs_ast_mapper.mapper) : exp =
                                transformed_jsx;
                              };
                          pexp_attributes = [];
+                         pexp_is_return = false;
                          pexp_loc = fn.pexp_loc;
                        }
                      | _ ->
                        Ast_compatible.app1 ~loc:fn.pexp_loc fn bounded_obj_arg));
             pexp_attributes = f.pexp_attributes;
+            pexp_is_return = false;
             pexp_loc = f.pexp_loc;
           })
     | _ -> Ast_compatible.app1 ~loc ~attrs:e.pexp_attributes f a)
diff --git a/compiler/frontend/ast_external_mk.ml b/compiler/frontend/ast_external_mk.ml
index 3ec65e1612..3f5e15d9cc 100644
--- a/compiler/frontend/ast_external_mk.ml
+++ b/compiler/frontend/ast_external_mk.ml
@@ -54,6 +54,7 @@ let local_external_apply loc ?(pval_attributes = []) ~(pval_prim : string list)
              Pexp_ident
                {txt = Ldot (Lident local_module_name, local_fun_name); loc};
            pexp_attributes = [];
+           pexp_is_return = false;
            pexp_loc = loc;
          }
           : Parsetree.expression)
@@ -90,6 +91,7 @@ let local_external_obj loc ?(pval_attributes = []) ~pval_prim ~pval_type
              Pexp_ident
                {txt = Ldot (Lident local_module_name, local_fun_name); loc};
            pexp_attributes = [];
+           pexp_is_return = false;
            pexp_loc = loc;
          }
           : Parsetree.expression)
@@ -126,5 +128,6 @@ let local_extern_cont_to_obj loc ?(pval_attributes = []) ~pval_prim ~pval_type
             Pexp_ident
               {txt = Ldot (Lident local_module_name, local_fun_name); loc};
           pexp_attributes = [];
+          pexp_is_return = false;
           pexp_loc = loc;
         } )
diff --git a/compiler/frontend/ast_open_cxt.ml b/compiler/frontend/ast_open_cxt.ml
index e9ebd7747b..abfb3c248f 100644
--- a/compiler/frontend/ast_open_cxt.ml
+++ b/compiler/frontend/ast_open_cxt.ml
@@ -56,6 +56,7 @@ let restore_exp (xs : Parsetree.expression) (qualifiers : t) :
         ({
            pexp_desc = Pexp_open (flag, lid, x);
            pexp_attributes = attrs;
+           pexp_is_return = false;
            pexp_loc = loc;
          }
           : Parsetree.expression))
diff --git a/compiler/ml/ast_helper.ml b/compiler/ml/ast_helper.ml
index 2fae640eb0..7300a9c324 100644
--- a/compiler/ml/ast_helper.ml
+++ b/compiler/ml/ast_helper.ml
@@ -154,7 +154,12 @@ end
 
 module Exp = struct
   let mk ?(loc = !default_loc) ?(attrs = []) d =
-    {pexp_desc = d; pexp_loc = loc; pexp_attributes = attrs}
+    {
+      pexp_desc = d;
+      pexp_loc = loc;
+      pexp_attributes = attrs;
+      pexp_is_return = false;
+    }
   let attr d a = {d with pexp_attributes = d.pexp_attributes @ [a]}
 
   let ident ?loc ?attrs a = mk ?loc ?attrs (Pexp_ident a)
diff --git a/compiler/ml/parsetree.ml b/compiler/ml/parsetree.ml
index 78c9899f74..1921ea95e0 100644
--- a/compiler/ml/parsetree.ml
+++ b/compiler/ml/parsetree.ml
@@ -212,6 +212,7 @@ and expression = {
   pexp_loc: Location.t;
   (* Hack: made pexp_attributes mutable for use in analysis exe. Please do not use elsewhere! *)
   mutable pexp_attributes: attributes; (* ... [@id1] [@id2] *)
+  mutable pexp_is_return: bool;
 }
 
 and expression_desc =
diff --git a/compiler/ml/typecore.ml b/compiler/ml/typecore.ml
index 90522f13d2..602fc3505f 100644
--- a/compiler/ml/typecore.ml
+++ b/compiler/ml/typecore.ml
@@ -273,7 +273,15 @@ let constant_or_raise env loc cst =
 let type_option ty = newty (Tconstr (Predef.path_option, [ty], ref Mnil))
 
 let mkexp exp_desc exp_type exp_loc exp_env =
-  {exp_desc; exp_type; exp_loc; exp_env; exp_extra = []; exp_attributes = []}
+  {
+    exp_desc;
+    exp_type;
+    exp_loc;
+    exp_env;
+    exp_extra = [];
+    exp_is_return = false;
+    exp_attributes = [];
+  }
 
 let option_none ty loc =
   let lid = Longident.Lident "None" and env = Env.initial_safe_string in
@@ -2273,6 +2281,7 @@ and type_expect ~context ?in_function ?recarg env sexp ty_expected =
     Builtin_attributes.warning_scope sexp.pexp_attributes (fun () ->
         type_expect_ ~context ?in_function ?recarg env sexp ty_expected)
   in
+  if sexp.pexp_is_return then exp.exp_is_return <- true;
   Cmt_format.set_saved_types
     (Cmt_format.Partial_expression exp :: previous_saved_types);
   exp
@@ -2323,6 +2332,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance env desc.val_type;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_constant cst ->
     let cst = constant_or_raise env loc cst in
@@ -2334,6 +2344,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = type_constant cst;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_let
       ( Nonrecursive,
@@ -2372,6 +2383,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = body.exp_type;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_fun
       {
@@ -2457,6 +2469,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
           exp_type = ty_res;
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
+          exp_is_return = false;
         }
     in
 
@@ -2511,6 +2524,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance env ty_expected;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_try (sbody, caselist) ->
     let body = type_expect ~context:None env sbody ty_expected in
@@ -2526,6 +2540,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = body.exp_type;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_tuple sexpl ->
     assert (List.length sexpl >= 2);
@@ -2546,6 +2561,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = newty (Ttuple (List.map (fun e -> e.exp_type) expl));
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_construct (lid, sarg) ->
     type_construct ~context env loc lid sarg ty_expected sexp.pexp_attributes
@@ -2572,6 +2588,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
               exp_type = ty_expected0;
               exp_attributes = sexp.pexp_attributes;
               exp_env = env;
+              exp_is_return = false;
             }
         | _ -> raise Not_found)
       | _ -> raise Not_found
@@ -2599,6 +2616,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
                  });
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
+          exp_is_return = false;
         })
   | Pexp_record (lid_sexp_list, None) ->
     let ty_record, opath, fields, repr_opt =
@@ -2696,6 +2714,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance env ty_expected;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_record (lid_sexp_list, Some sexp) ->
     assert (lid_sexp_list <> []);
@@ -2793,6 +2812,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance env ty_expected;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_field (srecord, lid) ->
     let record, label, _ = type_label_access env srecord lid in
@@ -2806,6 +2826,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = ty_arg;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_setfield (srecord, lid, snewval) ->
     let record, label, opath = type_label_access env srecord lid in
@@ -2827,6 +2848,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance_def Predef.type_unit;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_array sargl ->
     let ty = newgenvar () in
@@ -2845,6 +2867,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance env ty_expected;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_ifthenelse (scond, sifso, sifnot) -> (
     (* TODO(attributes) Unify the attribute handling in the parser and rest of the compiler. *)
@@ -2875,6 +2898,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
           exp_type = ifso.exp_type;
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
+          exp_is_return = false;
         }
     | Some sifnot ->
       let ifso = type_expect ~context:return_context env sifso ty_expected in
@@ -2889,6 +2913,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
           exp_type = ifso.exp_type;
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
+          exp_is_return = false;
         })
   | Pexp_sequence (sexp1, sexp2) ->
     let exp1 = type_statement ~context:None env sexp1 in
@@ -2901,6 +2926,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = exp2.exp_type;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_while (scond, sbody) ->
     let cond =
@@ -2915,6 +2941,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance_def Predef.type_unit;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_for (param, slow, shigh, dir, sbody) ->
     let low =
@@ -2947,6 +2974,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = instance_def Predef.type_unit;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_constraint (sarg, sty) ->
     let separate = true in
@@ -2969,6 +2997,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = ty';
         exp_attributes = arg.exp_attributes;
         exp_env = env;
+        exp_is_return = false;
         exp_extra =
           (Texp_constraint cty, loc, sexp.pexp_attributes) :: arg.exp_extra;
       }
@@ -3025,6 +3054,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = ty';
         exp_attributes = arg.exp_attributes;
         exp_env = env;
+        exp_is_return = false;
         exp_extra =
           (Texp_coerce cty', loc, sexp.pexp_attributes) :: arg.exp_extra;
       }
@@ -3056,6 +3086,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
           exp_type = typ;
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
+          exp_is_return = false;
         }
     with Unify _ ->
       let valid_methods =
@@ -3108,6 +3139,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = ty;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_letexception (cd, sbody) ->
     let cd, newenv = Typedecl.transl_exception env cd in
@@ -3120,6 +3152,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = body.exp_type;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_assert e ->
     let cond =
@@ -3138,6 +3171,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type;
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_newtype ({txt = name}, sbody) ->
     let ty = newvar () in
@@ -3211,6 +3245,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
         exp_type = newty (Tpackage (p, nl, tl'));
         exp_attributes = sexp.pexp_attributes;
         exp_env = env;
+        exp_is_return = false;
       }
   | Pexp_open (ovf, lid, e) ->
     let path, newenv = !type_open ovf env sexp.pexp_loc lid in
@@ -3245,6 +3280,7 @@ and type_expect_ ~context ?in_function ?(recarg = Rejected) env sexp ty_expected
           exp_type = instance_def Predef.type_extension_constructor;
           exp_attributes = sexp.pexp_attributes;
           exp_env = env;
+          exp_is_return = false;
         }
     | _ -> raise (Error (loc, env, Invalid_extension_constructor_payload)))
   | Pexp_extension ext ->
@@ -3322,6 +3358,7 @@ and type_function ?in_function ~arity ~async loc attrs env ty_expected_ l
       exp_type;
       exp_attributes = attrs;
       exp_env = env;
+      exp_is_return = false;
     }
 
 and type_label_access env srecord lid =
@@ -3822,6 +3859,7 @@ and type_construct ~context env loc lid sarg ty_expected attrs =
         exp_type = ty_res;
         exp_attributes = attrs;
         exp_env = env;
+        exp_is_return = false;
       }
   in
   (* Forward context if this is a Some constructor injected (meaning it's
diff --git a/compiler/ml/typedtree.ml b/compiler/ml/typedtree.ml
index 9ef328be4a..985f749f02 100644
--- a/compiler/ml/typedtree.ml
+++ b/compiler/ml/typedtree.ml
@@ -61,6 +61,7 @@ and expression = {
   exp_extra: (exp_extra * Location.t * attribute list) list;
   exp_type: type_expr;
   exp_env: Env.t;
+  mutable exp_is_return: bool;
   exp_attributes: attribute list;
 }
 
diff --git a/compiler/ml/typedtree.mli b/compiler/ml/typedtree.mli
index 6e6b1c5159..8d5459d8fc 100644
--- a/compiler/ml/typedtree.mli
+++ b/compiler/ml/typedtree.mli
@@ -104,6 +104,7 @@ and expression = {
   exp_extra: (exp_extra * Location.t * attributes) list;
   exp_type: type_expr;
   exp_env: Env.t;
+  mutable exp_is_return: bool;
   exp_attributes: attributes;
 }
 
diff --git a/compiler/syntax/src/jsx_v4.ml b/compiler/syntax/src/jsx_v4.ml
index f35e3014ab..11d647430c 100644
--- a/compiler/syntax/src/jsx_v4.ml
+++ b/compiler/syntax/src/jsx_v4.ml
@@ -1159,12 +1159,18 @@ let mk_record_from_props mapper (jsx_expr_loc : Location.t) (props : jsx_props)
   in
   match (record_fields, spread_props) with
   | [], Some spread_props ->
-    {pexp_desc = spread_props.pexp_desc; pexp_loc = loc; pexp_attributes = []}
+    {
+      pexp_desc = spread_props.pexp_desc;
+      pexp_loc = loc;
+      pexp_attributes = [];
+      pexp_is_return = false;
+    }
   | record_fields, spread_props ->
     {
       pexp_desc = Pexp_record (record_fields, spread_props);
       pexp_loc = loc;
       pexp_attributes = [];
+      pexp_is_return = false;
     }
 
 let try_find_key_prop (props : jsx_props) : (arg_label * expression) option =
diff --git a/compiler/syntax/src/res_driver.ml b/compiler/syntax/src/res_driver.ml
index 64039e7656..bdab030482 100644
--- a/compiler/syntax/src/res_driver.ml
+++ b/compiler/syntax/src/res_driver.ml
@@ -52,6 +52,7 @@ let parsing_engine =
       (fun ~for_printer ~filename ->
         let engine = setup ~filename ~for_printer () in
         let structure = Res_core.parse_implementation engine in
+        Res_return_marker.structure structure;
         let invalid, diagnostics =
           match engine.diagnostics with
           | [] as diagnostics -> (false, diagnostics)
@@ -69,6 +70,7 @@ let parsing_engine =
       (fun ~for_printer ~filename ->
         let engine = setup ~filename ~for_printer () in
         let signature = Res_core.parse_specification engine in
+        Res_return_marker.signature signature;
         let invalid, diagnostics =
           match engine.diagnostics with
           | [] as diagnostics -> (false, diagnostics)
@@ -90,6 +92,7 @@ let parsing_engine =
 let parse_implementation_from_source ~for_printer ~display_filename ~source =
   let engine = setup_from_source ~display_filename ~source ~for_printer () in
   let structure = Res_core.parse_implementation engine in
+  Res_return_marker.structure structure;
   let invalid, diagnostics =
     match engine.diagnostics with
     | [] as diagnostics -> (false, diagnostics)
@@ -107,6 +110,7 @@ let parse_implementation_from_source ~for_printer ~display_filename ~source =
 let parse_interface_from_source ~for_printer ~display_filename ~source =
   let engine = setup_from_source ~display_filename ~source ~for_printer () in
   let signature = Res_core.parse_specification engine in
+  Res_return_marker.signature signature;
   let invalid, diagnostics =
     match engine.diagnostics with
     | [] as diagnostics -> (false, diagnostics)
diff --git a/compiler/syntax/src/res_return_marker.ml b/compiler/syntax/src/res_return_marker.ml
new file mode 100644
index 0000000000..013b8b6b84
--- /dev/null
+++ b/compiler/syntax/src/res_return_marker.ml
@@ -0,0 +1,158 @@
+open Parsetree
+
+let rec mark_expression ~is_tail expr =
+  if is_tail then expr.pexp_is_return <- true;
+  match expr.pexp_desc with
+  | Pexp_ident _ | Pexp_constant _ -> ()
+  | Pexp_fun {default; rhs; _} ->
+    Option.iter (mark_expression ~is_tail:false) default;
+    mark_expression ~is_tail:true rhs
+  | Pexp_let (_, bindings, body) ->
+    List.iter mark_value_binding bindings;
+    mark_expression ~is_tail body
+  | Pexp_apply {funct; args; _} ->
+    mark_expression ~is_tail:false funct;
+    List.iter (fun (_, arg) -> mark_expression ~is_tail:false arg) args
+  | Pexp_match (scrutinee, cases) ->
+    mark_expression ~is_tail:false scrutinee;
+    List.iter (mark_case ~is_tail) cases
+  | Pexp_try (body, handlers) ->
+    mark_expression ~is_tail body;
+    List.iter (mark_case ~is_tail) handlers
+  | Pexp_tuple items | Pexp_array items ->
+    List.iter (mark_expression ~is_tail:false) items
+  | Pexp_construct (_, payload) | Pexp_variant (_, payload) ->
+    Option.iter (mark_expression ~is_tail:false) payload
+  | Pexp_record (fields, base) ->
+    List.iter (fun {x} -> mark_expression ~is_tail:false x) fields;
+    Option.iter (mark_expression ~is_tail:false) base
+  | Pexp_field (record, _) | Pexp_send (record, _) ->
+    mark_expression ~is_tail:false record
+  | Pexp_setfield (record, _, value) ->
+    mark_expression ~is_tail:false record;
+    mark_expression ~is_tail:false value
+  | Pexp_ifthenelse (cond, ifso, ifnot_opt) ->
+    mark_expression ~is_tail:false cond;
+    mark_expression ~is_tail ifso;
+    Option.iter (mark_expression ~is_tail) ifnot_opt
+  | Pexp_sequence (first, second) ->
+    mark_expression ~is_tail:false first;
+    mark_expression ~is_tail second
+  | Pexp_while (cond, body) ->
+    mark_expression ~is_tail:false cond;
+    mark_expression ~is_tail:false body
+  | Pexp_for (_pat, start, stop, _dir, body) ->
+    mark_expression ~is_tail:false start;
+    mark_expression ~is_tail:false stop;
+    mark_expression ~is_tail:false body
+  | Pexp_constraint (body, _)
+  | Pexp_coerce (body, _, _)
+  | Pexp_open (_, _, body)
+  | Pexp_newtype (_, body)
+  | Pexp_await body ->
+    mark_expression ~is_tail body
+  | Pexp_assert cond -> mark_expression ~is_tail:false cond
+  | Pexp_letmodule (_, me, body) ->
+    mark_module_expr me;
+    mark_expression ~is_tail body
+  | Pexp_letexception (_, body) -> mark_expression ~is_tail body
+  | Pexp_pack me -> mark_module_expr me
+  | Pexp_extension ext -> mark_extension ext
+  | Pexp_jsx_element element -> mark_jsx_element element
+
+and mark_case ~is_tail case =
+  Option.iter (mark_expression ~is_tail:false) case.pc_guard;
+  mark_expression ~is_tail case.pc_rhs
+
+and mark_value_binding vb = mark_expression ~is_tail:false vb.pvb_expr
+
+and mark_module_expr me =
+  match me.pmod_desc with
+  | Pmod_structure str -> mark_structure str
+  | Pmod_functor (_, param, body) ->
+    Option.iter mark_module_type param;
+    mark_module_expr body
+  | Pmod_apply (me1, me2) ->
+    mark_module_expr me1;
+    mark_module_expr me2
+  | Pmod_constraint (me', mt) ->
+    mark_module_expr me';
+    mark_module_type mt
+  | Pmod_unpack expr -> mark_expression ~is_tail:false expr
+  | Pmod_extension ext -> mark_extension ext
+  | Pmod_ident _ -> ()
+
+and mark_module_type mt =
+  match mt.pmty_desc with
+  | Pmty_signature sig_ -> mark_signature sig_
+  | Pmty_functor (_, param, res) ->
+    Option.iter mark_module_type param;
+    mark_module_type res
+  | Pmty_with (mt', _) -> mark_module_type mt'
+  | Pmty_typeof me -> mark_module_expr me
+  | Pmty_extension ext -> mark_extension ext
+  | Pmty_alias _ | Pmty_ident _ -> ()
+
+and mark_signature sig_items = List.iter mark_signature_item sig_items
+
+and mark_signature_item item =
+  match item.psig_desc with
+  | Psig_module md -> mark_module_type md.pmd_type
+  | Psig_recmodule mds -> List.iter (fun md -> mark_module_type md.pmd_type) mds
+  | Psig_modtype mtd -> Option.iter mark_module_type mtd.pmtd_type
+  | Psig_include incl -> mark_module_type incl.pincl_mod
+  | Psig_extension (ext, _) -> mark_extension ext
+  | Psig_attribute attr -> mark_attribute attr
+  | Psig_open _ | Psig_value _ | Psig_type _ | Psig_typext _ | Psig_exception _
+    ->
+    ()
+
+and mark_extension (_id, payload) = mark_payload payload
+
+and mark_attribute (_id, payload) = mark_payload payload
+
+and mark_payload = function
+  | PStr str -> mark_structure str
+  | PSig sig_ -> mark_signature sig_
+  | PTyp _ -> ()
+  | PPat (_, expr_opt) -> Option.iter (mark_expression ~is_tail:false) expr_opt
+
+and mark_structure_item item =
+  match item.pstr_desc with
+  | Pstr_value (_, bindings) -> List.iter mark_value_binding bindings
+  | Pstr_eval (expr, _) -> mark_expression ~is_tail:false expr
+  | Pstr_module mb -> mark_module_binding mb
+  | Pstr_recmodule mbs -> List.iter mark_module_binding mbs
+  | Pstr_include incl -> mark_module_expr incl.pincl_mod
+  | Pstr_extension (ext, _) -> mark_extension ext
+  | Pstr_attribute attr -> mark_attribute attr
+  | Pstr_exception _ | Pstr_primitive _ | Pstr_type _ | Pstr_typext _
+  | Pstr_open _ | Pstr_modtype _ ->
+    ()
+
+and mark_module_binding mb = mark_module_expr mb.pmb_expr
+
+and mark_structure items = List.iter mark_structure_item items
+
+and mark_jsx_element = function
+  | Jsx_fragment fragment ->
+    List.iter (mark_expression ~is_tail:false) fragment.jsx_fragment_children
+  | Jsx_unary_element element -> mark_jsx_props element.jsx_unary_element_props
+  | Jsx_container_element element ->
+    mark_jsx_props element.jsx_container_element_props;
+    List.iter
+      (mark_expression ~is_tail:false)
+      element.jsx_container_element_children
+
+and mark_jsx_props props =
+  List.iter
+    (function
+      | JSXPropPunning _ -> ()
+      | JSXPropValue (_, _, expr) -> mark_expression ~is_tail:false expr
+      | JSXPropSpreading (_, expr) -> mark_expression ~is_tail:false expr)
+    props
+
+let structure = mark_structure
+let signature = mark_signature
+
+let expression_is_return expr = expr.pexp_is_return
diff --git a/tests/ounit_tests/ounit_ast_return_tests.ml b/tests/ounit_tests/ounit_ast_return_tests.ml
new file mode 100644
index 0000000000..b0345acb4b
--- /dev/null
+++ b/tests/ounit_tests/ounit_ast_return_tests.ml
@@ -0,0 +1,99 @@
+open Parsetree
+
+let ( >:: ), ( >::: ) = OUnit.(( >:: ), ( >::: ))
+
+let parse_structure source =
+  let result =
+    Res_driver.parse_implementation_from_source ~for_printer:false
+      ~display_filename:"Test.res" ~source
+  in
+  result.parsetree
+
+let test_simple_return _ctx =
+  let structure = parse_structure "let f = x => x + 1\n" in
+  match structure with
+  | [
+   {
+     Parsetree.pstr_desc =
+       Pstr_value
+         ( _,
+           [{pvb_expr = {Parsetree.pexp_desc = Pexp_fun {rhs = body; _}; _}; _}]
+         );
+     _;
+   };
+  ] -> (
+    OUnit.assert_bool "function body should be marked as a return"
+      (Res_return_marker.expression_is_return body);
+    match body.pexp_desc with
+    | Pexp_apply {funct; args; _} ->
+      OUnit.assert_bool "callee should not be marked"
+        (not (Res_return_marker.expression_is_return funct));
+      List.iter
+        (fun (_, arg) ->
+          OUnit.assert_bool "arguments should not be marked"
+            (not (Res_return_marker.expression_is_return arg)))
+        args
+    | _ -> OUnit.assert_failure "expected application in function body")
+  | _ -> OUnit.assert_failure "unexpected structure for simple function"
+
+let test_sequence_return _ctx =
+  let structure = parse_structure "let f = x => {let y = x; y}\n" in
+  match structure with
+  | [
+   {
+     Parsetree.pstr_desc =
+       Pstr_value
+         ( _,
+           [{pvb_expr = {Parsetree.pexp_desc = Pexp_fun {rhs = body; _}; _}; _}]
+         );
+     _;
+   };
+  ] -> (
+    OUnit.assert_bool "function body should be marked"
+      (Res_return_marker.expression_is_return body);
+    match body.pexp_desc with
+    | Pexp_sequence (first, second) ->
+      OUnit.assert_bool "last expression in sequence should be marked"
+        (Res_return_marker.expression_is_return second);
+      OUnit.assert_bool "first expression in sequence should not be marked"
+        (not (Res_return_marker.expression_is_return first))
+    | Pexp_let (_, _, inner) ->
+      OUnit.assert_bool "inner expression should be marked"
+        (Res_return_marker.expression_is_return inner)
+    | _ -> OUnit.assert_failure "expected sequence or let in body")
+  | _ -> OUnit.assert_failure "unexpected structure for block body"
+
+let test_switch_returns _ctx =
+  let structure =
+    parse_structure "let f = x => switch x { | 0 => 1 | _ => x }\n"
+  in
+  match structure with
+  | [
+   {
+     Parsetree.pstr_desc =
+       Pstr_value
+         ( _,
+           [{pvb_expr = {Parsetree.pexp_desc = Pexp_fun {rhs = body; _}; _}; _}]
+         );
+     _;
+   };
+  ] -> (
+    OUnit.assert_bool "switch expression should be marked"
+      (Res_return_marker.expression_is_return body);
+    match body.pexp_desc with
+    | Pexp_match (_, cases) ->
+      List.iter
+        (fun {pc_rhs; _} ->
+          OUnit.assert_bool "case rhs should be marked"
+            (Res_return_marker.expression_is_return pc_rhs))
+        cases
+    | _ -> OUnit.assert_failure "expected switch in function body")
+  | _ -> OUnit.assert_failure "unexpected structure for switch body"
+
+let suites =
+  "return marker"
+  >::: [
+         "function body" >:: test_simple_return;
+         "sequence tail" >:: test_sequence_return;
+         "switch cases" >:: test_switch_returns;
+       ]
diff --git a/tests/ounit_tests/ounit_tests_main.ml b/tests/ounit_tests/ounit_tests_main.ml
index 37a1d7e597..302fac52bb 100644
--- a/tests/ounit_tests/ounit_tests_main.ml
+++ b/tests/ounit_tests/ounit_tests_main.ml
@@ -21,6 +21,7 @@ let suites =
       Ounit_bsb_regex_tests.suites;
       Ounit_bsb_pkg_tests.suites;
       Ounit_util_tests.suites;
+      Ounit_ast_return_tests.suites;
     ]
 
 let _ = OUnit.run_test_tt_main suites