cache Type{wrapper} for faster lookups (JuliaLang#44704)

mostly fixes JuliaLang#44402

Co-authored-by: Jameson Nash <[email protected]>

src/datatype.c

@@ -69,6 +69,7 @@ JL_DLLEXPORT jl_typename_t *jl_new_typename_in(jl_sym_t *name, jl_module_t *modu
  tn->name = name;
  tn->module = module;
  tn->wrapper = NULL;
+ jl_atomic_store_release(&tn->Typeofwrapper, NULL);
  jl_atomic_store_relaxed(&tn->cache, jl_emptysvec);
  jl_atomic_store_relaxed(&tn->linearcache, jl_emptysvec);
  tn->names = NULL;

src/dump.c

@@ -2010,6 +2010,7 @@ static jl_value_t *jl_deserialize_value_any(jl_serializer_state *s, uint8_t tag,
  jl_gc_wb(tn, tn->names);
  tn->wrapper = jl_deserialize_value(s, &tn->wrapper);
  jl_gc_wb(tn, tn->wrapper);
+ tn->Typeofwrapper = NULL;
  tn->mt = (jl_methtable_t*)jl_deserialize_value(s, (jl_value_t**)&tn->mt);
  jl_gc_wb(tn, tn->mt);
  ios_read(s->s, (char*)&tn->hash, sizeof(tn->hash));

src/jltypes.c

@@ -715,6 +715,12 @@ static ssize_t lookup_type_idx_linearvalue(jl_svec_t *cache, jl_value_t *key1, j
 static jl_value_t *lookup_type(jl_typename_t *tn JL_PROPAGATES_ROOT, jl_value_t **key, size_t n)
 {
  JL_TIMING(TYPE_CACHE_LOOKUP);
+ if (tn == jl_type_typename) {
+ assert(n == 1);
+ jl_value_t *uw = jl_unwrap_unionall(key[0]);
+ if (jl_is_datatype(uw) && key[0] == ((jl_datatype_t*)uw)->name->wrapper)
+ return jl_atomic_load_acquire(&((jl_datatype_t*)uw)->name->Typeofwrapper);
+ }
  unsigned hv = typekey_hash(tn, key, n, 0);
  if (hv) {
  jl_svec_t *cache = jl_atomic_load_relaxed(&tn->cache);
@@ -856,6 +862,16 @@ void jl_cache_type_(jl_datatype_t *type)
  assert(is_cacheable(type));
  jl_value_t **key = jl_svec_data(type->parameters);
  int n = jl_svec_len(type->parameters);
+ if (type->name == jl_type_typename) {
+ assert(n == 1);
+ jl_value_t *uw = jl_unwrap_unionall(key[0]);
+ if (jl_is_datatype(uw) && key[0] == ((jl_datatype_t*)uw)->name->wrapper) {
+ jl_typename_t *tn2 = ((jl_datatype_t*)uw)->name;
+ jl_atomic_store_release(&tn2->Typeofwrapper, (jl_value_t*)type);
+ jl_gc_wb(tn2, type);
+ return;
+ }
+ }
  unsigned hv = typekey_hash(type->name, key, n, 0);
  if (hv) {
  assert(hv == type->hash);
@@ -1371,7 +1387,7 @@ static jl_value_t *inst_datatype_inner(jl_datatype_t *dt, jl_svec_t *p, jl_value
  jl_typename_t *tn = dt->name;
  int istuple = (tn == jl_tuple_typename);
  int isnamedtuple = (tn == jl_namedtuple_typename);
- if (dt->name != jl_type_typename) {
+ if (tn != jl_type_typename) {
  size_t i;
  for (i = 0; i < ntp; i++)
  iparams[i] = normalize_unionalls(iparams[i]);
@@ -1415,6 +1431,11 @@ static jl_value_t *inst_datatype_inner(jl_datatype_t *dt, jl_svec_t *p, jl_value
  if (p) jl_gc_wb(p, tw);
  }
  }
+ if (tn == jl_type_typename && jl_is_datatype(iparams[0]) && ((jl_datatype_t*)iparams[0])->name == jl_type_typename &&
+ jl_tparam0(iparams[0]) == jl_bottom_type) {
+ // normalize Type{Type{Union{}}} to Type{TypeofBottom}
+ iparams[0] = (jl_value_t*)jl_typeofbottom_type;
+ }
  jl_value_t *lkup = (jl_value_t*)lookup_type(tn, iparams, ntp);
  if (lkup != NULL)
  return lkup;
@@ -2025,22 +2046,22 @@ void jl_init_types(void) JL_GC_DISABLED
  jl_typename_type->name->mt = jl_nonfunction_mt;
  jl_typename_type->super = jl_any_type;
  jl_typename_type->parameters = jl_emptysvec;
- jl_typename_type->name->n_uninitialized = 14 - 2;
- jl_typename_type->name->names = jl_perm_symsvec(14, "name", "module",
+ jl_typename_type->name->n_uninitialized = 15 - 2;
+ jl_typename_type->name->names = jl_perm_symsvec(15, "name", "module",
  "names", "atomicfields", "constfields",
- "wrapper", "cache", "linearcache",
+ "wrapper", "Typeofwrapper", "cache", "linearcache",
  "mt", "partial",
  "hash", "n_uninitialized",
  "flags", // "abstract", "mutable", "mayinlinealloc",
  "max_methods");
- jl_typename_type->types = jl_svec(14, jl_symbol_type, jl_any_type /*jl_module_type*/,
+ jl_typename_type->types = jl_svec(15, jl_symbol_type, jl_any_type /*jl_module_type*/,
  jl_simplevector_type, jl_any_type/*jl_voidpointer_type*/, jl_any_type/*jl_voidpointer_type*/,
- jl_type_type, jl_simplevector_type, jl_simplevector_type,
+ jl_type_type, jl_type_type, jl_simplevector_type, jl_simplevector_type,
  jl_methtable_type, jl_any_type,
  jl_any_type /*jl_long_type*/, jl_any_type /*jl_int32_type*/,
  jl_any_type /*jl_uint8_type*/,
  jl_any_type /*jl_uint8_type*/);
- const static uint32_t typename_constfields[1] = { 0x00001d3f }; // (1<<0)|(1<<1)|(1<<2)|(1<<3)|(1<<4)|(1<<5)|(1<<8)|(1<<10)|(1<<11)|(1<<12)
+ const static uint32_t typename_constfields[1] = { 0x00003a3f }; // (1<<0)|(1<<1)|(1<<2)|(1<<3)|(1<<4)|(1<<5)|(1<<9)|(1<<11)|(1<<12)|(1<<13)
  jl_typename_type->name->constfields = typename_constfields;
  jl_precompute_memoized_dt(jl_typename_type, 1);
 
@@ -2676,10 +2697,11 @@ void jl_init_types(void) JL_GC_DISABLED
  jl_svecset(jl_typename_type->types, 3, jl_voidpointer_type);
  jl_svecset(jl_typename_type->types, 4, jl_voidpointer_type);
  jl_svecset(jl_typename_type->types, 5, jl_type_type);
- jl_svecset(jl_typename_type->types, 10, jl_long_type);
- jl_svecset(jl_typename_type->types, 11, jl_int32_type);
- jl_svecset(jl_typename_type->types, 12, jl_uint8_type);
+ jl_svecset(jl_typename_type->types, 6, jl_type_type);
+ jl_svecset(jl_typename_type->types, 11, jl_long_type);
+ jl_svecset(jl_typename_type->types, 12, jl_int32_type);
  jl_svecset(jl_typename_type->types, 13, jl_uint8_type);
+ jl_svecset(jl_typename_type->types, 14, jl_uint8_type);
  jl_svecset(jl_methtable_type->types, 4, jl_long_type);
  jl_svecset(jl_methtable_type->types, 6, jl_module_type);
  jl_svecset(jl_methtable_type->types, 7, jl_array_any_type);

src/julia.h

@@ -464,6 +464,7 @@ typedef struct {
  // `wrapper` is either the only instantiation of the type (if no parameters)
  // or a UnionAll accepting parameters to make an instantiation.
  jl_value_t *wrapper;
+ _Atomic(jl_value_t*) Typeofwrapper; // cache for Type{wrapper}
  _Atomic(jl_svec_t*) cache; // sorted array
  _Atomic(jl_svec_t*) linearcache; // unsorted array
  struct _jl_methtable_t *mt;

src/staticdata.c

@@ -568,6 +568,7 @@ static void jl_serialize_value__(jl_serializer_state *s, jl_value_t *v, int recu
  jl_serialize_value(s, tn->module);
  jl_serialize_value(s, tn->names);
  jl_serialize_value(s, tn->wrapper);
+ jl_serialize_value(s, tn->Typeofwrapper);
  jl_serialize_value_(s, (jl_value_t*)tn->cache, 0);
  jl_serialize_value_(s, (jl_value_t*)tn->linearcache, 0);
  jl_serialize_value(s, tn->mt);