[WIP] - Preserve JSX by adding extra app flag

analysis/src/CompletionFrontEnd.ml

@@ -267,12 +267,14 @@ let rec exprToContextPathInner ~(inJsxContext : bool) (e : Parsetree.expression)
           };
         args =
           [(_, lhs); (_, {pexp_desc = Pexp_apply {funct = d; args; partial}})];
+        transformed_jsx;
       } ->
     (* Transform away pipe with apply call *)
     exprToContextPath ~inJsxContext
       {
         pexp_desc =
-          Pexp_apply {funct = d; args = (Nolabel, lhs) :: args; partial};
+          Pexp_apply
+            {funct = d; args = (Nolabel, lhs) :: args; partial; transformed_jsx};
         pexp_loc;
         pexp_attributes;
       }
@@ -284,6 +286,7 @@ let rec exprToContextPathInner ~(inJsxContext : bool) (e : Parsetree.expression)
             (_, lhs); (_, {pexp_desc = Pexp_ident id; pexp_loc; pexp_attributes});
           ];
         partial;
+        transformed_jsx;
       } ->
     (* Transform away pipe with identifier *)
     exprToContextPath ~inJsxContext
@@ -294,6 +297,7 @@ let rec exprToContextPathInner ~(inJsxContext : bool) (e : Parsetree.expression)
               funct = {pexp_desc = Pexp_ident id; pexp_loc; pexp_attributes};
               args = [(Nolabel, lhs)];
               partial;
+              transformed_jsx;
             };
         pexp_loc;
         pexp_attributes;

compiler/core/js_call_info.ml

@@ -33,10 +33,16 @@ type call_info =
    {[ fun x y -> (f x y) === f ]} when [f] is an atom
 *)
 
-type t = {call_info: call_info; arity: arity}
+type t = {
+  call_info: call_info;
+  arity: arity;
+  call_transformed_jsx: Parsetree.jsx_element option;
+}
 
-let dummy = {arity = NA; call_info = Call_na}
+let dummy = {arity = NA; call_info = Call_na; call_transformed_jsx = None}
 
-let builtin_runtime_call = {arity = Full; call_info = Call_builtin_runtime}
+let builtin_runtime_call =
+  {arity = Full; call_info = Call_builtin_runtime; call_transformed_jsx = None}
 
-let ml_full_call = {arity = Full; call_info = Call_ml}
+let ml_full_call =
+  {arity = Full; call_info = Call_ml; call_transformed_jsx = None}

compiler/core/js_call_info.mli

@@ -35,7 +35,11 @@ type call_info =
    {[ fun x y -> f x y === f ]} when [f] is an atom
 *)
 
-type t = {call_info: call_info; arity: arity}
+type t = {
+  call_info: call_info;
+  arity: arity;
+  call_transformed_jsx: Parsetree.jsx_element option;
+}
 
 val dummy : t
 

compiler/core/js_dump.ml

@@ -524,6 +524,10 @@ and expression_desc cxt ~(level : int) f x : cxt =
          when Ext_list.length_equal el i
        ]}
     *)
+  | Call (e, el, {call_transformed_jsx = Some jsx_element}) ->
+    (* The grand point would be to reconstruct the JSX here *)
+    P.string f "<meh />";
+    cxt
   | Call (e, el, info) ->
     P.cond_paren_group f (level > 15) (fun _ ->
         P.group f 0 (fun _ ->

compiler/core/lam.ml

@@ -81,7 +81,12 @@ module Types = struct
   *)
   and prim_info = {primitive: Lam_primitive.t; args: t list; loc: Location.t}
 
-  and apply = {ap_func: t; ap_args: t list; ap_info: ap_info}
+  and apply = {
+    ap_func: t;
+    ap_args: t list;
+    ap_info: ap_info;
+    ap_transformed_jsx: Parsetree.jsx_element option;
+  }
 
   and t =
     | Lvar of ident
@@ -121,7 +126,12 @@ module X = struct
     loc: Location.t;
   }
 
