Refactor

Don't recurse through tag hierarchy
FabricMC · Jun 23, 2023 · 7fd8410 · 7fd8410
1 parent b2682d4
commit 7fd8410
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 72 deletions.
diff --git a/...-client-tags-api-v1/src/client/java/net/fabricmc/fabric/api/tag/client/v1/ClientTags.java b/...-client-tags-api-v1/src/client/java/net/fabricmc/fabric/api/tag/client/v1/ClientTags.java
@@ -16,20 +16,17 @@
 
 package net.fabricmc.fabric.api.tag.client.v1;
 
-import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
 
-import net.minecraft.client.MinecraftClient;
-import net.minecraft.registry.Registries;
 import net.minecraft.registry.Registry;
 import net.minecraft.registry.RegistryKey;
 import net.minecraft.registry.entry.RegistryEntry;
 import net.minecraft.registry.tag.TagKey;
 import net.minecraft.util.Identifier;
 
+import net.fabricmc.fabric.impl.tag.client.ClientTagsImpl;
 import net.fabricmc.fabric.impl.tag.client.ClientTagsLoader;
 
 /**
@@ -45,8 +42,6 @@
  * even when connected to a vanilla server.
  */
 public final class ClientTags {
-	private static final Map<TagKey<?>, ClientTagsLoader.LoadedTag> LOCAL_TAG_HIERARCHY = new ConcurrentHashMap<>();
-
 	private ClientTags() {
 	}
 
@@ -57,7 +52,7 @@ private ClientTags() {
 	 * @return a set of {@code Identifier}s this tag contains
 	 */
 	public static Set<Identifier> getOrCreateLocalTag(TagKey<?> tagKey) {
-		return getOrCreatePartiallySyncedTag(tagKey).completeIds();
+		return ClientTagsImpl.getOrCreatePartiallySyncedTag(tagKey).completeIds();
 	}
 
 	/**
@@ -75,7 +70,7 @@ public static <T> boolean isInWithLocalFallback(TagKey<T> tagKey, T entry) {
 		Objects.requireNonNull(tagKey);
 		Objects.requireNonNull(entry);
 
-		return getRegistryEntry(tagKey, entry).map(re -> isInWithLocalFallback(tagKey, re)).orElse(false);
+		return ClientTagsImpl.getRegistryEntry(tagKey, entry).map(re -> isInWithLocalFallback(tagKey, re)).orElse(false);
 	}
 
 	/**
@@ -96,7 +91,7 @@ public static <T> boolean isInWithLocalFallback(TagKey<T> tagKey, RegistryEntry<
 		Objects.requireNonNull(registryEntry);
 
 		// Check if the tag exists in the dynamic registry first
-		Optional<? extends Registry<T>> maybeRegistry = getRegistry(tagKey);
+		Optional<? extends Registry<T>> maybeRegistry = ClientTagsImpl.getRegistry(tagKey);
 
 		if (maybeRegistry.isPresent()) {
 			// Check the synced tag exists and use that
@@ -110,20 +105,21 @@ public static <T> boolean isInWithLocalFallback(TagKey<T> tagKey, RegistryEntry<
 			return false;
 		}
 
-		// Recursively search the entries contained with the tag
-		ClientTagsLoader.LoadedTag wt = getOrCreatePartiallySyncedTag(tagKey);
-
-		if (wt.immediateChildIds().contains(registryEntry.getKey().get().getValue())) {
-			return true;
-		}
+		ClientTagsLoader.LoadedTag wt = ClientTagsImpl.getOrCreatePartiallySyncedTag(tagKey);
 
-		for (TagKey<?> key : wt.immediateChildTags()) {
-			if (isInWithLocalFallback((TagKey<T>) key, registryEntry)) {
-				return true;
+		// Check if child tags exist
+		for (TagKey<?> childTag : wt.completeChildTags()) {
+			if (maybeRegistry.isPresent()) {
+				// Check the synced tag exists and use that
+				if (maybeRegistry.get().getEntryList((TagKey<T>) childTag).isPresent()) {
+					if (registryEntry.isIn((TagKey<T>) childTag)) {
+						return true;
+					}
+				}
 			}
 		}
 
-		return false;
+		return wt.completeIds().contains(registryEntry.getKey().get().getValue());
 	}
 
 	/**
@@ -145,48 +141,4 @@ public static <T> boolean isInLocal(TagKey<T> tagKey, RegistryKey<T> registryKey
 
 		return false;
 	}
-
-	@SuppressWarnings("unchecked")
-	private static <T> Optional<? extends Registry<T>> getRegistry(TagKey<T> tagKey) {
-		Objects.requireNonNull(tagKey);
-
-		// Check if the tag represents a dynamic registry
-		if (MinecraftClient.getInstance() != null) {
-			if (MinecraftClient.getInstance().world != null) {
-				if (MinecraftClient.getInstance().world.getRegistryManager() != null) {
-					Optional<? extends Registry<T>> maybeRegistry = MinecraftClient.getInstance().world
-							.getRegistryManager().getOptional(tagKey.registry());
-					if (maybeRegistry.isPresent()) return maybeRegistry;
-				}
-			}
-		}
-
-		return (Optional<? extends Registry<T>>) Registries.REGISTRIES.getOrEmpty(tagKey.registry().getValue());
-	}
-
-	@SuppressWarnings("unchecked")
-	private static <T> Optional<RegistryEntry<T>> getRegistryEntry(TagKey<T> tagKey, T entry) {
-		Optional<? extends Registry<?>> maybeRegistry = getRegistry(tagKey);
-
-		if (maybeRegistry.isEmpty() || !tagKey.isOf(maybeRegistry.get().getKey())) {
-			return Optional.empty();
-		}
-
-		Registry<T> registry = (Registry<T>) maybeRegistry.get();
-
-		Optional<RegistryKey<T>> maybeKey = registry.getKey(entry);
-
-		return maybeKey.map(registry::entryOf);
-	}
-
-	private static ClientTagsLoader.LoadedTag getOrCreatePartiallySyncedTag(TagKey<?> tagKey) {
-		ClientTagsLoader.LoadedTag loadedTag = LOCAL_TAG_HIERARCHY.get(tagKey);
-
-		if (loadedTag == null) {
-			loadedTag = ClientTagsLoader.loadTag(tagKey);
-			LOCAL_TAG_HIERARCHY.put(tagKey, loadedTag);
-		}
-
-		return loadedTag;
-	}
 }
diff --git a/...lient-tags-api-v1/src/client/java/net/fabricmc/fabric/impl/tag/client/ClientTagsImpl.java b/...lient-tags-api-v1/src/client/java/net/fabricmc/fabric/impl/tag/client/ClientTagsImpl.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.tag.client;
+
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+
+import net.minecraft.client.MinecraftClient;
+import net.minecraft.registry.Registries;
+import net.minecraft.registry.Registry;
+import net.minecraft.registry.RegistryKey;
+import net.minecraft.registry.entry.RegistryEntry;
+import net.minecraft.registry.tag.TagKey;
+
+public class ClientTagsImpl {
+	private static final Map<TagKey<?>, ClientTagsLoader.LoadedTag> LOCAL_TAG_HIERARCHY = new ConcurrentHashMap<>();
+
+	@SuppressWarnings("unchecked")
+	public static <T> Optional<? extends Registry<T>> getRegistry(TagKey<T> tagKey) {
+		Objects.requireNonNull(tagKey);
+
+		// Check if the tag represents a dynamic registry
+		if (MinecraftClient.getInstance() != null) {
+			if (MinecraftClient.getInstance().world != null) {
+				if (MinecraftClient.getInstance().world.getRegistryManager() != null) {
+					Optional<? extends Registry<T>> maybeRegistry = MinecraftClient.getInstance().world
+							.getRegistryManager().getOptional(tagKey.registry());
+					if (maybeRegistry.isPresent()) return maybeRegistry;
+				}
+			}
+		}
+
+		return (Optional<? extends Registry<T>>) Registries.REGISTRIES.getOrEmpty(tagKey.registry().getValue());
+	}
+
+	@SuppressWarnings("unchecked")
+	public static <T> Optional<RegistryEntry<T>> getRegistryEntry(TagKey<T> tagKey, T entry) {
+		Optional<? extends Registry<?>> maybeRegistry = getRegistry(tagKey);
+
+		if (maybeRegistry.isEmpty() || !tagKey.isOf(maybeRegistry.get().getKey())) {
+			return Optional.empty();
+		}
+
+		Registry<T> registry = (Registry<T>) maybeRegistry.get();
+
+		Optional<RegistryKey<T>> maybeKey = registry.getKey(entry);
+
+		return maybeKey.map(registry::entryOf);
+	}
+
+	public static ClientTagsLoader.LoadedTag getOrCreatePartiallySyncedTag(TagKey<?> tagKey) {
+		ClientTagsLoader.LoadedTag loadedTag = LOCAL_TAG_HIERARCHY.get(tagKey);
+
+		if (loadedTag == null) {
+			loadedTag = ClientTagsLoader.loadTag(tagKey);
+			LOCAL_TAG_HIERARCHY.put(tagKey, loadedTag);
+		}
+
+		return loadedTag;
+	}
+}
diff --git a/...ent-tags-api-v1/src/client/java/net/fabricmc/fabric/impl/tag/client/ClientTagsLoader.java b/...ent-tags-api-v1/src/client/java/net/fabricmc/fabric/impl/tag/client/ClientTagsLoader.java
@@ -33,15 +33,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import net.minecraft.registry.Registry;
+import net.minecraft.registry.RegistryKey;
 import net.minecraft.registry.tag.TagEntry;
 import net.minecraft.registry.tag.TagFile;
 import net.minecraft.registry.tag.TagKey;
 import net.minecraft.registry.tag.TagManagerLoader;
 import net.minecraft.util.Identifier;
-import net.minecraft.registry.Registry;
-import net.minecraft.registry.RegistryKey;
 
-import net.fabricmc.fabric.api.tag.client.v1.ClientTags;
 import net.fabricmc.loader.api.FabricLoader;
 import net.fabricmc.loader.api.ModContainer;
 
@@ -73,7 +72,8 @@ public static LoadedTag loadTag(TagKey<?> tagKey) {
 			}
 		}
 
-		HashSet<Identifier> ids = new HashSet<>();
+		HashSet<Identifier> completeIds = new HashSet<>();
+		HashSet<TagKey<?>> completeTags = new HashSet<>();
 		HashSet<Identifier> immediateChildIds = new HashSet<>();
 		HashSet<TagKey<?>> immediateChildTags = new HashSet<>();
 
@@ -91,16 +91,23 @@ public Identifier direct(Identifier id) {
 				public Collection<Identifier> tag(Identifier id) {
 					TagKey<?> tag = TagKey.of(tagKey.registry(), id);
 					immediateChildTags.add(tag);
-					return ClientTags.getOrCreateLocalTag(tag);
+					LoadedTag localTag = ClientTagsImpl.getOrCreatePartiallySyncedTag(tag);
+					completeTags.addAll(localTag.completeChildTags);
+					completeTags.add(tag);
+					return localTag.completeIds;
 				}
-			}, ids::add);
+			}, completeIds::add);
 		}
 
-		return new LoadedTag(Collections.unmodifiableSet(ids), Collections.unmodifiableSet(immediateChildTags),
-				Collections.unmodifiableSet(immediateChildIds));
+		// Ensure that the tag does not refer to itself
+		completeTags.remove(tagKey);
+		immediateChildTags.remove(tagKey);
+
+		return new LoadedTag(Collections.unmodifiableSet(completeIds), Collections.unmodifiableSet(completeTags),
+				Collections.unmodifiableSet(immediateChildTags), Collections.unmodifiableSet(immediateChildIds));
 	}
 
-	public record LoadedTag(Set<Identifier> completeIds, Set<TagKey<?>> immediateChildTags, Set<Identifier> immediateChildIds) {
+	public record LoadedTag(Set<Identifier> completeIds, Set<TagKey<?>> completeChildTags, Set<TagKey<?>> immediateChildTags, Set<Identifier> immediateChildIds) {
 	}
 
 	/**