apply @react.component to fundef

Author: mununki

cli/reactjs_jsx_ppx_v3.ml

@@ -20,6 +20,8 @@ let optionIdent = Lident "option"
 
 let optionalAttr = [ ({ txt = "optional"; loc = Location.none }, PStr []) ]
 
+let reactComponentAttr = [ ({ txt = "react.component"; loc = Location.none }, PStr []) ]
+
 let constantString ~loc str = Ast_helper.Exp.constant ~loc (Pconst_string (str, None))
 
 let recordWithOnlyKey ~loc = Exp.record ~loc
@@ -279,7 +281,7 @@ let jsxMapper () =
       | _ -> false
     in
     (* check if record which goes to Foo.make({ ... } as record) empty or not
-      if empty then change it to {key: None} only for upper case jsx
+      if empty then change it to {key: @optional None} only for upper case jsx
       This would be redundant regarding PR progress https://github.com/rescript-lang/syntax/pull/299
     *)
     let props = if isEmptyRecord record then recordWithOnlyKey ~loc else record in
@@ -388,21 +390,21 @@ let jsxMapper () =
         recursivelyTransformNamedArgsForMake mapper expression
           ((arg, default, pattern, alias, pattern.ppat_loc, type_) :: args) newtypes
     | Pexp_fun (Nolabel, _, { ppat_desc = Ppat_construct ({ txt = Lident "()" }, _) | Ppat_any }, _expression) ->
-        (args, newtypes, None)
+        (args, newtypes)
     | Pexp_fun
         ( Nolabel,
           _,
-          { ppat_desc = Ppat_var { txt } | Ppat_constraint ({ ppat_desc = Ppat_var { txt } }, _) },
+          { ppat_desc = Ppat_var _ | Ppat_constraint ({ ppat_desc = Ppat_var _ }, _) },
           _expression ) ->
-        (args, newtypes, Some txt)
+        (args, newtypes)
     | Pexp_fun (Nolabel, _, pattern, _expression) ->
         Location.raise_errorf ~loc:pattern.ppat_loc
           "React: react.component refs only support plain arguments and type annotations."
     | Pexp_newtype (label, expression) ->
         recursivelyTransformNamedArgsForMake mapper expression args (label :: newtypes)
     | Pexp_constraint (expression, _typ) ->
         recursivelyTransformNamedArgsForMake mapper expression args newtypes
-    | _ -> (args, newtypes, None)
+    | _ -> (args, newtypes)
     [@@raises Invalid_argument]
   in
 
@@ -546,14 +548,14 @@ let jsxMapper () =
               in
               spelunkForFunExpression expression
             in
+            let wrapExpressionWithBinding expressionFn expression =
+              Vb.mk ~loc:bindingLoc
+                ~attrs:(List.filter otherAttrsPure binding.pvb_attributes)
+                (Pat.var ~loc:bindingPatLoc { loc = bindingPatLoc; txt = fnName })
+                (expressionFn expression)
+            in
             let modifiedBinding binding =
               let hasApplication = ref false in
-              let wrapExpressionWithBinding expressionFn expression =
-                Vb.mk ~loc:bindingLoc
-                  ~attrs:(List.filter otherAttrsPure binding.pvb_attributes)
-                  (Pat.var ~loc:bindingPatLoc { loc = bindingPatLoc; txt = fnName })
-                  (expressionFn expression)
-              in
               let expression = binding.pvb_expr in
               let unerasableIgnoreExp exp =
                 { exp with pexp_attributes = unerasableIgnore emptyLoc :: exp.pexp_attributes }
@@ -616,25 +618,10 @@ let jsxMapper () =
             in
             let bindingWrapper, _hasUnit, expression = modifiedBinding binding in
             (* do stuff here! *)
-            let namedArgList, _newtypes, _forwardRef =
+            let namedArgList, _newtypes =
               recursivelyTransformNamedArgsForMake mapper (modifiedBindingOld binding) [] []
             in
             let namedTypeList = List.fold_left argToType [] namedArgList in
