TRMC, reloaded

.depend

@@ -3531,6 +3531,7 @@ lambda/runtimedef.cmx : \
 lambda/runtimedef.cmi :
 lambda/simplif.cmo : \
     utils/warnings.cmi \
+    lambda/tmc.cmi \
     typing/primitive.cmi \
     parsing/location.cmi \
     lambda/lambda.cmi \
@@ -3541,6 +3542,7 @@ lambda/simplif.cmo : \
     lambda/simplif.cmi
 lambda/simplif.cmx : \
     utils/warnings.cmx \
+    lambda/tmc.cmx \
     typing/primitive.cmx \
     parsing/location.cmx \
     lambda/lambda.cmx \
@@ -3557,6 +3559,24 @@ lambda/switch.cmo : \
 lambda/switch.cmx : \
     lambda/switch.cmi
 lambda/switch.cmi :
+lambda/tmc.cmo : \
+    utils/warnings.cmi \
+    parsing/location.cmi \
+    lambda/lambda.cmi \
+    typing/ident.cmi \
+    lambda/debuginfo.cmi \
+    parsing/asttypes.cmi \
+    lambda/tmc.cmi
+lambda/tmc.cmx : \
+    utils/warnings.cmx \
+    parsing/location.cmx \
+    lambda/lambda.cmx \
+    typing/ident.cmx \
+    lambda/debuginfo.cmx \
+    parsing/asttypes.cmi \
+    lambda/tmc.cmi
+lambda/tmc.cmi : \
+    lambda/lambda.cmi
 lambda/translattribute.cmo : \
     utils/warnings.cmi \
     typing/typedtree.cmi \

Changes

@@ -10,6 +10,12 @@ Working version
   Allow 'object ... end # f', 'f object ... end', etc.
   (Yan Dong, review by Nicolás Ojeda Bär, Florian Angeletti and Gabriel Scherer)
 
+- #181, #9760: opt-in tail-modulo-cons (TMC) transformation
+    let[@tail_mod_cons] rec map f li = ...
+  (Frédéric Bour, Gabriel Scherer, Basile Clément,
+   review by Basile Clément and Pierre Chambart,
+   tested by Konstantin Romanov)
+
 ### Runtime system:
 
 * #9391, #9424: Fix failed assertion in runtime due to ephemerons *set_* and

compilerlibs/Makefile.compilerlibs

@@ -134,6 +134,7 @@ LAMBDA = \
   lambda/translcore.cmo \
   lambda/translclass.cmo \
   lambda/translmod.cmo \
+  lambda/tmc.cmo \
   lambda/simplif.cmo \
   lambda/runtimedef.cmo
 LAMBDA_CMI =

driver/main_args.ml

@@ -569,6 +569,10 @@ let mk_no_unboxed_types f =
   " unannotated unboxable types will not be unboxed (default)"
 ;;
 
+let mk_force_tmc f =
+  "-force-tmc", Arg.Unit f, " Rewrite all possible TMC calls"
+;;
+
 let mk_unsafe f =
   "-unsafe", Arg.Unit f,
   " Do not compile bounds checking on array and string access"
@@ -922,6 +926,7 @@ module type Common_options = sig
   val _no_strict_sequence : unit -> unit
   val _strict_formats : unit -> unit
   val _no_strict_formats : unit -> unit
+  val _force_tmc : unit -> unit
   val _unboxed_types : unit -> unit
   val _no_unboxed_types : unit -> unit
   val _unsafe_string : unit -> unit
@@ -1216,6 +1221,7 @@ struct
     mk_strict_formats F._strict_formats;
     mk_no_strict_formats F._no_strict_formats;
     mk_thread F._thread;
+    mk_force_tmc F._force_tmc;
     mk_unboxed_types F._unboxed_types;
     mk_no_unboxed_types F._no_unboxed_types;
     mk_unsafe F._unsafe;
@@ -1415,6 +1421,7 @@ struct
     mk_strict_formats F._strict_formats;
     mk_no_strict_formats F._no_strict_formats;
     mk_thread F._thread;
