Fix loss of LOC during AST0 translation

CHANGELOG.md

@@ -22,6 +22,8 @@
 
 - Fix partial application generalization for `...`. https://github.com/rescript-lang/rescript/pull/8343
 
+- Preserve JSX prop locations across the AST0 translation layer, fixing `0:0` editor diagnostics in PPX-related flows. https://github.com/rescript-lang/rescript/pull/8350
+
 #### :memo: Documentation
 
 #### :nail_care: Polish

compiler/bsc/rescript_compiler_main.ml

@@ -268,6 +268,10 @@ let command_line_flags : (string * Bsc_args.spec * string) array =
           Js_config.binary_ast := true;
           Js_config.syntax_only := true),
       "*internal* Generate binary .mli_ast and ml_ast and stop" );
+    ( "-bs-test-ast-conversion",
+      set Js_config.test_ast_conversion,
+      "*internal* Roundtrip the parsed AST through Parsetree0 before \
+       continuing" );
     ( "-bs-syntax-only",
       set Js_config.syntax_only,
       "*internal* Only check syntax" );

compiler/common/js_config.ml

@@ -43,6 +43,7 @@ let check_div_by_zero = ref true
 let get_check_div_by_zero () = !check_div_by_zero
 let syntax_only = ref false
 let binary_ast = ref false
+let test_ast_conversion = ref false
 let debug = ref false
 let cmi_only = ref false
 let cmj_only = ref false

compiler/common/js_config.mli

@@ -65,6 +65,8 @@ val syntax_only : bool ref
 
 val binary_ast : bool ref
 
+val test_ast_conversion : bool ref
+
 val debug : bool ref
 
 val cmi_only : bool ref

compiler/common/ml_binary.ml

@@ -52,6 +52,12 @@ let ast0_to_signature : ast0 -> Parsetree.signature = function
     Ast_mapper_from0.default_mapper.signature Ast_mapper_from0.default_mapper
       sig0
 
+let ast0_roundtrip : type a. a kind -> a -> a =
+ fun kind ast ->
+  match kind with
+  | Ml -> ast |> to_ast0 Ml |> ast0_to_structure
+  | Mli -> ast |> to_ast0 Mli |> ast0_to_signature
+
 let magic_of_kind : type a. a kind -> string = function
   | Ml -> Config.ast_impl_magic_number
   | Mli -> Config.ast_intf_magic_number

compiler/common/ml_binary.mli

@@ -32,5 +32,6 @@ type ast0 = Impl of Parsetree0.structure | Intf of Parsetree0.signature
 val magic_of_kind : 'a kind -> string
 val magic_of_ast0 : ast0 -> string
 val to_ast0 : 'a kind -> 'a -> ast0
+val ast0_roundtrip : 'a kind -> 'a -> 'a
 val ast0_to_structure : ast0 -> Parsetree.structure
 val ast0_to_signature : ast0 -> Parsetree.signature

compiler/core/cmd_ppx_apply.ml

@@ -95,7 +95,10 @@ let rewrite kind ppxs ast =
 
 let apply_rewriters_str ?(restore = true) ~tool_name ast =
   match !Clflags.all_ppx with
-  | [] -> ast
+  | [] ->
+    if !Js_config.test_ast_conversion then
+      Ml_binary.ast0_roundtrip Ml_binary.Ml ast
+    else ast
   | ppxs ->
     ast
     |> Ast_mapper.add_ppx_context_str ~tool_name
@@ -104,7 +107,10 @@ let apply_rewriters_str ?(restore = true) ~tool_name ast =
 
 let apply_rewriters_sig ?(restore = true) ~tool_name ast =
   match !Clflags.all_ppx with
-  | [] -> ast
+  | [] ->
+    if !Js_config.test_ast_conversion then
+      Ml_binary.ast0_roundtrip Ml_binary.Mli ast
+    else ast
   | ppxs ->
     ast
     |> Ast_mapper.add_ppx_context_sig ~tool_name

compiler/ml/ast_mapper_from0.ml

@@ -25,6 +25,19 @@ open Ast_helper
 open Location
 module Pt = Parsetree
 
