@Directives in functions (#6756)

zth · web-flow · commit b1d5a435707f · 2024-05-16T13:34:39.000+02:00
* base implementation for directives in functions

* prevent inlining of functions with directives

* changelog
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,10 @@
 
 # 12.0.0-alpha.1 (Unreleased)
 
+#### :rocket: New Feature
+
+- Allow `@directive` on functions for emitting function level directive code (`let serverAction = @directive("'use server'") (~name) => {...}`). https://github.com/rescript-lang/rescript-compiler/pull/6756
+
 #### :boom: Breaking Change
 
 - `lazy` syntax is no longer supported. If you're using it, use `Lazy` module or `React.lazy_` instead. https://github.com/rescript-lang/rescript-compiler/pull/6342
diff --git a/jscomp/core/j.ml b/jscomp/core/j.ml
@@ -135,6 +135,7 @@ and expression_desc =
       env : Js_fun_env.t;
       return_unit : bool;
       async : bool;
+      directive : string option;
     }
   | Str of { delim : delim; txt : string }
   (* A string is UTF-8 encoded, and may contain
diff --git a/jscomp/core/js_dump.ml b/jscomp/core/js_dump.ml
@@ -282,7 +282,7 @@ let rec try_optimize_curry cxt f len function_id =
   Curry_gen.pp_optimize_curry f len;
   P.paren_group f 1 (fun _ -> expression ~level:1 cxt f function_id)
 
-and pp_function ~return_unit ~async ~is_method cxt (f : P.t) ~fn_state
+and pp_function ~return_unit ~async ~is_method ?directive cxt (f : P.t) ~fn_state
     (l : Ident.t list) (b : J.block) (env : Js_fun_env.t) : cxt =
   match b with
   | [
@@ -363,13 +363,13 @@ and pp_function ~return_unit ~async ~is_method cxt (f : P.t) ~fn_state
                     if Js_fun_env.get_unused env 0 then cxt
                     else pp_var_assign_this cxt f this
                   in
-                  function_body ~return_unit cxt f b))
+                  function_body ?directive ~return_unit cxt f b))
         else
           let cxt =
             P.paren_group f 1 (fun _ -> formal_parameter_list inner_cxt f l)
           in
           P.space f;
-          P.brace_vgroup f 1 (fun _ -> function_body ~return_unit cxt f b)
+          P.brace_vgroup f 1 (fun _ -> function_body ?directive ~return_unit cxt f b)
       in
       let enclose () =
         let handle () =
@@ -483,9 +483,9 @@ and expression_desc cxt ~(level : int) f x : cxt =
           let cxt = expression ~level:0 cxt f e1 in
           comma_sp f;
           expression ~level:0 cxt f e2)
-  | Fun { is_method; params; body; env; return_unit; async } ->
+  | Fun { is_method; params; body; env; return_unit; async; directive } ->
       (* TODO: dump for comments *)
-      pp_function ~is_method cxt f ~fn_state:default_fn_exp_state params body
+      pp_function ?directive ~is_method cxt f ~fn_state:default_fn_exp_state params body
         env ~return_unit ~async
       (* TODO:
          when [e] is [Js_raw_code] with arity
@@ -515,10 +515,11 @@ and expression_desc cxt ~(level : int) f x : cxt =
                                env;
                                return_unit;
                                async;
+                               directive;
                              };
                        };
                       ] ->
-                          pp_function ~is_method ~return_unit ~async cxt f
+                          pp_function ?directive ~is_method ~return_unit ~async cxt f
                             ~fn_state:(No_name { single_arg = true })
                             params body env
                       | _ ->
