Use cmx reachability to decide locality of symbols (#1246)

Co-authored-by: Mark Shinwell <mshinwell@pm.me>

Author: Gbury

backend/cmm.ml

@@ -225,6 +225,11 @@ type kind_for_unboxing =
 
 type is_global = Global | Local
 
+let equal_is_global g g' =
+  match g, g' with
+  | Local, Local | Global, Global -> true
+  | Local, Global | Global, Local -> false
+
 type symbol =
   { sym_name : string;
     sym_global : is_global }

backend/cmm.mli

@@ -228,6 +228,7 @@ type kind_for_unboxing =
   | Boxed_float
 
 type is_global = Global | Local
+val equal_is_global : is_global -> is_global -> bool
 
 (* Symbols are marked with whether they are local or global,
    at both definition and use sites.

backend/cmm_helpers.ml

@@ -1208,7 +1208,7 @@ let apply_function_sym arity result mode =
   Compilenv.need_apply_fun arity result mode;
   global_symbol (apply_function_name arity result mode)
 
-let curry_function_sym function_kind arity result =
+let curry_function_sym_name function_kind arity result =
   match function_kind with
   | Lambda.Curried { nlocal } ->
     Compilenv.need_curry_fun function_kind arity result;
@@ -1233,6 +1233,11 @@ let curry_function_sym function_kind arity result =
     ^
     match result with [| Val |] -> "" | _ -> "_R" ^ machtype_identifier result)
 
+let curry_function_sym function_kind arity result =
+  { sym_name = curry_function_sym_name function_kind arity result;
+    sym_global = Global
+  }
+
 (* Big arrays *)
 
 let bigarray_elt_size_in_bytes : Lambda.bigarray_kind -> int = function
@@ -2809,7 +2814,7 @@ let final_curry_function nlocal arity result =
   let narity = List.length arity in
   let fun_name =
     global_symbol
-      (curry_function_sym (Lambda.Curried { nlocal }) arity result
+      (curry_function_sym_name (Lambda.Curried { nlocal }) arity result
       ^ "_"
       ^ Int.to_string (narity - 1))
   in
@@ -2828,7 +2833,9 @@ let final_curry_function nlocal arity result =
     }
 
 let intermediate_curry_functions ~nlocal ~arity result =
-  let name1 = curry_function_sym (Lambda.Curried { nlocal }) arity result in
+  let name1 =
+    curry_function_sym_name (Lambda.Curried { nlocal }) arity result
+  in
   let narity = List.length arity in
   let dbg = placeholder_dbg in
   let rec loop accumulated_args remaining_args num =
@@ -3756,9 +3763,8 @@ let emit_constant_closure symb fundecls clos_vars cont =
           in
           (Cint (infix_header pos) :: closure_symbol f2)
           @ Csymbol_address
-              (global_symbol
-                 (curry_function_sym arity.function_kind params_machtypes
-                    return_machtype))
+              (curry_function_sym arity.function_kind params_machtypes
+                 return_machtype)
             :: Cint (closure_info ~arity ~startenv:(startenv - pos) ~is_last)
             :: Csymbol_address
                  { sym_name = f2.label; sym_global = symb.sym_global }
@@ -3776,11 +3782,10 @@ let emit_constant_closure symb fundecls clos_vars cont =
       :: emit_others 3 remainder
     | arity ->
       Csymbol_address
-        (global_symbol
-           (curry_function_sym arity.function_kind
-              (List.map machtype_of_layout_changing_tagged_int_to_val
-                 arity.params_layout)
-              (machtype_of_layout_changing_tagged_int_to_val arity.return_layout)))
+        (curry_function_sym arity.function_kind
+           (List.map machtype_of_layout_changing_tagged_int_to_val
+              arity.params_layout)
+           (machtype_of_layout_changing_tagged_int_to_val arity.return_layout))
       :: Cint (closure_info ~arity ~startenv ~is_last)
       :: Csymbol_address { sym_name = f1.label; sym_global = symb.sym_global }
       :: emit_others 4 remainder)
@@ -3844,7 +3849,7 @@ let unit ~dbg = Cconst_int (1, dbg)
 
 let var v = Cvar v
 
-let symbol_from_string ~dbg sym = Cconst_symbol (global_symbol sym, dbg)
+let symbol ~dbg sym = Cconst_symbol (sym, dbg)
 
 let float ~dbg f = Cconst_float (f, dbg)
 
@@ -4136,7 +4141,7 @@ let gc_root_table syms =
   let table_symbol = make_symbol ?compilation_unit:None "gc_roots" in
   cdata
     (define_symbol { sym_name = table_symbol; sym_global = Global }
-    @ List.map (fun s -> symbol_address (global_symbol s)) syms
+    @ List.map symbol_address syms
     @ [cint 0n])
 
 let cmm_arith_size (e : Cmm.expression) =

