bpart: Skip implicit import reval if using'd export set is unchanged (JuliaLang#57421)

When loading a pkgimage, the new bpart validation code needs to check if
the export set of any using'd packages differs from what it would have
been during precompile. This could e.g. happen if somebody (or Revise)
eval'd a new `export` statement into a package that was `using`'d.
However, this case is somewhat rare, so let's optimize it by keeping a
bit in `Module` that keeps track of whether anything like that has
happened and if not skipping the revalidation. This slightly improves
pkgimage load time in the ordinary case. More optimizations to follow.

(cherry picked from commit 39255d4)

Author: Keno

src/julia.h

@@ -770,6 +770,8 @@ typedef struct _jl_module_t {
     int8_t infer;
     uint8_t istopmod;
     int8_t max_methods;
+    // If cleared no binding partition in this module has BINDING_FLAG_EXPORTED and min_world > jl_require_world.
+    _Atomic(int8_t) export_set_changed_since_require_world;
     jl_mutex_t lock;
     intptr_t hash;
 } jl_module_t;

src/julia_internal.h

@@ -923,6 +923,7 @@ extern JL_DLLEXPORT jl_module_t *jl_precompile_toplevel_module JL_GLOBALLY_ROOTE
 extern jl_genericmemory_t *jl_global_roots_list JL_GLOBALLY_ROOTED;
 extern jl_genericmemory_t *jl_global_roots_keyset JL_GLOBALLY_ROOTED;
 extern arraylist_t *jl_entrypoint_mis;
+JL_DLLEXPORT extern size_t jl_require_world;
 JL_DLLEXPORT int jl_is_globally_rooted(jl_value_t *val JL_MAYBE_UNROOTED) JL_NOTSAFEPOINT;
 JL_DLLEXPORT jl_value_t *jl_as_global_root(jl_value_t *val, int insert) JL_GLOBALLY_ROOTED;
 extern jl_svec_t *precompile_field_replace JL_GLOBALLY_ROOTED;
@@ -938,6 +939,14 @@ STATIC_INLINE int jl_bkind_is_some_import(enum jl_partition_kind kind) JL_NOTSAF
     return kind == BINDING_KIND_IMPLICIT || kind == BINDING_KIND_EXPLICIT || kind == BINDING_KIND_IMPORTED;
 }
 
+STATIC_INLINE int jl_bkind_is_some_guard(enum jl_partition_kind kind) JL_NOTSAFEPOINT {
+    return kind == BINDING_KIND_FAILED || kind == BINDING_KIND_GUARD;
+}
+
+STATIC_INLINE int jl_bkind_is_some_implicit(enum jl_partition_kind kind) JL_NOTSAFEPOINT {
+    return kind == BINDING_KIND_IMPLICIT || jl_bkind_is_some_guard(kind);
+}
+
 STATIC_INLINE int jl_bkind_is_some_constant(enum jl_partition_kind kind) JL_NOTSAFEPOINT {
     return kind == BINDING_KIND_CONST || kind == BINDING_KIND_CONST_IMPORT || kind == BINDING_KIND_UNDEF_CONST || kind == BINDING_KIND_BACKDATED_CONST;
 }
@@ -946,10 +955,6 @@ STATIC_INLINE int jl_bkind_is_defined_constant(enum jl_partition_kind kind) JL_N
     return kind == BINDING_KIND_CONST || kind == BINDING_KIND_CONST_IMPORT || kind == BINDING_KIND_BACKDATED_CONST;
 }
 
-STATIC_INLINE int jl_bkind_is_some_guard(enum jl_partition_kind kind) JL_NOTSAFEPOINT {
-    return kind == BINDING_KIND_FAILED || kind == BINDING_KIND_GUARD;
-}
-
 JL_DLLEXPORT jl_binding_partition_t *jl_get_binding_partition(jl_binding_t *b JL_PROPAGATES_ROOT, size_t world) JL_GLOBALLY_ROOTED;
 JL_DLLEXPORT jl_binding_partition_t *jl_get_binding_partition_all(jl_binding_t *b JL_PROPAGATES_ROOT, size_t min_world, size_t max_world) JL_GLOBALLY_ROOTED;
 

src/module.c