@@ -920,8 +921,8 @@ and variable_declaration top cxt f (variable : J.variable_declaration) : cxt =
           statement_desc top cxt f (J.Exp e)
       | _ -> (
           match e.expression_desc with
-          | Fun { is_method; params; body; env; return_unit; async } ->
-              pp_function ~is_method cxt f ~return_unit ~async
+          | Fun { is_method; params; body; env; return_unit; async; directive } ->
+              pp_function ?directive ~is_method cxt f ~return_unit ~async
                 ~fn_state:(if top then Name_top name else Name_non_top name)
                 params body env
           | _ ->
@@ -1124,9 +1125,9 @@ and statement_desc top cxt f (s : J.statement_desc) : cxt =
       cxt
   | Return e -> (
       match e.expression_desc with
-      | Fun { is_method; params; body; env; return_unit; async } ->
+      | Fun { is_method; params; body; env; return_unit; async; directive } ->
           let cxt =
-            pp_function ~return_unit ~is_method ~async cxt f ~fn_state:Is_return
+            pp_function ?directive ~return_unit ~is_method ~async cxt f ~fn_state:Is_return
               params body env
           in
           semi f;
@@ -1216,8 +1217,15 @@ and statement_desc top cxt f (s : J.statement_desc) : cxt =
                   P.string f L.finally;
                   P.space f;
                   brace_block cxt f b))
+                  
 
-and function_body (cxt : cxt) f ~return_unit (b : J.block) : unit =
+and function_body ?directive (cxt : cxt) f ~return_unit (b : J.block) : unit =
+  (match directive with 
+  | None -> () 
+  | Some directive -> 
+    P.newline f;
+    P.string f directive; P.string f ";";
+    P.newline f);
   match b with
   | [] -> ()
   | [ s ] -> (
diff --git a/jscomp/core/js_exp_make.ml b/jscomp/core/js_exp_make.ml
@@ -207,7 +207,7 @@ let unit : t = { expression_desc = Undefined {isUnit = true}; comment = None }
    [Js_fun_env.empty] is a mutable state ..
 *)
 
-let ocaml_fun ?comment ?immutable_mask ~return_unit ~async ~oneUnitArg params body : t =
+let ocaml_fun ?comment ?immutable_mask ~return_unit ~async ~oneUnitArg ?directive params body : t =
   let params = if oneUnitArg then [] else params in
   let len = List.length params in
   {
@@ -220,6 +220,7 @@ let ocaml_fun ?comment ?immutable_mask ~return_unit ~async ~oneUnitArg params bo
           env = Js_fun_env.make ?immutable_mask len;
           return_unit;
           async;
+          directive;
         };
     comment;
   }
@@ -236,6 +237,7 @@ let method_ ?comment ?immutable_mask ~return_unit params body : t =
           env = Js_fun_env.make ?immutable_mask len;
           return_unit;
           async = false;
+          directive = None;
         };
     comment;
   }
@@ -1301,6 +1303,7 @@ let of_block ?comment ?e block : t =
             env = Js_fun_env.make 0;
             return_unit;
             async = false;
+            directive = None;
           };
     }
     []
diff --git a/jscomp/core/js_exp_make.mli b/jscomp/core/js_exp_make.mli
@@ -91,6 +91,7 @@ val ocaml_fun :
   return_unit:bool ->
   async:bool ->
   oneUnitArg:bool ->
+  ?directive:string ->
   J.ident list ->
   J.block ->
   t
