From 11cc471c9390a6df1a241d8966bb784a58a82515 Mon Sep 17 00:00:00 2001
From: Zesen Qian <github@riaqn.org>
Date: Mon, 25 Nov 2024 08:51:09 -0500
Subject: [PATCH 1/6] some refactoring

---
 parsing/parser.mly | 43 +++++++++++++++++++------------------------
 1 file changed, 19 insertions(+), 24 deletions(-)

diff --git a/parsing/parser.mly b/parsing/parser.mly
index d6cf9bd64e0..b337fe44d62 100644
--- a/parsing/parser.mly
+++ b/parsing/parser.mly
@@ -196,20 +196,15 @@ let ghpat_with_modes ~loc ~pat ~cty ~modes =
   let pat = mkpat_with_modes ~loc ~pat ~cty ~modes in
   { pat with ppat_loc = { pat.ppat_loc with loc_ghost = true }}
 
-let mkexp_with_modes ~loc ~exp ~cty ~modes =
+let mkexp_constraint ~loc ~exp ~cty ~modes =
   match exp.pexp_desc with
   | Pexp_constraint (exp', cty', modes') ->
      begin match cty, cty' with
-     | Some _, None ->
+     | cty, None | None, cty ->
         { exp with
           pexp_desc = Pexp_constraint (exp', cty, modes @ modes');
           pexp_loc = make_loc loc
         }
-     | None, _ ->
-        { exp with
-          pexp_desc = Pexp_constraint (exp', cty', modes @ modes');
-          pexp_loc = make_loc loc
-        }
      | _ ->
         mkexp ~loc (Pexp_constraint (exp, cty, modes))
      end
@@ -219,8 +214,8 @@ let mkexp_with_modes ~loc ~exp ~cty ~modes =
      | cty, modes -> mkexp ~loc (Pexp_constraint (exp, cty, modes))
      end
 
-let ghexp_with_modes ~loc ~exp ~cty ~modes =
-  let exp = mkexp_with_modes ~loc ~exp ~cty ~modes in
+let ghexp_constraint ~loc ~exp ~cty ~modes =
+  let exp = mkexp_constraint ~loc ~exp ~cty ~modes in
   { exp with pexp_loc = { exp.pexp_loc with loc_ghost = true }}
 
 let exclave_ext_loc loc = mkloc "extension.exclave" loc
@@ -375,10 +370,10 @@ let removed_string_set loc =
 let not_expecting loc nonterm =
     raise Syntaxerr.(Error(Not_expecting(make_loc loc, nonterm)))
 
-let mkexp_type_constraint ?(ghost=false) ~loc ~modes e t =
+let mkexp_type_constraint_with_modes ?(ghost=false) ~loc ~modes e t =
   match t with
   | Pconstraint t ->
-     let mk = if ghost then ghexp_with_modes else mkexp_with_modes in
+     let mk = if ghost then ghexp_constraint else mkexp_constraint in
      mk ~loc ~exp:e ~cty:(Some t) ~modes
   | Pcoerce(t1, t2)  ->
      match modes with
@@ -387,9 +382,9 @@ let mkexp_type_constraint ?(ghost=false) ~loc ~modes e t =
       mk ~loc (Pexp_coerce(e, t1, t2))
      | _ :: _ -> not_expecting loc "mode annotations"
 
-let mkexp_opt_type_constraint ~loc ~modes e = function
+let mkexp_opt_type_constraint_with_modes ~loc ~modes e = function
   | None -> e
-  | Some c -> mkexp_type_constraint ~loc ~modes e c
+  | Some c -> mkexp_type_constraint_with_modes ~loc ~modes e c
 
 (* Helper functions for desugaring array indexing operators *)
 type paren_kind = Paren | Brace | Bracket
@@ -563,7 +558,7 @@ let mk_newtypes ~loc newtypes exp =
    in [let_binding_body_no_punning]. *)
 let wrap_type_annotation ~loc ?(typloc=loc) ~modes newtypes core_type body =
   let mk_newtypes = mk_newtypes ~loc in
-  let exp = mkexp_with_modes ~loc ~exp:body ~cty:(Some core_type) ~modes in
+  let exp = mkexp_constraint ~loc ~exp:body ~cty:(Some core_type) ~modes in
   let exp = mk_newtypes newtypes exp in
   let inner_type = Typ.varify_constructors (List.map fst newtypes) core_type in
   (exp, ghtyp ~loc:typloc (Ptyp_poly (newtypes, inner_type)))
@@ -756,7 +751,7 @@ let mkghost_newtype_function_body newtypes body_constraint body ~loc =
     | Some { type_constraint; mode_annotations } ->
         let {Location.loc_start; loc_end} = body.pexp_loc in
         let loc = loc_start, loc_end in
-        mkexp_type_constraint ~ghost:true ~loc ~modes:mode_annotations body type_constraint
+        mkexp_type_constraint_with_modes ~ghost:true ~loc ~modes:mode_annotations body type_constraint
   in
   mk_newtypes ~loc newtypes wrapped_body
 
@@ -1696,7 +1691,7 @@ paren_module_expr:
     e = expr
       { e }
   | e = expr COLON ty = package_type
-      { ghexp_with_modes ~loc:$loc ~exp:e ~cty:(Some ty) ~modes:[] }
+      { ghexp_constraint ~loc:$loc ~exp:e ~cty:(Some ty) ~modes:[] }
   | e = expr COLON ty1 = package_type COLONGREATER ty2 = package_type
       { ghexp ~loc:$loc (Pexp_coerce (e, Some ty1, ty2)) }
   | e = expr COLONGREATER ty2 = package_type
@@ -2358,7 +2353,7 @@ value:
       { ($4, $3, Cfk_concrete ($1, $6)), $2 }
   | override_flag attributes mutable_flag mkrhs(label) type_constraint
     EQUAL seq_expr
-      { let e = mkexp_type_constraint ~loc:$sloc ~modes:[] $7 $5 in
+      { let e = mkexp_type_constraint_with_modes ~loc:$sloc ~modes:[] $7 $5 in
         ($4, $3, Cfk_concrete ($1, e)), $2
       }
 ;
@@ -2822,7 +2817,7 @@ fun_expr:
   | UNDERSCORE
     { mkexp ~loc:$sloc Pexp_hole }
   | mode=mode_legacy exp=seq_expr
-     { mkexp_with_modes ~loc:$sloc ~exp ~cty:None ~modes:[mode] }
+     { mkexp_constraint ~loc:$sloc ~exp ~cty:None ~modes:[mode] }
   | EXCLAVE seq_expr
      { mkexp_exclave ~loc:$sloc ~kwd_loc:($loc($1)) $2 }
 ;
@@ -2897,7 +2892,7 @@ simple_expr:
       { unclosed "(" $loc($1) ")" $loc($3) }
   | LPAREN seq_expr type_constraint_with_modes RPAREN
       { let (t, m) = $3 in
-        mkexp_type_constraint ~ghost:true ~loc:$sloc ~modes:m $2 t }
+        mkexp_type_constraint_with_modes ~ghost:true ~loc:$sloc ~modes:m $2 t }
   | indexop_expr(DOT, seq_expr, { None })
       { mk_indexop_expr builtin_indexing_operators ~loc:$sloc $1 }
   (* Immutable array indexing is a regular operator, so it doesn't need its own
@@ -2962,7 +2957,7 @@ comprehension_clause_binding:
      over to the RHS of the binding, so we need everything to be visible. *)
   | attributes mode_legacy pattern IN expr
       { let expr =
-          mkexp_with_modes ~loc:$sloc ~exp:$5 ~cty:None ~modes:[$2]
+          mkexp_constraint ~loc:$sloc ~exp:$5 ~cty:None ~modes:[$2]
         in
         { pcomp_cb_pattern    = $3
         ; pcomp_cb_iterator   = Pcomp_in expr
@@ -3126,7 +3121,7 @@ labeled_simple_expr:
       { let loc = $loc(label) in
         (Labelled label, mkexpvar ~loc label) }
   | TILDE LPAREN label = LIDENT c = type_constraint RPAREN
-      { (Labelled label, mkexp_type_constraint ~loc:($startpos($2), $endpos) ~modes:[]
+      { (Labelled label, mkexp_type_constraint_with_modes ~loc:($startpos($2), $endpos) ~modes:[]
                            (mkexpvar ~loc:$loc(label) label) c) }
   | QUESTION label = LIDENT
       { let loc = $loc(label) in
@@ -3399,7 +3394,7 @@ fun_params:
        Some label, mkexpvar ~loc label }
   | TILDE LPAREN label = LIDENT c = type_constraint RPAREN %prec below_HASH
       { Some label,
-        mkexp_type_constraint
+        mkexp_type_constraint_with_modes
           ~loc:($startpos($2), $endpos) ~modes:[] (mkexpvar ~loc:$loc(label) label) c }
 ;
 reversed_labeled_tuple_body:
@@ -3426,7 +3421,7 @@ reversed_labeled_tuple_body:
   COMMA
   x2 = labeled_tuple_element
   { let x1 =
-      mkexp_type_constraint
+      mkexp_type_constraint_with_modes
         ~loc:($startpos($2), $endpos) ~modes:[] (mkexpvar ~loc:$loc(l1) l1) c
     in
     [ x2; Some l1, x1] }
@@ -3453,7 +3448,7 @@ record_expr_content:
           | Some e ->
               ($startpos(c), $endpos), label, e
         in
-        label, mkexp_opt_type_constraint ~loc:constraint_loc ~modes:[] e c }
+        label, mkexp_opt_type_constraint_with_modes ~loc:constraint_loc ~modes:[] e c }
 ;
 %inline object_expr_content:
   xs = separated_or_terminated_nonempty_list(SEMI, object_expr_field)

From 90940fd77b95a351b23d8dfafc6ad1b6db4879a9 Mon Sep 17 00:00:00 2001
From: Zesen Qian <github@riaqn.org>
Date: Fri, 29 Nov 2024 06:47:35 -0500
Subject: [PATCH 2/6] move things around

---
 .../tests/parsetree/source_jane_street.ml     | 354 +++++++++---------
 1 file changed, 177 insertions(+), 177 deletions(-)

diff --git a/testsuite/tests/parsetree/source_jane_street.ml b/testsuite/tests/parsetree/source_jane_street.ml
index e42376e9372..6db20a2185f 100644
--- a/testsuite/tests/parsetree/source_jane_street.ml
+++ b/testsuite/tests/parsetree/source_jane_street.ml
@@ -523,6 +523,183 @@ val f2 : local_ float -> (float -> float) @ once -> local_ t2 @ once @@
   global many = <fun>
 |}]
 
+
+module type S = sig
+  val bar : 'a -> 'a
+  module M : sig
+    val foo : 'a -> 'a
+  end
+end
+
+module type S' = sig
+  include [@no_recursive_modalities] S @@ portable
+end
+
+[%%expect{|
+module type S =
+  sig val bar : 'a -> 'a module M : sig val foo : 'a -> 'a end end
+module type S' =
+  sig
+    val bar : 'a -> 'a @@ portable
+    module M : sig val foo : 'a -> 'a end
+  end
+|}]
+
+(* Modules *)
+
+module type S = sig end
+module M = struct end
+[%%expect{|
+module type S = sig end
+module M : sig end
+|}]
+
+module F (X : S @@ portable) = struct
+end
+[%%expect{|
+Line 1, characters 19-27:
+1 | module F (X : S @@ portable) = struct
+                       ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module F (_ : S @@ portable) = struct
+end
+[%%expect{|
+Line 1, characters 19-27:
+1 | module F (_ : S @@ portable) = struct
+                       ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module M' = (M : S @@ portable)
+[%%expect{|
+Line 1, characters 22-30:
+1 | module M' = (M : S @@ portable)
+                          ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module F (M : S @@ portable) : S @@ portable = struct
+end
+[%%expect{|
+Line 1, characters 19-27:
+1 | module F (M : S @@ portable) : S @@ portable = struct
+                       ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module F (M : S @@ portable) @ portable = struct
+end
+[%%expect{|
+Line 1, characters 19-27:
+1 | module F (M : S @@ portable) @ portable = struct
+                       ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+
+
+(* CR zqian: the similar syntax for expressions are not allowed because @ might
+  be an binary operator *)
+module M' = (M @ portable)
+[%%expect{|
+Line 1, characters 17-25:
+1 | module M' = (M @ portable)
+                     ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module M' = (M : S @@ portable)
+[%%expect{|
+Line 1, characters 22-30:
+1 | module M' = (M : S @@ portable)
+                          ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module M @ portable = struct end
+[%%expect{|
+Line 1, characters 11-19:
+1 | module M @ portable = struct end
+               ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module M : S @@ portable = struct end
+[%%expect{|
+Line 1, characters 16-24:
+1 | module M : S @@ portable = struct end
+                    ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module type S' = functor () (M : S @@ portable) (_ : S @@ portable) -> S @ portable
+[%%expect{|
+Line 1, characters 38-46:
+1 | module type S' = functor () (M : S @@ portable) (_ : S @@ portable) -> S @ portable
+                                          ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+
+module type S' = () -> S @ portable -> S @ portable -> S @ portable
+[%%expect{|
+Line 1, characters 27-35:
+1 | module type S' = () -> S @ portable -> S @ portable -> S @ portable
+                               ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module (F @ portable) () = struct end
+[%%expect{|
+Line 1, characters 12-20:
+1 | module (F @ portable) () = struct end
+                ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module rec (F @ portable) () = struct end
+and (G @ portable) () = struct end
+[%%expect{|
+Line 1, characters 16-24:
+1 | module rec (F @ portable) () = struct end
+                    ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
+module type T = sig
+  module M : S @@ portable
+  module M = N @ portable
+  module (M @ portable) = N
+  module M0 (_ : S @@ portable) (X : S @@ portable) : S @@ portable
+  (* The above [@@ portable] is the return mode of the functor, not the modality on the
+  module declaration; To do that, one must write in the following ways. *)
+  module (M0 @ portable) (_ : S @@ portable) (X : S @@ portable) : S @@ portable
+  module M0 : functor (_ : S @@ portable) (X : S @@ portable) -> S @ portable @@ portable
+
+  module M0 : S @ portable -> S @ portable -> S @ portable @@ portable
+
+  module rec F : sig end @@ portable
+  and G : sig end @@ portable
+end
+[%%expect{|
+Line 3, characters 13-14:
+3 |   module M = N @ portable
+                 ^
+Error: Unbound module "N"
+|}]
+
+let foo () =
+  let module (F @ portable) () = struct end in
+  ()
+[%%expect{|
+Line 2, characters 18-26:
+2 |   let module (F @ portable) () = struct end in
+                      ^^^^^^^^
+Error: Mode annotations on modules are not supported yet.
+|}]
+
 (**********)
 (* unique *)
 
@@ -1023,183 +1200,6 @@ Line 9, characters 11-95:
 9 | module _ = Base(Name1)(Value1)(Name2)(Value2(Name2_1)(Value2_1)) [@jane.non_erasable.instances]
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 Error: Unbound module "Base[Name1:Value1][Name2:Value2[Name2_1:Value2_1]]"
-|}](*********)
-(* modes *)
-
-module type S = sig
-  val bar : 'a -> 'a
-  module M : sig
-    val foo : 'a -> 'a
-  end
-end
-
-module type S' = sig
-  include [@no_recursive_modalities] S @@ portable
-end
-
-[%%expect{|
-module type S =
-  sig val bar : 'a -> 'a module M : sig val foo : 'a -> 'a end end
-module type S' =
-  sig
-    val bar : 'a -> 'a @@ portable
-    module M : sig val foo : 'a -> 'a end
-  end
-|}]
-
-(* Modules *)
-
-module type S = sig end
-module M = struct end
-[%%expect{|
-module type S = sig end
-module M : sig end
-|}]
-
-module F (X : S @@ portable) = struct
-end
-[%%expect{|
-Line 1, characters 19-27:
-1 | module F (X : S @@ portable) = struct
-                       ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module F (_ : S @@ portable) = struct
-end
-[%%expect{|
-Line 1, characters 19-27:
-1 | module F (_ : S @@ portable) = struct
-                       ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module M' = (M : S @@ portable)
-[%%expect{|
-Line 1, characters 22-30:
-1 | module M' = (M : S @@ portable)
-                          ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module F (M : S @@ portable) : S @@ portable = struct
-end
-[%%expect{|
-Line 1, characters 19-27:
-1 | module F (M : S @@ portable) : S @@ portable = struct
-                       ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module F (M : S @@ portable) @ portable = struct
-end
-[%%expect{|
-Line 1, characters 19-27:
-1 | module F (M : S @@ portable) @ portable = struct
-                       ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-
-
-(* CR zqian: the similar syntax for expressions are not allowed because @ might
-  be an binary operator *)
-module M' = (M @ portable)
-[%%expect{|
-Line 1, characters 17-25:
-1 | module M' = (M @ portable)
-                     ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module M' = (M : S @@ portable)
-[%%expect{|
-Line 1, characters 22-30:
-1 | module M' = (M : S @@ portable)
-                          ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module M @ portable = struct end
-[%%expect{|
-Line 1, characters 11-19:
-1 | module M @ portable = struct end
-               ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module M : S @@ portable = struct end
-[%%expect{|
-Line 1, characters 16-24:
-1 | module M : S @@ portable = struct end
-                    ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module type S' = functor () (M : S @@ portable) (_ : S @@ portable) -> S @ portable
-[%%expect{|
-Line 1, characters 38-46:
-1 | module type S' = functor () (M : S @@ portable) (_ : S @@ portable) -> S @ portable
-                                          ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-
-module type S' = () -> S @ portable -> S @ portable -> S @ portable
-[%%expect{|
-Line 1, characters 27-35:
-1 | module type S' = () -> S @ portable -> S @ portable -> S @ portable
-                               ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module (F @ portable) () = struct end
-[%%expect{|
-Line 1, characters 12-20:
-1 | module (F @ portable) () = struct end
-                ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module rec (F @ portable) () = struct end
-and (G @ portable) () = struct end
-[%%expect{|
-Line 1, characters 16-24:
-1 | module rec (F @ portable) () = struct end
-                    ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
-|}]
-
-module type T = sig
-  module M : S @@ portable
-  module M = N @ portable
-  module (M @ portable) = N
-  module M0 (_ : S @@ portable) (X : S @@ portable) : S @@ portable
-  (* The above [@@ portable] is the return mode of the functor, not the modality on the
-  module declaration; To do that, one must write in the following ways. *)
-  module (M0 @ portable) (_ : S @@ portable) (X : S @@ portable) : S @@ portable
-  module M0 : functor (_ : S @@ portable) (X : S @@ portable) -> S @ portable @@ portable
-
-  module M0 : S @ portable -> S @ portable -> S @ portable @@ portable
-
-  module rec F : sig end @@ portable
-  and G : sig end @@ portable
-end
-[%%expect{|
-Line 3, characters 13-14:
-3 |   module M = N @ portable
-                 ^
-Error: Unbound module "N"
-|}]
-
-let foo () =
-  let module (F @ portable) () = struct end in
-  ()
-[%%expect{|
-Line 2, characters 18-26:
-2 |   let module (F @ portable) () = struct end in
-                      ^^^^^^^^
-Error: Mode annotations on modules are not supported yet.
 |}]
 
 (***************)

From f06da03258b3e27a5fae4d1da05f08aa65b4fadd Mon Sep 17 00:00:00 2001
From: Zesen Qian <github@riaqn.org>
Date: Fri, 29 Nov 2024 09:31:15 -0500
Subject: [PATCH 3/6] support modes on function body

---
 parsing/ast_helper.mli                        |  2 +-
 parsing/ast_iterator.ml                       | 13 ++-
 parsing/ast_mapper.ml                         |  7 +-
 parsing/depend.ml                             | 11 ++-
 parsing/parser.mly                            | 68 +++++++------
 parsing/parsetree.mli                         | 22 ++++-
 parsing/pprintast.ml                          | 28 +++---
 parsing/printast.ml                           |  9 +-
 printer/printast_with_mappings.ml             | 10 +-
 .../parsetree/modes_ast_mapper.reference      |  1 +
 .../tests/parsetree/source_jane_street.ml     | 25 +++++
 testsuite/tests/typing-modes/modes.ml         | 95 ++++++++++++++++++-
 typing/printtyped.ml                          |  9 +-
 typing/tast_iterator.ml                       |  3 +-
 typing/tast_mapper.ml                         |  5 +-
 typing/typeclass.ml                           |  2 +-
 typing/typecore.ml                            | 75 ++++++++++-----
 typing/typedtree.ml                           |  3 +-
 typing/typedtree.mli                          |  6 +-
 typing/untypeast.ml                           | 29 +++---
 20 files changed, 301 insertions(+), 122 deletions(-)

diff --git a/parsing/ast_helper.mli b/parsing/ast_helper.mli
index 7c82178c063..2ccefe541ea 100644
--- a/parsing/ast_helper.mli
+++ b/parsing/ast_helper.mli
@@ -152,7 +152,7 @@ module Exp:
     val let_: ?loc:loc -> ?attrs:attrs -> rec_flag -> value_binding list
               -> expression -> expression
     val function_ : ?loc:loc -> ?attrs:attrs -> function_param list
-                   -> function_constraint option -> function_body
+                   -> function_constraint -> function_body
                    -> expression
     val apply: ?loc:loc -> ?attrs:attrs -> expression
                -> (arg_label * expression) list -> expression
diff --git a/parsing/ast_iterator.ml b/parsing/ast_iterator.ml
index 08c6581aea5..c1ca9b78b48 100644
--- a/parsing/ast_iterator.ml
+++ b/parsing/ast_iterator.ml
@@ -432,14 +432,17 @@ module E = struct
   let iter_function_constraint sub : function_constraint -> _ =
     (* Enable warning 9 to ensure that the record pattern doesn't miss any
        field. *)
-    fun[@ocaml.warning "+9"] { mode_annotations; type_constraint } ->
+    fun[@ocaml.warning "+9"] { mode_annotations; ret_type_constraint; ret_mode_annotations } ->
       sub.modes sub mode_annotations;
-      match type_constraint with
-      | Pconstraint ty ->
+      begin match ret_type_constraint with
+      | Some (Pconstraint ty) ->
           sub.typ sub ty
-      | Pcoerce (ty1, ty2) ->
+      | Some (Pcoerce (ty1, ty2)) ->
           Option.iter (sub.typ sub) ty1;
           sub.typ sub ty2
+      | None -> ()
+      end;
+      sub.modes sub ret_mode_annotations
 
   let iter_function_body sub : function_body -> _ = function
     | Pfunction_body expr ->
@@ -462,7 +465,7 @@ module E = struct
         sub.expr sub e
     | Pexp_function (params, constraint_, body) ->
         List.iter (iter_function_param sub) params;
-        iter_opt (iter_function_constraint sub) constraint_;
+        iter_function_constraint sub constraint_;
         iter_function_body sub body
     | Pexp_apply (e, l) ->
         sub.expr sub e; List.iter (iter_snd (sub.expr sub)) l
diff --git a/parsing/ast_mapper.ml b/parsing/ast_mapper.ml
index 5bbedf1be5e..0eb886cde2f 100644
--- a/parsing/ast_mapper.ml
+++ b/parsing/ast_mapper.ml
@@ -487,9 +487,10 @@ module E = struct
     | Pcoerce (ty1, ty2) ->
         Pcoerce (Option.map (sub.typ sub) ty1, sub.typ sub ty2)
 
-  let map_function_constraint sub { mode_annotations; type_constraint } =
+  let map_function_constraint sub { mode_annotations; ret_type_constraint; ret_mode_annotations } =
     { mode_annotations = sub.modes sub mode_annotations;
-      type_constraint = map_type_constraint sub type_constraint;
+      ret_type_constraint = Option.map (map_type_constraint sub) ret_type_constraint;
+      ret_mode_annotations = sub.modes sub ret_mode_annotations
     }
 
   let map_iterator sub = function
@@ -534,7 +535,7 @@ module E = struct
     | Pexp_function (ps, c, b) ->
       function_ ~loc ~attrs
         (List.map (map_function_param sub) ps)
-        (map_opt (map_function_constraint sub) c)
+        (map_function_constraint sub c)
         (map_function_body sub b)
     | Pexp_apply (e, l) ->
         apply ~loc ~attrs (sub.expr sub e) (List.map (map_snd (sub.expr sub)) l)
diff --git a/parsing/depend.ml b/parsing/depend.ml
index f975418baed..be505a8ea9c 100644
--- a/parsing/depend.ml
+++ b/parsing/depend.ml
@@ -243,7 +243,7 @@ let rec add_expr bv exp =
       let bv = add_bindings rf bv pel in add_expr bv e
   | Pexp_function (params, constraint_, body) ->
       let bv = List.fold_left add_function_param bv params in
-      add_opt add_function_constraint bv constraint_;
+      add_function_constraint bv constraint_;
       add_function_body bv body
   | Pexp_apply(e, el) ->
       add_expr bv e; List.iter (fun (_,e) -> add_expr bv e) el
@@ -365,13 +365,14 @@ and add_function_body bv body =
   | Pfunction_cases (cases, _, _) ->
       add_cases bv cases
 
-and add_function_constraint bv { mode_annotations = _; type_constraint } =
-  match type_constraint with
-  | Pconstraint ty ->
+and add_function_constraint bv { mode_annotations = _; ret_type_constraint; ret_mode_annotations = _ } =
+  match ret_type_constraint with
+  | Some (Pconstraint ty) ->
       add_type bv ty
-  | Pcoerce (ty1, ty2) ->
+  | Some (Pcoerce (ty1, ty2)) ->
       add_opt add_type bv ty1;
       add_type bv ty2
+  | None -> ()
 
 and add_cases bv cases =
   List.iter (add_case bv) cases
diff --git a/parsing/parser.mly b/parsing/parser.mly
index b337fe44d62..4c60074d276 100644
--- a/parsing/parser.mly
+++ b/parsing/parser.mly
@@ -382,9 +382,9 @@ let mkexp_type_constraint_with_modes ?(ghost=false) ~loc ~modes e t =
       mk ~loc (Pexp_coerce(e, t1, t2))
      | _ :: _ -> not_expecting loc "mode annotations"
 
-let mkexp_opt_type_constraint_with_modes ~loc ~modes e = function
+let mkexp_opt_type_constraint_with_modes ?ghost ~loc ~modes e = function
   | None -> e
-  | Some c -> mkexp_type_constraint_with_modes ~loc ~modes e c
+  | Some c -> mkexp_type_constraint_with_modes ?ghost ~loc ~modes e c
 
 (* Helper functions for desugaring array indexing operators *)
 type paren_kind = Paren | Brace | Bracket
@@ -740,18 +740,22 @@ let all_params_as_newtypes =
     then Some (List.filter_map as_newtype params)
     else None
 
+let empty_body_constraint =
+  { ret_type_constraint = None; mode_annotations = []; ret_mode_annotations = []}
+
 (* Given a construct [fun (type a b c) : t -> e], we construct
    [Pexp_newtype(a, Pexp_newtype(b, Pexp_newtype(c, Pexp_constraint(e, t))))]
    rather than a [Pexp_function].
 *)
 let mkghost_newtype_function_body newtypes body_constraint body ~loc =
   let wrapped_body =
-    match body_constraint with
-    | None -> body
-    | Some { type_constraint; mode_annotations } ->
-        let {Location.loc_start; loc_end} = body.pexp_loc in
-        let loc = loc_start, loc_end in
-        mkexp_type_constraint_with_modes ~ghost:true ~loc ~modes:mode_annotations body type_constraint
+    let { ret_type_constraint; mode_annotations; ret_mode_annotations } =
+      body_constraint
+    in
+    let modes = mode_annotations @ ret_mode_annotations in
+    let {Location.loc_start; loc_end} = body.pexp_loc in
+    let loc = loc_start, loc_end in
+    mkexp_opt_type_constraint_with_modes ~ghost:true ~loc ~modes body ret_type_constraint
   in
   mk_newtypes ~loc newtypes wrapped_body
 
@@ -2578,7 +2582,7 @@ class_type_declarations:
            typechecking. For standalone function cases, we want the compiler to
            respect, e.g., [@inline] attributes.
          *)
-        mkfunction [] None (Pfunction_cases (cases, loc, [])) ~attrs:$2
+        mkfunction [] empty_body_constraint (Pfunction_cases (cases, loc, [])) ~attrs:$2
           ~loc:$sloc
       }
 ;
@@ -2772,6 +2776,21 @@ let_pattern_no_modes:
 
 %inline qualified_dotop: ioption(DOT mod_longident {$2}) DOTOP { $1, $2 };
 
+optional_atomic_constraint_:
+  | COLON atomic_type optional_atat_mode_expr {
+    { ret_type_constraint = Some (Pconstraint $2)
+    ; mode_annotations = []
+    ; ret_mode_annotations = $3
+    }
+   }
+  | at_mode_expr {
+    { ret_type_constraint = None
+    ; mode_annotations = []
+    ; ret_mode_annotations = $1
+    }
+  }
+  | { empty_body_constraint }
+
 fun_expr:
     simple_expr %prec below_HASH
       { $1 }
@@ -2779,18 +2798,9 @@ fun_expr:
       { let desc, attrs = $1 in
         mkexp_attrs ~loc:$sloc desc attrs }
     /* Cf #5939: we used to accept (fun p when e0 -> e) */
-  | FUN ext_attributes fun_params preceded(COLON, atomic_type)?
+  | FUN ext_attributes fun_params body_constraint = optional_atomic_constraint_
       MINUSGREATER fun_body
-      { let body_constraint =
-          Option.map
-            (fun x ->
-              { type_constraint = Pconstraint x
-              ; mode_annotations = []
-              })
-          $4
-        in
-        mkfunction $3 body_constraint $6 ~loc:$sloc ~attrs:$2
-      }
+      {  mkfunction $3 body_constraint $6 ~loc:$sloc ~attrs:$2 }
   | expr_
       { $1 }
   | let_bindings(ext) IN seq_expr
@@ -3273,14 +3283,16 @@ letop_bindings:
 strict_binding_modes:
     EQUAL seq_expr
       { fun _ -> $2 }
-  | fun_params type_constraint? EQUAL fun_body
-  (* CR zqian: The above [type_constraint] should be replaced by [constraint_]
-    to support mode annotation *)
+  | fun_params constraint_? EQUAL fun_body
     { fun mode_annotations ->
-        let constraint_ : function_constraint option =
-          match $2 with
-          | None -> None
-          | Some type_constraint -> Some { type_constraint; mode_annotations }
+        let constraint_ : function_constraint =
+          let ret_type_constraint, ret_mode_annotations =
+            match $2 with
+            | None -> None, []
+            | Some (ret_type_constraint, ret_mode_annotations) ->
+                ret_type_constraint, ret_mode_annotations
+          in
+          {mode_annotations; ret_type_constraint ; ret_mode_annotations }
         in
         let exp = mkfunction $1 constraint_ $4 ~loc:$sloc ~attrs:(None, []) in
         { exp with pexp_loc = { exp.pexp_loc with loc_ghost = true } }
@@ -3299,7 +3311,7 @@ fun_body:
         | Some _ ->
           (* function%foo extension nodes interrupt the arity *)
           let cases = Pfunction_cases ($3, make_loc $sloc, []) in
-          let function_ = mkfunction [] None cases ~loc:$sloc ~attrs:$2 in
+          let function_ = mkfunction [] empty_body_constraint cases ~loc:$sloc ~attrs:$2 in
           Pfunction_body function_
       }
   | fun_seq_expr
diff --git a/parsing/parsetree.mli b/parsing/parsetree.mli
index d00cb264a6c..74aba7388ed 100644
--- a/parsing/parsetree.mli
+++ b/parsing/parsetree.mli
@@ -368,7 +368,7 @@ and expression_desc =
                when [flag] is {{!Asttypes.rec_flag.Recursive}[Recursive]}.
          *)
   | Pexp_function of
-      function_param list * function_constraint option * function_body
+      function_param list * function_constraint * function_body
   (** [Pexp_function ([P1; ...; Pn], C, body)] represents any construct
       involving [fun] or [function], including:
       - [fun P1 ... Pn -> E]
@@ -595,11 +595,23 @@ and type_constraint =
 
 and function_constraint =
   { mode_annotations : modes;
-    (** The mode annotation placed on a function let-binding when the function
-            has a type constraint on the body, e.g.
-            [let local_ f x : int -> int = ...].
+    (** The mode annotation placed on a function let-binding, e.g.
+       [let local_ f x : int -> int = ...].
+       The [local_] syntax is parsed into two nodes: the field here, and [pvb_modes].
+       This field only affects the interpretation of [ret_type_constraint], while the
+       latter is translated in [typecore] to [Pexp_constraint] to contrain the mode of the
+       function.
+       (* CR zqian: This field is not failthful representation of the user syntax, and
+       complicates [pprintast]. It should be removed and their functionality should be
+       moved to [pvb_modes]. *)
     *)
-    type_constraint : type_constraint;
+    ret_mode_annotations : modes;
+    (** The mode annotation placed on a function's body, e.g.
+       [let f x : int -> int @@ local = ...].
+       This field constrains the mode of function's body.
+    *)
+    ret_type_constraint : type_constraint option;
+    (** The type constraint placed on a function's body. *)
   }
 (** See the comment on {{!expression_desc.Pexp_function}[Pexp_function]}. *)
 
diff --git a/parsing/pprintast.ml b/parsing/pprintast.ml
index 4705108989a..29d28928904 100644
--- a/parsing/pprintast.ml
+++ b/parsing/pprintast.ml
@@ -942,7 +942,7 @@ and expression ctxt f x =
     | Pexp_function (params, constraint_, body) ->
         begin match params, constraint_ with
           (* Omit [fun] if there are no params. *)
-          | [], None ->
+          | [], {ret_type_constraint = None; ret_mode_annotations = []; _} ->
               (* If function cases are a direct body of a function,
                  the function node should be wrapped in parens so
                  it doesn't become part of the enclosing function. *)
@@ -953,8 +953,8 @@ and expression ctxt f x =
               in
               let ctxt' = if should_paren then reset_ctxt else ctxt in
               pp f "@[<2>%a@]" (paren should_paren (function_body ctxt')) body
-          | [], Some constraint_ ->
-            pp f "@[<2>(%a@;%a)@]"
+          | [], constraint_ ->
+            pp f "@[<2>(%a%a)@]"
               (function_body ctxt) body
               (function_constraint ctxt) constraint_
           | _ :: _, _ ->
@@ -2237,21 +2237,21 @@ and function_body ctxt f x =
       (case_list ctxt) cases
 
 and function_constraint ctxt f x =
-  (* We don't currently print [x.alloc_mode]; this would need
-     to go on the enclosing [let] binding.
-  *)
+  (* We don't print [mode_annotations], which describes the whole function and goes on the
+     [let] binding. *)
   (* Enable warning 9 to ensure that the record pattern doesn't miss any field.
   *)
   match[@ocaml.warning "+9"] x with
-  | { type_constraint = Pconstraint ty; mode_annotations } ->
-    let _, modes = split_out_legacy_modes mode_annotations in
-    pp f ":@;%a%a" (core_type ctxt) ty optional_atat_modes modes
-  | { type_constraint = Pcoerce (ty1, ty2); mode_annotations } ->
-    let _, modes = split_out_legacy_modes mode_annotations in
-    pp f "%a:>@;%a%a"
+  | { ret_type_constraint = Some (Pconstraint ty); ret_mode_annotations; _ } ->
+
+    pp f "@;:@;%a%a" (core_type ctxt) ty optional_atat_modes ret_mode_annotations
+  | { ret_type_constraint = Some (Pcoerce (ty1, ty2)); ret_mode_annotations; _ } ->
+    pp f "@;%a:>@;%a%a"
       (option ~first:":@;" (core_type ctxt)) ty1
       (core_type ctxt) ty2
-      optional_atat_modes modes
+      optional_atat_modes ret_mode_annotations
+  | { ret_type_constraint = None; ret_mode_annotations; _} ->
+    pp f "%a" optional_at_modes ret_mode_annotations
 
 and function_params_then_body ctxt f params constraint_ body ~delimiter =
   let pp_params f =
@@ -2261,7 +2261,7 @@ and function_params_then_body ctxt f params constraint_ body ~delimiter =
   in
   pp f "%t%a%s@;%a"
     pp_params
-    (option (function_constraint ctxt) ~first:"@;") constraint_
+    (function_constraint ctxt) constraint_
     delimiter
     (function_body (under_functionrhs ctxt)) body
 
diff --git a/parsing/printast.ml b/parsing/printast.ml
index ba4bf6b2964..eb37dd1513a 100644
--- a/parsing/printast.ml
+++ b/parsing/printast.ml
@@ -311,7 +311,7 @@ and expression i ppf x =
   | Pexp_function (params, c, body) ->
       line i ppf "Pexp_function\n";
       list i function_param ppf params;
-      option i function_constraint ppf c;
+      function_constraint i ppf c;
       function_body i ppf body
   | Pexp_apply (e, l) ->
       line i ppf "Pexp_apply\n";
@@ -548,9 +548,10 @@ and type_constraint i ppf type_constraint =
       option (i+1) core_type ppf ty1;
       core_type (i+1) ppf ty2
 
-and function_constraint i ppf { type_constraint = c; mode_annotations } =
-  type_constraint i ppf c;
-  modes i ppf mode_annotations
+and function_constraint i ppf { ret_type_constraint = c; mode_annotations; ret_mode_annotations } =
+  option i type_constraint ppf c;
+  modes i ppf mode_annotations;
+  modes i ppf ret_mode_annotations
 
 and value_description i ppf x =
   line i ppf "value_description %a %a\n" fmt_string_loc
diff --git a/printer/printast_with_mappings.ml b/printer/printast_with_mappings.ml
index d40def98a8a..89e32c499cc 100644
--- a/printer/printast_with_mappings.ml
+++ b/printer/printast_with_mappings.ml
@@ -335,7 +335,7 @@ and expression i ppf x =
   | Pexp_function (params, c, body) ->
       line i ppf "Pexp_function\n";
       list i function_param ppf params;
-      option i function_constraint ppf c;
+      function_constraint i ppf c;
       function_body i ppf body
   | Pexp_apply (e, l) ->
       line i ppf "Pexp_apply\n";
@@ -574,9 +574,11 @@ and type_constraint i ppf type_constraint =
       option (i+1) core_type ppf ty1;
       core_type (i+1) ppf ty2
 
-and function_constraint i ppf { type_constraint = c; mode_annotations } =
-  type_constraint i ppf c;
-  modes i ppf mode_annotations
+and function_constraint i ppf
+  { ret_type_constraint; mode_annotations; ret_mode_annotations } =
+  modes i ppf mode_annotations;
+  option i type_constraint ppf ret_type_constraint;
+  modes i ppf ret_mode_annotations
 
 and value_description i ppf x =
   with_location_mapping ~loc:x.pval_loc ppf (fun () ->
diff --git a/testsuite/tests/parsetree/modes_ast_mapper.reference b/testsuite/tests/parsetree/modes_ast_mapper.reference
index 37d370402f5..1e44f592a40 100644
--- a/testsuite/tests/parsetree/modes_ast_mapper.reference
+++ b/testsuite/tests/parsetree/modes_ast_mapper.reference
@@ -1,5 +1,6 @@
 modes: local [File "_none_", line 1, characters 7-13]
 ------------------------------
+modes: unique [File "_none_", line 1, characters 4-11]
 modes: local [File "_none_", line 1, characters 15-21]
 modes: unique [File "_none_", line 1, characters 4-11]
 ------------------------------
diff --git a/testsuite/tests/parsetree/source_jane_street.ml b/testsuite/tests/parsetree/source_jane_street.ml
index 6db20a2185f..c6010930211 100644
--- a/testsuite/tests/parsetree/source_jane_street.ml
+++ b/testsuite/tests/parsetree/source_jane_street.ml
@@ -308,6 +308,11 @@ let g () =
   let local_ unique_ once_ f : int -> int = fun z -> z + z in
   let local_ f x: int -> int = x in
   let local_ f x y : int -> int = fun z -> x + y + z in
+
+  let foo a b @ local = exclave_ "hello" in
+  let foo = fun a b @ local -> exclave_ "hello" in
+  let foo a b : int @@ local = 42 in
+  let foo = fun a b : int @@ local -> 42 in
   ();;
 
 [%%expect{|
@@ -341,6 +346,26 @@ Line 7, characters 13-14:
                  ^
 Warning 26 [unused-var]: unused variable f.
 
+Line 9, characters 6-9:
+9 |   let foo a b @ local = exclave_ "hello" in
+          ^^^
+Warning 26 [unused-var]: unused variable foo.
+
+Line 10, characters 6-9:
+10 |   let foo = fun a b @ local -> exclave_ "hello" in
+           ^^^
+Warning 26 [unused-var]: unused variable foo.
+
+Line 11, characters 6-9:
+11 |   let foo a b : int @@ local = 42 in
+           ^^^
+Warning 26 [unused-var]: unused variable foo.
+
+Line 12, characters 6-9:
+12 |   let foo = fun a b : int @@ local -> 42 in
+           ^^^
+Warning 26 [unused-var]: unused variable foo.
+
 val g : unit -> unit @@ global many = <fun>
 |}]
 
diff --git a/testsuite/tests/typing-modes/modes.ml b/testsuite/tests/typing-modes/modes.ml
index af7c95803aa..31f51e17eab 100644
--- a/testsuite/tests/typing-modes/modes.ml
+++ b/testsuite/tests/typing-modes/modes.ml
@@ -66,9 +66,98 @@ Line 2, characters 18-23:
 Error: The locality axis has already been specified.
 |}]
 
-(* CR zqian: this should be supported *)
-(* let foo a b @ local = "hello"
-let foo a b : _ @@ local = "hello" *)
+let foo a b @ local = local_ "hello"
+[%%expect{|
+Line 1, characters 22-36:
+1 | let foo a b @ local = local_ "hello"
+                          ^^^^^^^^^^^^^^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+let foo = fun a b @ local -> local_ "hello"
+[%%expect{|
+Line 1, characters 29-43:
+1 | let foo = fun a b @ local -> local_ "hello"
+                                 ^^^^^^^^^^^^^^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+let foo a b @ local = exclave_ "hello"
+[%%expect{|
+val foo : 'a -> 'b -> local_ string = <fun>
+|}]
+
+let foo = fun a b @ local -> exclave_ "hello"
+[%%expect{|
+val foo : 'a -> 'b -> local_ string = <fun>
+|}]
+
+
+let foo a b @ local = local_ 42
+[%%expect{|
+Line 1, characters 22-31:
+1 | let foo a b @ local = local_ 42
+                          ^^^^^^^^^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+let foo = fun a b @ local -> local_ 42
+[%%expect{|
+Line 1, characters 29-38:
+1 | let foo = fun a b @ local -> local_ 42
+                                 ^^^^^^^^^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+let foo a b : int @@ local = local_ 42
+[%%expect{|
+val foo : 'a -> 'b -> local_ int = <fun>
+|}]
+
+let foo = fun a b : int @@ local -> local_ 42
+[%%expect{|
+val foo : 'a -> 'b -> local_ int = <fun>
+|}]
+
+let foo a b : _ @@ local = local_ 42
+[%%expect{|
+Line 1, characters 27-36:
+1 | let foo a b : _ @@ local = local_ 42
+                               ^^^^^^^^^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+let foo = fun a b : _ @@ local -> local_ 42
+[%%expect{|
+Line 1, characters 34-43:
+1 | let foo = fun a b : _ @@ local -> local_ 42
+                                      ^^^^^^^^^
+Error: This value escapes its region.
+  Hint: Cannot return a local value without an "exclave_" annotation.
+|}]
+
+
+(* the return mode annotation is used to interpret the return type annotation, giving the
+    expected currying behavior *)
+let foo a : string -> string -> unit @@ local = fun b c -> ()
+[%%expect{|
+val foo : 'a -> local_ (string -> string -> unit) = <fun>
+|}]
+
+(* the return mode annotation overrides the whole-function mode annotation, even when they
+    describe different axes. *)
+(* CR zqian: this should probably be improved somehow. *)
+let bar () = exclave_
+  let (foo @ local) a : string -> string -> unit @@ nonportable = fun b c -> () in
+  foo
+[%%expect{|
+val bar : unit -> local_ ('a -> (string -> string -> unit)) = <fun>
+|}]
 
 (* Expressions *)
 let foo = ("hello" : _ @@ local)
diff --git a/typing/printtyped.ml b/typing/printtyped.ml
index a71891ec1a4..af6d9915b4e 100644
--- a/typing/printtyped.ml
+++ b/typing/printtyped.ml
@@ -420,11 +420,10 @@ and function_body i ppf (body : function_body) =
 
 and expression_extra i ppf x attrs =
   match x with
-  | Texp_constraint (ct, m) ->
+  | Texp_constraint (ct) ->
       line i ppf "Texp_constraint\n";
       attributes i ppf attrs;
-      option i core_type ppf ct;
-      alloc_const_option_mode i ppf m;
+      core_type i ppf ct;
   | Texp_coerce (cto1, cto2) ->
       line i ppf "Texp_coerce\n";
       attributes i ppf attrs;
@@ -440,6 +439,10 @@ and expression_extra i ppf x attrs =
   | Texp_stack ->
       line i ppf "Texp_stack\n";
       attributes i ppf attrs
+  | Texp_mode m ->
+      line i ppf "Texp_mode\n";
+      attributes i ppf attrs;
+      alloc_const_option_mode i ppf m
 
 and alloc_mode_raw: type l r. _ -> _ -> (l * r) Mode.Alloc.t -> _
   = fun i ppf m -> line i ppf "alloc_mode %a\n" (Mode.Alloc.print ()) m
diff --git a/typing/tast_iterator.ml b/typing/tast_iterator.ml
index a2d8dbec1f1..bb57abd021a 100644
--- a/typing/tast_iterator.ml
+++ b/typing/tast_iterator.ml
@@ -279,13 +279,14 @@ let pat
       sub.pat sub p2
 
 let extra sub = function
-  | Texp_constraint (cty, _modes) -> Option.iter (sub.typ sub) cty
+  | Texp_constraint (cty) -> sub.typ sub cty
   | Texp_coerce (cty1, cty2) ->
       Option.iter (sub.typ sub) cty1;
       sub.typ sub cty2
   | Texp_newtype _ -> ()
   | Texp_poly cto -> Option.iter (sub.typ sub) cto
   | Texp_stack -> ()
+  | Texp_mode _ -> ()
 
 let function_param sub { fp_loc; fp_kind; fp_newtypes; _ } =
   sub.location sub fp_loc;
diff --git a/typing/tast_mapper.ml b/typing/tast_mapper.ml
index bc7dd19b6dc..8aa39ce3654 100644
--- a/typing/tast_mapper.ml
+++ b/typing/tast_mapper.ml
@@ -371,13 +371,14 @@ let function_param sub
   }
 
 let extra sub = function
-  | Texp_constraint (cty, modes) ->
-    Texp_constraint (Option.map (sub.typ sub) cty, modes)
+  | Texp_constraint cty ->
+    Texp_constraint (sub.typ sub cty)
   | Texp_coerce (cty1, cty2) ->
     Texp_coerce (Option.map (sub.typ sub) cty1, sub.typ sub cty2)
   | Texp_newtype _ as d -> d
   | Texp_poly cto -> Texp_poly (Option.map (sub.typ sub) cto)
   | Texp_stack as d -> d
+  | Texp_mode _ as d -> d
 
 let function_body sub body =
   match body with
diff --git a/typing/typeclass.ml b/typing/typeclass.ml
index ced9cde6ea2..8abeeab3e92 100644
--- a/typing/typeclass.ml
+++ b/typing/typeclass.ml
@@ -284,7 +284,7 @@ let make_method loc cl_num expr =
         pparam_loc = pat.ppat_loc;
       }
     ]
-    None (Pfunction_body expr)
+    {ret_type_constraint=None; ret_mode_annotations=[]; mode_annotations=[]} (Pfunction_body expr)
 
 (*******************************)
 
diff --git a/typing/typecore.ml b/typing/typecore.ml
index e5a7948cb17..773115fc68a 100644
--- a/typing/typecore.ml
+++ b/typing/typecore.ml
@@ -4381,10 +4381,10 @@ let type_approx_constraint ~loc env constraint_ ty_expected =
       ty_expected
 
 let type_approx_constraint_opt ~loc env constraint_ ty_expected =
-  match constraint_ with
+  let { ret_type_constraint; _} = constraint_ in
+  match ret_type_constraint with
   | None -> ty_expected
-  | Some { type_constraint; mode_annotations = _ } ->
-      type_approx_constraint ~loc env type_constraint ty_expected
+  | Some ty -> type_approx_constraint ~loc env ty ty_expected
 
 let type_approx_fun_one_param
     env loc label spato ty_expected ~first ~in_function
@@ -5036,7 +5036,7 @@ type split_function_ty =
 *)
 let split_function_ty
     env (expected_mode : expected_mode) ty_expected loc ~arg_label ~has_poly
-    ~mode_annots ~in_function ~is_first_val_param ~is_final_val_param
+    ~mode_annots ~ret_mode_annots ~in_function ~is_first_val_param ~is_final_val_param
   =
   let alloc_mode =
       (* Unlike most allocations which can be the highest mode allowed by
@@ -5073,6 +5073,7 @@ let split_function_ty
           generalize_structure ty_ret)
   in
   apply_mode_annots ~loc:loc_fun ~env mode_annots arg_mode;
+  apply_mode_annots ~loc:loc_fun ~env ret_mode_annots ret_mode;
   if not has_poly && not (tpoly_is_mono ty_arg) && !Clflags.principal
       && get_level ty_arg < Btype.generic_level then begin
     let snap = Btype.snapshot () in
@@ -6252,7 +6253,8 @@ and type_expect_
         exp_attributes = arg.exp_attributes;
         exp_env = env;
         exp_extra =
-          (Texp_constraint (Some extra_cty, modes),
+          (Texp_mode modes, loc, []) ::
+          (Texp_constraint extra_cty,
            loc,
            sexp.pexp_attributes) :: arg.exp_extra;
       }
@@ -6988,12 +6990,11 @@ and type_constraint env sty type_mode =
     @param loc_arg the location of the thing being constrained
 *)
 and type_constraint_expect
-  : type a. a constraint_arg -> _ -> _ -> _ -> loc_arg:_ -> _ -> _ -> a * _ * _
+  : type a. a constraint_arg -> _ -> _ -> _ -> loc_arg:_ -> _ -> _ -> _ -> a * _ * _
   =
-  fun constraint_arg env expected_mode loc ~loc_arg constraint_ ty_expected ->
+  fun constraint_arg env expected_mode loc ~loc_arg type_mode constraint_ ty_expected ->
   let ret, ty, exp_extra =
-    let { type_constraint = constraint_; mode_annotations } = constraint_ in
-    let type_mode = Typemode.transl_alloc_mode mode_annotations in
+    let type_mode = Alloc.Const.Option.value ~default:Alloc.Const.legacy type_mode in
     match constraint_ with
     | Pcoerce (ty_constrain, ty_coerce) ->
         type_coerce constraint_arg env expected_mode loc ty_constrain ty_coerce
@@ -7002,7 +7003,7 @@ and type_constraint_expect
         let ty, extra_cty = type_constraint env ty_constrain type_mode in
         constraint_arg.type_with_constraint env expected_mode ty,
         ty,
-        Texp_constraint (Some extra_cty, Mode.Alloc.Const.Option.none)
+        Texp_constraint extra_cty
   in
   unify_exp_types loc env ty (instance ty_expected);
   ret, ty, exp_extra
@@ -7142,6 +7143,16 @@ and type_function
                 | Pparam_newtype _ -> true)
               rest
       in
+      let ret_mode_annots =
+        match rest with
+        | [] ->
+          (* if this is the last parameter before the equal sign, then [ret_mode_annotation]
+            applies to the RHS of the equal sign. This is before [split_function_ty] which
+            enters new region. *)
+          let annots = Typemode.transl_mode_annots body_constraint.ret_mode_annotations in
+          annots
+        | _ :: _ -> Alloc.Const.Option.none
+      in
       let env,
           { filtered_arrow = { ty_arg; arg_mode; ty_ret; ret_mode };
             arg_sort; ret_sort;
@@ -7150,7 +7161,7 @@ and type_function
           } =
         split_function_ty env expected_mode ty_expected loc
           ~is_first_val_param:first ~is_final_val_param
-          ~arg_label:typed_arg_label ~in_function ~has_poly ~mode_annots
+          ~arg_label:typed_arg_label ~in_function ~has_poly ~mode_annots ~ret_mode_annots
       in
       (* [ty_arg_internal] is the type of the parameter viewed internally
          to the function. This is different than [ty_arg_mono] exactly for
@@ -7318,21 +7329,38 @@ and type_function
       }
   | [] ->
     let exp_type, body, fun_alloc_mode, ret_info =
+      let { ret_type_constraint; mode_annotations; ret_mode_annotations } =
+        body_constraint
+      in
+      let mode = Typemode.transl_mode_annots mode_annotations in
+      let ret_mode = Typemode.transl_mode_annots ret_mode_annotations in
+      let type_mode =
+        match ret_mode_annotations with
+        | _ :: _ ->
+            (* if return mode annotation is present, we use that to interpret the return
+               type annotation (currying mode behavior) *)
+            ret_mode
+        | [] ->
+            (* otherwise, we do not constrain the body mode, and we use the mode of the whole
+            function to interpret the return type, which is not precise and need to be fixed.
+            *)
+            mode
+      in
       match body with
       | Pfunction_body body ->
           let body =
-            match body_constraint with
+            match ret_type_constraint with
             | None -> type_expect env expected_mode body (mk_expected ty_expected)
             | Some constraint_ ->
-                let body_loc = body.pexp_loc in
-                let body, exp_type, exp_extra =
-                  type_constraint_expect (expression_constraint body)
-                    env expected_mode body_loc ~loc_arg:body_loc constraint_ ty_expected
-                in
-                { body with
-                    exp_extra = (exp_extra, body_loc, []) :: body.exp_extra;
-                    exp_type;
-                }
+            let body_loc = body.pexp_loc in
+            let body, exp_type, exp_extra =
+              type_constraint_expect (expression_constraint body)
+                env expected_mode body_loc ~loc_arg:body_loc type_mode constraint_ ty_expected
+            in
+            { body with
+                exp_extra = (exp_extra, body_loc, []) :: body.exp_extra;
+                exp_type;
+            }
           in
           body.exp_type, Tfunction_body body, None, None
       | Pfunction_cases (cases, _, attributes) ->
@@ -7342,7 +7370,7 @@ and type_function
               ~first
           in
           let (cases, exp_type, fun_alloc_mode, ret_info), exp_extra =
-            match body_constraint with
+            match ret_type_constraint with
             | None -> type_cases_expect env expected_mode ty_expected, None
             | Some constraint_ ->
               (* The typing of function case coercions/constraints is
@@ -7374,7 +7402,7 @@ and type_function
               in
               let (body, fun_alloc_mode, ret_info), exp_type, exp_extra =
                 type_constraint_expect function_cases_constraint_arg
-                  env expected_mode loc constraint_ ty_expected ~loc_arg:loc
+                  env expected_mode loc type_mode constraint_ ty_expected ~loc_arg:loc
               in
               (body, exp_type, fun_alloc_mode, ret_info), Some exp_extra
           in
@@ -8867,6 +8895,7 @@ and type_function_cases_expect
         } =
       split_function_ty env expected_mode ty_expected loc ~arg_label:Nolabel
         ~in_function ~has_poly:false ~mode_annots:Mode.Alloc.Const.Option.none
+        ~ret_mode_annots:Mode.Alloc.Const.Option.none
         ~is_first_val_param:first ~is_final_val_param:true
     in
     let cases, partial =
diff --git a/typing/typedtree.ml b/typing/typedtree.ml
index c1fc87d38f0..6263321b1b5 100644
--- a/typing/typedtree.ml
+++ b/typing/typedtree.ml
@@ -193,12 +193,13 @@ and expression =
    }
 
 and exp_extra =
-  | Texp_constraint of core_type option * Mode.Alloc.Const.Option.t
+  | Texp_constraint of core_type
   | Texp_coerce of core_type option * core_type
   | Texp_poly of core_type option
   | Texp_newtype of Ident.t * string loc *
                     Parsetree.jkind_annotation option * Uid.t
   | Texp_stack
+  | Texp_mode of Mode.Alloc.Const.Option.t
 
 and arg_label = Types.arg_label =
   | Nolabel
diff --git a/typing/typedtree.mli b/typing/typedtree.mli
index c4a2f2a52e6..d95ca7173a0 100644
--- a/typing/typedtree.mli
+++ b/typing/typedtree.mli
@@ -264,8 +264,8 @@ and expression =
    }
 
 and exp_extra =
-  | Texp_constraint of core_type option * Mode.Alloc.Const.Option.t
-        (** E : T @@ M *)
+  | Texp_constraint of core_type
+        (** E : T *)
   | Texp_coerce of core_type option * core_type
         (** E :> T           [Texp_coerce (None, T)]
             E : T0 :> T      [Texp_coerce (Some T0, T)]
@@ -282,6 +282,8 @@ and exp_extra =
         them here, as the cost of tracking this additional information is minimal. *)
   | Texp_stack
         (** stack_ E *)
+  | Texp_mode of Mode.Alloc.Const.Option.t
+        (** E : _ @@ M  *)
 
 and arg_label = Types.arg_label =
   | Nolabel
diff --git a/typing/untypeast.ml b/typing/untypeast.ml
index 63d9e580766..5da39d9a422 100644
--- a/typing/untypeast.ml
+++ b/typing/untypeast.ml
@@ -416,15 +416,13 @@ let exp_extra sub (extra, loc, attrs) sexp =
         Pexp_coerce (sexp,
                      Option.map (sub.typ sub) cty1,
                      sub.typ sub cty2)
-    | Texp_constraint (cty, modes) ->
-      Pexp_constraint
-        (sexp,
-         Option.map (sub.typ sub) cty,
-         Typemode.untransl_mode_annots ~loc modes)
+    | Texp_constraint (cty) ->
+        Pexp_constraint (sexp, Some (sub.typ sub cty), [])
     | Texp_poly cto -> Pexp_poly (sexp, Option.map (sub.typ sub) cto)
     | Texp_newtype (_, label_loc, jkind, _) ->
         Pexp_newtype (label_loc, jkind, sexp)
     | Texp_stack -> Pexp_stack sexp
+    | Texp_mode _ -> Misc.fatal_error "not supported Texp_mode"
   in
   Exp.mk ~loc ~attrs desc
 
@@ -497,30 +495,27 @@ let expression sub exp =
           | Tfunction_body body ->
               (* Unlike function cases, the [exp_extra] is placed on the body
                  itself. *)
-              Pfunction_body (sub.expr sub body), None
+              Pfunction_body (sub.expr sub body),
+              { mode_annotations = []; ret_type_constraint = None; ret_mode_annotations = []}
           | Tfunction_cases
               { fc_cases = cases; fc_loc = loc; fc_exp_extra = exp_extra;
                 fc_attributes = attributes; _ }
             ->
               let cases = List.map (sub.case sub) cases in
-              let constraint_ =
+              let ret_type_constraint =
                 match exp_extra with
                 | Some (Texp_coerce (ty1, ty2)) ->
                     Some
-                      (Pcoerce (Option.map (sub.typ sub) ty1, sub.typ sub ty2), [])
-                | Some (Texp_constraint (Some ty, modes)) ->
-                  Some (
-                    Pconstraint (sub.typ sub ty),
-                    Typemode.untransl_mode_annots ~loc modes
-                  )
-                | Some (Texp_poly _ | Texp_newtype _) | Some (Texp_constraint (None, _))
+                      (Pcoerce (Option.map (sub.typ sub) ty1, sub.typ sub ty2))
+                | Some (Texp_constraint ty) ->
+                  Some (Pconstraint (sub.typ sub ty))
+                | Some (Texp_mode _) -> Misc.fatal_error "not supported Texp_mode"
+                | Some (Texp_poly _ | Texp_newtype _)
                 | Some Texp_stack
                 | None -> None
               in
               let constraint_ =
-                Option.map
-                  (fun (type_constraint, mode_annotations) -> { mode_annotations; type_constraint })
-                  constraint_
+                { ret_type_constraint; mode_annotations=[]; ret_mode_annotations = [] }
               in
               Pfunction_cases (cases, loc, attributes), constraint_
         in

From 840be27532d331677e2436f685f0fe5a71c1a78c Mon Sep 17 00:00:00 2001
From: Zesen Qian <github@riaqn.org>
Date: Fri, 29 Nov 2024 09:41:14 -0500
Subject: [PATCH 4/6] update document

---
 jane/doc/extensions/modes/syntax.md | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/jane/doc/extensions/modes/syntax.md b/jane/doc/extensions/modes/syntax.md
index 7e86c6e8a39..f986c8c3581 100644
--- a/jane/doc/extensions/modes/syntax.md
+++ b/jane/doc/extensions/modes/syntax.md
@@ -62,7 +62,15 @@ let foo ((x, y) @ modes) = .. (error)
 let x @ modes = ..
 let x : ty @@ modes = ..
 let x : type a b. ty @@ modes = ..
+```
+
+## Functions and their return
+```ocaml
 let (foo @ modes) x y = ..
+let foo x y @ modes = ..
+let foo x y : ty @@ modes = ..
+fun foo x y @ modes -> ..
+fun foo x y : ty @@ modes -> ..
 ```
 
 ## Arrow types

From ac5d8f2d5ef32dea949eb4d917003abb8c2a005f Mon Sep 17 00:00:00 2001
From: Zesen Qian <github@riaqn.org>
Date: Fri, 6 Dec 2024 12:19:22 -0500
Subject: [PATCH 5/6] address comments

---
 parsing/printast.ml               | 5 ++---
 printer/printast_with_mappings.ml | 3 +--
 typing/typecore.ml                | 4 ++--
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/parsing/printast.ml b/parsing/printast.ml
index eb37dd1513a..cda06afc0fa 100644
--- a/parsing/printast.ml
+++ b/parsing/printast.ml
@@ -548,9 +548,8 @@ and type_constraint i ppf type_constraint =
       option (i+1) core_type ppf ty1;
       core_type (i+1) ppf ty2
 
-and function_constraint i ppf { ret_type_constraint = c; mode_annotations; ret_mode_annotations } =
-  option i type_constraint ppf c;
-  modes i ppf mode_annotations;
+and function_constraint i ppf { ret_type_constraint; ret_mode_annotations; mode_annotations = _ } =
+  option i type_constraint ppf ret_type_constraint;
   modes i ppf ret_mode_annotations
 
 and value_description i ppf x =
diff --git a/printer/printast_with_mappings.ml b/printer/printast_with_mappings.ml
index 89e32c499cc..4ac6d4e0941 100644
--- a/printer/printast_with_mappings.ml
+++ b/printer/printast_with_mappings.ml
@@ -575,8 +575,7 @@ and type_constraint i ppf type_constraint =
       core_type (i+1) ppf ty2
 
 and function_constraint i ppf
-  { ret_type_constraint; mode_annotations; ret_mode_annotations } =
-  modes i ppf mode_annotations;
+  { ret_type_constraint; ret_mode_annotations; mode_annotations = _ } =
   option i type_constraint ppf ret_type_constraint;
   modes i ppf ret_mode_annotations
 
diff --git a/typing/typecore.ml b/typing/typecore.ml
index 773115fc68a..0b4f242fc48 100644
--- a/typing/typecore.ml
+++ b/typing/typecore.ml
@@ -7342,8 +7342,8 @@ and type_function
             ret_mode
         | [] ->
             (* otherwise, we do not constrain the body mode, and we use the mode of the whole
-            function to interpret the return type, which is not precise and need to be fixed.
-            *)
+            function to interpret the return type *)
+            (* CR zqian: We should infer from [mode], instead of using directly. *)
             mode
       in
       match body with

From 8fd10b9e3cba6a322f09743bfc9c2b623b2cba00 Mon Sep 17 00:00:00 2001
From: Zesen Qian <github@riaqn.org>
Date: Tue, 10 Dec 2024 07:05:07 -0500
Subject: [PATCH 6/6] improve untypeast

---
 typing/untypeast.ml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/typing/untypeast.ml b/typing/untypeast.ml
index 5da39d9a422..007813a7deb 100644
--- a/typing/untypeast.ml
+++ b/typing/untypeast.ml
@@ -422,7 +422,8 @@ let exp_extra sub (extra, loc, attrs) sexp =
     | Texp_newtype (_, label_loc, jkind, _) ->
         Pexp_newtype (label_loc, jkind, sexp)
     | Texp_stack -> Pexp_stack sexp
-    | Texp_mode _ -> Misc.fatal_error "not supported Texp_mode"
+    | Texp_mode modes ->
+        Pexp_constraint (sexp, None, Typemode.untransl_mode_annots ~loc modes)
   in
   Exp.mk ~loc ~attrs desc
 
@@ -509,7 +510,7 @@ let expression sub exp =
                       (Pcoerce (Option.map (sub.typ sub) ty1, sub.typ sub ty2))
                 | Some (Texp_constraint ty) ->
                   Some (Pconstraint (sub.typ sub ty))
-                | Some (Texp_mode _) -> Misc.fatal_error "not supported Texp_mode"
+                | Some (Texp_mode _) (* CR zqian: [Texp_mode] should be possible here *)
                 | Some (Texp_poly _ | Texp_newtype _)
                 | Some Texp_stack
                 | None -> None