Optimize some local functions (#2143)

8 files changed, 383 insertions, 150 deletions

diff --git a/bytecomp/lambda.ml b/bytecomp/lambda.ml
index 974e9b0080..b998502768 100644
--- a/bytecomp/lambda.ml
+++ b/bytecomp/lambda.ml
@@ -221,7 +221,8 @@ let equal_inline_attribute x y =
   match x, y with
   | Always_inline, Always_inline
   | Never_inline, Never_inline
-  | Default_inline, Default_inline ->
+  | Default_inline, Default_inline
+    ->
     true
   | Unroll u, Unroll v ->
     u = v
@@ -242,6 +243,11 @@ let equal_specialise_attribute x y =
   | (Always_specialise | Never_specialise | Default_specialise), _ ->
     false
 
+type local_attribute =
+  | Always_local (* [@local] or [@local always] *)
+  | Never_local (* [@local never] *)
+  | Default_local (* [@local maybe] or no [@local] attribute *)
+
 type function_kind = Curried | Tupled
 
 type let_kind = Strict | Alias | StrictOpt | Variable
@@ -260,6 +266,7 @@ type shared_code = (int * int) list
 type function_attribute = {
   inline : inline_attribute;
   specialise : specialise_attribute;
+  local: local_attribute;
   is_a_functor: bool;
   stub: bool;
 }
@@ -336,6 +343,7 @@ let lambda_unit = Lconst const_unit
 let default_function_attribute = {
   inline = Default_inline;
   specialise = Default_specialise;
+  local = Default_local;
   is_a_functor = false;
   stub = false;
 }
@@ -456,7 +464,7 @@ let iter_opt f = function
   | None -> ()
   | Some e -> f e
 
-let iter_head_constructor f = function
+let shallow_iter ~tail ~non_tail:f = function
     Lvar _
   | Lconst _ -> ()
   | Lapply{ap_func = fn; ap_args = args} ->
@@ -464,31 +472,37 @@ let iter_head_constructor f = function
   | Lfunction{body} ->
       f body
   | Llet(_str, _k, _id, arg, body) ->
-      f arg; f body
+      f arg; tail body
   | Lletrec(decl, body) ->
-      f body;
+      tail body;
       List.iter (fun (_id, exp) -> f exp) decl
+  | Lprim (Pidentity, [l], _) ->
+      tail l
+  | Lprim (Psequand, [l1; l2], _)
+  | Lprim (Psequor, [l1; l2], _) ->
+      f l1;
+      tail l2
   | Lprim(_p, args, _loc) ->
       List.iter f args
   | Lswitch(arg, sw,_) ->
       f arg;
-      List.iter (fun (_key, case) -> f case) sw.sw_consts;
-      List.iter (fun (_key, case) -> f case) sw.sw_blocks;
-      iter_opt f sw.sw_failaction
+      List.iter (fun (_key, case) -> tail case) sw.sw_consts;
+      List.iter (fun (_key, case) -> tail case) sw.sw_blocks;
+      iter_opt tail sw.sw_failaction
   | Lstringswitch (arg,cases,default,_) ->
       f arg ;
-      List.iter (fun (_,act) -> f act) cases ;
-      iter_opt f default
+      List.iter (fun (_,act) -> tail act) cases ;
+      iter_opt tail default
   | Lstaticraise (_,args) ->
       List.iter f args
   | Lstaticcatch(e1, _, e2) ->
-      f e1; f e2
+      tail e1; tail e2
   | Ltrywith(e1, _, e2) ->
-      f e1; f e2
+      f e1; tail e2
   | Lifthenelse(e1, e2, e3) ->
-      f e1; f e2; f e3
+      f e1; tail e2; tail e3
   | Lsequence(e1, e2) ->
-      f e1; f e2
+      f e1; tail e2
   | Lwhile(e1, e2) ->
       f e1; f e2
   | Lfor(_v, e1, e2, _dir, e3) ->
@@ -497,10 +511,13 @@ let iter_head_constructor f = function
       f e
   | Lsend (_k, met, obj, args, _) ->
       List.iter f (met::obj::args)
-  | Levent (lam, _evt) ->
-      f lam
+  | Levent (e, _evt) ->
+      tail e
   | Lifused (_v, e) ->
-      f e
+      tail e
+
+let iter_head_constructor f l =
+  shallow_iter ~tail:f ~non_tail:f l
 
 let rec free_variables = function
   | Lvar id -> Ident.Set.singleton id
@@ -728,68 +745,68 @@ let rename idmap lam =
   let s = Ident.Map.map (fun new_id -> Lvar new_id) idmap in
   subst update_env s lam
 
-let rec map f lam =
-  let lam =
-    match lam with
-    | Lvar _ -> lam
-    | Lconst _ -> lam
-    | Lapply { ap_func; ap_args; ap_loc; ap_should_be_tailcall;
-          ap_inlined; ap_specialised } ->
-        Lapply {
-          ap_func = map f ap_func;
-          ap_args = List.map (map f) ap_args;
-          ap_loc;
-          ap_should_be_tailcall;
-          ap_inlined;
-          ap_specialised;
-        }
-    | Lfunction { kind; params; return; body; attr; loc; } ->
-        Lfunction { kind; params; return; body = map f body; attr; loc; }
-    | Llet (str, k, v, e1, e2) ->
-        Llet (str, k, v, map f e1, map f e2)
-    | Lletrec (idel, e2) ->
-        Lletrec (List.map (fun (v, e) -> (v, map f e)) idel, map f e2)
-    | Lprim (p, el, loc) ->
-        Lprim (p, List.map (map f) el, loc)
-    | Lswitch (e, sw, loc) ->
-        Lswitch (map f e,
-          { sw_numconsts = sw.sw_numconsts;
-            sw_consts = List.map (fun (n, e) -> (n, map f e)) sw.sw_consts;
-            sw_numblocks = sw.sw_numblocks;
-            sw_blocks = List.map (fun (n, e) -> (n, map f e)) sw.sw_blocks;
-            sw_failaction = Misc.may_map (map f) sw.sw_failaction;
-          },
-          loc)
-    | Lstringswitch (e, sw, default, loc) ->
-        Lstringswitch (
-          map f e,
-          List.map (fun (s, e) -> (s, map f e)) sw,
-          Misc.may_map (map f) default,
-          loc)
-    | Lstaticraise (i, args) ->
-        Lstaticraise (i, List.map (map f) args)
-    | Lstaticcatch (body, id, handler) ->
-        Lstaticcatch (map f body, id, map f handler)
-    | Ltrywith (e1, v, e2) ->
-        Ltrywith (map f e1, v, map f e2)
-    | Lifthenelse (e1, e2, e3) ->
-        Lifthenelse (map f e1, map f e2, map f e3)
-    | Lsequence (e1, e2) ->
-        Lsequence (map f e1, map f e2)
-    | Lwhile (e1, e2) ->
-        Lwhile (map f e1, map f e2)
-    | Lfor (v, e1, e2, dir, e3) ->
-        Lfor (v, map f e1, map f e2, dir, map f e3)
-    | Lassign (v, e) ->
-        Lassign (v, map f e)
-    | Lsend (k, m, o, el, loc) ->
-        Lsend (k, map f m, map f o, List.map (map f) el, loc)
-    | Levent (l, ev) ->
-        Levent (map f l, ev)
-    | Lifused (v, e) ->
-        Lifused (v, map f e)
-  in
-  f lam
+let shallow_map f = function
+  | Lvar _
+  | Lconst _ as lam -> lam
+  | Lapply { ap_func; ap_args; ap_loc; ap_should_be_tailcall;
+             ap_inlined; ap_specialised } ->
+      Lapply {
+        ap_func = f ap_func;
+        ap_args = List.map f ap_args;
+        ap_loc;
+        ap_should_be_tailcall;
+        ap_inlined;
+        ap_specialised;
+      }
+  | Lfunction { kind; params; return; body; attr; loc; } ->
+      Lfunction { kind; params; return; body = f body; attr; loc; }
+  | Llet (str, k, v, e1, e2) ->
+      Llet (str, k, v, f e1, f e2)
+  | Lletrec (idel, e2) ->
+      Lletrec (List.map (fun (v, e) -> (v, f e)) idel, f e2)
+  | Lprim (p, el, loc) ->
+      Lprim (p, List.map f el, loc)
+  | Lswitch (e, sw, loc) ->
+      Lswitch (f e,
+               { sw_numconsts = sw.sw_numconsts;
+                 sw_consts = List.map (fun (n, e) -> (n, f e)) sw.sw_consts;
+                 sw_numblocks = sw.sw_numblocks;
+                 sw_blocks = List.map (fun (n, e) -> (n, f e)) sw.sw_blocks;
+                 sw_failaction = Misc.may_map f sw.sw_failaction;
+               },
+               loc)
+  | Lstringswitch (e, sw, default, loc) ->
+      Lstringswitch (
+        f e,
+        List.map (fun (s, e) -> (s, f e)) sw,
+        Misc.may_map f default,
+        loc)
+  | Lstaticraise (i, args) ->
+      Lstaticraise (i, List.map f args)
+  | Lstaticcatch (body, id, handler) ->
+      Lstaticcatch (f body, id, f handler)
+  | Ltrywith (e1, v, e2) ->
+      Ltrywith (f e1, v, f e2)
+  | Lifthenelse (e1, e2, e3) ->
+      Lifthenelse (f e1, f e2, f e3)
+  | Lsequence (e1, e2) ->
+      Lsequence (f e1, f e2)
+  | Lwhile (e1, e2) ->
+      Lwhile (f e1, f e2)
+  | Lfor (v, e1, e2, dir, e3) ->
+      Lfor (v, f e1, f e2, dir, f e3)
+  | Lassign (v, e) ->
+      Lassign (v, f e)
+  | Lsend (k, m, o, el, loc) ->
+      Lsend (k, f m, f o, List.map f el, loc)
+  | Levent (l, ev) ->
+      Levent (f l, ev)
+  | Lifused (v, e) ->
+      Lifused (v, f e)
+
+let map f =
+  let rec g lam = f (shallow_map g lam) in
+  g
 
 (* To let-bind expressions to variables *)
 
diff --git a/bytecomp/lambda.mli b/bytecomp/lambda.mli
index c09943c7fc..190f9438cb 100644
--- a/bytecomp/lambda.mli
+++ b/bytecomp/lambda.mli
@@ -219,6 +219,11 @@ val equal_specialise_attribute
   -> specialise_attribute
   -> bool
 
+type local_attribute =
+  | Always_local (* [@local] or [@local always] *)
+  | Never_local (* [@local never] *)
+  | Default_local (* [@local maybe] or no [@local] attribute *)
+
 type function_kind = Curried | Tupled
 
 type let_kind = Strict | Alias | StrictOpt | Variable
@@ -242,6 +247,7 @@ type shared_code = (int * int) list     (* stack size -> code label *)
 type function_attribute = {
   inline : inline_attribute;
   specialise : specialise_attribute;
+  local: local_attribute;
   is_a_functor: bool;
   stub: bool;
 }
@@ -335,7 +341,16 @@ val name_lambda_list: lambda list -> (lambda list -> lambda) -> lambda
 val iter_head_constructor: (lambda -> unit) -> lambda -> unit
 (** [iter_head_constructor f lam] apply [f] to only the first level of
     sub expressions of [lam]. It does not recursively traverse the
-    expression. *)
+    expression.
+*)
+
+val shallow_iter:
+  tail:(lambda -> unit) ->
+  non_tail:(lambda -> unit) ->
+  lambda -> unit
+(** Same as [iter_head_constructor], but use a different callback for
+    sub-terms which are in tail position or not. *)
+
 
 val free_variables: lambda -> Ident.Set.t
 