+let jsx_prop_loc_attr = "res.jsxPropLoc"
+let jsx_spread_loc_attr = "res.jsxSpreadLoc"
+
+let extract_internal_loc_attr attr_name attrs =
+  let rec loop rev_acc = function
+    | [] -> (None, List.rev rev_acc)
+    | (({txt; loc}, payload) as attr) :: rest ->
+      if txt = attr_name && payload = PStr [] then
+        (Some loc, List.rev_append rev_acc rest)
+      else loop (attr :: rev_acc) rest
+  in
+  loop [] attrs
+
 type mapper = {
   attribute: mapper -> attribute -> Pt.attribute;
   attributes: mapper -> attribute list -> Pt.attribute list;
@@ -331,9 +344,18 @@ module E = struct
 
   let try_map_jsx_prop (sub : mapper) (lbl : Asttypes.Noloc.arg_label)
       (e : expression) : Parsetree.jsx_prop option =
+    let map_expr_with_loc_attr attr_name fallback make_prop =
+      let loc, attrs = extract_internal_loc_attr attr_name e.pexp_attributes in
+      let e = {e with pexp_attributes = attrs} in
+      let expr = sub.expr sub e in
+      make_prop (match loc with Some loc -> loc | None -> fallback expr) expr
+    in
     match (lbl, e) with
-    | Asttypes.Noloc.Labelled "_spreadProps", expr ->
-      Some (Parsetree.JSXPropSpreading (Location.none, sub.expr sub expr))
+    | Asttypes.Noloc.Labelled "_spreadProps", _expr ->
+      Some
+        (map_expr_with_loc_attr jsx_spread_loc_attr
+           (fun expr -> expr.pexp_loc)
+           (fun loc expr -> Parsetree.JSXPropSpreading (loc, expr)))
     | ( Asttypes.Noloc.Labelled name,
         {pexp_desc = Pexp_ident {txt = Longident.Lident v}; pexp_loc = name_loc}
       )
@@ -344,14 +366,18 @@ module E = struct
       )
       when name = v ->
       Some (Parsetree.JSXPropPunning (true, {txt = name; loc = name_loc}))
-    | Asttypes.Noloc.Labelled name, exp ->
+    | Asttypes.Noloc.Labelled name, _exp ->
       Some
-        (Parsetree.JSXPropValue
-           ({txt = name; loc = Location.none}, false, sub.expr sub exp))
-    | Asttypes.Noloc.Optional name, exp ->
+        (map_expr_with_loc_attr jsx_prop_loc_attr
+           (fun expr -> expr.pexp_loc)
+           (fun loc expr ->
+             Parsetree.JSXPropValue ({txt = name; loc}, false, expr)))
+    | Asttypes.Noloc.Optional name, _exp ->
       Some
-        (Parsetree.JSXPropValue
-           ({txt = name; loc = Location.none}, true, sub.expr sub exp))
+        (map_expr_with_loc_attr jsx_prop_loc_attr
+           (fun expr -> expr.pexp_loc)
+           (fun loc expr ->
+             Parsetree.JSXPropValue ({txt = name; loc}, true, expr)))
     | _ -> None
 
   let extract_props_and_children (sub : mapper) items =

compiler/ml/ast_mapper_to0.ml

@@ -25,6 +25,13 @@ open Ast_helper0
 open Location
 module Pt = Parsetree0
 
+let jsx_prop_loc_attr = "res.jsxPropLoc"
+let jsx_spread_loc_attr = "res.jsxSpreadLoc"
+
+let wrap_with_loc_attr attr_name loc (expr : Pt.expression) =
+  let attr : Pt.attribute = (Location.mkloc attr_name loc, Pt.PStr []) in
+  {expr with pexp_attributes = attr :: expr.pexp_attributes}
+
 type mapper = {
   attribute: mapper -> attribute -> Pt.attribute;
   attributes: mapper -> attribute list -> Pt.attribute list;
@@ -334,9 +341,12 @@ module E = struct
              if is_optional then Asttypes.Noloc.Optional name.txt
              else Asttypes.Noloc.Labelled name.txt
            in
-           (label, sub.expr sub value)
-         | JSXPropSpreading (_, value) ->
-           (Asttypes.Noloc.Labelled "_spreadProps", sub.expr sub value))
+           ( label,
+             sub.expr sub value |> wrap_with_loc_attr jsx_prop_loc_attr name.loc
+           )
+         | JSXPropSpreading (loc, value) ->
+           ( Asttypes.Noloc.Labelled "_spreadProps",
+             sub.expr sub value |> wrap_with_loc_attr jsx_spread_loc_attr loc ))
 
   let map_jsx_children sub loc children =
     match children with

tests/build_tests/super_errors/expected/jsx_invalid_prop_ast0_conversion.res.expected

@@ -0,0 +1,12 @@
+
+  We've found a bug for you!
+  /.../fixtures/jsx_invalid_prop_ast0_conversion.res:18:20-22
+
+  16 │ }
+  17 │ 
+  18 │ let _ = <Component bar="hello" />
+  19 │ 
+
+  The prop bar does not belong to the JSX component <Component /> 
+
+

tests/build_tests/super_errors/fixtures/jsx_invalid_prop_ast0_conversion.res

