JSX PPX: add Jsx.element return constraint

Author: cknitt

compiler/syntax/src/jsx_v4.ml

@@ -46,6 +46,38 @@ let ref_type loc =
     {loc; txt = Ldot (Ldot (Lident "Js", "Nullable"), "t")}
     [ref_type_var loc]
 
+let jsx_element_type ~loc =
+  Typ.constr ~loc {loc; txt = Ldot (Lident "Jsx", "element")} []
+
+let jsx_element_constraint ~loc expr =
+  Exp.constraint_ ~loc expr (jsx_element_type ~loc)
+
+let rec constrain_jsx_return ~loc expr =
+  match expr.pexp_desc with
+  | Pexp_fun ({rhs} as desc) ->
+    {
+      expr with
+      pexp_desc = Pexp_fun {desc with rhs = constrain_jsx_return ~loc rhs};
+    }
+  | Pexp_newtype (param, inner) ->
+    {
+      expr with
+      pexp_desc = Pexp_newtype (param, constrain_jsx_return ~loc inner);
+    }
+  | Pexp_constraint (inner, _) ->
+    jsx_element_constraint ~loc (constrain_jsx_return ~loc inner)
+  | Pexp_let (rec_flag, bindings, body) ->
+    {
+      expr with
+      pexp_desc = Pexp_let (rec_flag, bindings, constrain_jsx_return ~loc body);
+    }
+  | Pexp_sequence (first, second) ->
+    {
+      expr with
+      pexp_desc = Pexp_sequence (first, constrain_jsx_return ~loc second);
+    }
+  | _ -> jsx_element_constraint ~loc expr
+
 let merlin_focus = ({loc = Location.none; txt = "merlin.focus"}, PStr [])
 
 (* Helper method to filter out any attribute that isn't [@react.component] *)
@@ -713,6 +745,10 @@ 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 =
+      Exp.constraint_ ~loc:binding_loc expression
+        (jsx_element_type ~loc:binding_loc)
+    in
     (* (ref) => expr *)
     let expression =
       List.fold_left
@@ -784,6 +820,7 @@ let map_binding ~config ~empty_loc ~pstr_loc ~file_name ~rec_flag binding =
     (Some props_record_type, binding, new_binding))
   else if Jsx_common.has_attr_on_binding Jsx_common.has_attr_with_props binding
   then
+    let binding_loc = binding.pvb_loc in
     let modified_binding =
       {
         binding with
@@ -839,21 +876,28 @@ 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 =
+        Exp.constraint_ ~loc applied_expression (jsx_element_type ~loc)
+      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
@@ -874,18 +918,20 @@ let map_binding ~config ~empty_loc ~pstr_loc ~file_name ~rec_flag binding =
       | Recursive -> None
       | Nonrecursive ->
         Some
-          (make_new_binding ~loc:empty_loc ~full_module_name modified_binding)
+          (make_new_binding ~loc:binding_loc ~full_module_name modified_binding)
     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 = [];
-          };
+          ( binding.pvb_expr |> fun expr ->
+            {
+              expr with
+              (* moved to wrapper_expr *)
+              pexp_attributes = [];
+            }
+            |> constrain_jsx_return ~loc:binding_loc );
       },
       new_binding )
   else (None, binding, None)
@@ -934,7 +980,7 @@ let transform_structure_item ~config item =
             (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 = get_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 +1001,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 =
         {
@@ -1046,7 +1092,7 @@ let transform_signature_item ~config item =
             (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 = get_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 +1113,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 =
         {

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/CodeLens.res.txt

@@ -1,5 +1,8 @@
 Code Lens src/CodeLens.res
 [{
+        "range": {"start": {"line": 9, "character": 4}, "end": {"line": 9, "character": 8}},
+        "command": {"title": "Jsx.element", "command": ""}
+    }, {
         "range": {"start": {"line": 4, "character": 4}, "end": {"line": 4, "character": 6}},
         "command": {"title": "(~opt1: int=?, ~a: int, ~b: int, unit, ~opt2: int=?, unit, ~c: int) => int", "command": ""}
     }, {

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/Hover.res.txt

@@ -22,10 +22,10 @@ Hover src/Hover.res 33:4
 {"contents": {"kind": "markdown", "value": "```rescript\nunit => int\n```\n---\nDoc comment for functionWithTypeAnnotation"}}
 
 Hover src/Hover.res 37:13
-{"contents": {"kind": "markdown", "value": "```rescript\nstring\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/Hover.res 42:15
-{"contents": {"kind": "markdown", "value": "```rescript\nstring\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/Hover.res 46:10
 {"contents": {"kind": "markdown", "value": "```rescript\nint\n```"}}

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

@@ -35,6 +35,12 @@ Inlay Hint src/InlayHint.res 1:34
     "kind": 1,
     "paddingLeft": true,
     "paddingRight": false
+}, {
+    "position": {"line": 14, "character": 8},
+    "label": ": Jsx.element",
+    "kind": 1,
+    "paddingLeft": true,
+    "paddingRight": false
 }, {
     "position": {"line": 10, "character": 10},
     "label": ": (~xx: int) => int",

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```"}}