@@ -362,6 +377,12 @@ val rename : Ident.t Ident.Map.t -> lambda -> lambda
     idents. *)
 
 val map : (lambda -> lambda) -> lambda -> lambda
+  (** Bottom-up rewriting, applying the function on
+      each node from the leaves to the root. *)
+
+val shallow_map  : (lambda -> lambda) -> lambda -> lambda
+  (** Rewrite each immediate sub-term with the function. *)
+
 val bind : let_kind -> Ident.t -> lambda -> lambda -> lambda
 val bind_with_value_kind:
   let_kind -> (Ident.t * value_kind) -> lambda -> lambda -> lambda
diff --git a/bytecomp/printlambda.ml b/bytecomp/printlambda.ml
index bf57dfb72a..e4bb26a686 100644
--- a/bytecomp/printlambda.ml
+++ b/bytecomp/printlambda.ml
@@ -445,7 +445,7 @@ let name_of_primitive = function
   | Pint_as_pointer -> "Pint_as_pointer"
   | Popaque -> "Popaque"
 
-let function_attribute ppf { inline; specialise; is_a_functor; stub } =
+let function_attribute ppf { inline; specialise; local; is_a_functor; stub } =
   if is_a_functor then
     fprintf ppf "is_a_functor@ ";
   if stub then
