Further updates to middle_end/ only

middle_end/closure_conversion.ml

@@ -27,6 +27,51 @@ type t = {
   mutable imported_symbols : Symbol.Set.t;
 }
 
+let add_default_argument_wrappers lam =
+  (* CR-someday mshinwell: Temporary hack to mark default argument wrappers
+     as stubs.  Other possibilities:
+     1. Change Lambda.inline_attribute to add another ("stub") case;
+     2. Add a "stub" field to the Lfunction record. *)
+  let stubify body : Lambda.lambda =
+    let stub_prim =
+      Primitive.simple ~name:Closure_conversion_aux.stub_hack_prim_name
+        ~arity:1 ~alloc:false
+    in
+    Lprim (Pccall stub_prim, [body])
+  in
+  let defs_are_all_functions (defs : (_ * Lambda.lambda) list) =
+    List.for_all (function (_, Lambda.Lfunction _) -> true | _ -> false) defs
+  in
+  let f (lam : Lambda.lambda) : Lambda.lambda =
+    match lam with
+    | Llet (( Strict | Alias | StrictOpt), id,
+        Lfunction {kind; params; body = fbody; attr}, body) ->
+      begin match
+        Simplif.split_default_wrapper id kind params fbody attr
+          ~create_wrapper_body:stubify
+      with
+      | [fun_id, def] -> Llet (Alias, fun_id, def, body)
+      | [fun_id, def; inner_fun_id, def_inner] ->
+        Llet (Alias, inner_fun_id, def_inner, Llet (Alias, fun_id, def, body))
+      | _ -> assert false
+      end
+    | Lletrec (defs, body) as lam ->
+      if defs_are_all_functions defs then
+        let defs =
+          List.flatten
+            (List.map
+               (function
+                 | (id, Lambda.Lfunction {kind; params; body; attr}) ->
+                   Simplif.split_default_wrapper id kind params body attr
+                 | _ -> assert false)
+               defs)
+        in
+        Lletrec (defs, body)
+      else lam
+    | lam -> lam
+  in
+  Lambda.map f lam
+
 (** Generate a wrapper ("stub") function that accepts a tuple argument and
     calls another function with arguments extracted in the obvious
     manner from the tuple. *)
@@ -557,6 +602,7 @@ and close_let_bound_expression t ?let_rec_ident let_bound_var env
   | lam -> Expr (close t env lam)
 
 let lambda_to_flambda ~backend ~module_ident ~size lam : Flambda.program =
+  let lam = add_default_argument_wrappers lam in
   let module Backend = (val backend : Backend_intf.S) in
   let compilation_unit = Compilation_unit.get_current_exn () in
   let t =

middle_end/closure_conversion_aux.ml

@@ -80,6 +80,8 @@ module Env = struct
   let not_at_toplevel t = { t with at_toplevel = false; }
 end
 
+let stub_hack_prim_name = "*stub*"
+
 module Function_decls = struct
   module Function_decl = struct
     type t = {
@@ -119,10 +121,10 @@ module Function_decls = struct
     let inline t = t.inline
     let is_a_functor t = t.is_a_functor
 
-    (* CR-someday mshinwell: eliminate "*stub*" *)
     let primitive_wrapper t =
       match t.body with
-      | Lprim (Pccall { Primitive.prim_name = "*stub*" }, [body]) -> Some body
+      | Lprim (Pccall { Primitive. prim_name; }, [body])
+        when prim_name = stub_hack_prim_name -> Some body
       | _ -> None
   end
 

middle_end/closure_conversion_aux.mli

@@ -92,3 +92,5 @@ module Function_decls : sig
      It also contains the globals bindings of the provided environment. *)
   val closure_env_without_parameters : Env.t -> t -> Env.t
 end
+
+val stub_hack_prim_name : string

middle_end/find_recursive_functions.ml

@@ -16,7 +16,7 @@
 
 let in_function_declarations (function_decls : Flambda.function_declarations)
       ~backend =
-  let module VCC = Sort_connected_components.Make (Variable) in
+  let module VCC = Strongly_connected_components.Make (Variable) in
   let directed_graph =
     Flambda_utils.fun_vars_referenced_in_decls function_decls ~backend
   in

middle_end/lift_code.ml

@@ -99,7 +99,7 @@ and lift_lets_named _var (named:Flambda.named) ~toplevel : Flambda.named =
   | Project_var _ | Prim _ | Set_of_closures _ ->
     named
 
-module Sort_lets = Sort_connected_components.Make(Variable)
+module Sort_lets = Strongly_connected_components.Make (Variable)
 
 let rebuild_let_rec (defs:(Variable.t * Flambda.named) list) body =
   let map = Variable.Map.of_list defs in

middle_end/lift_constants.ml

@@ -552,7 +552,7 @@ let program_graph
       )
       effect_tbl graph_with_initialisation
   in
-  let module Symbol_SCC = Sort_connected_components.Make (Symbol) in
+  let module Symbol_SCC = Strongly_connected_components.Make (Symbol) in
   let components =
     Symbol_SCC.connected_components_sorted_from_roots_to_leaf
       graph
@@ -572,7 +572,7 @@ let add_definition_of_symbol constant_definitions
     assert(not (Symbol.Tbl.mem initialize_symbol_tbl sym));
     (sym, Symbol.Map.find sym constant_definitions)
   in
-  let module Symbol_SCC = Sort_connected_components.Make (Symbol) in
+  let module Symbol_SCC = Strongly_connected_components.Make (Symbol) in
   match component with
   | Symbol_SCC.Has_loop l ->
     let l = List.map symbol_declaration l in