diff --git a/jscomp/core/js_pass_tailcall_inline.ml b/jscomp/core/js_pass_tailcall_inline.ml
@@ -165,7 +165,7 @@ let subst (export_set : Set_ident.t) stats =
                      Some
                        {
                          expression_desc =
-                           Fun {is_method=false; params; body; env; async=false};
+                           Fun {is_method=false; params; body; env; async=false; directive=None};
                          comment = _;
                        };
                    (*TODO: don't inline method tail call yet,
@@ -200,7 +200,7 @@ let subst (export_set : Set_ident.t) stats =
                    Call
                      ( {
                          expression_desc =
-                           Fun {is_method=false; params; body; env; async=false};
+                           Fun {is_method=false; params; body; env; async=false; directive=None};
                        },
                        args,
                        _info );
diff --git a/jscomp/core/lam_analysis.ml b/jscomp/core/lam_analysis.ml
@@ -249,8 +249,8 @@ let destruct_pattern (body : Lam.t) params args =
   | _ -> false
 
 (* Async functions cannot be beta reduced *)
-let lfunction_can_be_beta_reduced (lfunction : Lam.lfunction) =
-  not lfunction.attr.async
+let lfunction_can_be_inlined (lfunction : Lam.lfunction) =
+  not lfunction.attr.async && lfunction.attr.directive = None
 
 (** Hints to inlining *)
 let ok_to_inline_fun_when_app (m : Lam.lfunction) (args : Lam.t list) =
diff --git a/jscomp/core/lam_analysis.mli b/jscomp/core/lam_analysis.mli
@@ -29,7 +29,7 @@ val no_side_effects : Lam.t -> bool
 
 val size : Lam.t -> int
 
-val lfunction_can_be_beta_reduced : Lam.lfunction -> bool
+val lfunction_can_be_inlined : Lam.lfunction -> bool
 
 val ok_to_inline_fun_when_app : Lam.lfunction -> Lam.t list -> bool
 
diff --git a/jscomp/core/lam_compile.ml b/jscomp/core/lam_compile.ml
@@ -272,7 +272,7 @@ and compile_external_field_apply ?(dynamic_import = false) (appinfo : Lam.apply)
   let ap_args = appinfo.ap_args in
   match ident_info.persistent_closed_lambda with
   | Some (Lfunction ({ params; body; _ } as lfunction))
-    when Ext_list.same_length params ap_args && Lam_analysis.lfunction_can_be_beta_reduced lfunction ->
+    when Ext_list.same_length params ap_args && Lam_analysis.lfunction_can_be_inlined lfunction ->
       (* TODO: serialize it when exporting to save compile time *)
       let _, param_map =
         Lam_closure.is_closed_with_map Set_ident.empty params body
@@ -319,7 +319,7 @@ and compile_external_field_apply ?(dynamic_import = false) (appinfo : Lam.apply)
 and compile_recursive_let ~all_bindings (cxt : Lam_compile_context.t)
     (id : Ident.t) (arg : Lam.t) : Js_output.t * initialization =
   match arg with
-  | Lfunction { params; body; attr = { return_unit; async; oneUnitArg } } ->
+  | Lfunction { params; body; attr = { return_unit; async; oneUnitArg; directive } } ->
       (* TODO: Think about recursive value
          {[
            let rec v = ref (fun _ ...
@@ -357,7 +357,7 @@ and compile_recursive_let ~all_bindings (cxt : Lam_compile_context.t)
              it will be renamed into [method]
              when it is detected by a primitive
           *)
-            ~return_unit ~async ~oneUnitArg ~immutable_mask:ret.immutable_mask
+            ~return_unit ~async ~oneUnitArg ?directive ~immutable_mask:ret.immutable_mask
             (Ext_list.map params (fun x ->
                  Map_ident.find_default ret.new_params x x))
             [
@@ -368,7 +368,7 @@ and compile_recursive_let ~all_bindings (cxt : Lam_compile_context.t)
             ]
         else
           (* TODO:  save computation of length several times *)
-          E.ocaml_fun params (Js_output.output_as_block output) ~return_unit ~async ~oneUnitArg
+          E.ocaml_fun params (Js_output.output_as_block output) ~return_unit ~async ~oneUnitArg ?directive
       in
       ( Js_output.output_of_expression
           (Declare (Alias, id))
@@ -1670,10 +1670,10 @@ and compile_prim (prim_info : Lam.prim_info)
 and compile_lambda (lambda_cxt : Lam_compile_context.t) (cur_lam : Lam.t) :
     Js_output.t =
   match cur_lam with
-  | Lfunction { params; body; attr = { return_unit; async; oneUnitArg } } ->
+  | Lfunction { params; body; attr = { return_unit; async; oneUnitArg; directive } } ->
       Js_output.output_of_expression lambda_cxt.continuation
         ~no_effects:no_effects_const
-        (E.ocaml_fun params ~return_unit ~async ~oneUnitArg
+        (E.ocaml_fun params ~return_unit ~async ~oneUnitArg ?directive
            (* Invariant:  jmp_table can not across function boundary,
               here we share env
            *)
diff --git a/jscomp/core/lam_pass_count.ml b/jscomp/core/lam_pass_count.ml
@@ -144,7 +144,7 @@ let collect_occurs lam : occ_tbl =
         count bv body
     (* Note there is a difference here when do beta reduction for *)
     | Lapply { ap_func = Lfunction ({ params; body } as lfunction); ap_args = args; _ }
-      when Ext_list.same_length params args && Lam_analysis.lfunction_can_be_beta_reduced lfunction ->
+      when Ext_list.same_length params args && Lam_analysis.lfunction_can_be_inlined lfunction ->
         count bv (Lam_beta_reduce.no_names_beta_reduce params body args)
     (* | Lapply{fn = Lfunction{function_kind = Tupled; params; body}; *)
     (*          args = [Lprim {primitive = Pmakeblock _;  args; _}]; _} *)
diff --git a/jscomp/core/lam_pass_lets_dce.ml b/jscomp/core/lam_pass_lets_dce.ml
@@ -147,7 +147,7 @@ let lets_helper (count_var : Ident.t -> Lam_pass_count.used_info) lam : Lam.t =
     | Lsequence(l1, l2) -> Lam.seq (simplif l1) (simplif l2)
 
     | Lapply{ap_func = Lfunction ({params; body} as lfunction);  ap_args = args; _}
-      when  Ext_list.same_length params args && Lam_analysis.lfunction_can_be_beta_reduced lfunction ->
+      when  Ext_list.same_length params args && Lam_analysis.lfunction_can_be_inlined lfunction ->
       simplif (Lam_beta_reduce.no_names_beta_reduce  params body args)
     (* | Lapply{ fn = Lfunction{function_kind = Tupled; params; body}; *)
     (*           args = [Lprim {primitive = Pmakeblock _;  args; _}]; _} *)
diff --git a/jscomp/core/lam_pass_remove_alias.ml b/jscomp/core/lam_pass_remove_alias.ml
@@ -82,7 +82,7 @@ let simplify_alias (meta : Lam_stats.t) (lam : Lam.t) : Lam.t =
                       rec_flag );
               })
           when Ext_list.same_length ap_args params
-               && Lam_analysis.lfunction_can_be_beta_reduced m
+               && Lam_analysis.lfunction_can_be_inlined m
                && Lam_analysis.ok_to_inline_fun_when_app m ap_args -> (
             let param_map =
               Lam_closure.is_closed_with_map meta.export_idents params body
@@ -184,7 +184,7 @@ let simplify_alias (meta : Lam_stats.t) (lam : Lam.t) : Lam.t =
                           | Some v -> v <> Parameter
                           | None -> true)
                       | _ -> true)
-               && Lam_analysis.lfunction_can_be_beta_reduced lfunction ->
+               && Lam_analysis.lfunction_can_be_inlined lfunction ->
             simpl (Lam_beta_reduce.propagate_beta_reduce meta params body args)
         | _ -> Lam.apply (simpl l1) (Ext_list.map args simpl) ap_info)
     (* Function inlining interact with other optimizations...
@@ -208,7 +208,7 @@ let simplify_alias (meta : Lam_stats.t) (lam : Lam.t) : Lam.t =
                     ( Lfunction ({ params; body; attr = { is_a_functor } } as m),
                       rec_flag );
               })
-          when Lam_analysis.lfunction_can_be_beta_reduced m ->
+          when Lam_analysis.lfunction_can_be_inlined m ->
             if Ext_list.same_length ap_args params then
               if
                 is_a_functor
@@ -266,7 +266,7 @@ let simplify_alias (meta : Lam_stats.t) (lam : Lam.t) : Lam.t =
           _;
         }
       when Ext_list.same_length params args