@@ -460,6 +460,11 @@ let function_attribute ppf { inline; specialise; is_a_functor; stub } =
   | Default_specialise -> ()
   | Always_specialise -> fprintf ppf "always_specialise@ "
   | Never_specialise -> fprintf ppf "never_specialise@ "
+  end;
+  begin match local with
+  | Default_local -> ()
+  | Always_local -> fprintf ppf "always_local@ "
+  | Never_local -> fprintf ppf "never_local@ "
   end
 
 let apply_tailcall_attribute ppf tailcall =
diff --git a/bytecomp/simplif.ml b/bytecomp/simplif.ml
index 03d503171f..8011953fb8 100644
--- a/bytecomp/simplif.ml
+++ b/bytecomp/simplif.ml
@@ -711,12 +711,137 @@ module Hooks = Misc.MakeHooks(struct
     type t = lambda
   end)
 
+(* Simplify local let-bound functions: if all occurrences are
+   fully-applied function calls in the same "tail scope", replace the
+   function by a staticcatch handler (on that scope).
+
+   This handles as a special case functions used exactly once (in any
+   scope) for a full application.
+*)
+
+type slot =
+  {
+    nargs: int;
+    mutable scope: lambda option;
+  }
+
+module LamTbl = Hashtbl.Make(struct
+    type t = lambda
+    let equal = (==)
+    let hash = Hashtbl.hash
+  end)
+
+let simplify_local_functions lam =
+  let slots = Hashtbl.create 16 in
+  let static_id = Hashtbl.create 16 in (* function id -> static id *)
+  let static = LamTbl.create 16 in (* scope -> static function on that scope *)
+  (* We keep track of the current "tail scope", identified
+     by the outermost lambda for which the the current lambda
+     is in tail position. *)
+  let current_scope = ref lam in
+  let check_static lf =
+    if lf.attr.local = Always_local then
+      Location.prerr_warning lf.loc
+        (Warnings.Inlining_impossible
+           "This function cannot be compiled into a static continuation")
+  in
+  let enabled = function
+    | {local = Always_local; _}
+    | {local = Default_local; inline = (Never_inline | Default_inline); _}
+      -> true
+    | {local = Default_local; inline = (Always_inline | Unroll _); _}
+    | {local = Never_local; _}
+      -> false
+  in
+  let rec tail = function
+    | Llet (_str, _kind, id, Lfunction lf, cont) when enabled lf.attr ->
+        let r = {nargs=List.length lf.params; scope=None} in
+        Hashtbl.add slots id r;
+        tail cont;
+        begin match Hashtbl.find_opt slots id with
+        | Some {scope = Some scope; _} ->
+            let st = next_raise_count () in
+            let sc =
+              (* Do not move higher than current lambda *)
+              if scope == !current_scope then cont
+              else scope
+            in
+            Hashtbl.add static_id id st;
+            LamTbl.add static sc (st, lf);
+            (* The body of the function will become an handler
+               in that "scope". *)
+            with_scope ~scope lf.body
+        | _ ->
+            check_static lf;
+            (* note: if scope = None, the function is unused *)
+            non_tail lf.body
+        end
+    | Lapply {ap_func = Lvar id; ap_args; _} ->
+        begin match Hashtbl.find_opt slots id with
+        | Some {nargs; _} when nargs <> List.length ap_args ->
+            (* Wrong arity *)
+            Hashtbl.remove slots id
+        | Some {scope = Some scope; _} when scope != !current_scope ->
+            (* Different "tail scope" *)
+            Hashtbl.remove slots id
+        | Some ({scope = None; _} as slot) ->
+            (* First use of the function: remember the current tail scope *)
+            slot.scope <- Some !current_scope
+        | _ ->
+            ()
+        end;
+        List.iter non_tail ap_args
+    | Lvar id ->
+        Hashtbl.remove slots id
+    | Lfunction lf as lam ->
+        check_static lf;
+        Lambda.shallow_iter ~tail ~non_tail lam
+    | lam ->
+        Lambda.shallow_iter ~tail ~non_tail lam
+  and non_tail lam =
+    with_scope ~scope:lam lam
+  and with_scope ~scope lam =
+    let old_scope = !current_scope in
+    current_scope := scope;
+    tail lam;
+    current_scope := old_scope
+  in
+  tail lam;
+  let rec rewrite lam0 =
+    let lam =
+      match lam0 with
+      | Llet (_, _, id, _, cont) when Hashtbl.mem static_id id ->
+          rewrite cont
+      | Lapply {ap_func = Lvar id; ap_args; _} when Hashtbl.mem static_id id ->
+          Lstaticraise (Hashtbl.find static_id id, List.map rewrite ap_args)
+      | lam ->
+          Lambda.shallow_map rewrite lam
+    in
+    List.fold_right
+      (fun (st, lf) lam ->
+         Lstaticcatch (lam, (st, lf.params), rewrite lf.body)
+      )
+      (LamTbl.find_all static lam0)
+      lam
+  in
+  if LamTbl.length static = 0 then
+    lam
+  else
+    rewrite lam
+
 (* The entry point:
    simplification + emission of tailcall annotations, if needed. *)
 
 let simplify_lambda sourcefile lam =
