Allow Module as type parameters (#47749)

The intended use case for this is generated functions that want to
generate some reference to a module-specific generic function. The
current solution is to duplicate the generated function into every
module (probably using a package-provided macro) or to have some
sort of registry system in the package providing the generated
function. Both of these seem a bit ugly and I don't think there's
any particularly good reason not to allow Modules to be type
parameters. Admittedly, modules are not part of the scope
contemplated by #33387 as they are mutable, but I think the mental
model of modules is that they're immutable references to a namespace
and what's actually mutable is the namespace itself (i.e. people
wouldn't expect two modules that happen to have the same content
be `==`). This makes me think it still fits the mental model.

Author: Keno

base/compiler/typeutils.jl

@@ -132,12 +132,12 @@ function valid_as_lattice(@nospecialize(x))
 end
 
 function valid_typeof_tparam(@nospecialize(t))
-    if t === Symbol || isbitstype(t)
+    if t === Symbol || t === Module || isbitstype(t)
         return true
     end
     isconcretetype(t) || return false
     if t <: NamedTuple
-        t = t.parameters[2]
+        t = t.parameters[2]::DataType
     end
     if t <: Tuple
         for p in t.parameters

src/builtins.c

@@ -198,7 +198,7 @@ static int egal_types(const jl_value_t *a, const jl_value_t *b, jl_typeenv_t *en
             return egal_types(vma->N, vmb->N, env, tvar_names);
         return !vma->N && !vmb->N;
     }
-    if (dt == jl_symbol_type)
+    if (dt == jl_symbol_type || dt == jl_module_type)
         return 0;
     assert(!dt->name->mutabl);
     return jl_egal__bits(a, b, dt);
@@ -414,6 +414,10 @@ static uintptr_t NOINLINE jl_object_id__cold(jl_datatype_t *dt, jl_value_t *v) J
         return memhash32_seed(jl_string_data(v), jl_string_len(v), 0xedc3b677);
 #endif
     }
+    if (dt == jl_module_type) {
+        jl_module_t *m = (jl_module_t*)v;
+        return m->hash;
+    }
     if (dt->name->mutabl)
         return inthash((uintptr_t)v);
     return immut_id_(dt, v, dt->hash);
@@ -1269,7 +1273,8 @@ static int is_nestable_type_param(jl_value_t *t)
         size_t i, l = jl_nparams(t);
         for (i = 0; i < l; i++) {
             jl_value_t *pi = jl_tparam(t, i);
-            if (!(pi == (jl_value_t*)jl_symbol_type || jl_isbits(pi) || is_nestable_type_param(pi)))
+            if (!(pi == (jl_value_t*)jl_symbol_type || jl_isbits(pi) || is_nestable_type_param(pi) ||
+        jl_is_module(pi)))
                 return 0;
         }
         return 1;
@@ -1284,7 +1289,8 @@ int jl_valid_type_param(jl_value_t *v)
     if (jl_is_vararg(v))
         return 0;
     // TODO: maybe more things
-    return jl_is_type(v) || jl_is_typevar(v) || jl_is_symbol(v) || jl_isbits(jl_typeof(v));
+    return jl_is_type(v) || jl_is_typevar(v) || jl_is_symbol(v) || jl_isbits(jl_typeof(v)) ||
+        jl_is_module(v);
 }
 
 JL_CALLABLE(jl_f_apply_type)