-  and apply = Types.apply = {ap_func: t; ap_args: t list; ap_info: ap_info}
+  and apply = Types.apply = {
+    ap_func: t;
+    ap_args: t list;
+    ap_info: ap_info;
+    ap_transformed_jsx: Parsetree.jsx_element option;
+  }
 
   and lfunction = Types.lfunction = {
     arity: int;
@@ -159,10 +169,10 @@ include Types
 let inner_map (l : t) (f : t -> X.t) : X.t =
   match l with
   | Lvar (_ : ident) | Lconst (_ : Lam_constant.t) -> ((* Obj.magic *) l : X.t)
-  | Lapply {ap_func; ap_args; ap_info} ->
+  | Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx} ->
     let ap_func = f ap_func in
     let ap_args = Ext_list.map ap_args f in
-    Lapply {ap_func; ap_args; ap_info}
+    Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx}
   | Lfunction {body; arity; params; attr} ->
     let body = f body in
     Lfunction {body; arity; params; attr}
@@ -279,7 +289,7 @@ let rec is_eta_conversion_exn params inner_args outer_args : t list =
   | _, _, _ -> raise_notrace Not_simple_form
 
 (** FIXME: more robust inlining check later, we should inline it before we add stub code*)
-let rec apply fn args (ap_info : ap_info) : t =
+let rec apply ?(ap_transformed_jsx = None) fn args (ap_info : ap_info) : t =
   match fn with
   | Lfunction
       {
@@ -300,15 +310,16 @@ let rec apply fn args (ap_info : ap_info) : t =
       Lprim
         {primitive = wrap; args = [Lprim {primitive_call with args; loc}]; loc}
     | exception Not_simple_form ->
-      Lapply {ap_func = fn; ap_args = args; ap_info})
+      Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx})
   | Lfunction
       {
         params;
         body = Lprim ({primitive = _; args = inner_args} as primitive_call);
       } -> (
     match is_eta_conversion_exn params inner_args args with
     | args -> Lprim {primitive_call with args; loc = ap_info.ap_loc}
-    | exception _ -> Lapply {ap_func = fn; ap_args = args; ap_info})
+    | exception _ ->
+      Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx})
   | Lfunction
       {
         params;
@@ -321,17 +332,17 @@ let rec apply fn args (ap_info : ap_info) : t =
     | args ->
       Lsequence (Lprim {primitive_call with args; loc = ap_info.ap_loc}, const)
     | exception _ ->
-      Lapply {ap_func = fn; ap_args = args; ap_info}
+      Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx}
       (* | Lfunction {params;body} when Ext_list.same_length params args ->
           Ext_list.fold_right2 (fun p arg acc ->
             Llet(Strict,p,arg,acc)
           ) params args body *)
       (* TODO: more rigirous analysis on [let_kind] *))
   | Llet (kind, id, e, (Lfunction _ as fn)) ->
-    Llet (kind, id, e, apply fn args ap_info)
+    Llet (kind, id, e, apply fn args ap_info ~ap_transformed_jsx)
   (* | Llet (kind0, id0, e0, Llet (kind,id, e, (Lfunction _ as fn))) ->
      Llet(kind0,id0,e0,Llet (kind, id, e, apply fn args loc status)) *)
-  | _ -> Lapply {ap_func = fn; ap_args = args; ap_info}
+  | _ -> Lapply {ap_func = fn; ap_args = args; ap_info; ap_transformed_jsx}
 
 let rec eq_approx (l1 : t) (l2 : t) =
   match l1 with

compiler/core/lam.mli

@@ -41,7 +41,12 @@ type lambda_switch = {
   sw_names: Ast_untagged_variants.switch_names option;
 }
 
