bpart: Allow inference/codegen to merge multiple partitions (#57602)

Keno · web-flow · commit a5157c0fc563 · 2025-03-02T15:50:03.000-05:00
When inference and codegen queries the partitions of a binding, they
often only care about a subset of the information in the partition. At
the moment, we always truncate world ranges to the most precise
partition that contains the relevant query. However, we can instead
expand the world range to cover all partitions that are equivalent for
the given query. To give a concrete example, both inference and codegen
never care about the exported flag in a binding partition, so we should
not unnecessarily truncate a world range just because an export was
introduced.

Further, this commit lays the ground work to stop invalidating code for
these same kinds of transitions, although the actual logic to do that
will come in a separate PR.
diff --git a/Compiler/src/abstractinterpretation.jl b/Compiler/src/abstractinterpretation.jl
diff --git a/Compiler/src/cicache.jl b/Compiler/src/cicache.jl
@@ -40,6 +40,14 @@ function intersect(a::WorldRange, b::WorldRange)
     return ret
 end
 
+function union(a::WorldRange, b::WorldRange)
+    if b.min_world < a.min_world
+        (b, a) = (a, b)
+    end
+    @assert a.max_world >= b.min_world - 1
+    return WorldRange(a.min_world, b.max_world)
+end
+
 """
     struct WorldView
 
diff --git a/Compiler/src/ssair/verify.jl b/Compiler/src/ssair/verify.jl
@@ -62,17 +62,14 @@ function check_op(ir::IRCode, domtree::DomTree, @nospecialize(op), use_bb::Int,
             raise_error()
         end
     elseif isa(op, GlobalRef)
-        bpart = lookup_binding_partition(min_world(ir.valid_worlds), op)
-        while is_some_imported(binding_kind(bpart)) && max_world(ir.valid_worlds) <= bpart.max_world
-            imported_binding = partition_restriction(bpart)::Core.Binding
-            bpart = lookup_binding_partition(min_world(ir.valid_worlds), imported_binding)
-        end
-        if (!is_defined_const_binding(binding_kind(bpart)) || (bpart.max_world < max_world(ir.valid_worlds))) &&
-                (op.mod !== Core) && (op.mod !== Base)
-            # Core and Base are excluded because the frontend uses them for intrinsics, etc.
-            # TODO: Decide which way to go with these.
-            @verify_error "Unbound or partitioned GlobalRef not allowed in value position"
-            raise_error()
+        if op.mod !== Core && op.mod !== Base
+            (valid_worlds, alldef) = scan_leaf_partitions(nothing, op, WorldWithRange(min_world(ir.valid_worlds), ir.valid_worlds)) do _, _, bpart
+                is_defined_const_binding(binding_kind(bpart))
+            end
+            if !alldef || max_world(valid_worlds) < max_world(ir.valid_worlds) || min_world(valid_worlds) > min_world(ir.valid_worlds)
+                @verify_error "Unbound or partitioned GlobalRef not allowed in value position"
+                raise_error()
+            end
         end
     elseif isa(op, Expr)
         # Only Expr(:boundscheck) is allowed in value position
diff --git a/Compiler/src/stmtinfo.jl b/Compiler/src/stmtinfo.jl
@@ -490,9 +490,7 @@ perform such accesses.
 """
 struct GlobalAccessInfo <: CallInfo
     b::Core.Binding
-    bpart::Core.BindingPartition
 end
-GlobalAccessInfo(::Core.Binding, ::Nothing) = NoCallInfo()
 function add_edges_impl(edges::Vector{Any}, info::GlobalAccessInfo)
     push!(edges, info.b)
 end
diff --git a/Compiler/test/effects.jl b/Compiler/test/effects.jl
@@ -400,6 +400,17 @@ let effects = Base.infer_effects(setglobal!_nothrow_undefinedyet2)
 end
 @test_throws TypeError setglobal!_nothrow_undefinedyet2()
 
+module ExportMutableGlobal
+    global mutable_global_for_setglobal_test::Int = 0
+    export mutable_global_for_setglobal_test
+end
+using .ExportMutableGlobal: mutable_global_for_setglobal_test
+f_assign_imported() = global mutable_global_for_setglobal_test = 42
+let effects = Base.infer_effects(f_assign_imported)
+    @test !Compiler.is_nothrow(effects)
+end
+@test_throws ErrorException f_assign_imported()
+
 # Nothrow for setfield!
 mutable struct SetfieldNothrow
     x::Int
diff --git a/base/Base.jl b/base/Base.jl
@@ -21,7 +21,6 @@ include(strcat(BUILDROOT, "version_git.jl")) # include($BUILDROOT/base/version_g
 
 # Initialize DL_LOAD_PATH as early as possible.  We are defining things here in
 # a slightly more verbose fashion than usual, because we're running so early.
-const DL_LOAD_PATH = String[]
 let os = ccall(:jl_get_UNAME, Any, ())
     if os === :Darwin || os === :Apple
         if Base.DARWIN_FRAMEWORK
diff --git a/base/Base_compiler.jl b/base/Base_compiler.jl
@@ -334,6 +334,7 @@ end
 
 BUILDROOT::String = ""
 DATAROOT::String = ""
+const DL_LOAD_PATH = String[]
 
 baremodule BuildSettings end
 
diff --git a/base/runtime_internals.jl b/base/runtime_internals.jl
@@ -225,6 +225,10 @@ function lookup_binding_partition(world::UInt, b::Core.Binding)
     ccall(:jl_get_binding_partition, Ref{Core.BindingPartition}, (Any, UInt), b, world)
 end
 
+function lookup_binding_partition(world::UInt, b::Core.Binding, previous_partition::Core.BindingPartition)
+    ccall(:jl_get_binding_partition_with_hint, Ref{Core.BindingPartition}, (Any, Any, UInt), b, previous_partition, world)
+end
+
 function convert(::Type{Core.Binding}, gr::Core.GlobalRef)
     if isdefined(gr, :binding)
         return gr.binding
diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -2928,14 +2928,13 @@ static jl_value_t *static_eval(jl_codectx_t &ctx, jl_value_t *ex)
 {
     if (jl_is_symbol(ex)) {
         jl_sym_t *sym = (jl_sym_t*)ex;
-        jl_binding_t *bnd = jl_get_module_binding(ctx.module, sym, 0);
-        jl_binding_partition_t *bpart = jl_get_binding_partition_all(bnd, ctx.min_world, ctx.max_world);
+        jl_binding_t *bnd = jl_get_module_binding(ctx.module, sym, 1);
         int possibly_deprecated = 0;
-        jl_walk_binding_inplace_all(&bnd, &bpart, &possibly_deprecated, ctx.min_world, ctx.max_world);
-        if (bpart && jl_bkind_is_some_constant(jl_binding_kind(bpart))) {
+        jl_value_t *cval = jl_get_binding_leaf_partitions_value_if_const(bnd, &possibly_deprecated, ctx.min_world, ctx.max_world);
+        if (cval) {
             if (possibly_deprecated)
                 emit_depwarn_check(ctx, bnd);
-            return bpart->restriction;
+            return cval;
         }
         return NULL;
     }
@@ -2957,13 +2956,9 @@ static jl_value_t *static_eval(jl_codectx_t &ctx, jl_value_t *ex)
     jl_sym_t *s = NULL;
     if (jl_is_globalref(ex)) {
         s = jl_globalref_name(ex);
-        jl_binding_t *bnd = jl_get_module_binding(jl_globalref_mod(ex), s, 0);
-        jl_binding_partition_t *bpart = jl_get_binding_partition_all(bnd, ctx.min_world, ctx.max_world);
+        jl_binding_t *bnd = jl_get_module_binding(jl_globalref_mod(ex), s, 1);
         int possibly_deprecated = 0;
-        jl_walk_binding_inplace_all(&bnd, &bpart, &possibly_deprecated, ctx.min_world, ctx.max_world);
-        jl_value_t *v = NULL;
-        if (bpart && jl_bkind_is_some_constant(jl_binding_kind(bpart)))
-            v = bpart->restriction;
+        jl_value_t *v = jl_get_binding_leaf_partitions_value_if_const(bnd, &possibly_deprecated, ctx.min_world, ctx.max_world);
         if (v) {
             if (possibly_deprecated)
                 emit_depwarn_check(ctx, bnd);
@@ -2985,13 +2980,9 @@ static jl_value_t *static_eval(jl_codectx_t &ctx, jl_value_t *ex)
                     // Assumes that the module is rooted somewhere.
                     s = (jl_sym_t*)static_eval(ctx, jl_exprarg(e, 2));
                     if (s && jl_is_symbol(s)) {
-                        jl_binding_t *bnd = jl_get_module_binding(m, s, 0);
-                        jl_binding_partition_t *bpart = jl_get_binding_partition_all(bnd, ctx.min_world, ctx.max_world);
+                        jl_binding_t *bnd = jl_get_module_binding(m, s, 1);
                         int possibly_deprecated = 0;
-                        jl_walk_binding_inplace_all(&bnd, &bpart, &possibly_deprecated, ctx.min_world, ctx.max_world);
-                        jl_value_t *v = NULL;
-                        if (bpart && jl_bkind_is_some_constant(jl_binding_kind(bpart)))
-                            v = bpart->restriction;
+                        jl_value_t *v = jl_get_binding_leaf_partitions_value_if_const(bnd, &possibly_deprecated, ctx.min_world, ctx.max_world);
                         if (v) {
                             if (possibly_deprecated)
                                 emit_depwarn_check(ctx, bnd);
@@ -3235,48 +3226,32 @@ static jl_cgval_t emit_globalref_runtime(jl_codectx_t &ctx, jl_binding_t *bnd, j
 static jl_cgval_t emit_globalref(jl_codectx_t &ctx, jl_module_t *mod, jl_sym_t *name, AtomicOrdering order)
 {
     jl_binding_t *bnd = jl_get_module_binding(mod, name, 1);
-    assert(bnd);
-    jl_binding_partition_t *bpart = jl_get_binding_partition_all(bnd, ctx.min_world, ctx.max_world);
-    if (!bpart) {
+    struct restriction_kind_pair rkp = { NULL, NULL, BINDING_KIND_GUARD, 0 };
+    if (!jl_get_binding_leaf_partitions_restriction_kind(bnd, &rkp, ctx.min_world, ctx.max_world)) {
         return emit_globalref_runtime(ctx, bnd, mod, name);
     }
-    int possibly_deprecated = 0;
-    int saw_explicit = 0;
-    while (bpart) {
-        if (!saw_explicit && (bpart->kind & BINDING_FLAG_DEPWARN))
-            possibly_deprecated = 1;
-        enum jl_partition_kind kind = jl_binding_kind(bpart);
-        if (!jl_bkind_is_some_import(kind))
-            break;
-        if (kind != BINDING_KIND_IMPLICIT)
-            saw_explicit = 1;
-        bnd = (jl_binding_t*)bpart->restriction;
-        bpart = jl_get_binding_partition_all(bnd, ctx.min_world, ctx.max_world);
-    }
-    Value *bp = NULL;
-    if (bpart) {
-        enum jl_partition_kind kind = jl_binding_kind(bpart);
-        if (jl_bkind_is_some_constant(kind) && kind != BINDING_KIND_BACKDATED_CONST) {
-            if (possibly_deprecated) {
-                bp = julia_binding_gv(ctx, bnd);
-                ctx.builder.CreateCall(prepare_call(jldepcheck_func), { bp });
-            }
-            jl_value_t *constval = bpart->restriction;
-            if (!constval) {
-                undef_var_error_ifnot(ctx, ConstantInt::get(getInt1Ty(ctx.builder.getContext()), 0), name, (jl_value_t*)mod);
-                return jl_cgval_t();
-            }
-            return mark_julia_const(ctx, constval);
+    if (jl_bkind_is_some_constant(rkp.kind) && rkp.kind != BINDING_KIND_BACKDATED_CONST) {
+        if (rkp.maybe_depwarn) {
+            Value *bp = julia_binding_gv(ctx, bnd);
+            ctx.builder.CreateCall(prepare_call(jldepcheck_func), { bp });
+        }
+        jl_value_t *constval = rkp.restriction;
+        if (!constval) {
+            undef_var_error_ifnot(ctx, ConstantInt::get(getInt1Ty(ctx.builder.getContext()), 0), name, (jl_value_t*)mod);
+            return jl_cgval_t();
         }
+        return mark_julia_const(ctx, constval);
     }
-    if (!bpart || jl_binding_kind(bpart) != BINDING_KIND_GLOBAL) {
+    if (rkp.kind != BINDING_KIND_GLOBAL) {
         return emit_globalref_runtime(ctx, bnd, mod, name);
     }
-    bp = julia_binding_gv(ctx, bnd);
-    if (possibly_deprecated) {
+    Value *bp = julia_binding_gv(ctx, bnd);
+    if (rkp.maybe_depwarn) {
         ctx.builder.CreateCall(prepare_call(jldepcheck_func), { bp });
     }
-    jl_value_t *ty = bpart->restriction;
+    if (bnd != rkp.binding_if_global)
+        bp = julia_binding_gv(ctx, rkp.binding_if_global);
+    jl_value_t *ty = rkp.restriction;
     Value *bpval = julia_binding_pvalue(ctx, bp);
     if (ty == nullptr)
         ty = (jl_value_t*)jl_any_type;
@@ -3844,27 +3819,27 @@ static jl_cgval_t emit_isdefinedglobal(jl_codectx_t &ctx, jl_module_t *modu, jl_
 {
     Value *isnull = NULL;
     jl_binding_t *bnd = allow_import ? jl_get_binding(modu, name) : jl_get_module_binding(modu, name, 0);
-    jl_binding_partition_t *bpart = jl_get_binding_partition_all(bnd, ctx.min_world, ctx.max_world);
-    enum jl_partition_kind kind = bpart ? jl_binding_kind(bpart) : BINDING_KIND_GUARD;
-    if (kind == BINDING_KIND_GLOBAL || jl_bkind_is_some_constant(kind)) {
-        if (jl_get_binding_value_if_const(bnd))
-            return mark_julia_const(ctx, jl_true);
-        Value *bp = julia_binding_gv(ctx, bnd);
-        bp = julia_binding_pvalue(ctx, bp);
-        LoadInst *v = ctx.builder.CreateAlignedLoad(ctx.types().T_prjlvalue, bp, Align(sizeof(void*)));
-        jl_aliasinfo_t ai = jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_binding);
-        ai.decorateInst(v);
-        v->setOrdering(get_llvm_atomic_order(order));
-        isnull = ctx.builder.CreateICmpNE(v, Constant::getNullValue(ctx.types().T_prjlvalue));
-    }
-    else {
-        Value *v = ctx.builder.CreateCall(prepare_call(jlboundp_func), {
-                literal_pointer_val(ctx, (jl_value_t*)modu),
-                literal_pointer_val(ctx, (jl_value_t*)name),
-                ConstantInt::get(getInt32Ty(ctx.builder.getContext()), allow_import)
-            });
-        isnull = ctx.builder.CreateICmpNE(v, ConstantInt::get(getInt32Ty(ctx.builder.getContext()), 0));
-    }
+    struct restriction_kind_pair rkp = { NULL, NULL, BINDING_KIND_GUARD, 0 };
+    if (allow_import && jl_get_binding_leaf_partitions_restriction_kind(bnd, &rkp, ctx.min_world, ctx.max_world)) {
+        if (jl_bkind_is_some_constant(rkp.kind))
+            return mark_julia_const(ctx, rkp.restriction);
+        if (rkp.kind == BINDING_KIND_GLOBAL) {
+            Value *bp = julia_binding_gv(ctx, rkp.binding_if_global);
+            bp = julia_binding_pvalue(ctx, bp);
+            LoadInst *v = ctx.builder.CreateAlignedLoad(ctx.types().T_prjlvalue, bp, Align(sizeof(void*)));
+            jl_aliasinfo_t ai = jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_binding);
+            ai.decorateInst(v);
+            v->setOrdering(get_llvm_atomic_order(order));
+            isnull = ctx.builder.CreateICmpNE(v, Constant::getNullValue(ctx.types().T_prjlvalue));
+            return mark_julia_type(ctx, isnull, false, jl_bool_type);
+        }
+    }
+    Value *v = ctx.builder.CreateCall(prepare_call(jlboundp_func), {
+            literal_pointer_val(ctx, (jl_value_t*)modu),
+            literal_pointer_val(ctx, (jl_value_t*)name),
+            ConstantInt::get(getInt32Ty(ctx.builder.getContext()), allow_import)
+        });
+    isnull = ctx.builder.CreateICmpNE(v, ConstantInt::get(getInt32Ty(ctx.builder.getContext()), 0));
     return mark_julia_type(ctx, isnull, false, jl_bool_type);
 }
 
diff --git a/src/dlload.c b/src/dlload.c
@@ -318,9 +318,9 @@ JL_DLLEXPORT void *jl_load_dynamic_library(const char *modname, unsigned flags,
       While these exist as OS concepts on Darwin, we want to use them on other platforms
       such as Windows, so we emulate them here.
     */
-    if (!abspath && !is_atpath && jl_base_module != NULL) {
+    if (!abspath && !is_atpath && jl_base_module != NULL && jl_typeinf_world != 1) {
         jl_binding_t *b = jl_get_module_binding(jl_base_module, jl_symbol("DL_LOAD_PATH"), 0);
-        jl_array_t *DL_LOAD_PATH = (jl_array_t*)(b ? jl_get_binding_value(b) : NULL);
+        jl_array_t *DL_LOAD_PATH = (jl_array_t*)(b ? jl_get_binding_value_in_world(b, jl_typeinf_world) : NULL);
         if (DL_LOAD_PATH != NULL) {
             size_t j;
             for (j = 0; j < jl_array_nrows(DL_LOAD_PATH); j++) {
diff --git a/src/julia.h b/src/julia.h
@@ -1907,6 +1907,7 @@ JL_DLLEXPORT jl_sym_t *jl_gensym(void);
 JL_DLLEXPORT jl_sym_t *jl_tagged_gensym(const char *str, size_t len);
 JL_DLLEXPORT jl_sym_t *jl_get_root_symbol(void);
 JL_DLLEXPORT jl_value_t *jl_get_binding_value(jl_binding_t *b JL_PROPAGATES_ROOT);
+JL_DLLEXPORT jl_value_t *jl_get_binding_value_in_world(jl_binding_t *b JL_PROPAGATES_ROOT, size_t world);
 JL_DLLEXPORT jl_value_t *jl_get_binding_value_if_const(jl_binding_t *b JL_PROPAGATES_ROOT);
 JL_DLLEXPORT jl_value_t *jl_get_binding_value_if_resolved(jl_binding_t *b JL_PROPAGATES_ROOT) JL_NOTSAFEPOINT;
 JL_DLLEXPORT jl_value_t *jl_get_binding_value_if_resolved_and_const(jl_binding_t *b JL_PROPAGATES_ROOT) JL_NOTSAFEPOINT;
diff --git a/src/julia_internal.h b/src/julia_internal.h
@@ -957,15 +957,26 @@ STATIC_INLINE int jl_bkind_is_defined_constant(enum jl_partition_kind kind) JL_N
 }
 
 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_with_hint(jl_binding_t *b JL_PROPAGATES_ROOT, jl_binding_partition_t *previous_part, 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;
 
+struct restriction_kind_pair {
+    jl_binding_t *binding_if_global;
+    jl_value_t *restriction;
+    enum jl_partition_kind kind;
+    int maybe_depwarn;
+};
+JL_DLLEXPORT int jl_get_binding_leaf_partitions_restriction_kind(jl_binding_t *b JL_PROPAGATES_ROOT, struct restriction_kind_pair *rkp, size_t min_world, size_t max_world) JL_GLOBALLY_ROOTED;
+JL_DLLEXPORT jl_value_t *jl_get_binding_leaf_partitions_value_if_const(jl_binding_t *b JL_PROPAGATES_ROOT, int *maybe_depwarn, size_t min_world, size_t max_world);
+
 EXTERN_INLINE_DECLARE uint8_t jl_bpart_get_kind(jl_binding_partition_t *bpart) JL_NOTSAFEPOINT {
     return (uint8_t)(bpart->kind & 0xf);
 }
 
 STATIC_INLINE void jl_walk_binding_inplace(jl_binding_t **bnd, jl_binding_partition_t **bpart JL_PROPAGATES_ROOT, size_t world) JL_NOTSAFEPOINT;
 STATIC_INLINE void jl_walk_binding_inplace_depwarn(jl_binding_t **bnd, jl_binding_partition_t **bpart, size_t world, int *depwarn) JL_NOTSAFEPOINT;
 STATIC_INLINE void jl_walk_binding_inplace_all(jl_binding_t **bnd, jl_binding_partition_t **bpart JL_PROPAGATES_ROOT, int *depwarn, size_t min_world, size_t max_world) JL_NOTSAFEPOINT;
+STATIC_INLINE void jl_walk_binding_inplace_worlds(jl_binding_t **bnd, jl_binding_partition_t **bpart, size_t *min_world, size_t *max_world, int *depwarn, size_t world) JL_NOTSAFEPOINT;
 
 #ifndef __clang_analyzer__
 STATIC_INLINE void jl_walk_binding_inplace(jl_binding_t **bnd, jl_binding_partition_t **bpart, size_t world) JL_NOTSAFEPOINT
@@ -1016,6 +1027,30 @@ STATIC_INLINE void jl_walk_binding_inplace_all(jl_binding_t **bnd, jl_binding_pa
         *bpart = jl_get_binding_partition_all(*bnd, min_world, max_world);
     }
 }
+
+STATIC_INLINE void jl_walk_binding_inplace_worlds(jl_binding_t **bnd, jl_binding_partition_t **bpart, size_t *min_world, size_t *max_world, int *depwarn, size_t world) JL_NOTSAFEPOINT
+{
+    int passed_explicit = 0;
+    while (*bpart) {
+        if (*min_world < (*bpart)->min_world)
+            *min_world = (*bpart)->min_world;
+        size_t bpart_max_world = jl_atomic_load_relaxed(&(*bpart)->max_world);
+        if (*max_world > bpart_max_world)
+            *max_world = bpart_max_world;
+        enum jl_partition_kind kind = jl_binding_kind(*bpart);
+        if (!jl_bkind_is_some_import(kind)) {
+            if (!passed_explicit && depwarn)
+                *depwarn |= (*bpart)->kind & BINDING_FLAG_DEPWARN;
+            return;
+        }
+        if (!passed_explicit && depwarn)
+            *depwarn |= (*bpart)->kind & BINDING_FLAG_DEPWARN;
+        if (kind != BINDING_KIND_IMPLICIT)
+            passed_explicit = 1;
+        *bnd = (jl_binding_t*)(*bpart)->restriction;
+        *bpart = jl_get_binding_partition(*bnd, world);
+    }
+}
 #endif
 
 STATIC_INLINE int is10digit(char c) JL_NOTSAFEPOINT
diff --git a/src/module.c b/src/module.c
diff --git a/stdlib/REPL/src/REPLCompletions.jl b/stdlib/REPL/src/REPLCompletions.jl
diff --git a/test/rebinding.jl b/test/rebinding.jl
