flambda-backend: Revert "Revert "Transform tail-recursive functions i…

…nto recursive continuations (ocaml-flambda#893)"" (ocaml-flambda#909)

This reverts commit 41454a0.

lambda/lambda.ml

@@ -386,6 +386,11 @@ type check_attribute =
   | Assert of property
   | Assume of property
 
+type loop_attribute =
+  | Always_loop (* [@loop] or [@loop always] *)
+  | Never_loop (* [@loop never] *)
+  | Default_loop (* no [@loop] attribute *)
+
 type function_kind = Curried of {nlocal: int} | Tupled
 
 type let_kind = Strict | Alias | StrictOpt
@@ -407,6 +412,7 @@ type function_attribute = {
   local: local_attribute;
   check : check_attribute;
   poll: poll_attribute;
+  loop: loop_attribute;
   is_a_functor: bool;
   stub: bool;
 }
@@ -540,6 +546,7 @@ let default_function_attribute = {
   local = Default_local;
   check = Default_check ;
   poll = Default_poll;
+  loop = Default_loop;
   is_a_functor = false;
   stub = false;
 }

lambda/lambda.mli

@@ -299,6 +299,11 @@ type poll_attribute =
   | Error_poll (* [@poll error] *)
   | Default_poll (* no [@poll] attribute *)
 
+type loop_attribute =
+  | Always_loop (* [@loop] or [@loop always] *)
+  | Never_loop (* [@loop never] *)
+  | Default_loop (* no [@loop] attribute *)
+
 type function_kind = Curried of {nlocal: int} | Tupled
 (* [nlocal] determines how many arguments may be partially applied
    before the resulting closure must be locally allocated.
@@ -327,6 +332,7 @@ type function_attribute = {
   local: local_attribute;
   check : check_attribute;
   poll: poll_attribute;
+  loop: loop_attribute;
   is_a_functor: bool;
   stub: bool;
 }

lambda/printlambda.ml

@@ -564,7 +564,7 @@ let check_attribute ppf check =
   | Assume p -> fprintf ppf "assume %s@ " (check_property p)
 
 let function_attribute ppf
-    { inline; specialise; check; local; is_a_functor; stub; poll } =
+    { inline; specialise; check; local; is_a_functor; stub; poll; loop } =
   if is_a_functor then
     fprintf ppf "is_a_functor@ ";
   if stub then
@@ -590,7 +590,12 @@ let function_attribute ppf
   | Default_poll -> ()
   | Error_poll -> fprintf ppf "error_poll@ "
   end;
-  check_attribute ppf check
+  check_attribute ppf check;
+  begin match loop with
+  | Default_loop -> ()
+  | Always_loop -> fprintf ppf "always_loop@ "
+  | Never_loop -> fprintf ppf "never_loop@ "
+  end
 
 let apply_tailcall_attribute ppf = function
   | Default_tailcall -> ()

lambda/translattribute.ml

@@ -47,6 +47,10 @@ let is_poll_attribute = function
   | {txt=("poll")} -> true
   | _ -> false
 
+let is_loop_attribute = function
+  | {txt=("loop"|"ocaml.loop")} -> true
+  | _ -> false
+
 let find_attribute p attributes =
   let inline_attribute, other_attributes =
     List.partition (fun a -> p a.Parsetree.attr_name) attributes
@@ -230,6 +234,19 @@ let parse_poll_attribute attr =
         ]
         payload
 
+let parse_loop_attribute attr =
+  match attr with
+  | None -> Default_loop
+  | Some {Parsetree.attr_name = {txt; loc}; attr_payload = payload} ->
+      parse_id_payload txt loc
+        ~default:Default_loop
+        ~empty:Always_loop
+        [
+          "never", Never_loop;
+          "always", Always_loop;
+        ]
+        payload
+
 let get_inline_attribute l =
   let attr, _ = find_attribute is_inline_attribute l in
   parse_inline_attribute attr
@@ -257,6 +274,10 @@ let get_poll_attribute l =
   let attr, _ = find_attribute is_poll_attribute l in
   parse_poll_attribute attr
 
+let get_loop_attribute l =
+  let attr, _ = find_attribute is_loop_attribute l in
+  parse_loop_attribute attr
+
 let check_local_inline loc attr =
   match attr.local, attr.inline with
   | Always_local, (Always_inline | Available_inline | Unroll _) ->
@@ -388,6 +409,23 @@ let add_poll_attribute expr loc attributes =
         (Warnings.Misplaced_attribute "error_poll");
       expr
 
+let add_loop_attribute expr loc attributes =
+  match expr, get_loop_attribute attributes with
+  | expr, Default_loop -> expr
+  | Lfunction({ attr = { stub = false } as attr } as funct), loop ->
+      begin match attr.loop with
+      | Default_loop -> ()
+      | Always_loop | Never_loop ->
+          Location.prerr_warning loc
+            (Warnings.Duplicated_attribute "loop")
+      end;
+      let attr = { attr with loop } in
+      Lfunction { funct with attr = attr }
+  | expr, (Always_loop | Never_loop) ->
+      Location.prerr_warning loc
+        (Warnings.Misplaced_attribute "loop");
+      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
@@ -504,6 +542,9 @@ let add_function_attributes lam loc attr =
   let lam =
     add_check_attribute lam loc attr
   in
+  let lam =
+    add_loop_attribute lam loc attr
+  in
   let lam =
     (* last because poll overrides inline and local *)
     add_poll_attribute lam loc attr

lambda/translattribute.mli

@@ -53,6 +53,16 @@ val get_local_attribute
    : Parsetree.attributes
   -> Lambda.local_attribute
 
+val add_loop_attribute
+  : Lambda.lambda
+  -> Location.t
+  -> Parsetree.attributes
+  -> Lambda.lambda
+
+val get_loop_attribute
+  : Parsetree.attributes
+  -> Lambda.loop_attribute
+
 val get_and_remove_inlined_attribute
    : Typedtree.expression
   -> Lambda.inlined_attribute * Typedtree.expression

lambda/translcore.ml

@@ -818,6 +818,7 @@ and transl_exp0 ~in_new_scope ~scopes e =
           specialise = Always_specialise;
           local = Never_local;
           check = Default_check;
+          loop = Never_loop;
           is_a_functor = false;
           stub = false;
           poll = Default_poll;

lambda/translmod.ml

@@ -548,6 +548,7 @@ let rec compile_functor ~scopes mexp coercion root_path loc =
       local = Default_local;
       poll = Default_poll;
       check = Default_check;
+      loop = Never_loop;
       is_a_functor = true;
       stub = false;
     };

