JSX PPX: add Jsx.element return constraint (#7939)

* JSX PPX: add Jsx.element return constraint

* Extract collect_prop_types helper

* Simplify collect_prop_types

* Cleanup

* Remove merlin.focus attribute

* CHANGELOG

# Conflicts:
#	CHANGELOG.md

* Update JSXV4 docs.

* Constrain return expression; do not use outer loc

* Revert binding loc change

Author: cknitt

CHANGELOG.md

@@ -25,6 +25,7 @@
 #### :nail_care: Polish
 
 - Improve circular dependency errors, and make sure they end up in the compiler log so the editor tooling can surface them. https://github.com/rescript-lang/rescript/pull/7940
+- JSX PPX: add Jsx.element return constraint. https://github.com/rescript-lang/rescript/pull/7939
 
 #### :house: Internal
 

compiler/syntax/src/jsx_v4.ml

@@ -46,7 +46,33 @@ let ref_type loc =
     {loc; txt = Ldot (Ldot (Lident "Js", "Nullable"), "t")}
     [ref_type_var loc]
 
-let merlin_focus = ({loc = Location.none; txt = "merlin.focus"}, PStr [])
+let jsx_element_type ~loc =
+  Typ.constr ~loc {loc; txt = Ldot (Lident "Jsx", "element")} []
+
+let jsx_element_constraint expr =
+  Exp.constraint_ expr (jsx_element_type ~loc:expr.pexp_loc)
+
+(* Traverse the component body and force every reachable return expression to
+   be annotated as `Jsx.element`. This walks through the wrapper constructs the
+   PPX introduces (fun/newtype/let/sequence) so that the constraint ends up on
+   the real return position even after we rewrite the function. *)
+let rec constrain_jsx_return expr =
+  match expr.pexp_desc with
+  | Pexp_fun ({rhs} as desc) ->
+    {expr with pexp_desc = Pexp_fun {desc with rhs = constrain_jsx_return rhs}}
+  | Pexp_newtype (param, inner) ->
+    {expr with pexp_desc = Pexp_newtype (param, constrain_jsx_return inner)}
+  | Pexp_constraint (inner, _) ->
+    let constrained_inner = constrain_jsx_return inner in
+    jsx_element_constraint constrained_inner
+  | Pexp_let (rec_flag, bindings, body) ->
+    {
+      expr with
+      pexp_desc = Pexp_let (rec_flag, bindings, constrain_jsx_return body);
+    }
+  | Pexp_sequence (first, second) ->
+    {expr with pexp_desc = Pexp_sequence (first, constrain_jsx_return second)}
+  | _ -> jsx_element_constraint expr
 
 (* Helper method to filter out any attribute that isn't [@react.component] *)
 let other_attrs_pure (loc, _) =
@@ -73,7 +99,7 @@ let make_new_binding binding expression new_name =
       pvb_pat =
         {pvb_pat with ppat_desc = Ppat_var {ppat_var with txt = new_name}};
       pvb_expr = expression;
-      pvb_attributes = [merlin_focus];
+      pvb_attributes = [];
     }
   | {pvb_loc} ->
     Jsx_common.raise_error ~loc:pvb_loc
@@ -713,6 +739,7 @@ let map_binding ~config ~empty_loc ~pstr_loc ~file_name ~rec_flag binding =
         vb_match_expr named_arg_list expression
       else expression
     in
+    let expression = constrain_jsx_return expression in
     (* (ref) => expr *)
     let expression =
       List.fold_left
@@ -839,21 +866,26 @@ let map_binding ~config ~empty_loc ~pstr_loc ~file_name ~rec_flag binding =
         | _ -> Pat.var {txt = "props"; loc}
       in
 
+      let applied_expression =
+        Exp.apply
+          (Exp.ident
+             {
+               txt =
+                 Lident
+                   (match rec_flag with
+                   | Recursive -> internal_fn_name
+                   | Nonrecursive -> fn_name);
+               loc;
+             })
+          [(Nolabel, Exp.ident {txt = Lident "props"; loc})]
+      in
+      let applied_expression =
+        Jsx_common.async_component ~async:is_async applied_expression
+      in
+      let applied_expression = constrain_jsx_return applied_expression in
       let wrapper_expr =
         Exp.fun_ ~arity:None Nolabel None props_pattern
-          ~attrs:binding.pvb_expr.pexp_attributes
-          (Jsx_common.async_component ~async:is_async
-             (Exp.apply
-                (Exp.ident
-                   {
-                     txt =
-                       Lident
-                         (match rec_flag with
-                         | Recursive -> internal_fn_name
-                         | Nonrecursive -> fn_name);
-                     loc;
-                   })
-                [(Nolabel, Exp.ident {txt = Lident "props"; loc})]))
+          ~attrs:binding.pvb_expr.pexp_attributes applied_expression
       in
 
       let wrapper_expr = Ast_uncurried.uncurried_fun ~arity:1 wrapper_expr in
@@ -876,20 +908,33 @@ let map_binding ~config ~empty_loc ~pstr_loc ~file_name ~rec_flag binding =
         Some
           (make_new_binding ~loc:empty_loc ~full_module_name modified_binding)
     in
+    let binding_expr =
+      {
+        binding.pvb_expr with
+        (* moved to wrapper_expr *)
+        pexp_attributes = [];
+      }
+    in
     ( None,
       {
         binding with
         pvb_attributes = binding.pvb_attributes |> List.filter other_attrs_pure;
-        pvb_expr =
-          {
-            binding.pvb_expr with
-            (* moved to wrapper_expr *)
-            pexp_attributes = [];
-          };
+        pvb_expr = binding_expr |> constrain_jsx_return;
       },
       new_binding )
   else (None, binding, None)
 
+let rec collect_prop_types types {ptyp_loc; ptyp_desc} =
+  match ptyp_desc with
+  | Ptyp_arrow {arg; ret = {ptyp_desc = Ptyp_arrow _} as rest}
+    when is_labelled arg.lbl || is_optional arg.lbl ->
+    collect_prop_types ((arg.lbl, arg.attrs, ptyp_loc, arg.typ) :: types) rest
+  | Ptyp_arrow {arg = {lbl = Nolabel}; ret} -> collect_prop_types types ret
+  | Ptyp_arrow {arg; ret = return_value}
+    when is_labelled arg.lbl || is_optional arg.lbl ->
+    (arg.lbl, arg.attrs, return_value.ptyp_loc, arg.typ) :: types
+  | _ -> types
+
 let transform_structure_item ~config item =
   match item with
   (* external *)
@@ -922,19 +967,7 @@ let transform_structure_item ~config item =
         |> Option.map Jsx_common.typ_vars_of_core_type
         |> Option.value ~default:[]
       in
-      let rec get_prop_types types ({ptyp_loc; ptyp_desc} as full_type) =
-        match ptyp_desc with
-        | Ptyp_arrow {arg; ret = {ptyp_desc = Ptyp_arrow _} as typ2}
-          when is_labelled arg.lbl || is_optional arg.lbl ->
-          get_prop_types ((arg.lbl, arg.attrs, ptyp_loc, arg.typ) :: types) typ2
-        | Ptyp_arrow {arg = {lbl = Nolabel}; ret} -> get_prop_types types ret
-        | Ptyp_arrow {arg; ret = return_value}
-          when is_labelled arg.lbl || is_optional arg.lbl ->
-          ( return_value,
-            (arg.lbl, arg.attrs, return_value.ptyp_loc, arg.typ) :: types )
-        | _ -> (full_type, types)
-      in
-      let inner_type, prop_types = get_prop_types [] pval_type in
+      let prop_types = collect_prop_types [] pval_type in
       let named_type_list = List.fold_left arg_to_concrete_type [] prop_types in
       let ret_props_type =
         Typ.constr ~loc:pstr_loc
@@ -955,7 +988,7 @@ let transform_structure_item ~config item =
       let new_external_type =
         Ptyp_constr
           ( {loc = pstr_loc; txt = module_access_name config "componentLike"},
-            [ret_props_type; inner_type] )
+            [ret_props_type; jsx_element_type ~loc:pstr_loc] )
       in
       let new_structure =
         {
@@ -1023,30 +1056,7 @@ let transform_signature_item ~config item =
         |> Option.map Jsx_common.typ_vars_of_core_type
         |> Option.value ~default:[]
       in
-      let rec get_prop_types types ({ptyp_loc; ptyp_desc} as full_type) =
-        match ptyp_desc with
-        | Ptyp_arrow {arg; ret = {ptyp_desc = Ptyp_arrow _} as rest}
-          when is_optional arg.lbl || is_labelled arg.lbl ->
-          get_prop_types ((arg.lbl, arg.attrs, ptyp_loc, arg.typ) :: types) rest
-        | Ptyp_arrow
-            {
-              arg =
-                {
-                  lbl = Nolabel;
-                  typ = {ptyp_desc = Ptyp_constr ({txt = Lident "unit"}, _)};
-                };
-              ret = rest;
-            } ->
-          get_prop_types types rest
-        | Ptyp_arrow {arg = {lbl = Nolabel}; ret = rest} ->
-          get_prop_types types rest
-        | Ptyp_arrow {arg; ret = return_value}
-          when is_optional arg.lbl || is_labelled arg.lbl ->
-          ( return_value,
-            (arg.lbl, arg.attrs, return_value.ptyp_loc, arg.typ) :: types )
-        | _ -> (full_type, types)
-      in
-      let inner_type, prop_types = get_prop_types [] pval_type in
+      let prop_types = collect_prop_types [] pval_type in
       let named_type_list = List.fold_left arg_to_concrete_type [] prop_types in
       let ret_props_type =
         Typ.constr
@@ -1067,7 +1077,7 @@ let transform_signature_item ~config item =
       let new_external_type =
         Ptyp_constr
           ( {loc = psig_loc; txt = module_access_name config "componentLike"},
-            [ret_props_type; inner_type] )
+            [ret_props_type; jsx_element_type ~loc:psig_loc] )
       in
       let new_structure =
         {

docs/JSXV4.md

@@ -272,7 +272,7 @@ is transformed to
 ```rescript
 type props<'x, 'y, 'z> = {x: 'x, y?: 'y, z?: 'z}
 
-({x, ?y, ?z}: props<_, _, _>) => {
+({x, ?y, ?z}: props<_, _, _>): Jsx.element => {
   let y = switch y {
   | None => 3 + x
   | Some(y) => y
@@ -281,7 +281,9 @@ type props<'x, 'y, 'z> = {x: 'x, y?: 'y, z?: 'z}
 }
 ```
 
-> Note: this implicit definition of type `props` means that there cannot be other type definitions of `props` in the same scope, or it will be a compiler error about multiple definitions of the type name.
+> Note:
+> - This implicit definition of type `props` means that there cannot be other type definitions of `props` in the same scope, or it will be a compiler error about multiple definitions of the type name.
+> - JSX V4 automatically adds a `Jsx.element` return constraint to all component functions. This ensures that components always return valid JSX elements and provides better type safety. If a component returns a non-JSX value, you'll get a helpful error message suggesting how to convert it (e.g., use `React.int` for integers, `React.string` for strings, etc.).
 
 ### Transformation for Component Application
 

tests/analysis_tests/tests/src/CompletionJsx.res

@@ -67,7 +67,7 @@ module MultiPropComp = {
   @react.component
   let make = (~name, ~age, ~time: time) => {
     ignore(time)
-    name ++ age
+    React.string(name ++ age)
   }
 }
 

tests/analysis_tests/tests/src/CompletionPattern.res

@@ -229,6 +229,7 @@ let make = (~thing: result<someVariant, unit>) => {
   //                     ^com
   | _ => ()
   }
+  React.null
 }
 
 type results = {

tests/analysis_tests/tests/src/CreateInterface.res

@@ -109,27 +109,27 @@ module type OptT = {
 
 module Opt = {
   @react.component
-  let withOpt1 = (~x=3, ~y) => x + y
+  let withOpt1 = (~x=3, ~y) => React.int(x + y)
 
   module Opt2 = {
     @react.component
     let withOpt2 = (~x: option<int>=?, ~y: int) =>
-      switch x {
+      React.int(switch x {
       | None => 0
       | Some(x) => x
       } +
-      y
+      y)
   }
   module type Opt2 = module type of Opt2
 
   module Opt3 = {
     @react.component
     let withOpt3 = (~x: option<int>, ~y: int) =>
-      switch x {
+      React.int(switch x {
       | None => 0
       | Some(x) => x
       } +
-      y
+      y)
   }
   module type Opt3 = module type of Opt3
 }

tests/analysis_tests/tests/src/expected/CompletionPattern.res.txt

@@ -1178,7 +1178,8 @@ Path res
   }]
 
 Complete src/CompletionPattern.res 227:25
-posCursor:[227:25] posNoWhite:[227:24] Found expr:[223:11->231:1]
+posCursor:[227:25] posNoWhite:[227:24] Found expr:[223:11->232:1]
+posCursor:[227:25] posNoWhite:[227:24] Found expr:[224:2->231:12]
 posCursor:[227:25] posNoWhite:[227:24] Found expr:[226:4->227:28]
 posCursor:[227:25] posNoWhite:[227:24] Found pattern:[227:18->227:27]
 Completable: Cpattern Value[r]->recordBody
@@ -1206,8 +1207,8 @@ Path r
     "documentation": {"kind": "markdown", "value": "```rescript\nnest: nestedRecord\n```\n\n```rescript\ntype someRecord = {first: int, second: (bool, option<someRecord>), optThird: option<[#first | #second(someRecord)]>, nest: nestedRecord}\n```"}
   }]
 
