Evalute :total function in the proper world

base/compiler/abstractinterpretation.jl

@@ -627,13 +627,14 @@ function is_all_const_arg((; fargs, argtypes)::ArgInfo)
     return true
 end
 
-function pure_eval_call_const_args(@nospecialize(f), argtypes::Vector{Any})
+function pure_eval_call_const_args(interp::AbstractInterpreter,
+        @nospecialize(f), argtypes::Vector{Any})
     args = Any[ (a = widenconditional(argtypes[i]);
         isa(a, Const) ? a.val :
         isconstType(a) ? (a::DataType).parameters[1] :
                          (a::DataType).instance) for i in 2:length(argtypes) ]
     try
-        value = Core._apply_pure(f, args)
+        value = Core._call_in_world_nonpure(get_world_counter(interp), f, args...)
         return Const(value)
     catch e
         return nothing
@@ -660,9 +661,9 @@ function abstract_call_method_with_const_args(interp::AbstractInterpreter, resul
         return nothing
     end
     if f !== nothing && result.edge !== nothing && is_total_or_error(result.edge_effects) && is_all_const_arg(arginfo)
-        rt = pure_eval_call_const_args(f, arginfo.argtypes) # TODO: Needs to be called in the correct world age
+        rt = pure_eval_call_const_args(interp, f, arginfo.argtypes)
         add_backedge!(result.edge, sv)
-        rt === nothing && return Union{}, false # The evaulation threw. By idempontency, we're guaranteed this would have happened at runtime
+        rt === nothing && return Union{}, false # The evaulation threw. By :consistent-cy, we're guaranteed this would have happened at runtime
         return rt, ConstResult(result.edge, rt.val)
     end
     mi = maybe_get_const_prop_profitable(interp, result, f, arginfo, match, sv)

src/builtin_proto.h

@@ -29,6 +29,7 @@ DECLARE_BUILTIN(arrayref);
 DECLARE_BUILTIN(arrayset);
 DECLARE_BUILTIN(arraysize);
 DECLARE_BUILTIN(_call_in_world);
+DECLARE_BUILTIN(_call_in_world_nonpure);
 DECLARE_BUILTIN(_call_latest);
 DECLARE_BUILTIN(replacefield);
 DECLARE_BUILTIN(const_arrayref);

src/builtins.c

@@ -778,6 +778,32 @@ JL_CALLABLE(jl_f__call_in_world)
     return ret;
 }
 
+JL_CALLABLE(jl_f__call_in_world_nonpure)
+{
+    JL_NARGSV(_apply_in_world_nonpure, 2);
+    JL_TYPECHK(_apply_in_world, ulong, args[0]);
+    jl_task_t *ct = jl_current_task;
+    int last_in = ct->ptls->in_pure_callback;
+    jl_value_t *ret = NULL;
+    size_t last_age = ct->world_age;
+    JL_TRY {
+        ct->ptls->in_pure_callback = 1;
+        size_t world = jl_unbox_ulong(args[0]);
+        ct->world_age = jl_atomic_load_acquire(&jl_world_counter);
+        if (ct->world_age > world)
+            ct->world_age = world;
+        ret = jl_apply(&args[1], nargs - 1);
+        ct->world_age = last_age;
+        ct->ptls->in_pure_callback = last_in;
+    }
+    JL_CATCH {
+        ct->world_age = last_age;
+        ct->ptls->in_pure_callback = last_in;
+        jl_rethrow();
+    }
+    return ret;
+}
+
 // tuples ---------------------------------------------------------------------
 
 JL_CALLABLE(jl_f_tuple)
@@ -1826,6 +1852,7 @@ void jl_init_primitives(void) JL_GC_DISABLED
     add_builtin_func("_apply_pure", jl_f__apply_pure);
     add_builtin_func("_call_latest", jl_f__call_latest);
     add_builtin_func("_call_in_world", jl_f__call_in_world);
+    add_builtin_func("_call_in_world_nonpure", jl_f__call_in_world_nonpure);
     add_builtin_func("_typevar", jl_f__typevar);
     add_builtin_func("_structtype", jl_f__structtype);
     add_builtin_func("_abstracttype", jl_f__abstracttype);

src/codegen.cpp

@@ -8028,6 +8028,7 @@ extern "C" void jl_init_llvm(void)
           { jl_f__apply_pure_addr,        new JuliaFunction{XSTR(jl_f__apply_pure), get_func_sig, get_func_attrs} },
           { jl_f__call_latest_addr,       new JuliaFunction{XSTR(jl_f__call_latest), get_func_sig, get_func_attrs} },
           { jl_f__call_in_world_addr,     new JuliaFunction{XSTR(jl_f__call_in_world), get_func_sig, get_func_attrs} },
+          { jl_f__call_in_world_nonpure_addr, new JuliaFunction{XSTR(jl_f__call_in_world_nonpure), get_func_sig, get_func_attrs} },
           { jl_f_throw_addr,              new JuliaFunction{XSTR(jl_f_throw), get_func_sig, get_func_attrs} },
           { jl_f_tuple_addr,              jltuple_func },
           { jl_f_svec_addr,               new JuliaFunction{XSTR(jl_f_svec), get_func_sig, get_func_attrs} },

src/staticdata.c

@@ -243,7 +243,7 @@ static htable_t field_replace;
 static const jl_fptr_args_t id_to_fptrs[] = {
     &jl_f_throw, &jl_f_is, &jl_f_typeof, &jl_f_issubtype, &jl_f_isa,
     &jl_f_typeassert, &jl_f__apply_iterate, &jl_f__apply_pure,
-    &jl_f__call_latest, &jl_f__call_in_world, &jl_f_isdefined,
+    &jl_f__call_latest, &jl_f__call_in_world, &jl_f__call_in_world_nonpure, &jl_f_isdefined,
     &jl_f_tuple, &jl_f_svec, &jl_f_intrinsic_call, &jl_f_invoke_kwsorter,
     &jl_f_getfield, &jl_f_setfield, &jl_f_swapfield, &jl_f_modifyfield,
     &jl_f_replacefield, &jl_f_fieldtype, &jl_f_nfields,

test/compiler/inference.jl

@@ -4009,3 +4009,12 @@ end
         end |> only === Union{}
     end
 end
+
+# Test that purity modeling doesn't accidentally introduce new world age issues
+f_redefine_me(x) = x+1
+f_call_redefine() = f_redefine_me(0)
+f_mk_opaque() = @Base.Experimental.opaque ()->Base.inferencebarrier(f_call_redefine)()
+const op_capture_world = f_mk_opaque()
+f_redefine_me(x) = x+2
+@test op_capture_world() == 1
+@test f_mk_opaque()() == 2