-and apply = private {ap_func: t; ap_args: t list; ap_info: ap_info}
+and apply = private {
+  ap_func: t;
+  ap_args: t list;
+  ap_info: ap_info;
+  ap_transformed_jsx: Parsetree.jsx_element option;
+}
 
 and lfunction = {
   arity: int;
@@ -103,7 +108,12 @@ val global_module : ?dynamic_import:bool -> ident -> t
 
 val const : Lam_constant.t -> t
 
-val apply : t -> t list -> ap_info -> t
+val apply :
+  ?ap_transformed_jsx:Parsetree.jsx_element option ->
+  t ->
+  t list ->
+  ap_info ->
+  t
 
 val function_ :
   attr:Lambda.function_attribute ->

compiler/core/lam_bounded_vars.ml

@@ -108,10 +108,10 @@ let rewrite (map : _ Hash_ident.t) (lam : Lam.t) : Lam.t =
       (* here it makes sure that global vars are not rebound *)
       Lam.prim ~primitive ~args:(Ext_list.map args aux) loc
     | Lglobal_module _ -> lam
-    | Lapply {ap_func; ap_args; ap_info} ->
+    | Lapply {ap_func; ap_args; ap_info; ap_transformed_jsx} ->
       let fn = aux ap_func in
       let args = Ext_list.map ap_args aux in
-      Lam.apply fn args ap_info
+      Lam.apply ~ap_transformed_jsx fn args ap_info
     | Lswitch
         ( l,
           {

compiler/core/lam_compile.ml

@@ -31,12 +31,13 @@ let args_either_function_or_const (args : Lam.t list) =
       | Lfunction _ | Lconst _ -> true
       | _ -> false)
 
-let call_info_of_ap_status (ap_status : Lam.apply_status) : Js_call_info.t =
+let call_info_of_ap_status call_transformed_jsx (ap_status : Lam.apply_status) :
+    Js_call_info.t =
   (* XXX *)
   match ap_status with
-  | App_infer_full -> {arity = Full; call_info = Call_ml}
-  | App_uncurry -> {arity = Full; call_info = Call_na}
-  | App_na -> {arity = NA; call_info = Call_ml}
+  | App_infer_full -> {arity = Full; call_info = Call_ml; call_transformed_jsx}
+  | App_uncurry -> {arity = Full; call_info = Call_na; call_transformed_jsx}
+  | App_na -> {arity = NA; call_info = Call_ml; call_transformed_jsx}
 
 let rec apply_with_arity_aux (fn : J.expression) (arity : int list)
     (args : E.t list) (len : int) : E.t =
@@ -49,7 +50,14 @@ let rec apply_with_arity_aux (fn : J.expression) (arity : int list)
       if len >= x then
         let first_part, continue = Ext_list.split_at args x in
         apply_with_arity_aux
-          (E.call ~info:{arity = Full; call_info = Call_ml} fn first_part)
+          (E.call
+             ~info:
+               {
+                 arity = Full;
+                 call_info = Call_ml;
+                 (* no clue if this is correct *) call_transformed_jsx = None;
+               }
+             fn first_part)
           rest continue (len - x)
       else if
         (* GPR #1423 *)
@@ -63,7 +71,13 @@ let rec apply_with_arity_aux (fn : J.expression) (arity : int list)
           [
             S.return_stmt
               (E.call
-                 ~info:{arity = Full; call_info = Call_ml}
+                 ~info:
+                   {
+                     arity = Full;
+                     call_info = Call_ml;
+                     (* no clue if this is correct *) call_transformed_jsx =
+                       None;
+                   }
                  fn
                  (Ext_list.append args @@ Ext_list.map params E.var));
           ]
@@ -306,7 +320,9 @@ let compile output_prefix =
       let expression =
         match appinfo.ap_info.ap_status with
         | (App_infer_full | App_uncurry) as ap_status ->
-          E.call ~info:(call_info_of_ap_status ap_status) fn args
+          E.call
+            ~info:(call_info_of_ap_status appinfo.ap_transformed_jsx ap_status)
+            fn args
         | App_na -> (
           match ident_info.arity with
           | Submodule _ | Single Arity_na ->
@@ -1439,14 +1455,17 @@ let compile output_prefix =
      ap_func =
        Lapply {ap_func; ap_args; ap_info = {ap_status = App_na; ap_inlined}};
      ap_info = {ap_status = App_na} as outer_ap_info;
+     ap_transformed_jsx;
     } ->
       (* After inlining, we can generate such code, see {!Ari_regress_test}*)
       let ap_info =
         if outer_ap_info.ap_inlined = ap_inlined then outer_ap_info
         else {outer_ap_info with ap_inlined}
       in
       compile_lambda lambda_cxt
-        (Lam.apply ap_func (Ext_list.append ap_args appinfo.ap_args) ap_info)
+        (Lam.apply ap_func
+           (Ext_list.append ap_args appinfo.ap_args)
+           ap_info ~ap_transformed_jsx)
     (* External function call: it can not be tailcall in this case*)
     | {
      ap_func =
@@ -1529,7 +1548,9 @@ let compile output_prefix =
         Js_output.output_of_block_and_expression lambda_cxt.continuation
           args_code
           (E.call
-             ~info:(call_info_of_ap_status appinfo.ap_info.ap_status)
+             ~info:
+               (call_info_of_ap_status appinfo.ap_transformed_jsx
+                  appinfo.ap_info.ap_status)
              fn_code args))
   and compile_prim (prim_info : Lam.prim_info)
       (lambda_cxt : Lam_compile_context.t) =

compiler/core/lam_compile_external_call.ml

@@ -287,7 +287,11 @@ let translate_ffi (cxt : Lam_compile_context.t) arg_types
     | _ ->
       let args, eff, dynamic = assemble_args_has_splice arg_types args in
       let args = if dynamic then E.variadic_args args else args in
-      add_eff eff (E.call ~info:{arity = Full; call_info = Call_na} fn args))
+      add_eff eff
+        (E.call
+           ~info:
+             {arity = Full; call_info = Call_na; call_transformed_jsx = None}
+           fn args))
   | Js_call
       {
         external_module_name = module_name;
@@ -302,20 +306,36 @@ let translate_ffi (cxt : Lam_compile_context.t) arg_types
     if splice then
       let args, eff, dynamic = assemble_args_has_splice arg_types args in
       let args = if dynamic then E.variadic_args args else args in
-      add_eff eff (E.call ~info:{arity = Full; call_info = Call_na} fn args)
+      add_eff eff
+        (E.call
+           ~info:
+             {arity = Full; call_info = Call_na; call_transformed_jsx = None}
+           fn args)
     else
       let args, eff = assemble_args_no_splice arg_types args in
-      add_eff eff @@ E.call ~info:{arity = Full; call_info = Call_na} fn args
+      add_eff eff
+      @@ E.call
+           ~info:
+             {arity = Full; call_info = Call_na; call_transformed_jsx = None}
+           fn args
   | Js_module_as_fn {external_module_name; splice} ->
     let fn = external_var external_module_name ~dynamic_import in
     if splice then
       let args, eff, dynamic = assemble_args_has_splice arg_types args in
       let args = if dynamic then E.variadic_args args else args in
-      add_eff eff (E.call ~info:{arity = Full; call_info = Call_na} fn args)
+      add_eff eff
+        (E.call
+           ~info:
+             {arity = Full; call_info = Call_na; call_transformed_jsx = None}
+           fn args)
     else
       let args, eff = assemble_args_no_splice arg_types args in
       (* TODO: fix in rest calling convention *)
-      add_eff eff (E.call ~info:{arity = Full; call_info = Call_na} fn args)
+      add_eff eff
+        (E.call
+           ~info:
+             {arity = Full; call_info = Call_na; call_transformed_jsx = None}
+           fn args)
   | Js_new {external_module_name = module_name; name = fn; splice; scopes} ->
     (* handle [@@new]*)
     (* This has some side effect, it will
@@ -362,14 +382,16 @@ let translate_ffi (cxt : Lam_compile_context.t) arg_types
         add_eff eff
           (let self = translate_scoped_access js_send_scopes self in
            E.call
-             ~info:{arity = Full; call_info = Call_na}
+             ~info:
+               {arity = Full; call_info = Call_na; call_transformed_jsx = None}
              (E.dot self name) args)
       else
         let args, eff = assemble_args_no_splice arg_types args in
         add_eff eff
           (let self = translate_scoped_access js_send_scopes self in
            E.call
-             ~info:{arity = Full; call_info = Call_na}
+             ~info:
+               {arity = Full; call_info = Call_na; call_transformed_jsx = None}
              (E.dot self name) args)
     | _ -> assert false)
   | Js_module_as_var module_name -> external_var module_name ~dynamic_import
@@ -384,7 +406,10 @@ let translate_ffi (cxt : Lam_compile_context.t) arg_types
         ~dynamic_import
     in
     if args = [] then e
-    else E.call ~info:{arity = Full; call_info = Call_na} e args
+    else
+      E.call
+        ~info:{arity = Full; call_info = Call_na; call_transformed_jsx = None}
+        e args
   | Js_module_as_class module_name ->
     let fn = external_var module_name ~dynamic_import in
     let args, eff = assemble_args_no_splice arg_types args in