-Complete src/CompletionPattern.res 242:33
-posCursor:[242:33] posNoWhite:[242:32] Found pattern:[242:7->242:35]
+Complete src/CompletionPattern.res 243:33
+posCursor:[243:33] posNoWhite:[243:32] Found pattern:[243:7->243:35]
 Completable: Cpattern Value[hitsUse](Nolabel)->recordBody
 Package opens Stdlib.place holder Pervasives.JsxModules.place holder
 Resolved opens 1 Stdlib

tests/analysis_tests/tests/src/expected/CreateInterface.res.txt

@@ -2,18 +2,18 @@ Create Interface src/CreateInterface.res
 type r = {name: string, age: int}
 let add: (~x: int, ~y: int) => int
 @react.component
-let make: (~name: string) => React.element
+let make: (~name: string) => Jsx.element
 module Other: {
   @react.component
-  let otherComponentName: (~name: string) => React.element
+  let otherComponentName: (~name: string) => Jsx.element
 }
 module Mod: {
   @react.component
-  let make: (~name: string) => React.element
+  let make: (~name: string) => Jsx.element
 }
 module type ModTyp = {
   @react.component
-  let make: (~name: string) => React.element
+  let make: (~name: string) => Jsx.element
 }
 @module("path") external dirname: string => string = "dirname"
 @module("path") @variadic