+    mk_force_tmc F._force_tmc;
     mk_unbox_closures F._unbox_closures;
     mk_unbox_closures_factor F._unbox_closures_factor;
     mk_inline_max_unroll F._inline_max_unroll;
@@ -1610,6 +1617,7 @@ struct
     mk_strict_formats F._strict_formats;
     mk_no_strict_formats F._no_strict_formats;
     mk_thread F._thread;
+    mk_force_tmc F._force_tmc;
     mk_unboxed_types F._unboxed_types;
     mk_no_unboxed_types F._no_unboxed_types;
     mk_unsafe_string F._unsafe_string;
@@ -1911,6 +1919,7 @@ module Default = struct
     let _noprompt = set noprompt
     let _nopromptcont = set nopromptcont
     let _stdin () = (* placeholder: file_argument ""*) ()
+    let _force_tmc = set force_tmc
     let _version () = print_version ()
     let _vnum () = print_version_num ()
     let _eval (_:string) = ()
@@ -1947,6 +1956,7 @@ module Default = struct
         "Profiling with \"gprof\" (option `-p') is only supported up to \
          OCaml 4.08.0"
     let _shared () = shared := true; dlcode := true
+    let _force_tmc = set force_tmc
     let _v () = Compenv.print_version_and_library "native-code compiler"
   end
 
@@ -1967,6 +1977,7 @@ module Default = struct
     let _pp s = Clflags.preprocessor := (Some s)
     let _ppx s = Clflags.all_ppx := (s :: (!Clflags.all_ppx))
     let _thread = set Clflags.use_threads
+    let _force_tmc = set force_tmc
     let _v () = Compenv.print_version_and_library "documentation generator"
     let _verbose = set Clflags.verbose
     let _version = Compenv.print_version_string
@@ -2000,6 +2011,7 @@ third-party libraries such as Lwt, but with a different API."
     let _output_complete_exe () =
       _output_complete_obj (); output_complete_executable := true
     let _output_obj () = output_c_object := true; custom_runtime := true
+    let _force_tmc = set force_tmc
     let _use_prims s = use_prims := s
     let _use_runtime s = use_runtime := s
     let _v () = Compenv.print_version_and_library "compiler"

driver/main_args.mli

@@ -40,6 +40,7 @@ module type Common_options = sig
   val _no_strict_sequence : unit -> unit
   val _strict_formats : unit -> unit
   val _no_strict_formats : unit -> unit
+  val _force_tmc : unit -> unit
   val _unboxed_types : unit -> unit
   val _no_unboxed_types : unit -> unit
   val _unsafe_string : unit -> unit

dune

@@ -68,7 +68,7 @@
    annot outcometree
 
    ;; lambda/
-   debuginfo lambda matching printlambda runtimedef simplif switch
+   debuginfo lambda matching printlambda runtimedef tmc simplif switch
    translattribute translclass translcore translmod translobj translprim
 
    ;; bytecomp/

lambda/lambda.ml

@@ -268,6 +268,7 @@ type function_attribute = {
   local: local_attribute;
   is_a_functor: bool;
   stub: bool;
+  tmc_candidate: bool;
 }
 
 type scoped_location = Debuginfo.Scoped_location.t
@@ -351,6 +352,7 @@ let default_function_attribute = {
   local = Default_local;
   is_a_functor = false;
   stub = false;
+  tmc_candidate = false;
 }
 
 let default_stub_attribute =
@@ -959,6 +961,25 @@ let function_is_curried func =
   | Curried -> true
   | Tupled -> false
 
