Further fixes/changes

TestMod/src/main/java/dhyces/testmod/TrimmedTestClient.java

@@ -4,17 +4,22 @@
 import dhyces.trimmed.impl.client.maps.ClientMapManager;
 import dhyces.trimmed.impl.client.maps.ClientRegistryMapKey;
 import dhyces.trimmed.impl.client.maps.ClientMapKey;
+import net.minecraft.client.Minecraft;
+import net.minecraft.core.Holder;
 import net.minecraft.core.registries.Registries;
+import net.minecraft.resources.ResourceKey;
 import net.minecraft.resources.ResourceLocation;
 import net.minecraft.world.item.Item;
 import net.minecraft.world.item.Items;
+import net.minecraft.world.level.biome.Biome;
 import net.minecraftforge.event.entity.player.PlayerEvent;
 import net.minecraftforge.eventbus.api.IEventBus;
 
 public class TrimmedTestClient {
 
     public static final ClientMapKey TEST_MAP = ClientMapKey.of(TrimmedTest.id("test_map"));
     public static final ClientRegistryMapKey<Item> TEST_ITEM_MAP = ClientRegistryMapKey.of(Registries.ITEM, TrimmedTest.id("checked_item_map"));
+    public static final ClientRegistryMapKey<Biome> TEST_BIOME_MAP = ClientRegistryMapKey.of(Registries.BIOME, TrimmedTest.id("checked_biome_map"));
 
     static void init(IEventBus forgeBus, IEventBus modBus) {
         forgeBus.addListener(TrimmedTestClient::loggedIn);
@@ -23,5 +28,6 @@ static void init(IEventBus forgeBus, IEventBus modBus) {
     private static void loggedIn(PlayerEvent.PlayerLoggedInEvent event) {
         ClientMapManager.getUnchecked(TEST_MAP).ifPresent(map -> TrimmedTest.LOGGER.info("Map present! " + map.get(new ResourceLocation(Trimmed.MODID, "not/a/real/place"))));
         ClientMapManager.getChecked(TEST_ITEM_MAP).ifPresent(itemStringMap -> TrimmedTest.LOGGER.info("Map present! " + itemStringMap.get(Items.IRON_INGOT)));
+        ClientMapManager.getDatapacked(TEST_BIOME_MAP).ifPresent(biomeStringMap -> biomeStringMap.forEach((biomeReference, s) -> TrimmedTest.LOGGER.info(biomeReference.key().toString())));
     }
 }

TestMod/src/main/resources/assets/testmod/maps/worldgen/biome/checked_biome_map.json

@@ -0,0 +1,5 @@
+{
+  "pairs": {
+    "minecraft:desert": "67"
+  }
+}

src/main/java/dhyces/trimmed/Trimmed.java

@@ -1,6 +1,7 @@
 package dhyces.trimmed;
 
 import net.minecraft.resources.ResourceLocation;
+import net.minecraftforge.common.MinecraftForge;
 import net.minecraftforge.fml.common.Mod;
 import net.minecraftforge.fml.javafmlmod.FMLJavaModLoadingContext;
 import net.minecraftforge.fml.loading.FMLLoader;
@@ -18,7 +19,7 @@ public static ResourceLocation id(String id) {
 
     public Trimmed() {
         if (FMLLoader.getDist().isClient()) {
-            TrimmedClient.init(FMLJavaModLoadingContext.get().getModEventBus());
+            TrimmedClient.init(MinecraftForge.EVENT_BUS, FMLJavaModLoadingContext.get().getModEventBus());
         }
     }
 }

src/main/java/dhyces/trimmed/TrimmedClient.java

@@ -9,15 +9,18 @@
 import net.minecraft.client.Minecraft;
 import net.minecraftforge.client.event.ModelEvent;
 import net.minecraftforge.client.event.RegisterClientReloadListenersEvent;
+import net.minecraftforge.event.TagsUpdatedEvent;
 import net.minecraftforge.eventbus.api.IEventBus;
 
 public class TrimmedClient {
 
-    static void init(IEventBus modBus) {
-        ItemOverrideProviderRegistry.init();
+    static void init(IEventBus forgeBus, IEventBus modBus) {
         TrimmedSpriteSourceTypes.bootstrap();
+        ItemOverrideProviderRegistry.init();
         modBus.addListener(TrimmedClient::registerClientReloadListener);
         modBus.addListener(TrimmedClient::addModels);
+
+        forgeBus.addListener(TrimmedClient::tagsSynced);
     }
 
     private static void registerClientReloadListener(final RegisterClientReloadListenersEvent event) {
@@ -29,4 +32,11 @@ private static void registerClientReloadListener(final RegisterClientReloadListe
     private static void addModels(final ModelEvent.RegisterAdditional event) {
         ItemOverrideReloadListener.getModelsToBake().forEach(event::register);
     }
+
+    private static void tagsSynced(final TagsUpdatedEvent event) {
+        if (event.shouldUpdateStaticData()) {
+            ClientTagManager.updateDatapacksSynced(event.getRegistryAccess());
+            ClientMapManager.updateDatapacksSynced(event.getRegistryAccess());
+        }
+    }
 }

src/main/java/dhyces/trimmed/api/codec/SetCodec.java

@@ -4,14 +4,13 @@
 import com.mojang.datafixers.util.Pair;
 import com.mojang.datafixers.util.Unit;
 import com.mojang.serialization.*;
+import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
 
+import java.util.Collections;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
-/**
- * Do not use for serialization, does not preserve order and will cause the cache to update every time datagen is run
- */
 public class SetCodec<A> implements Codec<Set<A>> {
     private final Codec<A> elementCodec;
 
@@ -22,20 +21,20 @@ public SetCodec(Codec<A> codec) {
     @Override
     public <T> DataResult<Pair<Set<A>, T>> decode(DynamicOps<T> ops, T input) {
         return ops.getList(input).setLifecycle(Lifecycle.stable()).flatMap(consumerConsumer -> {
-            ImmutableSet.Builder<A> builder = new ImmutableSet.Builder<>();
+            Set<A> linkedSet = new ObjectLinkedOpenHashSet<>(); // Preserve order
             Stream.Builder<T> failed = Stream.builder();
             AtomicReference<DataResult<Unit>> ref = new AtomicReference<>(DataResult.success(Unit.INSTANCE, Lifecycle.stable()));
 
             consumerConsumer.accept(t -> {
                 DataResult<Pair<A, T>> result = elementCodec.decode(ops, t);
                 result.error().ifPresent(e -> failed.add(t));
                 ref.setPlain(ref.getPlain().apply2stable((unit, o) -> {
-                    builder.add(o.getFirst());
+                    linkedSet.add(o.getFirst());
                     return unit;
                 }, result));
             });
 
-            Pair<Set<A>, T> pair = Pair.of(builder.build(), ops.createList(failed.build()));
+            Pair<Set<A>, T> pair = Pair.of(Collections.unmodifiableSet(linkedSet), ops.createList(failed.build()));
             return ref.getPlain().map(unit -> pair).setPartial(pair);
         });
     }

src/main/java/dhyces/trimmed/api/util/ResourcePath.java

@@ -62,7 +62,7 @@ public String getDirectoryStringFrom(String beginningDirectory) {
         for (int i = 0; i < path.length; i++) {
             String pathElement = path[i];
             if (pathElement.equals(beginningDirectory)) {
-                return ofElements(i+1, !path[path.length - 1].contains(".") ? path.length-1 : path.length-2);
+                return ofElements(i+1, !path[path.length - 1].contains(".") ? path.length : path.length-1);
             }
         }
         return "";
@@ -98,7 +98,10 @@ private String ofElements(int startIndexInclusive, int endIndexExclusive) {
         }
         StringBuilder builder = new StringBuilder();
         for (int i = startIndexInclusive; i < endIndexExclusive; i++) {
-            builder.append(path[i]).append("/");
+            builder.append(path[i]);
+            if (i+1 != endIndexExclusive) {
+                builder.append("/");
+            }
         }
         return builder.toString();
     }

src/main/java/dhyces/trimmed/impl/client/maps/ClientMapManager.java

@@ -8,9 +8,12 @@
 import com.mojang.serialization.codecs.UnboundedMapCodec;
 import dhyces.trimmed.Trimmed;
 import dhyces.trimmed.api.util.ResourcePath;
-import dhyces.trimmed.api.util.Utils;
 import net.minecraft.client.Minecraft;
+import net.minecraft.core.Holder;
 import net.minecraft.core.Registry;
+import net.minecraft.core.RegistryAccess;
+import net.minecraft.core.registries.BuiltInRegistries;
+import net.minecraft.network.chat.Component;
 import net.minecraft.resources.FileToIdConverter;
 import net.minecraft.resources.ResourceKey;
 import net.minecraft.resources.ResourceLocation;
@@ -20,6 +23,7 @@
 import net.minecraft.util.GsonHelper;
 import net.minecraft.util.Unit;
 import net.minecraft.util.profiling.ProfilerFiller;
+import net.minecraftforge.fml.ModList;
 import net.minecraftforge.registries.ForgeRegistry;
 import net.minecraftforge.registries.RegistryManager;
 
@@ -32,11 +36,14 @@
 public class ClientMapManager implements PreparableReloadListener {
     private static final Map<ClientMapKey, Map<ResourceLocation, String>> UNCHECKED_MAPS = new HashMap<>();
     private static final Map<ClientRegistryMapKey<?>, Map<?, String>> CHECKED_MAPS = new HashMap<>();
-    private static final Map<ClientRegistryMapKey<?>, Map<?, String>> LAZY_CHECKED_MAPS = new HashMap<>();
-    private static final Map<ClientRegistryMapKey<?>, Map<ResourceLocation, String>> LAZY_TO_CHECK_MAPS = new HashMap<>();
+    private static final Map<ClientRegistryMapKey<?>, Map<Holder.Reference<?>, String>> LAZY_CHECKED_MAPS = new HashMap<>();
+    private static final Map<ClientRegistryMapKey<?>, Map<ResourceKey<?>, String>> LAZY_TO_CHECK_MAPS = new HashMap<>();
 
     public static final UnboundedMapCodec<ResourceLocation, String> UNCHECKED_CODEC = Codec.unboundedMap(ResourceLocation.CODEC, Codec.STRING);
     private static final Map<ForgeRegistry<?>, UnboundedMapCodec<?, String>> USED_MAP_CODECS = new HashMap<>();
+    private static final Map<ResourceLocation, UnboundedMapCodec<ResourceKey<?>, String>> USED_DATAPACK_MAP_CODECS = new HashMap<>();
+
+    private static boolean datapacksSynced;
 
     private static final FileToIdConverter FILE_TO_ID_CONVERTER = FileToIdConverter.json("maps");
 
@@ -48,25 +55,43 @@ public static Optional<Map<ResourceLocation, String>> getUnchecked(ClientMapKey
         return Optional.ofNullable(UNCHECKED_MAPS.get(clientMapKey));
     }
 
-    public static <T> Optional<Map<T, String>> getChecked(ClientRegistryMapKey<T> checkedClientRegistryMapKey) {
+    public static <T> Optional<Map<T, String>> getChecked(ClientRegistryMapKey<T> clientRegistryMapKey) {
         if (CHECKED_MAPS.isEmpty()) {
-            Trimmed.LOGGER.error("Maps have not been loaded or are empty! Tried to get: " + checkedClientRegistryMapKey);
+            Trimmed.LOGGER.error("Maps have not been loaded or are empty! Tried to get: " + clientRegistryMapKey);
             return Optional.empty();
         }
-        if (RegistryManager.ACTIVE.getRegistry(checkedClientRegistryMapKey.getRegistryKey()) == null) {
-            if (LAZY_CHECKED_MAPS.isEmpty()) {
-                Trimmed.LOGGER.error("Datapacks are not yet loaded! Tried to get: " + checkedClientRegistryMapKey);
-                return Optional.empty();
-            }
-            // TODO: Probably move this to the datapack sync method
-            Optional<Registry<T>> registryOptional = Minecraft.getInstance().level.registryAccess().registry(checkedClientRegistryMapKey.getRegistryKey());
-            if (registryOptional.isEmpty() || registryOptional.get().keySet().isEmpty()) {
-                Trimmed.LOGGER.error("Datapack does not exist or is not synced to client! Tried to get: " + checkedClientRegistryMapKey);
-                return Optional.empty();
+        return Optional.ofNullable((Map<T, String>) CHECKED_MAPS.get(clientRegistryMapKey));
+    }
+
+    public static <T> Optional<Map<Holder.Reference<T>, String>> getDatapacked(ClientRegistryMapKey<T> clientRegistryMapKey) {
+        if (!datapacksSynced) {
+            Trimmed.LOGGER.error("Datapacks are not yet loaded! Tried to get: " + clientRegistryMapKey);
+            return Optional.empty();
+        }
+        return Optional.ofNullable(cast(LAZY_CHECKED_MAPS.get(clientRegistryMapKey)));
+    }
+
+    public static void updateDatapacksSynced(RegistryAccess registryAccess) {
+        if (datapacksSynced) {
+            Trimmed.LOGGER.info("Datapacks have been updated! Client may need to reload...");
+            Minecraft.getInstance().player.displayClientMessage(Component.translatable("trimmed.info.datapacksReloaded"), true);
+            return;
+        }
+
+        for (Map.Entry<ClientRegistryMapKey<?>, Map<ResourceKey<?>, String>> entry : LAZY_TO_CHECK_MAPS.entrySet()) {
+            Optional<? extends Registry<?>> datapackRegistryOptional = registryAccess.registry(entry.getKey().getRegistryKey());
+            if (datapackRegistryOptional.isEmpty()) {
+                Trimmed.LOGGER.error("Datapack registry " + entry.getKey().getRegistryKey().location() + " does not exist or is not synced to client!");
+            } else {
+                ImmutableMap.Builder<Object, String> mapBuilder = ImmutableMap.builder();
+                for (Map.Entry<ResourceKey<?>, String> entryTransform : entry.getValue().entrySet()) {
+                    datapackRegistryOptional.get().getHolder(cast(entryTransform.getKey())).ifPresent(o -> mapBuilder.put(o, entryTransform.getValue()));
+                }
+                LAZY_CHECKED_MAPS.put(entry.getKey(), cast(mapBuilder.build()));
             }
-            return Optional.ofNullable((Map<T, String>) LAZY_CHECKED_MAPS.get(checkedClientRegistryMapKey.getRegistryKey()));
         }
-        return Optional.ofNullable((Map<T, String>) CHECKED_MAPS.get(checkedClientRegistryMapKey));
+        LAZY_TO_CHECK_MAPS.clear();
+        datapacksSynced = true;
     }
 
     @Override
@@ -76,6 +101,7 @@ public CompletableFuture<Void> reload(PreparationBarrier pPreparationBarrier, Re
 
     @SuppressWarnings("UnstableApiUsage")
     private CompletableFuture<Unit> load(ResourceManager resourceManager, Executor backgroundExecutor) {
+        datapacksSynced = false;
         UNCHECKED_MAPS.clear();
         CHECKED_MAPS.clear();
         LAZY_CHECKED_MAPS.clear();
@@ -86,21 +112,27 @@ private CompletableFuture<Unit> load(ResourceManager resourceManager, Executor b
                 ClientMapKey mapKey = ClientMapKey.of(idPath.getFileNameOnly(5).asResourceLocation());
                 UNCHECKED_MAPS.put(mapKey, readMap);
             } else {
-                String[] registryDirectories = idPath.getDirectoryStringFrom("maps").split("/");
+                String registryDirectoryPath = idPath.getDirectoryStringFrom("maps");
+                String[] registryDirectories = registryDirectoryPath.split("/");
                 ResourceLocation registryId;
-                if (registryDirectories.length == 1) {
-                    registryId = new ResourceLocation(registryDirectories[0]);
+                final boolean isModded;
+                if (registryDirectories.length > 1 && ModList.get().isLoaded(registryDirectories[0])) {
+                    registryId = new ResourceLocation(registryDirectories[0], idPath.getDirectoryStringFrom(registryDirectories[0])); // TODO: test this
+                    isModded = true;
                 } else {
-                    registryId = new ResourceLocation(registryDirectories[0], registryDirectories[1]);
+                    registryId = new ResourceLocation(registryDirectoryPath);
+                    isModded = false;
                 }
-                if (RegistryManager.ACTIVE.getRegistry(registryId) != null) {
+
+                if (BuiltInRegistries.REGISTRY.get(registryId) != null || (isModded && RegistryManager.ACTIVE.getRegistry(registryId) != null)) {
                     ForgeRegistry<?> registry = RegistryManager.ACTIVE.getRegistry(registryId);
                     UnboundedMapCodec<?, String> mapCodec = USED_MAP_CODECS.computeIfAbsent(registry, reg -> Codec.unboundedMap(reg.getCodec(), Codec.STRING));
                     Map<?, String> readMap = readResources(entry.getValue(), mapCodec);
                     ClientRegistryMapKey<?> clientRegistryMapKey = ClientRegistryMapKey.of(registry.getRegistryKey(), idPath.getFileNameOnly(5).asResourceLocation());
                     CHECKED_MAPS.put(clientRegistryMapKey, readMap);
                 } else {
-                    Map<ResourceLocation, String> readMap = readResources(entry.getValue(), UNCHECKED_CODEC);
+                    UnboundedMapCodec<ResourceKey<?>, String> codec = USED_DATAPACK_MAP_CODECS.computeIfAbsent(registryId, resourceLocation -> Codec.unboundedMap(cast(ResourceKey.codec(ResourceKey.createRegistryKey(resourceLocation))), Codec.STRING));
+                    Map<ResourceKey<?>, String> readMap = readResources(entry.getValue(), codec);
                     ClientRegistryMapKey<?> clientRegistryMapKey = ClientRegistryMapKey.of(ResourceKey.createRegistryKey(registryId), idPath.getFileNameOnly(5).asResourceLocation());
                     LAZY_TO_CHECK_MAPS.put(clientRegistryMapKey, readMap);
                 }
@@ -127,4 +159,8 @@ private <T> Map<T, String> readResources(List<Resource> resourceStack, Codec<Map
         }
         return mapBuilder.build();
     }
+
+    private static <T> T cast(Object o) {
+        return (T) o;
+    }
 }