Cursed fixes for TagContext, adds support for hasTag() on fluids duri…

…ng recipes

common/src/main/java/dev/latvian/mods/kubejs/core/mixin/common/IngredientTagValueMixin.java

@@ -23,7 +23,7 @@ public abstract class IngredientTagValueMixin {
 	@Inject(method = "getItems", at = @At("HEAD"), cancellable = true)
 	private void kjs$getItems(CallbackInfoReturnable<Collection<ItemStack>> info) {
 		if (RecipesEventJS.instance != null) {
-			info.setReturnValue(TagContext.INSTANCE.getValue().patchTags(tag));
+			info.setReturnValue(TagContext.INSTANCE.getValue().patchIngredientTags(tag));
 		}
 	}
 }

common/src/main/java/dev/latvian/mods/kubejs/core/mixin/common/ReloadableServerResourcesMixin.java

@@ -19,6 +19,6 @@ private void init(RegistryAccess.Frozen frozen, Commands.CommandSelection comman
 
 	@Inject(method = "updateRegistryTags(Lnet/minecraft/core/RegistryAccess;)V", at = @At("RETURN"))
 	public void updateRegistryTags(RegistryAccess registryAccess, CallbackInfo ci) {
-		TagContext.INSTANCE.setValue(TagContext.REGISTRY);
+		TagContext.INSTANCE.setValue(TagContext.usingRegistry(registryAccess));
 	}
 }

common/src/main/java/dev/latvian/mods/kubejs/fluid/FluidStackJS.java

@@ -4,6 +4,7 @@
 import com.google.gson.JsonObject;
 import dev.architectury.fluid.FluidStack;
 import dev.architectury.registry.registries.Registries;
+import dev.latvian.mods.kubejs.item.ingredient.TagContext;
 import dev.latvian.mods.kubejs.recipe.RecipeExceptionJS;
 import dev.latvian.mods.kubejs.registry.KubeJSRegistries;
 import dev.latvian.mods.kubejs.util.MapJS;
