Optimize registry access

Author: mezz

Common/src/main/java/mezz/jei/common/util/RegistryUtil.java

@@ -5,19 +5,45 @@
 import net.minecraft.core.Registry;
 import net.minecraft.core.RegistryAccess;
 import net.minecraft.resources.ResourceKey;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
 
 public class RegistryUtil {
+	private static final Map<ResourceKey<? extends Registry<?>>, Registry<?>> REGISTRY_CACHE = new HashMap<>();
+	private static @Nullable RegistryAccess REGISTRY_ACCESS;
+
 	public static <T> Registry<T> getRegistry(ResourceKey<? extends Registry<T>> key) {
+		Registry<?> registry = REGISTRY_CACHE.get(key);
+		if (registry == null) {
+			registry = getRegistryUncached(key);
+			REGISTRY_CACHE.put(key, registry);
+		}
+		@SuppressWarnings("unchecked")
+		Registry<T> castRegistry = (Registry<T>) registry;
+		return castRegistry;
+	}
+
+	private static Registry<?> getRegistryUncached(ResourceKey<? extends Registry<?>> key) {
 		RegistryAccess registryAccess = getRegistryAccess();
 		return registryAccess.registryOrThrow(key);
 	}
 
 	public static RegistryAccess getRegistryAccess() {
-		Minecraft minecraft = Minecraft.getInstance();
-		ClientLevel level = minecraft.level;
-		if (level == null) {
-			throw new IllegalStateException("Could not get registry, registry access is unavailable because the level is currently null");
+		if (REGISTRY_ACCESS == null) {
+			Minecraft minecraft = Minecraft.getInstance();
+			ClientLevel level = minecraft.level;
+			if (level == null) {
+				throw new IllegalStateException("Could not get registry, registry access is unavailable because the level is currently null");
+			}
+			REGISTRY_ACCESS = level.registryAccess();
 		}
-		return level.registryAccess();
+		return REGISTRY_ACCESS;
+	}
+
+	public static void setRegistryAccess(@Nullable RegistryAccess registryAccess) {
+		REGISTRY_ACCESS = registryAccess;
+		REGISTRY_CACHE.clear();
 	}
 }

Library/src/main/java/mezz/jei/library/startup/JeiStarter.java

@@ -16,6 +16,7 @@
 import mezz.jei.common.config.file.IConfigSchemaBuilder;
 import mezz.jei.common.platform.Services;
 import mezz.jei.common.util.ErrorUtil;
+import mezz.jei.common.util.RegistryUtil;
 import mezz.jei.core.util.LoggedTimer;
 import mezz.jei.library.color.ColorHelper;
 import mezz.jei.library.config.ColorNameConfig;
@@ -100,6 +101,7 @@ public void start() {
 			return;
 		}
 		RegistryAccess registryAccess = minecraft.level.registryAccess();
+		RegistryUtil.setRegistryAccess(registryAccess);
 
 		LoggedTimer totalTime = new LoggedTimer();
 		totalTime.start("Starting JEI");
@@ -178,5 +180,6 @@ public void stop() {
 		List<IModPlugin> plugins = data.plugins();
 		PluginCaller.callOnPlugins("Sending Runtime Unavailable", plugins, IModPlugin::onRuntimeUnavailable);
 		Internal.setRuntime(null);
+		RegistryUtil.setRegistryAccess(null);
 	}
 }