support type annotations on global variables (#43671)

add type declarations to global bindings

replaces #43455
closes #964

Co-authored-by: Miguel Raz Guzmán Macedo <[email protected]>

Author: simeonschaub

NEWS.md

@@ -20,6 +20,7 @@ New language features
   them after construction, providing for greater clarity and optimization
   ability of these objects ([#43305]).
 * Empty n-dimensional arrays can now be created using multiple semicolons inside square brackets, i.e. `[;;;]` creates a 0×0×0 `Array`. ([#41618])
+* Type annotations can now be added to global variables to make accessing the variable type stable. ([#43671])
 
 Language changes
 ----------------
@@ -35,6 +36,10 @@ Language changes
   to mitigate the ["trojan source"](https://www.trojansource.codes) vulnerability ([#42918]).
 * `Base.ifelse` is now defined as a generic function rather than a builtin one, allowing packages to
   extend its definition ([#37343]).
+* Every assignment to a global variable now first goes through a call to `convert(Any, x)` (or `convert(T, x)`
+  respectively if a type `T` has already been declared for the global). This means great care should be taken
+  to ensure the invariant `convert(Any, x) === x` always holds, as this change could otherwise lead to
+  unexpected behavior. ([#43671])
 
 Compiler/Runtime improvements
 -----------------------------

base/Base.jl

@@ -57,7 +57,7 @@ modifyproperty!(x, f::Symbol, op, v, order::Symbol=:notatomic) =
 replaceproperty!(x, f::Symbol, expected, desired, success_order::Symbol=:notatomic, fail_order::Symbol=success_order) =
     (@inline; Core.replacefield!(x, f, expected, convert(fieldtype(typeof(x), f), desired), success_order, fail_order))
 
-
+convert(::Type{Any}, Core.@nospecialize x) = x
 include("coreio.jl")
 
 eval(x) = Core.eval(Base, x)

base/compiler/abstractinterpretation.jl

@@ -1793,7 +1793,9 @@ function abstract_eval_global(M::Module, s::Symbol)
     if isdefined(M,s) && isconst(M,s)
         return Const(getfield(M,s))
     end
-    return Any
+    ty = ccall(:jl_binding_type, Any, (Any, Any), M, s)
+    ty === nothing && return Any
+    return ty
 end
 
 function abstract_eval_ssavalue(s::SSAValue, src::CodeInfo)

base/compiler/compiler.jl

@@ -28,6 +28,8 @@ include(mod, x) = Core.include(mod, x)
 macro inline()   Expr(:meta, :inline)   end
 macro noinline() Expr(:meta, :noinline) end
 
+convert(::Type{Any}, Core.@nospecialize x) = x
+
 # essential files and libraries
 include("essentials.jl")
 include("ctypes.jl")

base/compiler/optimize.jl

@@ -156,7 +156,7 @@ const _PURE_BUILTINS = Any[tuple, svec, ===, typeof, nfields]
 const _PURE_OR_ERROR_BUILTINS = [
     fieldtype, apply_type, isa, UnionAll,
     getfield, arrayref, const_arrayref, arraysize, isdefined, Core.sizeof,
-    Core.kwfunc, Core.ifelse, Core._typevar, (<:)
+    Core.kwfunc, Core.ifelse, Core._typevar, (<:),
 ]
 
 const TOP_TUPLE = GlobalRef(Core, :tuple)
@@ -219,6 +219,12 @@ function stmt_effect_free(@nospecialize(stmt), @nospecialize(rt), src::Union{IRC
                         Any[argextype(args[i], src) for i = 2:length(args)])
             end
             contains_is(_PURE_BUILTINS, f) && return true
+            # `get_binding_type` sets the type to Any if the binding doesn't exist yet
+            if f === Core.get_binding_type
+                length(args) == 3 || return false
+                M, s = argextype(args[2], src), argextype(args[3], src)
+                return get_binding_type_effect_free(M, s)
+            end
             contains_is(_PURE_OR_ERROR_BUILTINS, f) || return false
             rt === Bottom && return false
             return _builtin_nothrow(f, Any[argextype(args[i], src) for i = 2:length(args)], rt)

base/compiler/ssair/inlining.jl

@@ -1404,6 +1404,11 @@ function early_inline_special_case(
             if _builtin_nothrow(f, argtypes[2:end], type)
                 return SomeCase(quoted(val))
             end
+        elseif f === Core.get_binding_type
+            length(argtypes) == 3 || return nothing
+            if get_binding_type_effect_free(argtypes[2], argtypes[3])
+                return SomeCase(quoted(val))
+            end
         end
     end
     return nothing

base/compiler/tfuncs.jl

@@ -1695,6 +1695,8 @@ function _builtin_nothrow(@nospecialize(f), argtypes::Array{Any,1}, @nospecializ
             return true
         end
         return false
+    elseif f === Core.get_binding_type
+        return length(argtypes) == 2
     end
     return false
 end
@@ -1898,4 +1900,20 @@ function typename_static(@nospecialize(t))
     return isType(t) ? _typename(t.parameters[1]) : Core.TypeName
 end
 
+function get_binding_type_effect_free(@nospecialize(M), @nospecialize(s))
+    if M isa Const && widenconst(M) === Module &&
+        s isa Const && widenconst(s) === Symbol
+        return ccall(:jl_binding_type, Any, (Any, Any), M.val, s.val) !== nothing
+    end
+    return false
+end
+function get_binding_type_tfunc(@nospecialize(M), @nospecialize(s))
+    if get_binding_type_effect_free(M, s)
+        @assert M isa Const && s isa Const
+        return Const(Core.get_binding_type(M.val, s.val))
+    end
+    return Type
+end
+add_tfunc(Core.get_binding_type, 2, 2, get_binding_type_tfunc, 0)
+
 @specialize

base/essentials.jl

@@ -211,7 +211,6 @@ See also: [`round`](@ref), [`trunc`](@ref), [`oftype`](@ref), [`reinterpret`](@r
 function convert end
 
 convert(::Type{Union{}}, @nospecialize x) = throw(MethodError(convert, (Union{}, x)))
-convert(::Type{Any}, @nospecialize x) = x
 convert(::Type{T}, x::T) where {T} = x
 convert(::Type{Type}, x::Type) = x # the ssair optimizer is strongly dependent on this method existing to avoid over-specialization
                                    # in the absence of inlining-enabled

doc/src/manual/performance-tips.md

@@ -11,9 +11,9 @@ The use of functions is not only important for performance: functions are more r
 
 The functions should take arguments, instead of operating directly on global variables, see the next point.
 
-## Avoid global variables
+## Avoid untyped global variables
 
-A global variable might have its value, and therefore possibly its type, changed at any point. This makes
+An untyped global variable might have its value, and therefore possibly its type, changed at any point. This makes
 it difficult for the compiler to optimize code using global variables. This also applies to type-valued variables,
 i.e. type aliases on the global level. Variables should be local, or passed as arguments to functions, whenever possible.
 
@@ -24,7 +24,9 @@ performance:
 const DEFAULT_VAL = 0
 ```
 
-Uses of non-constant globals can be optimized by annotating their types at the point of use:
+If a global is known to always be of the same type, [the type should be annotated](@ref man-typed-globals).
+
+Uses of untyped globals can be optimized by annotating their types at the point of use:
 
 ```julia
 global x = rand(1000)

doc/src/manual/variables-and-scoping.md

@@ -799,3 +799,58 @@ WARNING: redefinition of constant x. This may fail, cause incorrect answers, or
 julia> f()
 1
 ```
+
+## [Typed Globals](@id man-typed-globals)
+
+!!! compat "Julia 1.8"
+    Support for typed globals was added in Julia 1.8
+
+Similar to being declared as constants, global bindings can also be declared to always be of a
+constant type. This can either be done without assigning an actual value using the syntax
+`global x::T` or upon assignment as `x::T = 123`.
+
+```jldoctest
+julia> x::Float64 = 2.718
+2.718
+
+julia> f() = x
+f (generic function with 1 method)
+
+julia> Base.return_types(f)
+1-element Vector{Any}:
+ Float64
+```
+
+For any assignment to a global, Julia will first try to convert it to the appropriate type using
+[`convert`](@ref):
+
+```jldoctest
+julia> global y::Int
+
+julia> y = 1.0
+1.0
+
+julia> y
+1
+
+julia> y = 3.14
+ERROR: InexactError: Int64(3.14)
+Stacktrace:
+[...]
+```
+
+The type does not need to be concrete, but annotations with abstract types typically have little
+performance benefit.
+
+Once a global has either been assigned to or its type has been set, the binding type is not allowed
+to change:
+
+```jldoctest
+julia> x = 1
+1
+
+julia> global x::Int
+ERROR: cannot set type for global x. It already has a value or is already set to a different type.
+Stacktrace:
+[...]
+```

src/builtin_proto.h

@@ -65,6 +65,8 @@ JL_CALLABLE(jl_f__abstracttype);
 JL_CALLABLE(jl_f__primitivetype);
 JL_CALLABLE(jl_f__setsuper);
 JL_CALLABLE(jl_f__equiv_typedef);
+JL_CALLABLE(jl_f_get_binding_type);
+JL_CALLABLE(jl_f_set_binding_type);
 
 #ifdef __cplusplus
 }

src/builtins.c

@@ -1655,6 +1655,44 @@ JL_CALLABLE(jl_f__equiv_typedef)
     return equiv_type(args[0], args[1]) ? jl_true : jl_false;
 }
 
+JL_CALLABLE(jl_f_get_binding_type)
+{
+    JL_NARGS(get_binding_type, 2, 2);
+    JL_TYPECHK(get_binding_type, module, args[0]);
+    JL_TYPECHK(get_binding_type, symbol, args[1]);
+    jl_module_t *mod = (jl_module_t*)args[0];
+    jl_sym_t *sym = (jl_sym_t*)args[1];
+    jl_value_t *ty = jl_binding_type(mod, sym);
+    if (ty == (jl_value_t*)jl_nothing) {
+        jl_binding_t *b = jl_get_binding_wr(mod, sym, 0);
+        if (b) {
+            jl_value_t *old_ty = NULL;
+            jl_atomic_cmpswap_relaxed(&b->ty, &old_ty, (jl_value_t*)jl_any_type);
+            return jl_atomic_load_relaxed(&b->ty);
+        }
+        return (jl_value_t*)jl_any_type;
+    }
+    return ty;
+}
+
+JL_CALLABLE(jl_f_set_binding_type)
+{
+    JL_NARGS(set_binding_type!, 2, 3);
+    JL_TYPECHK(set_binding_type!, module, args[0]);
+    JL_TYPECHK(set_binding_type!, symbol, args[1]);
+    jl_value_t *ty = nargs == 2 ? (jl_value_t*)jl_any_type : args[2];
+    JL_TYPECHK(set_binding_type!, type, ty);
+    jl_binding_t *b = jl_get_binding_wr((jl_module_t*)args[0], (jl_sym_t*)args[1], 1);
+    jl_value_t *old_ty = NULL;
+    if (!jl_atomic_cmpswap_relaxed(&b->ty, &old_ty, ty) && ty != old_ty) {
+        if (nargs == 2)
+            return jl_nothing;
+        jl_errorf("cannot set type for global %s. It already has a value or is already set to a different type.",
+                  jl_symbol_name(b->name));
+    }
+    return jl_nothing;
+}
+
 // IntrinsicFunctions ---------------------------------------------------------
 
 static void (*runtime_fp[num_intrinsics])(void);
@@ -1834,6 +1872,8 @@ void jl_init_primitives(void) JL_GC_DISABLED
     add_builtin_func("_setsuper!", jl_f__setsuper);
     jl_builtin__typebody = add_builtin_func("_typebody!", jl_f__typebody);
     add_builtin_func("_equiv_typedef", jl_f__equiv_typedef);
+    add_builtin_func("get_binding_type", jl_f_get_binding_type);
+    add_builtin_func("set_binding_type!", jl_f_set_binding_type);
 
     // builtin types
     add_builtin("Any", (jl_value_t*)jl_any_type);

src/codegen.cpp

@@ -2385,7 +2385,7 @@ static jl_cgval_t emit_globalref(jl_codectx_t &ctx, jl_module_t *mod, jl_sym_t *
         LoadInst *v = ctx.builder.CreateAlignedLoad(ctx.types().T_prjlvalue, bp, Align(sizeof(void*)));
         v->setOrdering(AtomicOrdering::Unordered);
         tbaa_decorate(ctx.tbaa().tbaa_binding, v);
-        return mark_julia_type(ctx, v, true, (jl_value_t*)jl_any_type);
+        return mark_julia_type(ctx, v, true, bnd->ty);
     }
     // todo: use type info to avoid undef check
     return emit_checked_var(ctx, bp, name, false, ctx.tbaa().tbaa_binding);

src/dump.c

@@ -387,6 +387,7 @@ static void jl_serialize_module(jl_serializer_state *s, jl_module_t *m)
                 jl_serialize_value(s, e);
             jl_serialize_value(s, jl_atomic_load_relaxed(&b->globalref));
             jl_serialize_value(s, b->owner);
+            jl_serialize_value(s, jl_atomic_load_relaxed(&b->ty));
             write_int8(s->s, (b->deprecated<<3) | (b->constp<<2) | (b->exportp<<1) | (b->imported));
         }
     }
@@ -1708,6 +1709,8 @@ static jl_value_t *jl_deserialize_value_module(jl_serializer_state *s) JL_GC_DIS
         if (bglobalref != NULL) jl_gc_wb(m, bglobalref);
         b->owner = (jl_module_t*)jl_deserialize_value(s, (jl_value_t**)&b->owner);
         if (b->owner != NULL) jl_gc_wb(m, b->owner);
+        jl_value_t *bty = jl_deserialize_value(s, (jl_value_t**)&b->ty);
+        *(jl_value_t**)&b->ty = bty;
         int8_t flags = read_int8(s->s);
         b->deprecated = (flags>>3) & 1;
         b->constp = (flags>>2) & 1;