testsuite/tests/functors/functors.compilers.reference

@@ -2,14 +2,14 @@
   (let
     (O =
        (module-defn(O) Functors functors.ml(12):184-279
-         (function X is_a_functor always_inline
+         (function X is_a_functor always_inline never_loop
            (let
              (cow = (function x[int] : int (apply (field 0 X) x))
               sheep = (function x[int] : int (+ 1 (apply cow x))))
              (makeblock 0 cow sheep))))
      F =
        (module-defn(F) Functors functors.ml(17):281-392
-         (function X Y is_a_functor always_inline
+         (function X Y is_a_functor always_inline never_loop
            (let
              (cow =
                 (function x[int] : int
@@ -18,7 +18,7 @@
              (makeblock 0 cow sheep))))
      F1 =
        (module-defn(F1) Functors functors.ml(31):516-632
-         (function X Y is_a_functor always_inline
+         (function X Y is_a_functor always_inline never_loop
            (let
              (cow =
                 (function x[int] : int
@@ -27,7 +27,7 @@
              (makeblock 0 sheep))))
      F2 =
        (module-defn(F2) Functors functors.ml(36):634-784
-         (function X Y is_a_functor always_inline
+         (function X Y is_a_functor always_inline never_loop
            (let
              (X =a (makeblock 0 (field 1 X))
               Y =a (makeblock 0 (field 1 Y))
@@ -41,7 +41,7 @@
          (let
            (F =
               (module-defn(F) Functors.M functors.ml(44):849-966
-                (function X Y is_a_functor always_inline
+                (function X Y is_a_functor always_inline never_loop
                   (let
                     (cow =
                        (function x[int] : int