-           && Lam_analysis.lfunction_can_be_beta_reduced lfunction ->
+           && Lam_analysis.lfunction_can_be_inlined lfunction ->
         simpl (Lam_beta_reduce.propagate_beta_reduce meta params body args)
     (* | Lapply{ fn = Lfunction{function_kind =  Tupled;  params; body};  *)
     (*          args = [Lprim {primitive = Pmakeblock _; args; _}]; _} *)
diff --git a/jscomp/ml/lambda.ml b/jscomp/ml/lambda.ml
@@ -323,6 +323,7 @@ type function_attribute = {
   stub: bool;
   return_unit : bool;
   async : bool;
+  directive : string option;
   oneUnitArg : bool;
 }
 
@@ -394,6 +395,7 @@ let default_function_attribute = {
   return_unit = false;
   async = false;
   oneUnitArg = false;
+  directive = None;
 }
 
 let default_stub_attribute =
diff --git a/jscomp/ml/lambda.mli b/jscomp/ml/lambda.mli
@@ -293,6 +293,7 @@ type function_attribute = {
   stub: bool;
   return_unit : bool;
   async : bool;
+  directive : string option;
   oneUnitArg : bool;
 }
 
diff --git a/jscomp/ml/translcore.ml b/jscomp/ml/translcore.ml
@@ -717,6 +717,11 @@ let try_ids = Hashtbl.create 8
 
 let has_async_attribute exp = exp.exp_attributes |> List.exists (fun ({txt}, _payload) -> txt = "res.async")
 
