@directives in functions

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

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 ] -> (

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;
           };
     }
     []

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

jscomp/core/lam_compile.ml

@@ -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
            *)

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 =

jscomp/ml/lambda.mli

@@ -293,6 +293,7 @@ type function_attribute = {
   stub: bool;
   return_unit : bool;
   async : bool;
+  directive : string option;
   oneUnitArg : bool;
 }
 

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

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;

jscomp/test/function_directives.res

@@ -0,0 +1 @@
+let testFnWithDirective = @directive("'use memo'") (name: string) => "Hello " ++ name