-  let res = simplify_lets (simplify_exits lam) in
-  let res = Hooks.apply_hooks { Misc.sourcefile } res in
+  let lam =
+    lam
+    |> (if !Clflags.native_code || not !Clflags.debug
+        then simplify_local_functions else Fun.id
+       )
+    |> simplify_exits
+    |> simplify_lets
+    |> Hooks.apply_hooks { Misc.sourcefile }
+  in
   if !Clflags.annotations || Warnings.is_active Warnings.Expect_tailcall
-    then emit_tail_infos true res;
-  res
+    then emit_tail_infos true lam;
+  lam
diff --git a/bytecomp/translattribute.ml b/bytecomp/translattribute.ml
index 712e42b956..1520a3b41f 100644
--- a/bytecomp/translattribute.ml
+++ b/bytecomp/translattribute.ml
@@ -34,6 +34,10 @@ let is_specialised_attribute = function
   | {txt=("specialised"|"ocaml.specialised")} when Config.flambda -> true
   | _ -> false
 
+let is_local_attribute = function
+  | {txt=("local"|"ocaml.local")} -> true
+  | _ -> false
+
 let find_attribute p attributes =
   let inline_attribute, other_attributes =
     List.partition (fun a -> p a.Parsetree.attr_name) attributes
@@ -53,6 +57,37 @@ let is_unrolled = function
   | {txt="inline"|"ocaml.inline"|"inlined"|"ocaml.inlined"} -> false
   | _ -> assert false
 
