Allow metadata-only code in environments

lthls · lthls · commit ed18bcace5a1 · 2022-02-16T09:53:03.000+01:00
diff --git a/middle_end/flambda2/from_lambda/closure_conversion.ml b/middle_end/flambda2/from_lambda/closure_conversion.ml
@@ -1606,21 +1606,33 @@ let close_program ~symbol_for_global ~big_endian ~module_ident
       ~handler_params:load_fields_handler_param ~handler:load_fields_body ~body
       ~is_exn_handler:false
   in
+  let all_code =
+    Code_id.Map.mapi (fun code_id code_or_metadata ->
+        match (code_or_metadata : Code_or_metadata.t) with
+        | Code_present code -> code
+        | Metadata_only _ -> Misc.fatal_errorf "Metadata only for code ID %a" Code_id.print code_id)
+      (Acc.code acc)
+  in
   let acc, body =
-    Code_id.Map.fold
-      (fun code_id code (acc, body) ->
-        let bound_symbols =
-          Bound_symbols.singleton (Bound_symbols.Pattern.code code_id)
-        in
-        let static_const = Static_const_or_code.create_code code in
-        let defining_expr =
-          Static_const_group.create [static_const] |> Named.create_static_consts
-        in
-        Let_with_acc.create acc
-          (Bound_pattern.symbols bound_symbols)
-          defining_expr ~body
-        |> Expr_with_acc.create_let)
-      (Acc.code acc) (acc, body)
+    let bound_symbols, static_consts =
+      Code_id.Map.fold
+        (fun code_id code (bound_symbols, static_consts) ->
+           let bound_symbols =
+             (Bound_symbols.Pattern.code code_id) :: bound_symbols
+           in
+           let static_consts =
+             (Static_const_or_code.create_code code) :: static_consts
+           in
+           (bound_symbols, static_consts))
+        all_code ([], [])
+    in
+    let defining_expr =
+      Static_const_group.create static_consts |> Named.create_static_consts
+    in
+    Let_with_acc.create acc
+      (Bound_pattern.symbols (Bound_symbols.create bound_symbols))
+      defining_expr ~body
+    |> Expr_with_acc.create_let
   in
   (* We must make sure there is always an outer [Let_symbol] binding so that
      lifted constants not in the scope of any other [Let_symbol] binding get put
@@ -1670,5 +1682,5 @@ let close_program ~symbol_for_global ~big_endian ~module_ident
       ~module_symbol ~used_closure_vars:Unknown,
     Exported_code.add_code
       ~keep_code:(fun _ -> false)
-      (Acc.code acc) Exported_code.empty,
+      all_code Exported_code.empty,
     exported_offsets )