@@ -49,21 +49,21 @@ module RFS: {
 module Functor: () =>
 {
   @react.component
-  let make: unit => React.element
+  let make: unit => Jsx.element
 }
 module type FT = {
   module Functor: (
     X: {
       let a: int
       @react.component
-      let make: (~name: string) => React.element
+      let make: (~name: string) => Jsx.element
       let b: int
     },
     Y: ModTyp,
   ) =>
   {
     @react.component
-    let make: (~name: string) => React.element
+    let make: (~name: string) => Jsx.element
   }
 }
 module NormaList = List
@@ -73,34 +73,34 @@ module rec RM: ModTyp
 and D: ModTyp
 module type OptT = {
   @react.component
-  let withOpt1: (~x: int=?, ~y: int) => int
+  let withOpt1: (~x: int=?, ~y: int) => Jsx.element
   module type Opt2 = {
     @react.component
-    let withOpt2: (~x: int=?, ~y: int) => int
+    let withOpt2: (~x: int=?, ~y: int) => Jsx.element
   }
   module type Opt3 = {
     @react.component
-    let withOpt3: (~x: option<int>, ~y: int) => int
+    let withOpt3: (~x: option<int>, ~y: int) => Jsx.element
   }
 }
 module Opt: {
   @react.component
-  let withOpt1: (~x: int=?, ~y: int) => int
+  let withOpt1: (~x: int=?, ~y: int) => Jsx.element
   module Opt2: {
     @react.component
-    let withOpt2: (~x: int=?, ~y: int) => int
+    let withOpt2: (~x: int=?, ~y: int) => Jsx.element
   }
   module type Opt2 = {
     @react.component
-    let withOpt2: (~x: int=?, ~y: int) => int
+    let withOpt2: (~x: int=?, ~y: int) => Jsx.element
   }
   module Opt3: {
     @react.component
-    let withOpt3: (~x: option<int>, ~y: int) => int
+    let withOpt3: (~x: option<int>, ~y: int) => Jsx.element
   }
   module type Opt3 = {
     @react.component
-    let withOpt3: (~x: option<int>, ~y: int) => int
+    let withOpt3: (~x: option<int>, ~y: int) => Jsx.element
   }
 }
 module Opt2: OptT

tests/analysis_tests/tests/src/expected/Jsx2.resi.txt

@@ -1,5 +1,5 @@
 Hover src/Jsx2.resi 1:4
-{"contents": {"kind": "markdown", "value": "```rescript\nprops<string>\n```\n\n---\n\n```\n \n```\n```rescript\ntype props<'first> = {first: 'first}\n```\nGo to: [Type definition](command:rescript-vscode.go_to_location?%5B%22Jsx2.resi%22%2C0%2C0%5D)\n"}}
+{"contents": {"kind": "markdown", "value": "```rescript\nJsx.element\n```\n\n---\n\n```\n \n```\n```rescript\ntype Jsx.element\n```\nGo to: [Type definition](command:rescript-vscode.go_to_location?%5B%22Jsx.res%22%2C25%2C0%5D)\n"}}
 
 Hover src/Jsx2.resi 4:4
 {"contents": {"kind": "markdown", "value": "```rescript\nint\n```"}}