+let get_id_payload =
+  let open Parsetree in
+  function
+  | PStr [] -> Some ""
+  | PStr [{pstr_desc = Pstr_eval ({pexp_desc},[])}] ->
+      begin match pexp_desc with
+      | Pexp_ident { txt = Longident.Lident id } -> Some id
+      | _ -> None
+      end
+  | _ -> None
+
+let parse_id_payload txt loc ~default ~empty cases payload =
+  let[@local] warn () =
+    let ( %> ) f g x = g (f x) in
+    let msg =
+      cases
+      |> List.map (fst %> Printf.sprintf "'%s'")
+      |> String.concat ", "
+      |> Printf.sprintf "It must be either %s or empty"
+    in
+    Location.prerr_warning loc (Warnings.Attribute_payload (txt, msg));
+    default
+  in
+  match get_id_payload payload with
+  | Some "" -> empty
+  | None -> warn ()
+  | Some id ->
+      match List.assoc_opt id cases with
+      | Some r -> r
+      | None -> warn ()
+
 let parse_inline_attribute attr =
   match attr with
   | None -> Default_inline
@@ -80,58 +115,42 @@ let parse_inline_attribute attr =
       | _ ->
         Location.prerr_warning loc (warning txt);
         Default_inline
-    end else begin
-      (* the 'inline' and 'inlined' attributes can be used as
-         [@inline], [@inline never] or [@inline always].
-         [@inline] is equivalent to [@inline always] *)
-      let warning txt =
-        Warnings.Attribute_payload
-          (txt, "It must be either empty, 'always' or 'never'")
-      in
-      match payload with
-      | PStr [] -> Always_inline
-      | PStr [{pstr_desc = Pstr_eval ({pexp_desc},[])}] -> begin
-          match pexp_desc with
-          | Pexp_ident { txt = Longident.Lident "never" } ->
-            Never_inline
-          | Pexp_ident { txt = Longident.Lident "always" } ->
-            Always_inline
-          | _ ->
-            Location.prerr_warning loc (warning txt);
-            Default_inline
-        end
-      | _ ->
-        Location.prerr_warning loc (warning txt);
-        Default_inline
-    end
+    end else
+      parse_id_payload txt loc
+        ~default:Default_inline
+        ~empty:Always_inline
+        [
+          "never", Never_inline;
+          "always", Always_inline;
+        ]
+        payload
 
 let parse_specialise_attribute attr =
   match attr with
   | None -> Default_specialise
   | Some {Parsetree.attr_name = {txt; loc}; attr_payload = payload} ->