@@ -0,0 +1,18 @@
+@@config({flags: ["-bs-test-ast-conversion"]})
+
+module React = {
+  type element = Jsx.element
+  @val external null: element = "null"
+  type componentLike<'props, 'return> = Jsx.componentLike<'props, 'return>
+  type component<'props> = Jsx.component<'props>
+  external component: componentLike<'props, element> => component<'props> = "%identity"
+  @module("react/jsx-runtime")
+  external jsx: (component<'props>, 'props) => element = "jsx"
+}
+
+module Component = {
+  @react.component
+  let make = (~foo: string) => React.null
+}
+
+let _ = <Component bar="hello" />

tests/ounit_tests/ounit_jsx_loc_tests.ml

@@ -0,0 +1,91 @@
+let ( >:: ), ( >::: ) = OUnit.(( >:: ), ( >::: ))
+let assert_equal = OUnit.assert_equal
+let assert_failure = OUnit.assert_failure
+
+let parse_structure source =
+  Res_driver.parse_implementation_from_source ~for_printer:false
+    ~display_filename:"JsxLocTest.res" ~source
+  |> fun result -> result.parsetree
+
+let roundtrip_structure source =
+  parse_structure source |> Ml_binary.to_ast0 Ml_binary.Ml
+  |> Ml_binary.ast0_to_structure
+
+let get_jsx_props structure =
+  let from_expr = function
+    | {
+        Parsetree.pexp_desc =
+          Pexp_jsx_element (Jsx_unary_element {jsx_unary_element_props = props});
+        _;
+      } ->
+      Some props
+    | {
+        Parsetree.pexp_desc =
+          Pexp_jsx_element
+            (Jsx_container_element {jsx_container_element_props = props});
+        _;
+      } ->
+      Some props
+    | _ -> None
+  in
+  let from_item = function
+    | {Parsetree.pstr_desc = Pstr_value (_, bindings); _} ->
+      bindings
+      |> List.find_map (fun {Parsetree.pvb_expr; _} -> from_expr pvb_expr)
+    | {pstr_desc = Pstr_eval (expr, _); _} -> from_expr expr
+    | _ -> None
+  in
+  match List.find_map from_item structure with
+  | Some props -> props
+  | None ->
+    assert_failure
+      (Format.asprintf "Expected a JSX unary element binding in:@.%a"
+         Pprintast.structure structure)
+
+let get_roundtripped_props source =
+  let original = parse_structure source |> get_jsx_props in
+  let roundtripped = roundtrip_structure source |> get_jsx_props in
+  (original, roundtripped)
+
+let assert_same_loc expected actual =
+  let to_tuple (loc : Location.t) =
+    ( loc.loc_start.pos_lnum,
+      loc.loc_start.pos_bol,
+      loc.loc_start.pos_cnum,
+      loc.loc_end.pos_lnum,
+      loc.loc_end.pos_bol,
+      loc.loc_end.pos_cnum )
+  in
+  assert_equal ~printer:(fun loc ->
+      let sl, sb, sc, el, eb, ec = to_tuple loc in
+      Printf.sprintf "(%d,%d,%d)-(%d,%d,%d)" sl sb sc el eb ec)
+    expected actual
+
+let test_jsx_prop_value_loc_roundtrip _ =
+  let source = {|
+let _ = <Comp foo=bar />
+|} in
+  let original, roundtripped = get_roundtripped_props source in
+  match (original, roundtripped) with
+  | [Parsetree.JSXPropValue (original_name, _, _)], [JSXPropValue (name, _, _)]
+    ->
+    assert_same_loc original_name.loc name.loc
+  | _ -> assert_failure "Expected one JSX prop value"
+
+let test_jsx_spread_loc_roundtrip _ =
+  let source = {|
+let _ = <Comp {...props} foo=bar />
+|} in
+  let original, roundtripped = get_roundtripped_props source in
+  match (original, roundtripped) with
+  | Parsetree.JSXPropSpreading (original_loc, _) :: _,
+    JSXPropSpreading (loc, _) :: _ ->
+    assert_same_loc original_loc loc
+  | _ -> assert_failure "Expected a leading JSX spread prop"
+
+let suites =
+  __FILE__
+  >::: [
+         "prop_value_roundtrip" >:: test_jsx_prop_value_loc_roundtrip;
+         "spread_roundtrip" >:: test_jsx_spread_loc_roundtrip;
+       ]

tests/ounit_tests/ounit_tests_main.ml

@@ -18,6 +18,7 @@ let suites =
       Ounit_utf8_test.suites;
       Ounit_unicode_tests.suites;
       Ounit_util_tests.suites;
+      Ounit_jsx_loc_tests.suites;
     ]
 
 let _ = OUnit.run_test_tt_main suites