Add basic infrastructure for binding replacement

Keno · Keno · commit 0102c8c3b982 · 2024-10-18T14:21:42.000Z
Now that I've had a few months to recover from the slog of adding `BindingPartition`, it's time to renew my quest to finish #54654. This adds the basic infrastructure for having multiple partitions, including making the lookup respect the `world` argument - on-demand allocation of missing partitions, `Base.delete_binding` and the `@world` macro. Not included is any inference or invalidation support, or any support for the runtime to create partitions itself (only `Base.delete_binding` does that for now), which will come in subsequent PRs.
diff --git a/base/essentials.jl b/base/essentials.jl
@@ -1250,6 +1250,50 @@ function isiterable(T)::Bool
     return hasmethod(iterate, Tuple{T})
 end
 
+"""
+    @world(sym, world)
+
+Resolve the binding `sym` in world `world`. See [`invoke_in_world`](@ref) for running
+arbitrary code in fixed worlds. `world` may be `UnitRange`, in which case the macro
+will error unless the binding is valid and has the same value across the entire world
+range.
+
+The `@world` macro is primarily used in the priniting of bindings that are no longer available
+in the current world.
+
+## Example
+```
+julia> struct Foo; a::Int; end
+Foo
+
+julia> fold = Foo(1)
+
+julia> Int(Base.get_world_counter())
+26866
+
+julia> struct Foo; a::Int; b::Int end
+Foo
+
+julia> fold
+@world(Foo, 26866)(1)
+```
+
+!!! compat "Julia 1.12"
+    This functionality requires at least Julia 1.12.
+"""
+macro world(sym, world)
+    if isa(sym, Symbol)
+        return :($(_resolve_in_world)($world, $(QuoteNode(GlobalRef(__module__, sym)))))
+    elseif isa(sym, GlobalRef)
+        return :($(_resolve_in_world)($world, $(QuoteNode(sym))))
+    else
+        error("`@world` requires a symbol or GlobalRef")
+    end
+end
+
+_resolve_in_world(world::Integer, gr::GlobalRef) =
+    invoke_in_world(UInt(world), Core.getglobal, gr.mod, gr.name)
+
 # Special constprop heuristics for various binary opes
 typename(typeof(function + end)).constprop_heuristic  = Core.SAMETYPE_HEURISTIC
 typename(typeof(function - end)).constprop_heuristic  = Core.SAMETYPE_HEURISTIC
diff --git a/base/range.jl b/base/range.jl
@@ -1680,3 +1680,14 @@ function show(io::IO, r::LogRange{T}) where {T}
     show(io, length(r))
     print(io, ')')
 end
+
+# Implementation detail of @world
+# The rest of this is defined in essentials.jl, but UnitRange is not available
+function _resolve_in_world(worlds::UnitRange, gr::GlobalRef)
+    # Validate that this binding's reference covers the entire world range
+    bpart = lookup_binding_partition(first(worlds), gr)
+    if bpart.max_world < last(world)
+        error("Binding does not cover the full world range")
+    end
+    _resolve_in_world(last(world), gr)
+end
diff --git a/base/runtime_internals.jl b/base/runtime_internals.jl
@@ -218,16 +218,19 @@ function _fieldnames(@nospecialize t)
     return t.name.names
 end
 
-const BINDING_KIND_GLOBAL       = 0x0
-const BINDING_KIND_CONST        = 0x1
-const BINDING_KIND_CONST_IMPORT = 0x2
+# N.B.: Needs to be synced with julia.h
+const BINDING_KIND_CONST        = 0x0
+const BINDING_KIND_CONST_IMPORT = 0x1
+const BINDING_KIND_GLOBAL       = 0x2
 const BINDING_KIND_IMPLICIT     = 0x3
 const BINDING_KIND_EXPLICIT     = 0x4
 const BINDING_KIND_IMPORTED     = 0x5
 const BINDING_KIND_FAILED       = 0x6
 const BINDING_KIND_DECLARED     = 0x7
 const BINDING_KIND_GUARD        = 0x8
 
+is_some_const_binding(kind::UInt8) = (kind == BINDING_KIND_CONST || kind == BINDING_KIND_CONST_IMPORT)
+
 function lookup_binding_partition(world::UInt, b::Core.Binding)
     ccall(:jl_get_binding_partition, Ref{Core.BindingPartition}, (Any, UInt), b, world)
 end
