From 9510702aa3a3a8263b9d3c9d437630a5549ba951 Mon Sep 17 00:00:00 2001 From: pietro-lopes <97140255+pietro-lopes@users.noreply.github.com> Date: Fri, 16 Aug 2024 00:37:39 -0300 Subject: [PATCH] added modded registry scanner --- .../kubejs/core/mixin/ResourceKeyMixin.java | 24 ++++ .../mods/kubejs/registry/RegistryType.java | 128 +++++++++++++++--- .../dev/latvian/mods/kubejs/util/UtilsJS.java | 11 ++ 3 files changed, 141 insertions(+), 22 deletions(-) diff --git a/src/main/java/dev/latvian/mods/kubejs/core/mixin/ResourceKeyMixin.java b/src/main/java/dev/latvian/mods/kubejs/core/mixin/ResourceKeyMixin.java index f3d728d72..1eb7927d6 100644 --- a/src/main/java/dev/latvian/mods/kubejs/core/mixin/ResourceKeyMixin.java +++ b/src/main/java/dev/latvian/mods/kubejs/core/mixin/ResourceKeyMixin.java @@ -1,14 +1,22 @@ package dev.latvian.mods.kubejs.core.mixin; +import dev.latvian.mods.kubejs.KubeJS; import dev.latvian.mods.rhino.Context; import dev.latvian.mods.rhino.util.RemapPrefixForJS; import dev.latvian.mods.rhino.util.SpecialEquality; +import net.minecraft.Util; +import net.minecraft.core.registries.Registries; import net.minecraft.resources.ResourceKey; import net.minecraft.resources.ResourceLocation; import org.spongepowered.asm.mixin.Final; import org.spongepowered.asm.mixin.Mixin; import org.spongepowered.asm.mixin.Shadow; import org.spongepowered.asm.mixin.Unique; +import org.spongepowered.asm.mixin.injection.At; +import org.spongepowered.asm.mixin.injection.Inject; +import org.spongepowered.asm.mixin.injection.callback.CallbackInfo; + +import static dev.latvian.mods.kubejs.registry.RegistryType.Scanner; @RemapPrefixForJS("kjs$") @Mixin(value = ResourceKey.class, priority = 1001) @@ -39,4 +47,20 @@ public boolean specialEquals(Context cx, Object o, boolean shallow) { return location.toString().equals(String.valueOf(o)); } } + + @Inject(method = "", at = @At(value = "RETURN")) + private void kjs$getKeyStackTraces(ResourceLocation registryName, ResourceLocation location, CallbackInfo ci){ + if (Scanner.isFrozen()) return; + if (!registryName.equals(Registries.ROOT_REGISTRY_NAME)) return; + if (Scanner.shouldSkipNamespace(location.getNamespace())) return; + var startTime = Util.getNanos(); + var stack = Thread.currentThread().getStackTrace(); + for (StackTraceElement stackTraceElement : stack) { + if (Scanner.shouldSkipModule(stackTraceElement.getModuleName())) continue; + var className = stackTraceElement.getClassName(); + if (Scanner.contains(className)) continue; + Scanner.add(className); + } + KubeJS.LOGGER.debug("Took {} ms to grab stacktrace classes.", (int)((Util.getNanos() - startTime)/1_000_000)); + } } diff --git a/src/main/java/dev/latvian/mods/kubejs/registry/RegistryType.java b/src/main/java/dev/latvian/mods/kubejs/registry/RegistryType.java index 62808ed8c..208dc469a 100644 --- a/src/main/java/dev/latvian/mods/kubejs/registry/RegistryType.java +++ b/src/main/java/dev/latvian/mods/kubejs/registry/RegistryType.java @@ -2,19 +2,27 @@ import dev.latvian.mods.kubejs.DevProperties; import dev.latvian.mods.kubejs.KubeJS; +import dev.latvian.mods.kubejs.util.UtilsJS; import dev.latvian.mods.rhino.type.TypeInfo; import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap; +import net.minecraft.Util; import net.minecraft.core.Registry; import net.minecraft.core.registries.Registries; import net.minecraft.resources.ResourceKey; + +import net.neoforged.neoforge.common.util.Lazy; +import net.neoforged.neoforge.registries.NeoForgeRegistries; import org.jetbrains.annotations.Nullable; import java.lang.reflect.Modifier; import java.lang.reflect.ParameterizedType; import java.util.ArrayList; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.stream.Stream; public record RegistryType(ResourceKey> key, Class baseClass, TypeInfo type) { private static final Map, RegistryType> KEY_MAP = new Reference2ObjectOpenHashMap<>(); @@ -23,23 +31,7 @@ public record RegistryType(ResourceKey> key, Class baseClass, // This is cursed, but it's better than manually registering every type public static synchronized void init() { - try { - for (var field : Registries.class.getDeclaredFields()) { - if (field.getType() == ResourceKey.class - && Modifier.isPublic(field.getModifiers()) - && Modifier.isStatic(field.getModifiers()) - && field.getGenericType() instanceof ParameterizedType t1 - && t1.getActualTypeArguments()[0] instanceof ParameterizedType t2 - ) { - var key = (ResourceKey) field.get(null); - var type = t2.getActualTypeArguments()[0]; - var typeInfo = TypeInfo.of(type); - register(key, typeInfo); - } - } - } catch (Exception ex) { - ex.printStackTrace(); - } + Scanner.processClass(Stream.of(Registries.class, NeoForgeRegistries.Keys.class)); } public static synchronized void register(ResourceKey> key, TypeInfo type) { @@ -55,22 +47,32 @@ public static synchronized void register(ResourceKey> key, TypeI @Nullable public static synchronized RegistryType ofKey(ResourceKey key) { - return KEY_MAP.get(key); + return (RegistryType) of(key); } @Nullable public static synchronized RegistryType ofType(TypeInfo typeInfo) { - return TYPE_MAP.get(typeInfo); + return (RegistryType) of(typeInfo); } @Nullable public static synchronized RegistryType ofClass(Class type) { - var list = CLASS_MAP.get(type); - return list != null && list.size() == 1 ? list.getFirst() : null; + var regList = ((List>) of(type)); + return regList != null && regList.size() == 1 ? regList.getFirst() : null; } public static synchronized List> allOfClass(Class type) { - return CLASS_MAP.getOrDefault(type, List.of()); + return (List>) of(type); + } + + private static synchronized Object of(Object obj){ + Scanner.startIfNotFrozen(); + return switch (obj) { + case ResourceKey key -> KEY_MAP.get(key); + case Class clazz -> CLASS_MAP.getOrDefault(clazz, List.of()); + case TypeInfo info -> TYPE_MAP.get(info); + default -> List.of(); + }; } @Nullable @@ -94,4 +96,86 @@ public static synchronized RegistryType lookup(TypeInfo target) { public String toString() { return key.location() + "=" + type; } + + public static class Scanner { + private static final Lazy>> VALID_TYPES = Lazy.of(() -> { + Set> set = new HashSet<>(); + set.add(ResourceKey.class); + set.add(Registry.class); + var registrar = UtilsJS.tryLoadClass("dev.architectury.registry.registries.Registrar"); + if (registrar != null) set.add(registrar); + return set; + }); + private static final Set CLASSES_TO_SCAN = new HashSet<>(); + private static final Set MODULES_TO_SKIP = Set.of("java.base","neoforge","fml_loader","kubejs"); + private static final Set NAMESPACES_TO_SKIP = Set.of("neoforge", "minecraft"); + + private static boolean frozen = false; + + private static synchronized void startIfNotFrozen(){ + if (isFrozen()) return; + frozen = true; + var startTime = Util.getNanos(); + Stream> classStream = CLASSES_TO_SCAN.stream().map(UtilsJS::tryLoadClass); + processClass(classStream); + CLASSES_TO_SCAN.clear(); + KubeJS.LOGGER.debug("Took {} ms to discover registry classes.", (int)((Util.getNanos() - startTime)/1_000_000)); + } + + public static synchronized boolean isFrozen(){ + return frozen; + } + + public static synchronized boolean shouldSkipModule(String moduleName){ + return MODULES_TO_SKIP.contains(moduleName); + } + + public static synchronized boolean shouldSkipNamespace(String namespace){ + return NAMESPACES_TO_SKIP.contains(namespace); + } + + public static synchronized void add(String className){ + CLASSES_TO_SCAN.add(className); + } + + public static synchronized boolean contains(String className){ + return CLASSES_TO_SCAN.contains(className); + } + + private static synchronized void processClass(Stream> classStream){ + classStream.map(Class::getDeclaredFields) + .flatMap(Stream::of) + .forEach(field -> { + try { + if (!VALID_TYPES.get().contains(field.getType())) return; + if (!Modifier.isPublic(field.getModifiers())) field.setAccessible(true); + var value = field.get(null); + if (value instanceof ResourceKey key) { + if (field.getGenericType() instanceof ParameterizedType t1 + && t1.getActualTypeArguments()[0] instanceof ParameterizedType t2) { + processKey(key, t2, false); + } + } else if (value instanceof Registry registry) { + if (field.getGenericType() instanceof ParameterizedType t1){ + processKey(registry.key(), t1, true); + } + } else if (field.getType().getName().equals("dev.architectury.registry.registries.Registrar")){ + if (field.getGenericType() instanceof ParameterizedType t1){ + var method = value.getClass().getDeclaredMethod("key"); + processKey((ResourceKey) method.invoke(value), t1, true); + } + } + } catch (Exception ex) { + KubeJS.LOGGER.error("Error while trying to get registry from field " + field.getName() + " from class " + field.getType().getName(), ex); + } + }); + } + + private static synchronized void processKey(ResourceKey key, ParameterizedType paramType, boolean checkIfContains){ + if (checkIfContains && RegistryType.ofKey(key) != null) return; + var type = paramType.getActualTypeArguments()[0]; + var typeInfo = TypeInfo.of(type); + register(key, typeInfo); + } + } } diff --git a/src/main/java/dev/latvian/mods/kubejs/util/UtilsJS.java b/src/main/java/dev/latvian/mods/kubejs/util/UtilsJS.java index 271e70963..8ea4c2178 100644 --- a/src/main/java/dev/latvian/mods/kubejs/util/UtilsJS.java +++ b/src/main/java/dev/latvian/mods/kubejs/util/UtilsJS.java @@ -15,6 +15,7 @@ import dev.latvian.mods.rhino.Wrapper; import dev.latvian.mods.rhino.type.TypeInfo; import dev.latvian.mods.rhino.type.TypeUtils; +import dev.latvian.mods.rhino.util.HideFromJS; import net.minecraft.advancements.critereon.MinMaxBounds; import net.minecraft.commands.arguments.selector.EntitySelector; import net.minecraft.commands.arguments.selector.EntitySelectorParser; @@ -407,4 +408,14 @@ public static CreativeModeTab findCreativeTab(ResourceLocation id) { public static T makeFunctionProxy(Context cx, TypeInfo targetClass, BaseFunction function) { return Cast.to(cx.createInterfaceAdapter(targetClass, function)); } + + @Nullable + @HideFromJS + public static Class tryLoadClass(String className){ + Class clazz = null; + try { + clazz = Class.forName(className); + } catch (Exception ignored) {} + return clazz; + } } \ No newline at end of file