+let extract_directive_for_fn exp = 
+  exp.exp_attributes |> List.find_map (
+    fun ({txt}, payload) -> if txt = "directive" then Ast_payload.is_single_string payload else None)
+
+
 let rec transl_exp e =
   List.iter (Translattribute.check_attribute e) e.exp_attributes;
   transl_exp0 e
@@ -732,6 +737,11 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
       transl_let rec_flag pat_expr_list (transl_exp body)
   | Texp_function { arg_label = _; param; cases; partial } ->
       let async = has_async_attribute e in
+      let directive = (
+        match extract_directive_for_fn e with 
+        | None -> None 
+        | Some (directive, _) -> Some directive
+      ) in
       let params, body, return_unit =
         let pl = push_defaults e.exp_loc [] cases partial in
         transl_function e.exp_loc partial param pl
@@ -742,6 +752,7 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
           inline = Translattribute.get_inline_attribute e.exp_attributes;
           async;
           return_unit;
+          directive;
         }
       in
       let loc = e.exp_loc in
diff --git a/jscomp/ml/translmod.ml b/jscomp/ml/translmod.ml
@@ -278,6 +278,7 @@ let rec compile_functor mexp coercion root_path loc =
           return_unit = false;
           async = false;
           oneUnitArg = false;
+          directive = None;
         };
       loc;
       body;
diff --git a/jscomp/test/build.ninja b/jscomp/test/build.ninja
diff --git a/jscomp/test/function_directives.js b/jscomp/test/function_directives.js
diff --git a/jscomp/test/function_directives.res b/jscomp/test/function_directives.res
@@ -0,0 +1 @@
+let testFnWithDirective = @directive("'use memo'") (name: string) => "Hello " ++ name
diff --git a/jscomp/test/function_directives_no_inline.js b/jscomp/test/function_directives_no_inline.js
diff --git a/jscomp/test/function_directives_no_inline.res b/jscomp/test/function_directives_no_inline.res
@@ -0,0 +1,3 @@
+let testFnWithDirective = @directive("'use server'") (name: string) => "Hello " ++ name
+
+let x = testFnWithDirective("test")

Original file line number	Diff line number	Diff line change
`@@ -135,6 +135,7 @@ and expression_desc =`
`135`	`135`	`env : Js_fun_env.t;`
`136`	`136`	`return_unit : bool;`
`137`	`137`	`async : bool;`
	`138`	`+ directive : string option;`
`138`	`139`	`}`
`139`	`140`	`\| Str of { delim : delim; txt : string }`
`140`	`141`	`(* A string is UTF-8 encoded, and may contain`
Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,7 @@ val ocaml_fun :`
`91`	`91`	`return_unit:bool ->`
`92`	`92`	`async:bool ->`
`93`	`93`	`oneUnitArg:bool ->`
	`94`	`+ ?directive:string ->`
`94`	`95`	`J.ident list ->`
`95`	`96`	`J.block ->`
`96`	`97`	`t`