@@ -236,9 +239,27 @@ function lookup_binding_partition(world::UInt, gr::Core.GlobalRef)
     ccall(:jl_get_globalref_partition, Ref{Core.BindingPartition}, (Any, UInt), gr, world)
 end
 
+partition_restriction(bpart::Core.BindingPartition) = ccall(:jl_bpart_get_restriction_value, Any, (Any,), bpart)
+
 binding_kind(bpart::Core.BindingPartition) = ccall(:jl_bpart_get_kind, UInt8, (Any,), bpart)
 binding_kind(m::Module, s::Symbol) = binding_kind(lookup_binding_partition(tls_world_age(), GlobalRef(m, s)))
 
+"""
+    delete_binding(mod::Module, sym::Symbol)
+
+Force the binding `mod.sym` to be undefined again, allowing it be redefined.
+Note that this operation is very expensive, requirinig a full scan of all code in the system,
+as well as potential recompilation of any methods that (may) have used binding
+information.
+
+!!! warning
+    The implementation of this functionality is currently incomplete. Do not use
+    this method on versions that contain this disclaimer except for testing.
+"""
+function delete_binding(mod::Module, sym::Symbol)
+    ccall(:jl_disable_binding, Cvoid, (Any,), GlobalRef(mod, sym))
+end
+
 """
     fieldname(x::DataType, i::Integer)
 
diff --git a/base/show.jl b/base/show.jl
@@ -1032,6 +1032,21 @@ function is_global_function(tn::Core.TypeName, globname::Union{Symbol,Nothing})
     return false
 end
 
+function check_world_bounded(tn::Core.TypeName)
+    bnd = ccall(:jl_get_module_binding, Ref{Core.Binding}, (Any, Any, Cint), tn.module, tn.name, true)
+    isdefined(bnd, :partitions) || return nothing
+    partition = @atomic bnd.partitions
+    while true
+        if is_some_const_binding(binding_kind(partition)) && partition_restriction(partition) <: tn.wrapper
+            max_world = @atomic partition.max_world
+            max_world == typemax(UInt) && return nothing
+            return Int(partition.min_world):Int(max_world)
+        end
+        isdefined(partition, :next) || return nothing
+        partition = @atomic partition.next
+    end
+end
+
 function show_type_name(io::IO, tn::Core.TypeName)
     if tn === UnionAll.name
         # by coincidence, `typeof(Type)` is a valid representation of the UnionAll type.
@@ -1060,7 +1075,10 @@ function show_type_name(io::IO, tn::Core.TypeName)
             end
         end
     end
+    world = check_world_bounded(tn)
+    world !== nothing && print(io, "@world(")
     show_sym(io, sym)
+    world !== nothing && print(io, ", ", world, ")")
     quo      && print(io, ")")
     globfunc && print(io, ")")
     nothing
diff --git a/src/clangsa/GCChecker.cpp b/src/clangsa/GCChecker.cpp
@@ -819,6 +819,7 @@ bool GCChecker::isGCTrackedType(QualType QT) {
                    Name.ends_with_insensitive("jl_tupletype_t") ||
                    Name.ends_with_insensitive("jl_gc_tracked_buffer_t") ||
                    Name.ends_with_insensitive("jl_binding_t") ||
+                   Name.ends_with_insensitive("jl_binding_partition_t") ||
                    Name.ends_with_insensitive("jl_ordereddict_t") ||
                    Name.ends_with_insensitive("jl_tvar_t") ||
                    Name.ends_with_insensitive("jl_typemap_t") ||
@@ -842,6 +843,7 @@ bool GCChecker::isGCTrackedType(QualType QT) {
                    Name.ends_with_insensitive("jl_stenv_t") ||
                    Name.ends_with_insensitive("jl_varbinding_t") ||
                    Name.ends_with_insensitive("set_world") ||
+                   Name.ends_with_insensitive("jl_ptr_kind_union_t") ||
                    Name.ends_with_insensitive("jl_codectx_t")) {
                  return true;
                }
diff --git a/src/julia.h b/src/julia.h
@@ -620,6 +620,7 @@ typedef struct _jl_weakref_t {
     jl_value_t *value;
 } jl_weakref_t;
 
+// N.B: Needs to be synced with runtime_internals.jl
 enum jl_partition_kind {
     // Constant: This binding partition is a constant declared using `const`
     //  ->restriction holds the constant value
@@ -684,7 +685,7 @@ typedef struct __attribute__((aligned(8))) _jl_binding_partition_t {
     _Atomic(jl_ptr_kind_union_t) restriction;
     size_t min_world;
     _Atomic(size_t) max_world;
-    _Atomic(struct _jl_binding_partition_t*) next;
+    _Atomic(struct _jl_binding_partition_t *) next;
     size_t reserved; // Reserved for ->kind. Currently this holds the low bits of ->restriction during serialization
 } jl_binding_partition_t;
 
@@ -1839,8 +1840,8 @@ JL_DLLEXPORT jl_sym_t *jl_symbol_n(const char *str, size_t len) JL_NOTSAFEPOINT;
 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_NOTSAFEPOINT;
-JL_DLLEXPORT jl_value_t *jl_get_binding_value_if_const(jl_binding_t *b JL_PROPAGATES_ROOT) JL_NOTSAFEPOINT;
+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_if_const(jl_binding_t *b JL_PROPAGATES_ROOT);
 JL_DLLEXPORT jl_value_t *jl_declare_const_gf(jl_binding_t *b, jl_module_t *mod, jl_sym_t *name);
 JL_DLLEXPORT jl_method_t *jl_method_def(jl_svec_t *argdata, jl_methtable_t *mt, jl_code_info_t *f, jl_module_t *module);
 JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo, size_t world, jl_code_instance_t **cache);
diff --git a/src/julia_internal.h b/src/julia_internal.h
@@ -886,13 +886,10 @@ EXTERN_INLINE_DECLARE enum jl_partition_kind decode_restriction_kind(jl_ptr_kind
 #endif
 }
 
-STATIC_INLINE jl_value_t *decode_restriction_value(jl_ptr_kind_union_t pku) JL_NOTSAFEPOINT
+STATIC_INLINE jl_value_t *decode_restriction_value(jl_ptr_kind_union_t JL_PROPAGATES_ROOT pku) JL_NOTSAFEPOINT
 {
 #ifdef _P64
     jl_value_t *val = (jl_value_t*)(pku & ~0x7);
-    // This is a little bit of a lie at the moment - it is one of the things that
-    // can go wrong with binding replacement.
-    JL_GC_PROMISE_ROOTED(val);
     return val;
 #else
     return pku.val;
@@ -926,14 +923,8 @@ STATIC_INLINE int jl_bkind_is_some_guard(enum jl_partition_kind kind) JL_NOTSAFE
     return kind == BINDING_KIND_FAILED || kind == BINDING_KIND_GUARD || kind == BINDING_KIND_DECLARED;
 }
 
-EXTERN_INLINE_DECLARE jl_binding_partition_t *jl_get_binding_partition(jl_binding_t *b, size_t world) JL_NOTSAFEPOINT {
-    if (!b)
-        return NULL;
-    assert(jl_is_binding(b));
-    return jl_atomic_load_relaxed(&b->partitions);
-}
-
-JL_DLLEXPORT jl_binding_partition_t *jl_get_globalref_partition(jl_globalref_t *gr, size_t world);
+JL_DLLEXPORT jl_binding_partition_t *jl_get_binding_partition(jl_binding_t *b JL_PROPAGATES_ROOT, size_t world);
+JL_DLLEXPORT jl_binding_partition_t *jl_get_globalref_partition(jl_globalref_t *gr JL_PROPAGATES_ROOT, size_t world);
 
 EXTERN_INLINE_DECLARE uint8_t jl_bpart_get_kind(jl_binding_partition_t *bpart) JL_NOTSAFEPOINT {
     return decode_restriction_kind(jl_atomic_load_relaxed(&bpart->restriction));
diff --git a/src/module.c b/src/module.c
@@ -13,10 +13,51 @@ extern "C" {
 #endif
 
 // In this translation unit and this translation unit only emit this symbol `extern` for use by julia
-EXTERN_INLINE_DEFINE jl_binding_partition_t *jl_get_binding_partition(jl_binding_t *b, size_t world) JL_NOTSAFEPOINT;
 EXTERN_INLINE_DEFINE uint8_t jl_bpart_get_kind(jl_binding_partition_t *bpart) JL_NOTSAFEPOINT;
 extern inline enum jl_partition_kind decode_restriction_kind(jl_ptr_kind_union_t pku) JL_NOTSAFEPOINT;
 
+static jl_binding_partition_t *new_binding_partition(void)
+{
+    jl_binding_partition_t *bpart = (jl_binding_partition_t*)jl_gc_alloc(jl_current_task->ptls, sizeof(jl_binding_partition_t), jl_binding_partition_type);
+    jl_atomic_store_relaxed(&bpart->restriction, encode_restriction(NULL, BINDING_KIND_GUARD));
+    bpart->min_world = 0;
+    jl_atomic_store_relaxed(&bpart->max_world, (size_t)-1);
+    jl_atomic_store_relaxed(&bpart->next, NULL);
+#ifdef _P64
+    bpart->reserved = 0;
+#endif
+    return bpart;
+}
+
+jl_binding_partition_t *jl_get_binding_partition(jl_binding_t *b, size_t world) {
+    if (!b)
+        return NULL;
+    assert(jl_is_binding(b));
+    jl_value_t *parent = (jl_value_t*)b;
+    _Atomic(jl_binding_partition_t *)*insert = &b->partitions;
+    jl_binding_partition_t *bpart = jl_atomic_load_relaxed(insert);
+    size_t max_world = (size_t)-1;
+    while (1) {
+        while (bpart && world < bpart->min_world) {
+            insert = &bpart->next;
+            max_world = bpart->min_world - 1;
+            parent = (jl_value_t *)bpart;
+            bpart = jl_atomic_load_relaxed(&bpart->next);
+        }
+        if (bpart && world <= jl_atomic_load_relaxed(&bpart->max_world))
+            return bpart;
+        jl_binding_partition_t *new_bpart = new_binding_partition();
+        jl_atomic_store_relaxed(&new_bpart->next, bpart);
+        if (bpart)
+            new_bpart->min_world = jl_atomic_load_relaxed(&bpart->max_world) + 1;
+        jl_atomic_store_relaxed(&new_bpart->max_world, max_world);
+        if (jl_atomic_cmpswap(insert, &bpart, new_bpart)) {
+            jl_gc_wb(parent, new_bpart);
+            return new_bpart;
+        }
+    }
+}
+
 JL_DLLEXPORT jl_binding_partition_t *jl_get_globalref_partition(jl_globalref_t *gr, size_t world)
 {
     if (!gr)
@@ -188,19 +229,6 @@ static jl_globalref_t *jl_new_globalref(jl_module_t *mod, jl_sym_t *name, jl_bin
     return g;
 }
 
-static jl_binding_partition_t *new_binding_partition(void)
-{
-    jl_binding_partition_t *bpart = (jl_binding_partition_t*)jl_gc_alloc(jl_current_task->ptls, sizeof(jl_binding_partition_t), jl_binding_partition_type);
-    jl_atomic_store_relaxed(&bpart->restriction, encode_restriction(NULL, BINDING_KIND_GUARD));
-    bpart->min_world = 0;
-    jl_atomic_store_relaxed(&bpart->max_world, (size_t)-1);
-    jl_atomic_store_relaxed(&bpart->next, NULL);
-#ifdef _P64
-    bpart->reserved = 0;
-#endif
-    return bpart;
-}
-
 static jl_binding_t *new_binding(jl_module_t *mod, jl_sym_t *name)
 {
     jl_task_t *ct = jl_current_task;
@@ -215,9 +243,7 @@ static jl_binding_t *new_binding(jl_module_t *mod, jl_sym_t *name)
     JL_GC_PUSH1(&b);
     b->globalref = jl_new_globalref(mod, name, b);
     jl_gc_wb(b, b->globalref);
-    jl_binding_partition_t *bpart = new_binding_partition();
-    jl_atomic_store_relaxed(&b->partitions, bpart);
-    jl_gc_wb(b, bpart);
+    jl_atomic_store_relaxed(&b->partitions, NULL);
     JL_GC_POP();
     return b;
 }
@@ -324,6 +350,12 @@ JL_DLLEXPORT jl_value_t *jl_get_binding_value_if_const(jl_binding_t *b)
     return decode_restriction_value(pku);
 }
 
+JL_DLLEXPORT jl_value_t *jl_bpart_get_restriction_value(jl_binding_partition_t *bpart)
+{
+    jl_ptr_kind_union_t pku = jl_atomic_load_relaxed(&bpart->restriction);
+    return decode_restriction_value(pku);
+}
+
 typedef struct _modstack_t {
     jl_module_t *m;
     jl_sym_t *var;
@@ -952,6 +984,28 @@ JL_DLLEXPORT void jl_set_const(jl_module_t *m JL_ROOTING_ARGUMENT, jl_sym_t *var
     jl_gc_wb(bpart, val);
 }
 
+extern jl_mutex_t world_counter_lock;
+JL_DLLEXPORT void jl_disable_binding(jl_globalref_t *gr)
+{
+    jl_binding_t *b = gr->binding;
+    b = jl_resolve_owner(b, gr->mod, gr->name, NULL);
+    jl_binding_partition_t *bpart = jl_get_binding_partition(b, jl_current_task->world_age);
+
+    if (decode_restriction_kind(jl_atomic_load_relaxed(&bpart->restriction)) == BINDING_KIND_GUARD) {
+        // Already guard
+        return;
+    }
+
+    JL_LOCK(&world_counter_lock);
+    jl_task_t *ct = jl_current_task;
+    size_t new_max_world = jl_atomic_load_acquire(&jl_world_counter);
+    // TODO: Trigger invalidation here
+    (void)ct;
+    jl_atomic_store_release(&bpart->max_world, new_max_world);
+    jl_atomic_store_release(&jl_world_counter, new_max_world + 1);
+    JL_UNLOCK(&world_counter_lock);
+}
+
 JL_DLLEXPORT int jl_globalref_is_const(jl_globalref_t *gr)
 {
     jl_binding_t *b = gr->binding;
@@ -1023,13 +1077,17 @@ void jl_binding_deprecation_warning(jl_module_t *m, jl_sym_t *s, jl_binding_t *b
 
 jl_value_t *jl_check_binding_wr(jl_binding_t *b JL_PROPAGATES_ROOT, jl_module_t *mod, jl_sym_t *var, jl_value_t *rhs JL_MAYBE_UNROOTED, int reassign)
 {
+    JL_GC_PUSH1(&rhs); // callee-rooted
     jl_binding_partition_t *bpart = jl_get_binding_partition(b, jl_current_task->world_age);
     jl_ptr_kind_union_t pku = jl_atomic_load_relaxed(&bpart->restriction);
     assert(!jl_bkind_is_some_guard(decode_restriction_kind(pku)) && !jl_bkind_is_some_import(decode_restriction_kind(pku)));
     if (jl_bkind_is_some_constant(decode_restriction_kind(pku))) {
         jl_value_t *old = decode_restriction_value(pku);
-        if (jl_egal(rhs, old))
+        JL_GC_PROMISE_ROOTED(old);
+        if (jl_egal(rhs, old)) {
+            JL_GC_POP();
             return NULL;
+        }
         if (jl_typeof(rhs) == jl_typeof(old))
             jl_errorf("invalid redefinition of constant %s.%s. This redefinition may be permitted using the `const` keyword.",
                         jl_symbol_name(mod->name), jl_symbol_name(var));
@@ -1038,13 +1096,13 @@ jl_value_t *jl_check_binding_wr(jl_binding_t *b JL_PROPAGATES_ROOT, jl_module_t
                 jl_symbol_name(mod->name), jl_symbol_name(var));
     }
     jl_value_t *old_ty = decode_restriction_value(pku);
+    JL_GC_PROMISE_ROOTED(old_ty);
     if (old_ty != (jl_value_t*)jl_any_type && jl_typeof(rhs) != old_ty) {
-        JL_GC_PUSH1(&rhs); // callee-rooted
         if (!jl_isa(rhs, old_ty))
             jl_errorf("cannot assign an incompatible value to the global %s.%s.",
                         jl_symbol_name(mod->name), jl_symbol_name(var));
-        JL_GC_POP();
     }
+    JL_GC_POP();
     return old_ty;
 }
 
@@ -1081,6 +1139,7 @@ JL_DLLEXPORT jl_value_t *jl_checked_modify(jl_binding_t *b, jl_module_t *mod, jl
         jl_errorf("invalid redefinition of constant %s.%s",
                   jl_symbol_name(mod->name), jl_symbol_name(var));
     jl_value_t *ty = decode_restriction_value(pku);
+    JL_GC_PROMISE_ROOTED(ty);
     return modify_value(ty, &b->value, (jl_value_t*)b, op, rhs, 1, mod, var);
 }
 
diff --git a/src/staticdata.c b/src/staticdata.c
diff --git a/test/choosetests.jl b/test/choosetests.jl
diff --git a/test/rebinding.jl b/test/rebinding.jl