backend/cmm_helpers.mli

@@ -475,7 +475,7 @@ val machtype_identifier : machtype -> string
 (** Get the symbol for the generic currying or tuplifying wrapper with [n]
     arguments, and ensure its presence in the set of defined symbols. *)
 val curry_function_sym :
-  Lambda.function_kind -> machtype list -> machtype -> string
+  Lambda.function_kind -> machtype list -> machtype -> Cmm.symbol
 
 (** Bigarrays *)
 
@@ -958,9 +958,8 @@ val unit : dbg:Debuginfo.t -> Cmm.expression
 (** Create an expression from a variable. *)
 val var : Backend_var.t -> Cmm.expression
 
-(** Create an expression that gives the value of an object file symbol, such
-    symbol's name being given by a string. *)
-val symbol_from_string : dbg:Debuginfo.t -> string -> Cmm.expression
+(** Create an expression that gives the value of an object file symbol. *)
+val symbol : dbg:Debuginfo.t -> Cmm.symbol -> Cmm.expression
 
 (** Create a constant float expression. *)
 val float : dbg:Debuginfo.t -> float -> expression
@@ -1244,7 +1243,7 @@ val cfunction : fundecl -> phrase
 val cdata : data_item list -> phrase
 
 (** Create the gc root table from a list of root symbols. *)
-val gc_root_table : string list -> phrase
+val gc_root_table : Cmm.symbol list -> phrase
 
 (* An estimate of the number of arithmetic instructions in a Cmm expression.
    This is currently used in Flambda 2 to determine whether untagging an

backend/cmmgen.ml

@@ -481,7 +481,7 @@ let rec transl env e =
                 transl_fundecls (pos + 3) rem
               | arity ->
                 Cconst_symbol
-                  (global_symbol (curry_function_sym
+                  ((curry_function_sym
                      arity.function_kind
                      (List.map machtype_of_layout_changing_tagged_int_to_val
                        arity.params_layout)

backend/cmmgen_state.ml

@@ -76,6 +76,10 @@ let is_local_function name =
 let clear_function_names () =
   Hashtbl.clear state.function_names
 
+let add_structured_constant (sym : Cmm.symbol) cst =
+  if not (Hashtbl.mem state.structured_constants sym.sym_name) then
+    Hashtbl.replace state.structured_constants sym.sym_name (sym.sym_global, cst)
+
 let set_local_structured_constants l =
   Hashtbl.clear state.structured_constants;
   List.iter

backend/cmmgen_state.mli

@@ -41,6 +41,8 @@ val is_local_function : Clambda.function_label -> bool
 
 val clear_function_names : unit -> unit
 
+val add_structured_constant : Cmm.symbol -> Clambda.ustructured_constant -> unit
+
 val set_local_structured_constants : Clambda.preallocated_constant list -> unit
 
 val add_global_structured_constant : string -> Clambda.ustructured_constant -> unit

backend/printcmm.ml

@@ -63,6 +63,13 @@ let extcall_signature ppf (ty_res, ty_args) =
       fprintf ppf "->%a" machtype ty_res
   end
 
+let is_global ppf = function
+  | Global -> fprintf ppf "G"
+  | Local -> fprintf ppf "L"
+
+let symbol ppf s =
+  fprintf ppf "%a:\"%s\"" is_global s.sym_global s.sym_name
+
 let integer_comparison = function
   | Ceq -> "=="
   | Cne -> "!="
@@ -230,12 +237,13 @@ let operation d = function
   | Cbeginregion -> "beginregion"
   | Cendregion -> "endregion"
 
+
 let rec expr ppf = function
   | Cconst_int (n, _dbg) -> fprintf ppf "%i" n
   | Cconst_natint (n, _dbg) ->
     fprintf ppf "%s" (Nativeint.to_string n)
   | Cconst_float (n, _dbg) -> fprintf ppf "%F" n
-  | Cconst_symbol (s, _dbg) -> fprintf ppf "\"%s\"" s.sym_name
+  | Cconst_symbol (s, _dbg) -> fprintf ppf "%a:\"%s\"" is_global s.sym_global s.sym_name
   | Cvar id -> V.print ppf id
   | Clet(id, def, (Clet(_, _, _) as body)) ->
       let print_binding id ppf def =
@@ -401,7 +409,7 @@ let data_item ppf = function
   | Cint n -> fprintf ppf "int %s" (Nativeint.to_string n)
   | Csingle f -> fprintf ppf "single %F" f
   | Cdouble f -> fprintf ppf "double %F" f
-  | Csymbol_address s -> fprintf ppf "addr \"%s\"" s.sym_name
+  | Csymbol_address s -> fprintf ppf "addr %a:\"%s\"" is_global s.sym_global s.sym_name
   | Cstring s -> fprintf ppf "string \"%s\"" s
   | Cskip n -> fprintf ppf "skip %i" n
   | Calign n -> fprintf ppf "align %i" n

backend/printcmm.mli

@@ -17,6 +17,7 @@
 
 open Format
 
+val symbol : formatter -> Cmm.symbol -> unit
 val rec_flag : formatter -> Cmm.rec_flag -> unit
 val machtype_component : formatter -> Cmm.machtype_component -> unit
 val machtype : formatter -> Cmm.machtype -> unit

middle_end/flambda2/cmx/flambda_cmx.ml

@@ -221,15 +221,19 @@ let prepare_cmx ~module_symbol create_typing_env ~free_names_of_name
     |> Exported_offsets.reexport_value_slots
          (Name_occurrences.all_value_slots slots_used_in_typing_env)
   in
-  Some
-    (Flambda_cmx_format.create ~final_typing_env ~all_code ~exported_offsets
-       ~used_value_slots)
+  let cmx =
+    Flambda_cmx_format.create ~final_typing_env ~all_code ~exported_offsets
+      ~used_value_slots
+  in
+  reachable_names, Some cmx
 
 let prepare_cmx_file_contents ~final_typing_env ~module_symbol ~used_value_slots
     ~exported_offsets all_code =
   match final_typing_env with
-  | None -> None
-  | Some _ when Flambda_features.opaque () -> None
+  | None ->
+    Name_occurrences.singleton_symbol module_symbol Name_mode.normal, None
+  | Some _ when Flambda_features.opaque () ->
+    Name_occurrences.singleton_symbol module_symbol Name_mode.normal, None
   | Some final_typing_env ->
     let typing_env, canonicalise =
       TE.Pre_serializable.create final_typing_env ~used_value_slots
@@ -247,7 +251,7 @@ let prepare_cmx_file_contents ~final_typing_env ~module_symbol ~used_value_slots
 let prepare_cmx_from_approx ~approxs ~module_symbol ~exported_offsets
     ~used_value_slots all_code =
   if Flambda_features.opaque ()
-  then None
+  then Name_occurrences.singleton_symbol module_symbol Name_mode.normal, None
   else
     let create_typing_env reachable_names =
       let approxs =

middle_end/flambda2/cmx/flambda_cmx.mli

@@ -40,12 +40,12 @@ val prepare_cmx_file_contents :
   used_value_slots:Value_slot.Set.t ->
   exported_offsets:Exported_offsets.t ->
   Exported_code.t ->
-  Flambda_cmx_format.t option
+  Name_occurrences.t * Flambda_cmx_format.t option
 
 val prepare_cmx_from_approx :
   approxs:Code_or_metadata.t Value_approximation.t Symbol.Map.t ->
   module_symbol:Symbol.t ->
   exported_offsets:Exported_offsets.t ->
   used_value_slots:Value_slot.Set.t ->
   Exported_code.t ->
-  Flambda_cmx_format.t option
+  Name_occurrences.t * Flambda_cmx_format.t option

middle_end/flambda2/flambda2.ml

@@ -136,20 +136,25 @@ let lambda_to_cmm ~ppf_dump:ppf ~prefixname ~filename:_ ~keep_symbol_tables
     in
     Compiler_hooks.execute Raw_flambda2 raw_flambda;
     print_rawflambda ppf raw_flambda;
-    let flambda, offsets, cmx, all_code =
+    let flambda, offsets, reachable_names, cmx, all_code =
       match mode, close_program_metadata with
-      | Classic, Classic (code, cmx, offsets) ->
+      | Classic, Classic (code, reachable_names, cmx, offsets) ->
         (if Flambda_features.inlining_report ()
         then
           let output_prefix = prefixname ^ ".cps_conv" in
           let inlining_tree =
             Inlining_report.output_then_forget_decisions ~output_prefix
           in
           Compiler_hooks.execute Inlining_tree inlining_tree);
-        raw_flambda, offsets, cmx, code
+        raw_flambda, offsets, reachable_names, cmx, code
       | Normal, Normal ->
         let round = 0 in
-        let { Simplify.unit = flambda; exported_offsets; cmx; all_code } =
+        let { Simplify.unit = flambda;
+              exported_offsets;
+              cmx;
+              all_code;
+              reachable_names
+            } =
           Profile.record_call ~accumulate:true "simplify" (fun () ->
               Simplify.run ~cmx_loader ~round raw_flambda)
         in
@@ -163,13 +168,15 @@ let lambda_to_cmm ~ppf_dump:ppf ~prefixname ~filename:_ ~keep_symbol_tables
         Compiler_hooks.execute Flambda2 flambda;
         print_flambda "simplify" ppf flambda;
         print_flexpect "simplify" ppf ~raw_flambda flambda;
-        flambda, exported_offsets, cmx, all_code
+        flambda, exported_offsets, reachable_names, cmx, all_code
     in
     (match cmx with
     | None ->
       () (* Either opaque was passed, or there is no need to export offsets *)
     | Some cmx -> Compilenv.flambda2_set_export_info cmx);
-    let cmm = Flambda2_to_cmm.To_cmm.unit flambda ~all_code ~offsets in
+    let cmm =
+      Flambda2_to_cmm.To_cmm.unit flambda ~all_code ~offsets ~reachable_names
+    in
     if not keep_symbol_tables
     then (
       Compilenv.reset_info_tables ();

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -35,7 +35,10 @@ module VB = Bound_var
 type 'a close_program_metadata =
   | Normal : [`Normal] close_program_metadata
   | Classic :
-      (Exported_code.t * Flambda_cmx_format.t option * Exported_offsets.t)
+      (Exported_code.t
+      * Name_occurrences.t
+      * Flambda_cmx_format.t option
+      * Exported_offsets.t)
       -> [`Classic] close_program_metadata
 
 type 'a close_program_result = Flambda_unit.t * 'a close_program_metadata
@@ -2580,7 +2583,7 @@ let close_program (type mode) ~(mode : mode Flambda_features.mode) ~big_endian
       Slot_offsets.finalize_offsets (Acc.slot_offsets acc) ~get_code_metadata
         ~used_slots
     in
-    let cmx =
+    let reachable_names, cmx =
       Flambda_cmx.prepare_cmx_from_approx ~approxs:symbols_approximations
         ~module_symbol ~exported_offsets ~used_value_slots all_code
     in
@@ -2589,4 +2592,4 @@ let close_program (type mode) ~(mode : mode Flambda_features.mode) ~big_endian
         ~toplevel_my_region ~body ~module_symbol
         ~used_value_slots:(Known used_value_slots)
     in
-    unit, Classic (all_code, cmx, exported_offsets)
+    unit, Classic (all_code, reachable_names, cmx, exported_offsets)

middle_end/flambda2/from_lambda/closure_conversion.mli

@@ -73,7 +73,10 @@ val close_switch :
 type 'a close_program_metadata =
   | Normal : [`Normal] close_program_metadata
   | Classic :
-      (Exported_code.t * Flambda_cmx_format.t option * Exported_offsets.t)
+      (Exported_code.t
+      * Name_occurrences.t
+      * Flambda_cmx_format.t option
+      * Exported_offsets.t)
       -> [`Classic] close_program_metadata
 
 type 'a close_program_result = Flambda_unit.t * 'a close_program_metadata

middle_end/flambda2/simplify/simplify.ml

@@ -20,7 +20,8 @@ type simplify_result =
   { cmx : Flambda_cmx_format.t option;
     unit : Flambda_unit.t;
     all_code : Exported_code.t;
-    exported_offsets : Exported_offsets.t
+    exported_offsets : Exported_offsets.t;
+    reachable_names : Name_occurrences.t
   }
 
 let run ~cmx_loader ~round unit =
@@ -97,12 +98,12 @@ let run ~cmx_loader ~round unit =
       in
       Slot_offsets.finalize_offsets slot_offsets ~get_code_metadata ~used_slots
   in
-  let cmx =
+  let reachable_names, cmx =
     Flambda_cmx.prepare_cmx_file_contents ~final_typing_env ~module_symbol
       ~used_value_slots ~exported_offsets all_code
   in
   let unit =
     FU.create ~return_continuation ~exn_continuation ~toplevel_my_region
       ~module_symbol ~body ~used_value_slots:(Known used_value_slots)
   in
-  { cmx; unit; all_code; exported_offsets }
+  { cmx; unit; all_code; exported_offsets; reachable_names }

middle_end/flambda2/simplify/simplify.mli

@@ -24,7 +24,8 @@ type simplify_result = private
   { cmx : Flambda_cmx_format.t option;
     unit : Flambda_unit.t;
     all_code : Exported_code.t;
-    exported_offsets : Exported_offsets.t
+    exported_offsets : Exported_offsets.t;
+    reachable_names : Name_occurrences.t
   }
 
 val run :