@@ -909,10 +909,7 @@ static void module_import_(jl_module_t *to, jl_module_t *from, jl_sym_t *asname,
     size_t new_world = jl_atomic_load_acquire(&jl_world_counter)+1;
     jl_binding_partition_t *btopart = jl_get_binding_partition(bto, new_world);
     enum jl_partition_kind btokind = jl_binding_kind(btopart);
-    if (btokind == BINDING_KIND_GUARD ||
-        btokind == BINDING_KIND_IMPLICIT ||
-        btokind == BINDING_KIND_FAILED) {
-
+    if (jl_bkind_is_some_implicit(btokind)) {
         jl_binding_partition_t *new_bpart = jl_replace_binding_locked(bto, btopart, (jl_value_t*)b, (explici != 0) ? BINDING_KIND_IMPORTED : BINDING_KIND_EXPLICIT, new_world);
         if (jl_atomic_load_relaxed(&new_bpart->max_world) == ~(size_t)0)
             jl_add_binding_backedge(b, (jl_value_t*)bto);
@@ -1016,7 +1013,7 @@ JL_DLLEXPORT void jl_module_using(jl_module_t *to, jl_module_t *from)
             if (tob) {
                 jl_binding_partition_t *tobpart = jl_get_binding_partition(tob, new_world);
                 enum jl_partition_kind kind = jl_binding_kind(tobpart);
-                if (kind == BINDING_KIND_IMPLICIT || jl_bkind_is_some_guard(kind)) {
+                if (jl_bkind_is_some_implicit(kind)) {
                     jl_replace_binding_locked(tob, tobpart, NULL, BINDING_KIND_IMPLICIT_RECOMPUTE, new_world);
                 }
             }
@@ -1294,6 +1291,10 @@ JL_DLLEXPORT jl_binding_partition_t *jl_replace_binding_locked2(jl_binding_t *b,
     jl_atomic_store_relaxed(&new_bpart->next, old_bpart);
     jl_gc_wb_fresh(new_bpart, old_bpart);
 
+    if (((old_bpart->kind & BINDING_FLAG_EXPORTED) || (kind & BINDING_FLAG_EXPORTED)) && jl_require_world != ~(size_t)0) {
+        jl_atomic_store_release(&b->globalref->mod->export_set_changed_since_require_world, 1);
+    }
+
     jl_atomic_store_release(&b->partitions, new_bpart);
     jl_gc_wb(b, new_bpart);
     JL_GC_POP();

src/staticdata.c

@@ -1394,6 +1394,10 @@ static void jl_write_module(jl_serializer_state *s, uintptr_t item, jl_module_t
         }
     }
     assert(ios_pos(s->s) - reloc_offset == tot);
+
+    // After reload, everything that has happened in this process happened semantically at
+    // (for .incremental) or before jl_require_world, so reset this flag.
+    jl_atomic_store_relaxed(&newm->export_set_changed_since_require_world, 0);
 }
 
 static void record_memoryref(jl_serializer_state *s, size_t reloc_offset, jl_genericmemoryref_t ref) {
@@ -3537,32 +3541,31 @@ extern void export_jl_small_typeof(void);
 int IMAGE_NATIVE_CODE_TAINTED = 0;
 
 // TODO: This should possibly be in Julia
-static void jl_validate_binding_partition(jl_binding_t *b, jl_binding_partition_t *bpart, size_t mod_idx)
+static int jl_validate_binding_partition(jl_binding_t *b, jl_binding_partition_t *bpart, size_t mod_idx, int unchanged_implicit)
 {
-
     if (jl_atomic_load_relaxed(&bpart->max_world) != ~(size_t)0)
-        return;
+        return 1;
     size_t raw_kind = bpart->kind;
     enum jl_partition_kind kind = (enum jl_partition_kind)(raw_kind & 0x0f);
-    if (!jl_bkind_is_some_import(kind))
-        return;
-    jl_binding_t *imported_binding = (jl_binding_t*)bpart->restriction;
-    jl_binding_partition_t *latest_imported_bpart = jl_atomic_load_relaxed(&imported_binding->partitions);
-    if (!latest_imported_bpart)
-        return;
-    if (kind == BINDING_KIND_IMPLICIT || kind == BINDING_KIND_FAILED) {
+    if (!unchanged_implicit && jl_bkind_is_some_implicit(kind)) {
         jl_check_new_binding_implicit(bpart, b, NULL, jl_atomic_load_relaxed(&jl_world_counter));
         bpart->kind |= (raw_kind & 0xf0);
         if (bpart->min_world > jl_require_world)
             goto invalidated;
     }
+    if (!jl_bkind_is_some_import(kind))
+        return 1;
+    jl_binding_t *imported_binding = (jl_binding_t*)bpart->restriction;
+    jl_binding_partition_t *latest_imported_bpart = jl_atomic_load_relaxed(&imported_binding->partitions);
+    if (!latest_imported_bpart)
+        return 1;
     if (latest_imported_bpart->min_world <= bpart->min_world) {
         // Imported binding is still valid
         if ((kind == BINDING_KIND_EXPLICIT || kind == BINDING_KIND_IMPORTED) &&
                 external_blob_index((jl_value_t*)imported_binding) != mod_idx) {
             jl_add_binding_backedge(imported_binding, (jl_value_t*)b);
         }
-        return;
+        return 1;
     }
     else {
         // Binding partition was invalidated
@@ -3580,13 +3583,14 @@ static void jl_validate_binding_partition(jl_binding_t *b, jl_binding_partition_
             jl_binding_t *bedge = (jl_binding_t*)edge;
             if (!jl_atomic_load_relaxed(&bedge->partitions))
                 continue;
-            jl_validate_binding_partition(bedge, jl_atomic_load_relaxed(&bedge->partitions), mod_idx);
+            jl_validate_binding_partition(bedge, jl_atomic_load_relaxed(&bedge->partitions), mod_idx, 0);
         }
     }
     if (bpart->kind & BINDING_FLAG_EXPORTED) {
         jl_module_t *mod = b->globalref->mod;
         jl_sym_t *name = b->globalref->name;
         JL_LOCK(&mod->lock);
+        jl_atomic_store_release(&mod->export_set_changed_since_require_world, 1);
         if (mod->usings_backedges) {
             for (size_t i = 0; i < jl_array_len(mod->usings_backedges); i++) {
                 jl_module_t *edge = (jl_module_t*)jl_array_ptr_ref(mod->usings_backedges, i);
@@ -3596,12 +3600,24 @@ static void jl_validate_binding_partition(jl_binding_t *b, jl_binding_partition_
                 if (!jl_atomic_load_relaxed(&importee->partitions))
                     continue;
                 JL_UNLOCK(&mod->lock);
-                jl_validate_binding_partition(importee, jl_atomic_load_relaxed(&importee->partitions), mod_idx);
+                jl_validate_binding_partition(importee, jl_atomic_load_relaxed(&importee->partitions), mod_idx, 0);
                 JL_LOCK(&mod->lock);
             }
         }
         JL_UNLOCK(&mod->lock);
+        return 0;
     }
+    return 1;
+}
+
+static int all_usings_unchanged_implicit(jl_module_t *mod)
+{
+    int unchanged_implicit = 1;
+    for (size_t i = 0; unchanged_implicit && i < module_usings_length(mod); i++) {
+        jl_module_t *usee = module_usings_getmod(mod, i);
+        unchanged_implicit &= !jl_atomic_load_acquire(&usee->export_set_changed_since_require_world);
+    }
+    return unchanged_implicit;
 }
 
 static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl_array_t *depmods, uint64_t checksum,
@@ -4063,12 +4079,15 @@ static void jl_restore_system_image_from_stream_(ios_t *f, jl_image_t *image, jl
                 jl_module_t *mod = (jl_module_t*)obj;
                 size_t mod_idx = external_blob_index((jl_value_t*)mod);
                 jl_svec_t *table = jl_atomic_load_relaxed(&mod->bindings);
+                int unchanged_implicit = all_usings_unchanged_implicit(mod);
                 for (size_t i = 0; i < jl_svec_len(table); i++) {
                     jl_binding_t *b = (jl_binding_t*)jl_svecref(table, i);
                     if ((jl_value_t*)b == jl_nothing)
                         continue;
                     jl_binding_partition_t *bpart = jl_atomic_load_relaxed(&b->partitions);
-                    jl_validate_binding_partition(b, bpart, mod_idx);
+                    if (!jl_validate_binding_partition(b, bpart, mod_idx, unchanged_implicit)) {
+                        unchanged_implicit = all_usings_unchanged_implicit(mod);
+                    }
                 }
             }
         }