-            let vbIgnoreUnusedRef = Vb.mk (Pat.any ()) (Exp.ident (Location.mknoloc (Lident "ref"))) in
-            let namedArgWithDefaultValueList = List.filter_map argWithDefaultValue namedArgList in
-            let vbMatch ((label, default)) =
-              Vb.mk (Pat.var (Location.mknoloc label))
-              (Exp.match_ (Exp.ident { txt = Lident label; loc=Location.none })
-              [
-                Exp.case
-                  (Pat.construct (Location.mknoloc @@ Lident "Some") (Some (Pat.var ( Location.mknoloc label))))
-                  (Exp.ident (Location.mknoloc @@ Lident label));
-                Exp.case
-                  (Pat.construct (Location.mknoloc @@ Lident "None") None)
-                  default
-              ])
-            in
-            let vbMatchList = List.map vbMatch namedArgWithDefaultValueList in
             (* type props = { ... } *)
             let propsRecordType =
               makePropsRecordType "props" emptyLoc
@@ -662,29 +649,8 @@ let jsxMapper () =
                     [ Vb.mk ~loc:emptyLoc (Pat.var ~loc:emptyLoc { loc = emptyLoc; txt }) fullExpression ]
                     (Exp.ident ~loc:emptyLoc { loc = emptyLoc; txt = Lident txt })
             in
-            let rec returnedExpression patterns ({ pexp_desc } as expr) =
-              match pexp_desc with
-              | Pexp_fun (_arg_label, _default, { ppat_desc = Ppat_construct ({ txt = Lident "()" }, _) | Ppat_any }, expr) ->
-                (patterns, expr)
-              | Pexp_fun (arg_label, _default, { ppat_loc }, expr) ->
-                returnedExpression (({loc = ppat_loc; txt = Lident (getLabel arg_label)}, Pat.var { txt = getLabel arg_label; loc = ppat_loc}) :: patterns) expr
-              | _ -> (patterns, expr)
-            in
-            let patternsWithLid, expression = returnedExpression [] expression in
-            let pattern = (Pat.record ((List.rev patternsWithLid) @ [(Location.mknoloc (Lident "ref"), Pat.var (Location.mknoloc "ref"))]) Closed)
-            in
-            (* add patttern matching for optional prop value *)
-            let expression = if List.length vbMatchList = 0 then expression else (Exp.let_ Nonrecursive vbMatchList expression) in
-            (* add let _ = ref to ignore unused warning *)
-            let expression = Exp.let_ Nonrecursive [ vbIgnoreUnusedRef ] expression in
-            let expression = Exp.fun_ Nolabel None
-              begin
-              Pat.constraint_ pattern
-                (Typ.constr ~loc:emptyLoc { txt = Lident "props"; loc=emptyLoc }
-                  (makePropsTypeParams namedTypeList))
-              end
-              expression
-            in 
+            (* add @react.component attr *)
+            let expression = { expression with pexp_attributes = reactComponentAttr } in
             (* let make = ({id, name, ...}: props<'id, 'name, ...>) => { ... } *)
             let bindings, newBinding =
               match recFlag with
@@ -699,7 +665,7 @@ let jsxMapper () =
                       (Exp.ident { loc = emptyLoc; txt = Lident fnName }));
                 ], None)
               | Nonrecursive ->
-                ([ { binding with pvb_expr = expression; pvb_attributes = [] } ], Some (bindingWrapper fullExpression))
+                ([ bindingWrapper expression ], Some (Vb.mk (Pat.var { loc = emptyLoc; txt = fnName }) fullExpression))
             in
             (Some propsRecordType, bindings, newBinding)
           else (None, [ binding ], None)
@@ -821,7 +787,58 @@ let jsxMapper () =
     match structure with structures -> default_mapper.structure mapper @@ reactComponentTransform mapper structures
     [@@raises Invalid_argument]
   in
