Merge pull request #19846 from JuliaLang/jb/fix18838

fix #18838, populate `source` field of generated functions

Author: JeffBezanson

src/alloc.c

@@ -484,7 +484,7 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo)
     jl_expr_t *ex = NULL;
     jl_value_t *linenum = NULL;
     jl_svec_t *sparam_vals = env;
-    jl_method_instance_t *generator = linfo->def->unspecialized;
+    jl_method_instance_t *generator = linfo->def->generator;
     assert(linfo != generator);
     assert(linfo->def->isstaged);
     jl_code_info_t *func = NULL;
@@ -509,7 +509,7 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *linfo)
         jl_array_t *argnames = jl_alloc_vec_any(nargs);
         jl_array_ptr_set(ex->args, 0, argnames);
         for (i = 0; i < nargs; i++)
-            jl_array_ptr_set(argnames, i, jl_array_ptr_ref(((jl_code_info_t*)generator->inferred)->slotnames, i));
+            jl_array_ptr_set(argnames, i, jl_array_ptr_ref(linfo->def->source->slotnames, i));
 
         jl_expr_t *scopeblock = jl_exprn(jl_symbol("scope-block"), 1);
         jl_array_ptr_set(ex->args, 1, scopeblock);
@@ -642,6 +642,7 @@ JL_DLLEXPORT jl_method_t *jl_new_method_uninit(void)
     m->module = ptls->current_module;
     m->source = NULL;
     m->unspecialized = NULL;
+    m->generator = NULL;
     m->name = NULL;
     m->file = empty_sym;
     m->line = 0;