@@ -115,7 +116,7 @@ public Collection<ResourceLocation> getTags() {
 	}
 
 	public boolean hasTag(ResourceLocation tag) {
-		return getFluid().is(Tags.fluid(tag));
+		return TagContext.INSTANCE.getValue().contains(Tags.fluid(tag), getFluid());
 	}
 
 	public Fluid getFluid() {

common/src/main/java/dev/latvian/mods/kubejs/item/ingredient/TagContext.java

@@ -1,84 +1,122 @@
 package dev.latvian.mods.kubejs.item.ingredient;
 
 import com.google.common.collect.Iterables;
+import dev.architectury.extensions.injected.InjectedRegistryEntryExtension;
 import dev.latvian.mods.kubejs.KubeJS;
+import dev.latvian.mods.kubejs.util.ConsoleJS;
+import dev.latvian.mods.kubejs.util.UtilsJS;
 import net.minecraft.core.Holder;
 import net.minecraft.core.Registry;
+import net.minecraft.core.RegistryAccess;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.resources.ResourceLocation;
 import net.minecraft.tags.TagKey;
 import net.minecraft.tags.TagManager;
 import net.minecraft.world.item.Item;
 import net.minecraft.world.item.ItemStack;
 import org.apache.commons.lang3.mutable.MutableObject;
+import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 public interface TagContext {
 	TagContext EMPTY = new TagContext() {
 		@Override
-		public boolean isEmpty(TagKey<Item> tag) {
+		public <T> boolean isEmpty(TagKey<T> tag) {
 			return true;
 		}
 
 		@Override
-		public boolean areTagsBound() {
-			return false;
-		}
-
-		@Override
-		public Iterable<Holder<Item>> getTag(TagKey<Item> tag) {
+		public <T> Iterable<Holder<T>> getTag(TagKey<T> tag) {
 			KubeJS.LOGGER.warn("Tried to get tag {} from an empty tag context!", tag.location());
 			return List.of();
 		}
 	};
 
-	TagContext REGISTRY = new TagContext() {
-		@Override
-		public boolean isEmpty(TagKey<Item> tag) {
-			return Registry.ITEM.getTag(tag).isEmpty();
-		}
-
-		@Override
-		public boolean areTagsBound() {
-			return true;
-		}
-
-		@Override
-		public Iterable<Holder<Item>> getTag(TagKey<Item> tag) {
-			return Registry.ITEM.getTagOrEmpty(tag);
-		}
-	};
+	static TagContext usingRegistry(RegistryAccess registryAccess) {
+		return new TagContext() {
+			@NotNull
+			private <T> Registry<T> registry(TagKey<T> tag) {
+				return registryAccess.registryOrThrow(tag.registry());
+			}
 
-	MutableObject<TagContext> INSTANCE = new MutableObject<>(TagContext.EMPTY);
+			@Override
+			public boolean areTagsBound() {
+				return true;
+			}
 
-	static TagContext usingResult(TagManager.LoadResult<Item> manager) {
-		return new TagContext() {
 			@Override
-			public boolean isEmpty(TagKey<Item> tag) {
-				return Iterables.isEmpty(getTag(tag));
+			public <T> Iterable<Holder<T>> getTag(TagKey<T> tag) {
+				return registry(tag).getTagOrEmpty(tag);
 			}
 
 			@Override
-			public boolean areTagsBound() {
-				return false;
+			public <T> boolean contains(TagKey<T> tag, T value) {
+				if (value instanceof InjectedRegistryEntryExtension<?> ext) {
+					Holder<T> holder = UtilsJS.cast(ext.arch$holder());
+					return holder.is(tag);
+				}
+
+				// cursed reverse holder lookup using the registry, and fallback to super
+				var reg = registry(tag);
+				return reg.getResourceKey(value).flatMap(reg::getHolder).map(holder -> holder.is(tag)).orElseGet(() -> TagContext.super.contains(tag, value));
 			}
+		};
+	}
+
+	MutableObject<TagContext> INSTANCE = new MutableObject<>(TagContext.EMPTY);
 
+	static TagContext fromLoadResult(List<TagManager.LoadResult<?>> results) {
+		final Map<ResourceKey<? extends Registry<?>>, Map<ResourceLocation, Collection<Holder<?>>>> tags = results.stream()
+				.collect(Collectors.toMap(result -> UtilsJS.cast(result.key()), result -> UtilsJS.cast(result.tags())));
+
+		if (!tags.containsKey(Registry.ITEM_REGISTRY)) {
+			ConsoleJS.getCurrent(ConsoleJS.SERVER).warn("Failed to load item tags during recipe event! Using replaceInput etc. will not work!");
+			return TagContext.EMPTY;
+		}
+
+		return new TagContext() {
 			@Override
-			public Iterable<Holder<Item>> getTag(TagKey<Item> tag) {
-				return manager.tags().getOrDefault(tag.location(), Set.of());
+			public <T> Iterable<Holder<T>> getTag(TagKey<T> tag) {
+				return UtilsJS.cast(tags.get(tag.registry()).getOrDefault(tag.location(), Set.of()));
 			}
 		};
 	}
 
-	boolean isEmpty(TagKey<Item> tag);
+	default <T> boolean isEmpty(TagKey<T> tag) {
+		return Iterables.isEmpty(getTag(tag));
+	}
+
+	default <T> boolean contains(TagKey<T> tag, T value) {
+		if (isEmpty(tag)) {
+			return false;
+		}
+
+		for (var holder : getTag(tag)) {
+			// if a mod doesn't do proper equality checks
+			// on their content i'm gonna die i guess
+			if (holder.value().equals(value)) {
+				return true;
+			}
+		}
+
+		return false;
+	}
+
+	default boolean areTagsBound() {
+		return false;
+	}
 
-	boolean areTagsBound();
+	;
 
-	Iterable<Holder<Item>> getTag(TagKey<Item> tag);
+	<T> Iterable<Holder<T>> getTag(TagKey<T> tag);
 
-	default Collection<ItemStack> patchTags(TagKey<Item> tag) {
+	default Collection<ItemStack> patchIngredientTags(TagKey<Item> tag) {
 		var c = getTag(tag);
 
 		var stacks = new ArrayList<ItemStack>(c instanceof Collection<?> cl ? cl.size() : 3);

common/src/main/java/dev/latvian/mods/kubejs/recipe/RecipesEventJS.java

@@ -29,7 +29,6 @@
 import dev.latvian.mods.kubejs.util.UtilsJS;
 import dev.latvian.mods.rhino.mod.util.JsonUtils;
 import dev.latvian.mods.rhino.util.HideFromJS;
-import net.minecraft.core.Registry;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.util.GsonHelper;
 import net.minecraft.world.item.crafting.Recipe;
@@ -147,15 +146,7 @@ public RecipesEventJS() {
 	public void post(RecipeManager recipeManager, Map<ResourceLocation, JsonObject> datapackRecipeMap) {
 		RecipeJS.itemErrors = false;
 
-		TagContext.INSTANCE.setValue(KubeJSReloadListener.resources.tagManager.getResult()
-				.stream()
-				.filter(result -> result.key() == Registry.ITEM_REGISTRY)
-				.findFirst()
-				.map(result -> TagContext.usingResult(UtilsJS.cast(result)))
-				.orElseGet(() -> {
-					ConsoleJS.SERVER.warn("Failed to load item tags during recipe event! Using replaceInput etc. will not work!");
-					return TagContext.EMPTY;
-				}));
+		TagContext.INSTANCE.setValue(TagContext.fromLoadResult(KubeJSReloadListener.resources.tagManager.getResult()));
 
 		var timer = Stopwatch.createStarted();
 

common/src/main/java/dev/latvian/mods/kubejs/util/Tags.java

@@ -1,5 +1,6 @@
 package dev.latvian.mods.kubejs.util;
 
+import dev.latvian.mods.kubejs.item.ingredient.TagContext;
 import net.minecraft.core.Holder;
 import net.minecraft.core.Registry;
 import net.minecraft.resources.ResourceKey;
@@ -66,6 +67,7 @@ public static Stream<TagKey<EntityType<?>>> byEntityType(EntityType<?> entityTyp
 	}
 
 	public static <T> Stream<TagKey<T>> forType(T object, Registry<T> registry) {
+		warnIfUnbound();
 		return registry.getResourceKey(object)
 				.flatMap(registry::getHolder)
 				.stream()
@@ -77,6 +79,13 @@ private static <T> TagKey<T> generic(ResourceLocation id, ResourceKey<Registry<T
 	}
 
 	private static <T> Stream<TagKey<T>> forHolder(Holder.Reference<T> registryHolder) {
+		warnIfUnbound();
 		return registryHolder.tags();
 	}
+
+	private static void warnIfUnbound() {
+		if (!TagContext.INSTANCE.getValue().areTagsBound()) {
+			ConsoleJS.getCurrent(ConsoleJS.STARTUP).warn("Tags have not been bound to registry yet! The values returned by this method may be outdated!", new Throwable());
+		}
+	}
 }