diff --git a/middle_end/flambda2/from_lambda/closure_conversion_aux.ml b/middle_end/flambda2/from_lambda/closure_conversion_aux.ml
@@ -270,7 +270,7 @@ module Acc = struct
   type t =
     { declared_symbols : (Symbol.t * Static_const.t) list;
       shareable_constants : Symbol.t Static_const.Map.t;
-      code : Code.t Code_id.Map.t;
+      code : Code_or_metadata.t Code_id.Map.t;
       free_names : Name_occurrences.t;
       continuation_applications : continuation_application Continuation.Map.t;
       cost_metrics : Cost_metrics.t;
@@ -325,7 +325,7 @@ module Acc = struct
     { t with shareable_constants }
 
   let add_code ~code_id ~code t =
-    { t with code = Code_id.Map.add code_id code t.code }
+    { t with code = Code_id.Map.add code_id (Code_or_metadata.create code) t.code }
 
   let add_free_names free_names t =
     { t with free_names = Name_occurrences.union free_names t.free_names }
@@ -629,8 +629,9 @@ module Let_with_acc = struct
         Cost_metrics.set_of_closures
           ~find_code_characteristics:(fun code_id ->
             let code = Code_id.Map.find code_id code_mapping in
-            { cost_metrics = Code.cost_metrics code;
-              params_arity = List.length (Code.params_arity code)
+            let code_metadata = Code_or_metadata.code_metadata code in
+            { cost_metrics = Code_metadata.cost_metrics code_metadata;
+              params_arity = List.length (Code_metadata.params_arity code_metadata)
             })
           set_of_closures
       | Rec_info _ -> Cost_metrics.zero
diff --git a/middle_end/flambda2/from_lambda/closure_conversion_aux.mli b/middle_end/flambda2/from_lambda/closure_conversion_aux.mli
@@ -162,7 +162,7 @@ module Acc : sig
 
   val shareable_constants : t -> Symbol.t Static_const.Map.t
 
-  val code : t -> Code.t Code_id.Map.t
+  val code : t -> Code_or_metadata.t Code_id.Map.t
 
   val free_names : t -> Name_occurrences.t
 
diff --git a/middle_end/flambda2/simplify/env/downwards_env.ml b/middle_end/flambda2/simplify/env/downwards_env.ml
@@ -44,7 +44,7 @@ type t =
     do_not_rebuild_terms : bool;
     closure_info : Closure_info.t;
     get_imported_code : unit -> Exported_code.t;
-    all_code : Code.t Code_id.Map.t
+    all_code : Code_or_metadata.t Code_id.Map.t
   }
 
 let print_debuginfo ppf dbg =
@@ -90,7 +90,7 @@ let [@ocamlformat "disable"] print ppf { round; typing_env;
     CSE.print cse
     do_not_rebuild_terms
     Closure_info.print closure_info
-    (Code_id.Map.print Code.print) all_code
+    (Code_id.Map.print Code_or_metadata.print) all_code
 
 let create ~round ~(resolver : resolver)
     ~(get_imported_names : get_imported_names)
@@ -394,7 +394,7 @@ let mem_code t id =
 
 let find_code_exn t id =
   match Code_id.Map.find id t.all_code with
-  | code -> Code_or_metadata.create code
+  | code_or_metadata -> code_or_metadata
   | exception Not_found -> (
     (* This [find_exn] call doesn't load any .cmx files, but in the majority of
        cases will succeed. *)
@@ -423,9 +423,18 @@ let define_code t ~code_id ~code =
     TE.add_to_code_age_relation t.typing_env ~new_code_id:code_id
       ~old_code_id:(Code.newer_version_of code)
   in
-  let all_code = Code_id.Map.add code_id code t.all_code in
+  let all_code = Code_id.Map.add code_id (Code_or_metadata.create code) t.all_code in
   { t with typing_env; all_code }
 
+let define_code_metadata_only t ~code_id ~code =
+  let code_or_metadata =
+    Code_or_metadata.remember_only_metadata (Code_or_metadata.create code)
+  in
+  let all_code =
+    Code_id.Map.add code_id code_or_metadata t.all_code
+  in
+  { t with all_code }
+
 let set_inlined_debuginfo t dbg = { t with inlined_debuginfo = dbg }
 
 let get_inlined_debuginfo t = t.inlined_debuginfo
diff --git a/middle_end/flambda2/simplify/env/downwards_env.mli b/middle_end/flambda2/simplify/env/downwards_env.mli
@@ -41,7 +41,7 @@ val create :
   unit_toplevel_return_continuation:Continuation.t ->
   t
 
-val all_code : t -> Code.t Code_id.Map.t
+val all_code : t -> Code_or_metadata.t Code_id.Map.t
 
 val resolver : t -> Compilation_unit.t -> Flambda2_types.Typing_env.t option
 
@@ -153,6 +153,8 @@ val check_simple_is_bound : t -> Simple.t -> unit
 
 val define_code : t -> code_id:Code_id.t -> code:Code.t -> t
 
+val define_code_metadata_only : t -> code_id:Code_id.t -> code:Code.t -> t
+
 val mem_code : t -> Code_id.t -> bool
 
 (** This function raises if the code ID is unbound. *)
diff --git a/middle_end/flambda2/simplify_shared/closure_offsets.ml b/middle_end/flambda2/simplify_shared/closure_offsets.ml
@@ -497,7 +497,7 @@ module Greedy = struct
           if Compilation_unit.is_current (Closure_id.get_compilation_unit c)
           then
             let code_metadata =
-              Code_id.Map.find code_id all_code |> Code.code_metadata
+              Code_id.Map.find code_id all_code |> Code_or_metadata.code_metadata
             in
             let module CM = Code_metadata in
             let is_tupled = CM.is_tupled code_metadata in
diff --git a/middle_end/flambda2/simplify_shared/closure_offsets.mli b/middle_end/flambda2/simplify_shared/closure_offsets.mli
@@ -43,7 +43,7 @@ val create : unit -> t
 val add_set_of_closures :
   t ->
   is_phantom:bool ->
-  all_code:Code.t Code_id.Map.t ->
+  all_code:Code_or_metadata.t Code_id.Map.t ->
   Set_of_closures.t ->
   t
 

Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ val create : unit -> t`
`43`	`43`	`val add_set_of_closures :`
`44`	`44`	`t ->`
`45`	`45`	`is_phantom:bool ->`
`46`		`- all_code:Code.t Code_id.Map.t ->`
	`46`	`+ all_code:Code_or_metadata.t Code_id.Map.t ->`
`47`	`47`	`Set_of_closures.t ->`
`48`	`48`	`t`
`49`	`49`