@@ -1896,7 +1902,7 @@ void jl_init_intrinsic_properties(void) JL_GC_DISABLED
 
 void jl_init_intrinsic_functions(void) JL_GC_DISABLED
 {
-    jl_module_t *inm = jl_new_module(jl_symbol("Intrinsics"));
+    jl_module_t *inm = jl_new_module(jl_symbol("Intrinsics"), NULL);
     inm->parent = jl_core_module;
     jl_set_const(jl_core_module, jl_symbol("Intrinsics"), (jl_value_t*)inm);
     jl_mk_builtin_func(jl_intrinsic_type, "IntrinsicFunction", jl_f_intrinsic_call);

src/init.c

@@ -821,7 +821,7 @@ static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_
     jl_init_serializer();
 
     if (!jl_options.image_file) {
-        jl_core_module = jl_new_module(jl_symbol("Core"));
+        jl_core_module = jl_new_module(jl_symbol("Core"), NULL);
         jl_core_module->parent = jl_core_module;
         jl_type_typename->mt->module = jl_core_module;
         jl_top_module = jl_core_module;

src/julia.h

@@ -603,6 +603,7 @@ typedef struct _jl_module_t {
     uint8_t istopmod;
     int8_t max_methods;
     jl_mutex_t lock;
+    intptr_t hash;
 } jl_module_t;
 
 typedef struct {
@@ -1615,7 +1616,7 @@ extern JL_DLLEXPORT jl_module_t *jl_main_module JL_GLOBALLY_ROOTED;
 extern JL_DLLEXPORT jl_module_t *jl_core_module JL_GLOBALLY_ROOTED;
 extern JL_DLLEXPORT jl_module_t *jl_base_module JL_GLOBALLY_ROOTED;
 extern JL_DLLEXPORT jl_module_t *jl_top_module JL_GLOBALLY_ROOTED;
-JL_DLLEXPORT jl_module_t *jl_new_module(jl_sym_t *name);
+JL_DLLEXPORT jl_module_t *jl_new_module(jl_sym_t *name, jl_module_t *parent);
 JL_DLLEXPORT void jl_set_module_nospecialize(jl_module_t *self, int on);
 JL_DLLEXPORT void jl_set_module_optlevel(jl_module_t *self, int lvl);
 JL_DLLEXPORT int jl_get_module_optlevel(jl_module_t *m);

src/module.c

@@ -11,15 +11,15 @@
 extern "C" {
 #endif
 
-JL_DLLEXPORT jl_module_t *jl_new_module_(jl_sym_t *name, uint8_t default_names)
+JL_DLLEXPORT jl_module_t *jl_new_module_(jl_sym_t *name, jl_module_t *parent, uint8_t default_names)
 {
     jl_task_t *ct = jl_current_task;
     const jl_uuid_t uuid_zero = {0, 0};
     jl_module_t *m = (jl_module_t*)jl_gc_alloc(ct->ptls, sizeof(jl_module_t),
                                                jl_module_type);
     assert(jl_is_symbol(name));
     m->name = name;
-    m->parent = NULL;
+    m->parent = parent;
     m->istopmod = 0;
     m->uuid = uuid_zero;
     static unsigned int mcounter; // simple counter backup, in case hrtime is not incrementing
@@ -34,6 +34,8 @@ JL_DLLEXPORT jl_module_t *jl_new_module_(jl_sym_t *name, uint8_t default_names)
     m->compile = -1;
     m->infer = -1;
     m->max_methods = -1;
+    m->hash = parent == NULL ? bitmix(name->hash, jl_module_type->hash) :
+        bitmix(name->hash, parent->hash);
     JL_MUTEX_INIT(&m->lock);
     htable_new(&m->bindings, 0);
     arraylist_new(&m->usings, 0);
@@ -50,9 +52,9 @@ JL_DLLEXPORT jl_module_t *jl_new_module_(jl_sym_t *name, uint8_t default_names)
     return m;
 }
 
-JL_DLLEXPORT jl_module_t *jl_new_module(jl_sym_t *name)
+JL_DLLEXPORT jl_module_t *jl_new_module(jl_sym_t *name, jl_module_t *parent)
 {
-    return jl_new_module_(name, 1);
+    return jl_new_module_(name, parent, 1);
 }
 
 uint32_t jl_module_next_counter(jl_module_t *m)
@@ -63,10 +65,9 @@ uint32_t jl_module_next_counter(jl_module_t *m)
 JL_DLLEXPORT jl_value_t *jl_f_new_module(jl_sym_t *name, uint8_t std_imports, uint8_t default_names)
 {
     // TODO: should we prohibit this during incremental compilation?
-    jl_module_t *m = jl_new_module_(name, default_names);
+    // TODO: the parent module is a lie
+    jl_module_t *m = jl_new_module_(name, jl_main_module, default_names);
     JL_GC_PUSH1(&m);
-    m->parent = jl_main_module; // TODO: this is a lie
-    jl_gc_wb(m, m->parent);
     if (std_imports)
         jl_add_standard_imports(m);
     JL_GC_POP();

src/toplevel.c

@@ -48,7 +48,7 @@ JL_DLLEXPORT void jl_add_standard_imports(jl_module_t *m)
 void jl_init_main_module(void)
 {
     assert(jl_main_module == NULL);
-    jl_main_module = jl_new_module(jl_symbol("Main"));
+    jl_main_module = jl_new_module(jl_symbol("Main"), NULL);
     jl_main_module->parent = jl_main_module;
     jl_set_const(jl_main_module, jl_symbol("Core"),
                  (jl_value_t*)jl_core_module);
@@ -134,7 +134,8 @@ static jl_value_t *jl_eval_module_expr(jl_module_t *parent_module, jl_expr_t *ex
         jl_type_error("module", (jl_value_t*)jl_symbol_type, (jl_value_t*)name);
     }
 
-    jl_module_t *newm = jl_new_module(name);
+    int is_parent__toplevel__ = jl_is__toplevel__mod(parent_module);
+    jl_module_t *newm = jl_new_module(name, is_parent__toplevel__ ? NULL : parent_module);
     jl_value_t *form = (jl_value_t*)newm;
     JL_GC_PUSH1(&form);
     JL_LOCK(&jl_modules_mutex);
@@ -145,15 +146,14 @@ static jl_value_t *jl_eval_module_expr(jl_module_t *parent_module, jl_expr_t *ex
 
     // copy parent environment info into submodule
     newm->uuid = parent_module->uuid;
-    if (jl_is__toplevel__mod(parent_module)) {
+    if (is_parent__toplevel__) {
         newm->parent = newm;
         jl_register_root_module(newm);
         if (jl_options.incremental) {
             jl_precompile_toplevel_module = newm;
         }
     }
     else {
-        newm->parent = parent_module;
         jl_binding_t *b = jl_get_binding_wr(parent_module, name, 1);
         jl_declare_constant(b);
         jl_value_t *old = NULL;

test/core.jl

@@ -7873,3 +7873,16 @@ let # https://github.com/JuliaLang/julia/issues/46918
     @test isempty(String(take!(stderr))) # make sure no error has happened
     @test String(take!(stdout)) == "nothing IO IO"
 end
+
+# Modules allowed as type parameters and usable in generated functions
+module ModTparamTest
+    foo_test_mod_tparam() = 1
+end
+foo_test_mod_tparam() = 2
+
+struct ModTparamTestStruct{M}; end
+@generated function ModTparamTestStruct{M}() where {M}
+    return :($(GlobalRef(M, :foo_test_mod_tparam))())
+end
+@test ModTparamTestStruct{@__MODULE__}() == 2
+@test ModTparamTestStruct{ModTparamTest}() == 1

test/precompile.jl

@@ -1580,6 +1580,23 @@ end
     @test which(f46778, Tuple{Any,DataType}).specializations[1].cache.invoke != C_NULL
 end
 
+
+precompile_test_harness("Module tparams") do load_path
+    write(joinpath(load_path, "ModuleTparams.jl"),
+        """
+        module ModuleTparams
+            module TheTParam
+            end
+
+            struct ParamStruct{T}; end
+            const the_struct = ParamStruct{TheTParam}()
+        end
+        """)
+    Base.compilecache(Base.PkgId("ModuleTparams"))
+    (@eval (using ModuleTparams))
+    @test ModuleTparams.the_struct === Base.invokelatest(ModuleTparams.ParamStruct{ModuleTparams.TheTParam})
+end
+
 empty!(Base.DEPOT_PATH)
 append!(Base.DEPOT_PATH, original_depot_path)
 empty!(Base.LOAD_PATH)