From dc8618cec03a6d849a0d267bfe655ab955587067 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Tue, 18 Oct 2022 18:30:30 -0400
Subject: [PATCH] Use uniquing target table to map gvars to uniqued values

---
 src/staticdata.c | 41 +++++++++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/staticdata.c b/src/staticdata.c
index 8e1c392b1cd578..8d6fdfda9ca564 100644
--- a/src/staticdata.c
+++ b/src/staticdata.c
@@ -39,6 +39,7 @@
 
     `uniquing_list` also holds the serialized location of external DataTypes, MethodInstances, and singletons
     in the serialized blob (i.e., new-at-the-time-of-serialization specializations).
+    `uniquing_target` is a hash table for which `uniquing_target[obj] -> chosen_target`.
 
     Most of step 2 is handled by `jl_write_values`, followed by special handling of the dedicated parallel streams.
 
@@ -1864,7 +1865,7 @@ static uint32_t write_gvars(jl_serializer_state *s, arraylist_t *globals, arrayl
 }
 
 // Pointer relocation for native-code referenced global variables
-static void jl_update_all_gvars(jl_serializer_state *s, uint32_t external_fns_begin)
+static void jl_update_all_gvars(jl_serializer_state *s, uint32_t external_fns_begin, htable_t *uniquing_target)
 {
     if (sysimg_gvars_base == NULL)
         return;
@@ -1880,7 +1881,11 @@ static void jl_update_all_gvars(jl_serializer_state *s, uint32_t external_fns_be
         if (offset) {
             if (gvname_index < (external_fns_begin-1)) { // external_fns_begin is 1-indexed, gvname_index is 0-indexed
                 uintptr_t v = get_item_for_reloc(s, base, size, offset, s->link_ids_gvars, &gvar_link_index);
-                // TODO(required): if (incremental) v = jl_as_global_root((jl_value_t*)v);
+                uintptr_t unique_value = (uintptr_t)ptrhash_get(uniquing_target, (void*)v);
+                if ((void*)unique_value != HT_NOTFOUND)
+                    v = unique_value;
+                // if (s->incremental)
+                //     v = (uintptr_t) jl_as_global_root((jl_value_t*)v);
                 *sysimg_gvars(sysimg_gvars_base, gvname_index) = (uintptr_t)v;
             }
             else {
@@ -2855,21 +2860,10 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_array_t *depmods,
             *pfld = (uintptr_t)newobj | GC_OLD;
         else
             *pfld = (uintptr_t)newobj;
+        ptrhash_put(&uniquing_target, (void*)obj, (void*)newobj);
         assert(!(image_base < (char*)newobj && (char*)newobj <= image_base + sizeof_sysimg + sizeof(uintptr_t)));
         assert(jl_typeis(obj, otyp));
     }
-    // Write junk in place of the source data we used during uniquing, to catch inadvertent references to
-    // it from elsewhere.
-    for (size_t i = 0; i < s.uniquing_list.len; i++) {
-        void *item = s.uniquing_list.items[i];
-        jl_taggedvalue_t *o = jl_astaggedvalue(item);
-        if (o->type == (jl_value_t*)jl_method_instance_type)
-            memset(o, 0xba, sizeof(jl_value_t*) + sizeof(jl_method_instance_t));
-        else if (o->type == (jl_value_t*)jl_datatype_type)
-            memset(o, 0xba, sizeof(jl_value_t*) + sizeof(jl_datatype_t));
-        else
-            memset(o, 0xba, sizeof(jl_value_t*));
-    }
     // Perform fixups: things like updating world ages, inserting methods & specializations, etc.
     size_t world = jl_atomic_load_acquire(&jl_world_counter);
     for (size_t i = 0; i < s.fixup_list.len; i++) {
@@ -2948,15 +2942,30 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_array_t *depmods,
             r->bnd_cache = b && b->value ? b : NULL;
         }
     }
-    arraylist_free(&s.uniquing_list);
     arraylist_free(&s.fixup_list);
 
     // s.link_ids_gvars will be processed in `jl_update_all_gvars`
     // s.link_ids_external_fnvars will be processed in `jl_update_all_gvars`
     ios_close(&relocs);
     ios_close(&const_data);
-    jl_update_all_gvars(&s, external_fns_begin); // gvars relocs
+    jl_update_all_gvars(&s, external_fns_begin, &uniquing_target); // gvars relocs
     ios_close(&gvar_record);
+
+    // Write junk in place of the source data we used during uniquing, to catch inadvertent references to
+    // it from elsewhere. This needs to be delayed after update_all_gvars so that we can rewire gvars as well.
+    for (size_t i = 0; i < s.uniquing_list.len; i++) {
+        void *item = s.uniquing_list.items[i];
+        jl_taggedvalue_t *o = jl_astaggedvalue(item);
+        if (o->type == (jl_value_t*)jl_method_instance_type)
+            memset(o, 0xba, sizeof(jl_value_t*) + sizeof(jl_method_instance_t));
+        else if (o->type == (jl_value_t*)jl_datatype_type)
+            memset(o, 0xba, sizeof(jl_value_t*) + sizeof(jl_datatype_t));
+        else
+            memset(o, 0xba, sizeof(jl_value_t*));
+    }
+    arraylist_free(&s.uniquing_list);
+    htable_free(&uniquing_target);
+
     s.s = NULL;
 
     if (0) {