+let find_exact_application kind ~arity args =
+  match kind with
+  | Curried ->
+      if arity <> List.length args
+      then None
+      else Some args
+  | Tupled ->
+      begin match args with
+      | [Lprim(Pmakeblock _, tupled_args, _)] ->
+          if arity <> List.length tupled_args
+          then None
+          else Some tupled_args
+      | [Lconst(Const_block (_, const_args))] ->
+          if arity <> List.length const_args
+          then None
+          else Some (List.map (fun cst -> Lconst cst) const_args)
+      | _ -> None
+      end
+
 let max_arity () =
   if !Clflags.native_code then 126 else max_int
   (* 126 = 127 (the maximal number of parameters supported in C--)

lambda/lambda.mli

@@ -254,6 +254,7 @@ type function_attribute = {
   local: local_attribute;
   is_a_functor: bool;
   stub: bool;
+  tmc_candidate: bool;
 }
 
 type scoped_location = Debuginfo.Scoped_location.t
@@ -424,6 +425,8 @@ val default_function_attribute : function_attribute
 val default_stub_attribute : function_attribute
 
 val function_is_curried : lfunction -> bool
+val find_exact_application :
+  function_kind -> arity:int -> lambda list -> lambda list option
 
 val max_arity : unit -> int
   (** Maximal number of parameters for a function, or in other words,

lambda/printlambda.ml

@@ -444,28 +444,30 @@ let name_of_primitive = function
   | Pint_as_pointer -> "Pint_as_pointer"
   | Popaque -> "Popaque"
 
-let function_attribute ppf { inline; specialise; local; is_a_functor; stub } =
-  if is_a_functor then
+let function_attribute ppf t =
+  if t.is_a_functor then
     fprintf ppf "is_a_functor@ ";
-  if stub then
+  if t.stub then
     fprintf ppf "stub@ ";
-  begin match inline with
+  begin match t.inline with
   | Default_inline -> ()
   | Always_inline -> fprintf ppf "always_inline@ "
   | Hint_inline -> fprintf ppf "hint_inline@ "
   | Never_inline -> fprintf ppf "never_inline@ "
   | Unroll i -> fprintf ppf "unroll(%i)@ " i
   end;
-  begin match specialise with
+  begin match t.specialise with
   | Default_specialise -> ()
   | Always_specialise -> fprintf ppf "always_specialise@ "
   | Never_specialise -> fprintf ppf "never_specialise@ "
   end;
-  begin match local with
+  begin match t.local with
   | Default_local -> ()
   | Always_local -> fprintf ppf "always_local@ "
   | Never_local -> fprintf ppf "never_local@ "
-  end
+  end;
+  if t.tmc_candidate then
+    fprintf ppf "tail_mod_cons@ "
 
 let apply_tailcall_attribute ppf = function
   | Default_tailcall -> ()

lambda/simplif.ml

@@ -333,23 +333,8 @@ let simplify_exits lam =
 *)
 
 let exact_application {kind; params; _} args =
-  match kind with
-  | Curried ->
-      if List.length params <> List.length args
-      then None
-      else Some args
-  | Tupled ->
-      begin match args with
-      | [Lprim(Pmakeblock _, tupled_args, _)] ->
-          if List.length params <> List.length tupled_args
-          then None
-          else Some tupled_args
-      | [Lconst(Const_block (_, const_args))] ->
-          if List.length params <> List.length const_args
-          then None
-          else Some (List.map (fun cst -> Lconst cst) const_args)
-      | _ -> None
-      end
+  let arity = List.length params in
+  Lambda.find_exact_application kind ~arity args
 
 let beta_reduce params body args =
   List.fold_left2 (fun l (param, kind) arg -> Llet(Strict, kind, param, arg, l))
@@ -904,7 +889,10 @@ let simplify_local_functions lam =
     rewrite lam
 
 (* The entry point:
-   simplification + emission of tailcall annotations, if needed. *)
+   simplification
+   + rewriting of tail-modulo-cons calls
+   + emission of tailcall annotations, if needed
+*)
 
 let simplify_lambda lam =
   let lam =
@@ -914,6 +902,7 @@ let simplify_lambda lam =
        )
     |> simplify_exits
     |> simplify_lets
+    |> Tmc.rewrite
   in
   if !Clflags.annotations
      || Warnings.is_active (Warnings.Wrong_tailcall_expectation true)