-    let open Parsetree in
-    let warning txt =
-      Warnings.Attribute_payload
-      (txt, "It must be either empty, 'always' or 'never'")
-    in
-    match payload with
-    | PStr [] -> Always_specialise
-    | PStr [{pstr_desc = Pstr_eval ({pexp_desc},[])}] -> begin
-        (* the 'specialise' and 'specialised' attributes can be used as
-           [@specialise], [@specialise never] or [@specialise always].
-           [@specialise] is equivalent to [@specialise always] *)
-        match pexp_desc with
-        | Pexp_ident { txt = Longident.Lident "never" } ->
-          Never_specialise
-        | Pexp_ident { txt = Longident.Lident "always" } ->
-          Always_specialise
-        | _ ->
-          Location.prerr_warning loc (warning txt);
-          Default_specialise
-      end
-    | _ ->
-      Location.prerr_warning loc (warning txt);
-      Default_specialise
+      parse_id_payload txt loc
+        ~default:Default_specialise
+        ~empty:Always_specialise
+        [
+          "never", Never_specialise;
+          "always", Always_specialise;
+        ]
+        payload
+
+let parse_local_attribute attr =
+  match attr with
+  | None -> Default_local
+  | Some {Parsetree.attr_name = {txt; loc}; attr_payload = payload} ->
+      parse_id_payload txt loc
+        ~default:Default_local
+        ~empty:Always_local
+        [
+          "never", Never_local;
+          "always", Always_local;
+          "maybe", Default_local;
+        ]
+        payload
 
 let get_inline_attribute l =
   let attr, _ = find_attribute is_inline_attribute l in
@@ -141,6 +160,18 @@ let get_specialise_attribute l =
   let attr, _ = find_attribute is_specialise_attribute l in
   parse_specialise_attribute attr
 
+let get_local_attribute l =
+  let attr, _ = find_attribute is_local_attribute l in
+  parse_local_attribute attr
+
+let check_local_inline loc attr =
+  match attr.local, attr.inline with
+  | Always_local, (Always_inline | Unroll _) ->
+      Location.prerr_warning loc
+        (Warnings.Duplicated_attribute "local/inline")
+  | _ ->
+      ()
+
 let add_inline_attribute expr loc attributes =
   match expr, get_inline_attribute attributes with
   | expr, Default_inline -> expr
@@ -152,6 +183,7 @@ let add_inline_attribute expr loc attributes =
             (Warnings.Duplicated_attribute "inline")
       end;
       let attr = { attr with inline } in
+      check_local_inline loc attr;
       Lfunction { funct with attr = attr }
   | expr, (Always_inline | Never_inline | Unroll _) ->
       Location.prerr_warning loc
@@ -175,6 +207,24 @@ let add_specialise_attribute expr loc attributes =
         (Warnings.Misplaced_attribute "specialise");
       expr
 