-
+  
+  let transformComponentExpression mapper expression attrs =
+    match expression with
+    | { pexp_desc = Pexp_fun _ } ->
+      (* remove @react.component to prevent infinite loop *)
+      let expression = { expression with pexp_attributes = attrs } in
+      let emptyLoc = Location.none in
+      let namedArgList, _newtypes = recursivelyTransformNamedArgsForMake mapper expression [] [] in
+      let vbIgnoreUnusedRef = Vb.mk (Pat.any ()) (Exp.ident (Location.mknoloc (Lident "ref"))) in
+      let namedArgWithDefaultValueList = List.filter_map argWithDefaultValue namedArgList in
+      let namedTypeList = List.fold_left argToType [] namedArgList in
+      let vbMatch ((label, default)) =
+        Vb.mk (Pat.var (Location.mknoloc label))
+        (Exp.match_ (Exp.ident { txt = Lident label; loc=Location.none })
+        [
+          Exp.case
+            (Pat.construct (Location.mknoloc @@ Lident "Some") (Some (Pat.var ( Location.mknoloc label))))
+            (Exp.ident (Location.mknoloc @@ Lident label));
+          Exp.case
+            (Pat.construct (Location.mknoloc @@ Lident "None") None)
+            default
+        ])
+      in
+      let vbMatchList = List.map vbMatch namedArgWithDefaultValueList in
+      let rec returnedExpression patterns ({ pexp_desc } as expr) =
+        match pexp_desc with
+        | Pexp_fun (_arg_label, _default, { ppat_desc = Ppat_construct ({ txt = Lident "()" }, _) | Ppat_any }, expr) ->
+          (patterns, expr)
+        | Pexp_fun (arg_label, _default, { ppat_loc }, expr) when isLabelled arg_label || isOptional arg_label ->
+          returnedExpression (({loc = ppat_loc; txt = Lident (getLabel arg_label)}, Pat.var { txt = getLabel arg_label; loc = ppat_loc}) :: patterns) expr
+        | _ -> (patterns, expr)
+      in
+      let patternsWithLid, expression = returnedExpression [] expression in
+      let expression = mapper.expr mapper expression in
+      let pattern = (Pat.record ((List.rev patternsWithLid) @ [(Location.mknoloc (Lident "ref"), Pat.var (Location.mknoloc "ref"))]) Closed)
+      in
+      (* add patttern matching for optional prop value *)
+      let expression = if List.length vbMatchList = 0 then expression else (Exp.let_ Nonrecursive vbMatchList expression) in
+      (* add let _ = ref to ignore unused warning *)
+      let expression = Exp.let_ Nonrecursive [ vbIgnoreUnusedRef ] expression in
+      let expression = Exp.fun_ Nolabel None
+        begin
+        Pat.constraint_ pattern
+          (Typ.constr ~loc:emptyLoc { txt = Lident "props"; loc=emptyLoc }
+            (makePropsTypeParams namedTypeList))
+        end
+        expression
+      in
+      expression
+    | _ -> expression
+  in
+  
   let expr mapper expression =
     match expression with
     (* Does the function application have the @JSX attribute? *)
@@ -857,6 +874,14 @@ let jsxMapper () =
               (* ReactDOMRe.createElement *)
               (Exp.ident ~loc { loc; txt = Ldot (Lident "ReactDOMRe", "createElement") })
               args )
+    | { pexp_desc = Pexp_fun _; pexp_attributes } ->
+      let reactComponentAttribute, nonReactComponentAttributes = List.partition hasAttr pexp_attributes in
+      begin
+      match (reactComponentAttribute, nonReactComponentAttributes) with
+      | [], _ -> default_mapper.expr mapper expression
+      | _, nonReactComponentAttributes ->
+        transformComponentExpression mapper expression nonReactComponentAttributes
+      end
     (* Delegate to the default mapper, a deep identity traversal *)
     | e -> default_mapper.expr mapper e
     [@@raises Invalid_argument]