Add roots field to CodeInstance

This is designed to eliminate confusion in the serialization by ensuring
that offsets are relative to a "private" roots table. This may allow
more extensive caching of inference results, because it should eliminate
root-indexing conflicts between different instances of the same method.

Author: timholy

base/compiler/ssair/inlining.jl

@@ -702,7 +702,7 @@ function analyze_method!(idx::Int, sig::Signature, @nospecialize(metharg), meths
         return spec_lambda(atype_unlimited, sv, invoke_data)
     end
 
-    isconst, src = find_inferred(mi, atypes, sv, stmttyp)
+    isconst, src, roots = find_inferred(mi, atypes, sv, stmttyp)
     if isconst
         add_backedge!(mi, sv)
         return ConstantCase(src, method, Any[methsp...], metharg)
@@ -722,7 +722,7 @@ function analyze_method!(idx::Int, sig::Signature, @nospecialize(metharg), meths
     add_backedge!(mi, sv)
 
     if !isa(src, CodeInfo)
-        src = ccall(:jl_uncompress_ast, Any, (Any, Ptr{Cvoid}, Any), method, C_NULL, src::Vector{UInt8})::CodeInfo
+        src = ccall(:jl_uncompress_ast, Any, (Any, Ptr{Cvoid}, Any, Any), method, C_NULL, src::Vector{UInt8}, roots)::CodeInfo
     end
 
     @timeit "inline IR inflation" begin
@@ -1294,10 +1294,10 @@ function find_inferred(mi::MethodInstance, @nospecialize(atypes), sv::Optimizati
     if isa(inf_result, InferenceResult)
         let inferred_src = inf_result.src
             if isa(inferred_src, CodeInfo)
-                return svec(false, inferred_src)
+                return svec(false, inferred_src, nothing)
             end
             if isa(inferred_src, Const) && is_inlineable_constant(inferred_src.val)
-                return svec(true, quoted(inferred_src.val),)
+                return svec(true, quoted(inferred_src.val), nothing)
             end
         end
     end
@@ -1306,9 +1306,9 @@ function find_inferred(mi::MethodInstance, @nospecialize(atypes), sv::Optimizati
     if linfo isa CodeInstance
         if invoke_api(linfo) == 2
             # in this case function can be inlined to a constant
-            return svec(true, quoted(linfo.rettype_const))
+            return svec(true, quoted(linfo.rettype_const), linfo.localroots)
         end
-        return svec(false, linfo.inferred)
+        return svec(false, linfo.inferred, linfo.localroots)
     end
-    return svec(false, nothing)
+    return svec(false, nothing, nothing)
 end

base/compiler/typeinfer.jl

@@ -83,7 +83,7 @@ end
 # update the MethodInstance and notify the edges
 function cache_result(result::InferenceResult, min_valid::UInt, max_valid::UInt)
     def = result.linfo.def
-    toplevel = !isa(result.linfo.def, Method)
+    toplevel = !isa(def, Method)
     if toplevel
         min_valid = UInt(0)
         max_valid = UInt(0)
@@ -129,12 +129,19 @@ function cache_result(result::InferenceResult, min_valid::UInt, max_valid::UInt)
                 end
             end
         end
-        if !isa(inferred_result, Union{CodeInfo, Vector{UInt8}})
-            inferred_result = nothing
+        if isa(inferred_result, CodeInfo)
+            ccall(:jl_set_method_inferred, Ref{CodeInstance}, (Any, Any, Any, Any, Any, Int32, UInt, UInt),
+                result.linfo, widenconst(result.result), rettype_const, inferred_result, C_NULL,
+                const_flags, min_valid, max_valid)
+        elseif isa(inferred_result, SimpleVector)
+            ccall(:jl_set_method_inferred, Ref{CodeInstance}, (Any, Any, Any, Any, Any, Int32, UInt, UInt),
+                result.linfo, widenconst(result.result), rettype_const, inferred_result[1], inferred_result[2],
+                const_flags, min_valid, max_valid)
+        else
+            ccall(:jl_set_method_inferred, Ref{CodeInstance}, (Any, Any, Any, Any, Any, Int32, UInt, UInt),
+                result.linfo, widenconst(result.result), rettype_const, nothing, C_NULL,
+                const_flags, min_valid, max_valid)
         end
-        ccall(:jl_set_method_inferred, Ref{CodeInstance}, (Any, Any, Any, Any, Int32, UInt, UInt),
-            result.linfo, widenconst(result.result), rettype_const, inferred_result,
-            const_flags, min_valid, max_valid)
     end
     result.linfo.inInference = false
     nothing

base/compiler/utilities.jl

@@ -105,7 +105,7 @@ function retrieve_code_info(linfo::MethodInstance)
     if c === nothing && isdefined(m, :source)
         src = m.source
         if isa(src, Array{UInt8,1})
-            c = ccall(:jl_uncompress_ast, Any, (Any, Ptr{Cvoid}, Any), m, C_NULL, src)
+            c = ccall(:jl_uncompress_ast, Any, (Any, Ptr{Cvoid}, Any, Any), m, C_NULL, src, m.roots)
         else
             c = copy(src::CodeInfo)
         end

base/reflection.jl

@@ -912,8 +912,8 @@ uncompressed_ast(m::Method) = isdefined(m, :source) ? _uncompressed_ast(m, m.sou
                               isdefined(m, :generator) ? error("Method is @generated; try `code_lowered` instead.") :
                               error("Code for this Method is not available.")
 _uncompressed_ast(m::Method, s::CodeInfo) = copy(s)
-_uncompressed_ast(m::Method, s::Array{UInt8,1}) = ccall(:jl_uncompress_ast, Any, (Any, Ptr{Cvoid}, Any), m, C_NULL, s)::CodeInfo
-_uncompressed_ast(ci::Core.CodeInstance, s::Array{UInt8,1}) = ccall(:jl_uncompress_ast, Any, (Any, Any, Any), ci.def.def::Method, ci, s)::CodeInfo
+_uncompressed_ast(m::Method, s::Array{UInt8,1}) = ccall(:jl_uncompress_ast, Any, (Any, Ptr{Cvoid}, Any, Any), m, C_NULL, s, m.roots)::CodeInfo
+_uncompressed_ast(ci::Core.CodeInstance, s::Array{UInt8,1}) = ccall(:jl_uncompress_ast, Any, (Any, Any, Any, Any), ci.def.def::Method, ci, s, ci.localroots)::CodeInfo
 
 function method_instances(@nospecialize(f), @nospecialize(t), world::UInt = typemax(UInt))
     tt = signature_type(f, t)

src/codegen.cpp

@@ -1125,7 +1125,7 @@ jl_code_instance_t *jl_compile_linfo(jl_method_instance_t *mi, jl_code_info_t *s
             }
             // found inferred code, prep it for codegen
             if ((jl_value_t*)src != jl_nothing)
-                src = jl_uncompress_ast(mi->def.method, codeinst, (jl_array_t*)src);
+                src = jl_uncompress_ast(mi->def.method, codeinst, (jl_array_t*)src, (jl_array_t*)(codeinst->localroots));
             if (!jl_is_code_info(src)) {
                 src = jl_type_infer(mi, world, 0);
                 if (!src)
@@ -1157,7 +1157,7 @@ jl_code_instance_t *jl_compile_linfo(jl_method_instance_t *mi, jl_code_info_t *s
                 }
             }
             assert((jl_value_t*)src != jl_nothing);
-            src = jl_uncompress_ast(mi->def.method, NULL, (jl_array_t*)src);
+            src = jl_uncompress_ast(mi->def.method, NULL, (jl_array_t*)src, mi->def.method->roots);
         }
         assert(jl_is_code_info(src));
 
@@ -1175,6 +1175,7 @@ jl_code_instance_t *jl_compile_linfo(jl_method_instance_t *mi, jl_code_info_t *s
             uncached->functionObjectsDecls.functionObject = NULL;
             uncached->functionObjectsDecls.specFunctionObject = NULL;
             uncached->inferred = jl_nothing;
+            uncached->localroots = jl_alloc_vec_any(0);
             if (uncached->invoke != jl_fptr_const_return)
                 uncached->invoke = NULL;
             uncached->specptr.fptr = NULL;
@@ -1240,8 +1241,11 @@ jl_code_instance_t *jl_compile_linfo(jl_method_instance_t *mi, jl_code_info_t *s
                 jl_options.debug_level > 1) {
                 // update the stored code
                 if (codeinst->inferred != (jl_value_t*)src) {
-                    if (jl_is_method(mi->def.method))
-                        src = (jl_code_info_t*)jl_compress_ast(mi->def.method, src);
+                    if (jl_is_method(mi->def.method)) {
+                        jl_svec_t *srcroots = jl_compress_ast(mi->def.method, src);
+                        src = (jl_code_info_t*) jl_svecref(srcroots, 0);
+                        codeinst->localroots = (jl_array_t*) jl_svecref(srcroots, 1);
+                    }
                     codeinst->inferred = (jl_value_t*)src;
                     jl_gc_wb(codeinst, src);
                 }
@@ -1504,7 +1508,7 @@ void *jl_get_llvmf_defn(jl_method_instance_t *mi, size_t world, bool getwrapper,
         jl_code_instance_t *codeinst = (jl_code_instance_t*)ci;
         src = (jl_code_info_t*)codeinst->inferred;
         if ((jl_value_t*)src != jl_nothing && !jl_is_code_info(src) && jl_is_method(mi->def.method))
-            src = jl_uncompress_ast(mi->def.method, codeinst, (jl_array_t*)src);
+            src = jl_uncompress_ast(mi->def.method, codeinst, (jl_array_t*)src, (jl_array_t*)(codeinst->localroots));
         jlrettype = codeinst->rettype;
     }
     if (!src || (jl_value_t*)src == jl_nothing) {
@@ -1514,7 +1518,7 @@ void *jl_get_llvmf_defn(jl_method_instance_t *mi, size_t world, bool getwrapper,
         if (!src && jl_is_method(mi->def.method)) {
             src = mi->def.method->generator ? jl_code_for_staged(mi) : (jl_code_info_t*)mi->def.method->source;
             if (src && !jl_is_code_info(src) && jl_is_method(mi->def.method))
-                src = jl_uncompress_ast(mi->def.method, NULL, (jl_array_t*)src);
+                src = jl_uncompress_ast(mi->def.method, NULL, (jl_array_t*)src, mi->def.method->roots);
         }
     }
     if (src == NULL || (jl_value_t*)src == jl_nothing || !jl_is_code_info(src))

src/dump.c

@@ -141,7 +141,8 @@ typedef struct {
     ios_t *s;
     DUMP_MODES mode;
     // method we're compressing for in MODE_IR
-    jl_method_t *method;
+    jl_module_t *module;
+    jl_array_t *roots;
     jl_ptls_t ptls;
     jl_array_t *loaded_modules_array;
 } jl_serializer_state;
@@ -484,7 +485,7 @@ static int is_ast_node(jl_value_t *v)
 
 static int literal_val_id(jl_serializer_state *s, jl_value_t *v) JL_GC_DISABLED
 {
-    jl_array_t *rs = s->method->roots;
+    jl_array_t *rs = s->roots;
     int i, l = jl_array_len(rs);
     if (jl_is_symbol(v) || jl_is_concrete_type(v)) {
         for (i = 0; i < l; i++) {
@@ -535,7 +536,7 @@ static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v, int as_li
     }
 
     if (s->mode == MODE_IR) {
-        if (v == (jl_value_t*)s->method->module) {
+        if (v == (jl_value_t*)s->module) {
             write_uint8(s->s, TAG_NEARBYMODULE);
             return;
         }
@@ -640,7 +641,7 @@ static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v, int as_li
         ios_write(s->s, jl_symbol_name((jl_sym_t*)v), l);
     }
     else if (jl_is_globalref(v)) {
-        if (s->mode == MODE_IR && jl_globalref_mod(v) == s->method->module) {
+        if (s->mode == MODE_IR && jl_globalref_mod(v) == s->module) {
             write_uint8(s->s, TAG_NEARBYGLOBAL);
             jl_serialize_value(s, jl_globalref_name(v));
         }
@@ -885,13 +886,15 @@ static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v, int as_li
         jl_serialize_value(s, (jl_value_t*)codeinst->def);
         if (validate || codeinst->min_world == 0) {
             jl_serialize_value(s, codeinst->inferred);
+            jl_serialize_value(s, codeinst->localroots);
             jl_serialize_value(s, codeinst->rettype_const);
             jl_serialize_value(s, codeinst->rettype);
         }
         else {
             // skip storing useless data
             jl_serialize_value(s, NULL);
             jl_serialize_value(s, NULL);
+            jl_serialize_value(s, NULL);
             jl_serialize_value(s, jl_any_type);
         }
         jl_serialize_value(s, codeinst->next);
@@ -1387,7 +1390,7 @@ static jl_value_t *jl_deserialize_datatype(jl_serializer_state *s, int pos, jl_v
         assert(0 && "corrupt deserialization state");
         abort();
     }
-    assert(s->method == NULL && s->mode != MODE_IR && "no new data-types expected during MODE_IR");
+    assert(s->roots == NULL && s->mode != MODE_IR && "no new data-types expected during MODE_IR");
     assert(pos == backref_list.len - 1 && "nothing should have been deserialized since assigning pos");
     backref_list.items[pos] = dt;
     dt->size = size;
@@ -1762,7 +1765,11 @@ static jl_value_t *jl_deserialize_value_code_instance(jl_serializer_state *s, jl
     codeinst->def = (jl_method_instance_t*)jl_deserialize_value(s, (jl_value_t**)&codeinst->def);
     jl_gc_wb(codeinst, codeinst->def);
     codeinst->inferred = jl_deserialize_value(s, &codeinst->inferred);
-    jl_gc_wb(codeinst, codeinst->inferred);
+    if (codeinst->inferred)
+        jl_gc_wb(codeinst, codeinst->inferred);
+    codeinst->localroots = (jl_array_t*) jl_deserialize_value(s, (jl_value_t**) &(codeinst->localroots));
+    if (codeinst->localroots)
+        jl_gc_wb(codeinst, codeinst->localroots);
     codeinst->rettype_const = jl_deserialize_value(s, &codeinst->rettype_const);
     if (codeinst->rettype_const)
         jl_gc_wb(codeinst, codeinst->rettype_const);
@@ -2023,7 +2030,7 @@ static jl_value_t *jl_deserialize_value(jl_serializer_state *s, jl_value_t **loc
         tag = read_uint8(s->s);
         return deser_tag[tag];
     case TAG_BACKREF: JL_FALLTHROUGH; case TAG_SHORT_BACKREF:
-        assert(s->method == NULL && s->mode != MODE_IR);
+        assert(s->roots == NULL && s->mode != MODE_IR);
         uintptr_t offs = (tag == TAG_BACKREF) ? read_int32(s->s) : read_uint16(s->s);
         int isflagref = 0;
         isflagref = !!(offs & 1);
@@ -2042,9 +2049,9 @@ static jl_value_t *jl_deserialize_value(jl_serializer_state *s, jl_value_t **loc
         }
         return (jl_value_t*)bp;
     case TAG_METHODROOT:
-        return jl_array_ptr_ref(s->method->roots, read_uint8(s->s));
+        return jl_array_ptr_ref(s->roots, read_uint8(s->s));
     case TAG_LONG_METHODROOT:
-        return jl_array_ptr_ref(s->method->roots, read_uint16(s->s));
+        return jl_array_ptr_ref(s->roots, read_uint16(s->s));
     case TAG_SVEC: JL_FALLTHROUGH; case TAG_LONG_SVEC:
         return jl_deserialize_value_svec(s, tag);
     case TAG_COMMONSYM:
@@ -2152,12 +2159,12 @@ static jl_value_t *jl_deserialize_value(jl_serializer_state *s, jl_value_t **loc
     case TAG_UINT8:
         return jl_box_uint8(read_uint8(s->s));
     case TAG_NEARBYGLOBAL:
-        assert(s->method != NULL);
+        assert(s->module != NULL);
         v = jl_deserialize_value(s, NULL);
-        return jl_module_globalref(s->method->module, (jl_sym_t*)v);
+        return jl_module_globalref(s->module, (jl_sym_t*)v);
     case TAG_NEARBYMODULE:
-        assert(s->method != NULL);
-        return (jl_value_t*)s->method->module;
+        assert(s->module != NULL);
+        return (jl_value_t*)s->module;
     case TAG_GLOBALREF:
         return jl_deserialize_value_globalref(s);
     case TAG_SINGLETON:
@@ -2479,7 +2486,7 @@ JL_DLLEXPORT void jl_init_restored_modules(jl_array_t *init_order)
 
 // --- entry points ---
 
-JL_DLLEXPORT jl_array_t *jl_compress_ast(jl_method_t *m, jl_code_info_t *code)
+JL_DLLEXPORT jl_svec_t *jl_compress_ast(jl_method_t *m, jl_code_info_t *code)
 {
     JL_TIMING(AST_COMPRESS);
     JL_LOCK(&m->writelock); // protect the roots array (Might GC)
@@ -2490,13 +2497,11 @@ JL_DLLEXPORT jl_array_t *jl_compress_ast(jl_method_t *m, jl_code_info_t *code)
     int en = jl_gc_enable(0); // Might GC
     size_t i;
 
-    if (m->roots == NULL) {
-        m->roots = jl_alloc_vec_any(0);
-        jl_gc_wb(m, m->roots);
-    }
+    jl_array_t *localroots = jl_alloc_vec_any(0);
     jl_serializer_state s = {
         &dest, MODE_IR,
-        m,
+        m->module,
+        localroots,
         jl_get_ptls_states(),
         NULL
     };
@@ -2552,26 +2557,24 @@ JL_DLLEXPORT jl_array_t *jl_compress_ast(jl_method_t *m, jl_code_info_t *code)
     }
 
     ios_flush(s.s);
-    jl_array_t *v = jl_take_buffer(&dest);
+    jl_svec_t *v = jl_svec2(jl_take_buffer(&dest), localroots);
     ios_close(s.s);
-    if (jl_array_len(m->roots) == 0) {
-        m->roots = NULL;
-    }
     JL_GC_PUSH1(&v);
     jl_gc_enable(en);
     JL_UNLOCK(&m->writelock); // Might GC
     JL_GC_POP();
     return v;
 }
 
-JL_DLLEXPORT jl_code_info_t *jl_uncompress_ast(jl_method_t *m, jl_code_instance_t *metadata, jl_array_t *data)
+JL_DLLEXPORT jl_code_info_t *jl_uncompress_ast(jl_method_t *m, jl_code_instance_t *metadata, jl_array_t *data, jl_array_t *localroots)
 {
     if (jl_is_code_info(data))
         return (jl_code_info_t*)data;
     JL_TIMING(AST_UNCOMPRESS);
     JL_LOCK(&m->writelock); // protect the roots array (Might GC)
     assert(jl_is_method(m));
     assert(jl_typeis(data, jl_array_uint8_type));
+    assert(localroots);
     size_t i;
     ios_t src;
     ios_mem(&src, 0);
@@ -2580,7 +2583,8 @@ JL_DLLEXPORT jl_code_info_t *jl_uncompress_ast(jl_method_t *m, jl_code_instance_
     int en = jl_gc_enable(0); // Might GC
     jl_serializer_state s = {
         &src, MODE_IR,
-        m,
+        m->module,
+        localroots,
         jl_get_ptls_states(),
         NULL
     };
@@ -2817,6 +2821,7 @@ JL_DLLEXPORT int jl_save_incremental(const char *fname, jl_array_t *worklist)
     jl_serializer_state s = {
         &f, MODE_MODULE,
         NULL,
+        NULL,
         jl_get_ptls_states(),
         mod_array
     };
@@ -3198,6 +3203,7 @@ static jl_value_t *_jl_restore_incremental(ios_t *f, jl_array_t *mod_array)
     jl_serializer_state s = {
         f, MODE_MODULE,
         NULL,
+        NULL,
         ptls,
         mod_array
     };

src/gf.c

@@ -131,6 +131,7 @@ JL_DLLEXPORT jl_method_t *jl_new_method_uninit(jl_module_t*);
 JL_DLLEXPORT jl_code_instance_t* jl_set_method_inferred(
         jl_method_instance_t *mi, jl_value_t *rettype,
         jl_value_t *inferred_const, jl_value_t *inferred,
+        jl_array_t *localroots,
         int32_t const_flags, size_t min_world, size_t max_world);
 
 jl_datatype_t *jl_mk_builtin_func(jl_datatype_t *dt, const char *name, jl_fptr_args_t fptr) JL_GC_DISABLED
@@ -155,7 +156,7 @@ jl_datatype_t *jl_mk_builtin_func(jl_datatype_t *dt, const char *name, jl_fptr_a
     m->unspecialized = mi;
     jl_gc_wb(m, mi);
 
-    jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, jl_nothing, jl_nothing,
+    jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, jl_nothing, jl_nothing, NULL,
         0, 1, ~(size_t)0);
     codeinst->specptr.fptr1 = fptr;
     codeinst->invoke = jl_fptr_args;
@@ -266,13 +267,14 @@ JL_DLLEXPORT jl_code_instance_t *jl_get_method_inferred(
         codeinst = codeinst->next;
     }
     return jl_set_method_inferred(
-        mi, rettype, NULL, NULL,
+        mi, rettype, NULL, NULL, NULL,
         0, min_world, max_world);
 }
 
 JL_DLLEXPORT jl_code_instance_t *jl_set_method_inferred(
         jl_method_instance_t *mi, jl_value_t *rettype,
         jl_value_t *inferred_const, jl_value_t *inferred,
+        jl_array_t *localroots,
         int32_t const_flags, size_t min_world, size_t max_world
         /*, jl_array_t *edges, int absolute_max*/)
 {
@@ -288,6 +290,7 @@ JL_DLLEXPORT jl_code_instance_t *jl_set_method_inferred(
     codeinst->functionObjectsDecls.specFunctionObject = NULL;
     codeinst->rettype = rettype;
     codeinst->inferred = inferred;
+    codeinst->localroots = localroots;
     //codeinst->edges = NULL;
     if ((const_flags & 2) == 0)
         inferred_const = NULL;
@@ -1738,7 +1741,7 @@ jl_code_instance_t *jl_compile_method_internal(jl_method_instance_t *mi, size_t
         if (jl_is_method(def) && def->unspecialized) {
             jl_code_instance_t *unspec = def->unspecialized->cache;
             if (unspec && unspec->invoke != NULL) {
-                jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, NULL, NULL,
+                jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, NULL, NULL, NULL,
                     0, 1, ~(size_t)0);
                 codeinst->specptr = unspec->specptr;
                 codeinst->rettype_const = unspec->rettype_const;
@@ -1748,7 +1751,7 @@ jl_code_instance_t *jl_compile_method_internal(jl_method_instance_t *mi, size_t
         }
         jl_code_info_t *src = jl_code_for_interpreter(mi);
         if (!jl_code_requires_compiler(src)) {
-            jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, NULL, NULL,
+            jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, NULL, NULL, NULL,
                 0, 1, ~(size_t)0);
             jl_atomic_store_release(&codeinst->invoke, jl_fptr_interpret_call);
             return codeinst;
@@ -1786,7 +1789,7 @@ jl_code_instance_t *jl_compile_method_internal(jl_method_instance_t *mi, size_t
                 ucache->invoke != jl_fptr_interpret_call) {
                 return ucache;
             }
-            jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, NULL, NULL,
+            jl_code_instance_t *codeinst = jl_set_method_inferred(mi, (jl_value_t*)jl_any_type, NULL, NULL, NULL,
                 0, 1, ~(size_t)0);
             codeinst->specptr = ucache->specptr;
             codeinst->rettype_const = ucache->rettype_const;

src/interpreter.c

@@ -790,22 +790,20 @@ jl_code_info_t *jl_code_for_interpreter(jl_method_instance_t *mi)
     if (jl_is_method(mi->def.value)) {
         if (!src || (jl_value_t*)src == jl_nothing) {
             if (mi->def.method->source) {
-                src = (jl_code_info_t*)mi->def.method->source;
+                src = jl_uncompress_ast(mi->def.method, NULL, mi->def.method->source,
+                                        mi->def.method->roots);
             }
             else {
                 assert(mi->def.method->generator);
                 src = jl_code_for_staged(mi);
             }
         }
-        if (src && (jl_value_t*)src != jl_nothing) {
-            JL_GC_PUSH1(&src);
-            src = jl_uncompress_ast(mi->def.method, NULL, (jl_array_t*)src);
-            mi->uninferred = (jl_value_t*)src;
-            jl_gc_wb(mi, src);
-            JL_GC_POP();
+        if (!src || !jl_is_code_info(src)) {
+            jl_error("source missing for method called in interpreter");
         }
-    }
-    if (!src || !jl_is_code_info(src)) {
+        mi->uninferred = (jl_value_t*)src;
+        jl_gc_wb(mi, src);
+    } else if (!src || !jl_is_code_info(src)) {
         jl_error("source missing for method called in interpreter");
     }
     return src;

src/jltypes.c

@@ -2181,26 +2181,28 @@ void jl_init_types(void) JL_GC_DISABLED
     jl_code_instance_type =
         jl_new_datatype(jl_symbol("CodeInstance"), core,
                         jl_any_type, jl_emptysvec,
-                        jl_perm_symsvec(11,
+                        jl_perm_symsvec(12,
                             "def",
                             "next",
                             "min_world",
                             "max_world",
                             "rettype",
                             "rettype_const",
                             "inferred",
+                            "localroots",
                             //"edges",
                             //"absolute_max",
                             "invoke", "specptr",
                             "", ""), // function object decls
-                        jl_svec(11,
+                        jl_svec(12,
                             jl_method_instance_type,
                             jl_any_type,
                             jl_ulong_type,
                             jl_ulong_type,
                             jl_any_type,
                             jl_any_type,
                             jl_any_type,
+                            jl_any_type,
                             //jl_any_type,
                             //jl_bool_type,
                             jl_any_type, jl_any_type, // fptrs
@@ -2316,10 +2318,10 @@ void jl_init_types(void) JL_GC_DISABLED
     jl_svecset(jl_methtable_type->types, 10, jl_uint8_type);
     jl_svecset(jl_method_type->types, 11, jl_method_instance_type);
     jl_svecset(jl_method_instance_type->types, 5, jl_code_instance_type);
-    jl_svecset(jl_code_instance_type->types, 7, jl_voidpointer_type);
     jl_svecset(jl_code_instance_type->types, 8, jl_voidpointer_type);
     jl_svecset(jl_code_instance_type->types, 9, jl_voidpointer_type);
     jl_svecset(jl_code_instance_type->types, 10, jl_voidpointer_type);
+    jl_svecset(jl_code_instance_type->types, 11, jl_voidpointer_type);
 
     jl_compute_field_offsets(jl_datatype_type);
     jl_compute_field_offsets(jl_typename_type);

src/julia.h

@@ -341,6 +341,7 @@ typedef struct _jl_code_instance_t {
     jl_value_t *rettype; // return type for fptr
     jl_value_t *rettype_const; // inferred constant return value, or null
     jl_value_t *inferred; // inferred jl_code_info_t, or jl_nothing, or null
+    jl_array_t *localroots;   // roots for this specific MethodInstance
     //TODO: jl_array_t *edges; // stored information about edges from this object
     //TODO: uint8_t absolute_max; // whether true max world is unknown
 
@@ -1592,8 +1593,8 @@ JL_DLLEXPORT void jl_register_newmeth_tracer(void (*callback)(jl_method_t *trace
 // AST access
 JL_DLLEXPORT jl_value_t *jl_copy_ast(jl_value_t *expr JL_MAYBE_UNROOTED);
 
-JL_DLLEXPORT jl_array_t *jl_compress_ast(jl_method_t *m, jl_code_info_t *code);
-JL_DLLEXPORT jl_code_info_t *jl_uncompress_ast(jl_method_t *m, jl_code_instance_t *metadata, jl_array_t *data);
+JL_DLLEXPORT jl_svec_t *jl_compress_ast(jl_method_t *m, jl_code_info_t *code);
+JL_DLLEXPORT jl_code_info_t *jl_uncompress_ast(jl_method_t *m, jl_code_instance_t *metadata, jl_array_t *data, jl_array_t *roots);
 JL_DLLEXPORT uint8_t jl_ast_flag_inferred(jl_array_t *data) JL_NOTSAFEPOINT;
 JL_DLLEXPORT uint8_t jl_ast_flag_inlineable(jl_array_t *data) JL_NOTSAFEPOINT;
 JL_DLLEXPORT uint8_t jl_ast_flag_pure(jl_array_t *data) JL_NOTSAFEPOINT;

src/method.c

@@ -556,10 +556,14 @@ static void jl_method_set_source(jl_method_t *m, jl_code_info_t *src)
     jl_gc_wb(src, copy);
     m->slot_syms = jl_compress_argnames(src->slotnames);
     jl_gc_wb(m, m->slot_syms);
-    if (gen_only)
+    if (gen_only) {
         m->source = NULL;
-    else
-        m->source = (jl_value_t*)jl_compress_ast(m, src);
+    }
+    else {
+        jl_svec_t *srcroots = jl_compress_ast(m, src);
+        m->source = jl_svecref(srcroots, 0);
+        m->roots = (jl_array_t*) jl_svecref(srcroots, 1);
+    }
     jl_gc_wb(m, m->source);
     JL_GC_POP();
 }