+let add_local_attribute expr loc attributes =
+  match expr, get_local_attribute attributes with
+  | expr, Default_local -> expr
+  | Lfunction({ attr = { stub = false } as attr } as funct), local ->
+      begin match attr.local with
+      | Default_local -> ()
+      | Always_local | Never_local ->
+          Location.prerr_warning loc
+            (Warnings.Duplicated_attribute "local")
+      end;
+      let attr = { attr with local } in
+      check_local_inline loc attr;
+      Lfunction { funct with attr }
+  | expr, (Always_local | Never_local) ->
+      Location.prerr_warning loc
+        (Warnings.Misplaced_attribute "local");
+      expr
+
 (* Get the [@inlined] attribute payload (or default if not present).
    It also returns the expression without this attribute. This is
    used to ensure that this attribute is not misplaced: If it
@@ -268,3 +318,15 @@ let check_attribute_on_module e {Parsetree.attr_name = { txt; loc }; _} =
       Location.prerr_warning loc
         (Warnings.Misplaced_attribute txt)
   | _ -> ()
+
+let add_function_attributes lam loc attr =
+  let lam =
+    add_inline_attribute lam loc attr
+  in
+  let lam =
+    add_specialise_attribute lam loc attr
+  in
+  let lam =
+    add_local_attribute lam loc attr
+  in
+  lam
diff --git a/bytecomp/translattribute.mli b/bytecomp/translattribute.mli
index 4b5840ede8..bf22fd1c5d 100644
--- a/bytecomp/translattribute.mli
+++ b/bytecomp/translattribute.mli
@@ -43,6 +43,16 @@ val get_specialise_attribute
    : Parsetree.attributes
   -> Lambda.specialise_attribute
 
+val add_local_attribute
+   : Lambda.lambda
+  -> Location.t
+  -> Parsetree.attributes
+  -> Lambda.lambda
+
+val get_local_attribute
+   : Parsetree.attributes
+  -> Lambda.local_attribute
+
 val get_and_remove_inlined_attribute
    : Typedtree.expression
   -> Lambda.inline_attribute * Typedtree.expression
@@ -58,3 +68,9 @@ val get_and_remove_specialised_attribute
 val get_tailcall_attribute
    : Typedtree.expression
   -> bool * Typedtree.expression
+
+val add_function_attributes
+  : Lambda.lambda
+  -> Location.t
+  -> Parsetree.attributes
+  -> Lambda.lambda
diff --git a/bytecomp/translcore.ml b/bytecomp/translcore.ml
index 6cd6caf4ba..e9098a2fb3 100644
--- a/bytecomp/translcore.ml
+++ b/bytecomp/translcore.ml
@@ -235,14 +235,10 @@ and transl_exp0 e =
             transl_function e.exp_loc return_kind !Clflags.native_code repr
               partial param pl)
       in
-      let attr = {
-        default_function_attribute with
-        inline = Translattribute.get_inline_attribute e.exp_attributes;
-        specialise = Translattribute.get_specialise_attribute e.exp_attributes;
-      }
-      in
+      let attr = default_function_attribute in
       let loc = e.exp_loc in
-      Lfunction{kind; params; return; body; attr; loc}
+      let lam = Lfunction{kind; params; return; body; attr; loc} in
+      Translattribute.add_function_attributes lam loc e.exp_attributes
   | Texp_apply({ exp_desc = Texp_ident(path, _, {val_kind = Val_prim p});
                 exp_type = prim_type } as funct, oargs)
     when List.length oargs >= p.prim_arity
@@ -769,12 +765,7 @@ and transl_let rec_flag pat_expr_list =
           fun body -> body
       | {vb_pat=pat; vb_expr=expr; vb_attributes=attr; vb_loc} :: rem ->
           let lam = transl_exp expr in
-          let lam =
-            Translattribute.add_inline_attribute lam vb_loc attr
-          in
-          let lam =
-            Translattribute.add_specialise_attribute lam vb_loc attr
-          in
+          let lam = Translattribute.add_function_attributes lam vb_loc attr in
           let mk_body = transl rem in
           fun body -> Matching.for_let pat.pat_loc lam pat (mk_body body)
       in transl pat_expr_list
@@ -789,12 +780,7 @@ and transl_let rec_flag pat_expr_list =
       let transl_case {vb_expr=expr; vb_attributes; vb_loc} id =
         let lam = transl_exp expr in
         let lam =
-          Translattribute.add_inline_attribute lam vb_loc
-            vb_attributes
-        in
-        let lam =
-          Translattribute.add_specialise_attribute lam vb_loc
-            vb_attributes
+          Translattribute.add_function_attributes lam vb_loc vb_attributes
         in
         (id, lam) in
       let lam_bds = List.map2 transl_case pat_expr_list idlist in
diff --git a/bytecomp/translmod.ml b/bytecomp/translmod.ml
index 64b84cbf23..9dc77da68f 100644
--- a/bytecomp/translmod.ml
+++ b/bytecomp/translmod.ml
@@ -445,6 +445,7 @@ let rec compile_functor mexp coercion root_path loc =
     attr = {
       inline = inline_attribute;
       specialise = Default_specialise;
+      local = Default_local;
       is_a_functor = true;
       stub = false;
     };