@@ -689,12 +690,10 @@ jl_method_t *jl_new_method(jl_code_info_t *definition,
     root = (jl_value_t*)m;
     jl_method_set_source(m, definition);
     if (isstaged) {
-        // remove the code from `->source` (since generic source isn't present)
-        // and use the `->unspecialized` field to be the source generator
-        m->unspecialized = jl_get_specialized(m, jl_anytuple_type, jl_emptysvec);
-        jl_gc_wb(m, m->unspecialized);
-        m->unspecialized->inferred = (jl_value_t*)m->source;
-        m->source = NULL;
+        // create and store generator for generated functions
+        m->generator = jl_get_specialized(m, jl_anytuple_type, jl_emptysvec);
+        jl_gc_wb(m, m->generator);
+        m->generator->inferred = (jl_value_t*)m->source;
     }
 
 #ifdef RECORD_METHOD_ORDER

src/codegen.cpp

@@ -1297,9 +1297,7 @@ void jl_extern_c(jl_function_t *f, jl_value_t *rt, jl_value_t *argt, char *name)
 extern "C" JL_DLLEXPORT
 void *jl_get_llvmf_defn(jl_method_instance_t *linfo, size_t world, bool getwrapper, bool optimize, const jl_cgparams_t params)
 {
-    // `source` is `NULL` for generated functions.
-    // The `isstaged` check can be removed if that is not the case anymore.
-    if (linfo->def && linfo->def->source == NULL && !linfo->def->isstaged) {
+    if (linfo->def && linfo->def->source == NULL) {
         // not a generic function
         return NULL;
     }
@@ -1376,9 +1374,7 @@ void *jl_get_llvmf_defn(jl_method_instance_t *linfo, size_t world, bool getwrapp
 extern "C" JL_DLLEXPORT
 void *jl_get_llvmf_decl(jl_method_instance_t *linfo, size_t world, bool getwrapper, const jl_cgparams_t params)
 {
-    // `source` is `NULL` for generated functions.
-    // The `isstaged` check can be removed if that is not the case anymore.
-    if (linfo->def && linfo->def->source == NULL && !linfo->def->isstaged) {
+    if (linfo->def && linfo->def->source == NULL) {
         // not a generic function
         return NULL;
     }

src/dump.c

@@ -920,6 +920,7 @@ static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v)
         jl_serialize_value(s, (jl_value_t*)m->roots);
         jl_serialize_value(s, (jl_value_t*)m->source);
         jl_serialize_value(s, (jl_value_t*)m->unspecialized);
+        jl_serialize_value(s, (jl_value_t*)m->generator);
         jl_serialize_value(s, (jl_value_t*)m->invokes.unknown);
         write_int8(s->s, m->needs_sparam_vals_ducttape);
     }
@@ -1625,6 +1626,9 @@ static jl_value_t *jl_deserialize_value_method(jl_serializer_state *s, jl_value_
     m->unspecialized = (jl_method_instance_t*)jl_deserialize_value(s, (jl_value_t**)&m->unspecialized);
     if (m->unspecialized)
         jl_gc_wb(m, m->unspecialized);
+    m->generator = (jl_method_instance_t*)jl_deserialize_value(s, (jl_value_t**)&m->generator);
+    if (m->generator)
+        jl_gc_wb(m, m->generator);
     m->invokes.unknown = jl_deserialize_value(s, (jl_value_t**)&m->invokes);
     jl_gc_wb(m, m->invokes.unknown);
     m->needs_sparam_vals_ducttape = read_int8(s->s);

src/jltypes.c

@@ -3893,7 +3893,7 @@ void jl_init_types(void)
     jl_method_type =
         jl_new_datatype(jl_symbol("Method"),
                         jl_any_type, jl_emptysvec,
-                        jl_svec(20,
+                        jl_svec(21,
                                 jl_symbol("name"),
                                 jl_symbol("module"),
                                 jl_symbol("file"),
@@ -3907,14 +3907,15 @@ void jl_init_types(void)
                                 jl_symbol("sparam_syms"),
                                 jl_symbol("source"),
                                 jl_symbol("unspecialized"),
+                                jl_symbol("generator"),
                                 jl_symbol("roots"),
                                 jl_symbol("invokes"),
                                 jl_symbol("nargs"),
                                 jl_symbol("called"),
                                 jl_symbol("isva"),
                                 jl_symbol("isstaged"),
                                 jl_symbol("needs_sparam_vals_ducttape")),
-                        jl_svec(20,
+                        jl_svec(21,
                                 jl_sym_type,
                                 jl_module_type,
                                 jl_sym_type,
@@ -3928,6 +3929,7 @@ void jl_init_types(void)
                                 jl_simplevector_type,
                                 jl_code_info_type,
                                 jl_any_type, // jl_method_instance_type
+                                jl_any_type, // jl_method_instance_type
                                 jl_array_any_type,
                                 jl_any_type,
                                 jl_int32_type,
@@ -4037,6 +4039,7 @@ void jl_init_types(void)
 #endif
     jl_svecset(jl_methtable_type->types, 8, jl_int32_type); // uint32_t
     jl_svecset(jl_method_type->types, 12, jl_method_instance_type);
+    jl_svecset(jl_method_type->types, 13, jl_method_instance_type);
     jl_svecset(jl_method_instance_type->types, 12, jl_voidpointer_type);
     jl_svecset(jl_method_instance_type->types, 13, jl_voidpointer_type);
     jl_svecset(jl_method_instance_type->types, 14, jl_voidpointer_type);

src/julia.h

@@ -241,12 +241,10 @@ typedef struct _jl_method_t {
     // table of all argument types for which we've inferred or compiled this code
     union jl_typemap_t specializations;
 
-    // sparams are the symbols in the tvars vector
-    jl_svec_t *sparam_syms;
-    // the code AST template
-    jl_code_info_t *source; // null for builtins and staged functions
-    // unspecialized executable thunk (for isstaged, code for the generator), or null
-    struct _jl_method_instance_t *unspecialized;
+    jl_svec_t *sparam_syms;  // symbols corresponding to the tvars vector
+    jl_code_info_t *source;  // original code template, null for builtins
+    struct _jl_method_instance_t *unspecialized;  // unspecialized executable method instance, or null
+    struct _jl_method_instance_t *generator;  // executable code-generating function if isstaged
     jl_array_t *roots;  // pointers in generated code (shared to reduce memory), or null
 
     // cache of specializations of this method for invoke(), i.e.

src/toplevel.c

@@ -717,7 +717,7 @@ void print_func_loc(JL_STREAM *s, jl_method_t *m);
 
 void jl_check_static_parameter_conflicts(jl_method_t *m, jl_svec_t *t)
 {
-    jl_code_info_t *src = m->isstaged ? (jl_code_info_t*)m->unspecialized->inferred : m->source;
+    jl_code_info_t *src = m->source;
     size_t nvars = jl_array_len(src->slotnames);
 
     size_t i, n = jl_svec_len(t);