Fabric implementation of REAs

RegistryEntryAttachment.java

package com.williambl.elysium.data;

import com.google.gson.*;
import com.mojang.datafixers.util.Either;
import com.mojang.datafixers.util.Pair;
import com.mojang.serialization.Codec;
import com.mojang.serialization.DataResult;
import com.mojang.serialization.JsonOps;
import net.fabricmc.fabric.api.client.networking.v1.ClientPlayNetworking;
import net.fabricmc.fabric.api.event.lifecycle.v1.ServerLifecycleEvents;
import net.fabricmc.fabric.api.networking.v1.PacketByteBufs;
import net.fabricmc.fabric.api.networking.v1.PacketSender;
import net.fabricmc.fabric.api.networking.v1.ServerPlayConnectionEvents;
import net.fabricmc.fabric.api.networking.v1.ServerPlayNetworking;
import net.fabricmc.fabric.api.resource.ResourceManagerHelper;
import net.fabricmc.fabric.api.resource.SimpleResourceReloadListener;
import net.minecraft.ResourceLocationException;
import net.minecraft.core.Holder;
import net.minecraft.core.Registry;
import net.minecraft.network.FriendlyByteBuf;
import net.minecraft.network.protocol.Packet;
import net.minecraft.resources.ResourceKey;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.MinecraftServer;
import net.minecraft.server.network.ServerGamePacketListenerImpl;
import net.minecraft.server.packs.PackType;
import net.minecraft.server.packs.resources.CloseableResourceManager;
import net.minecraft.server.packs.resources.ResourceManager;
import net.minecraft.tags.TagKey;
import net.minecraft.util.GsonHelper;
import net.minecraft.util.profiling.ProfilerFiller;
import org.jetbrains.annotations.Nullable;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Collectors;

// provide your own
import static YourMod.LOGGER;
import static YourMod.id;

// doesn't work on dynamic registries because I don't need it to. sorry!
public class RegistryEntryAttachment<K, V> {
	private static Loader loader;

	private final Registry<K> registry;
	private final ResourceLocation name;
	private final Class<V> vClass;
	private final Codec<V> codec;

	private Map<K, V> map;

	public RegistryEntryAttachment(Registry<K> registry, ResourceLocation name, Class<V> vClass, Codec<V> codec) {
		this.registry = registry;
		this.name = name;
		this.vClass = vClass;
		this.codec = codec;
	}

	public Optional<V> get(K key) {
		return Optional.ofNullable(this.map.get(key));
	}

	public static <K, V> Builder<K, V> builder(Registry<K> registry, ResourceLocation name, Class<V> vClass, Codec<V> codec) {
		return new Builder<>(registry, name, vClass, codec);
	}

	public static void init() {
		loader = new Loader();
		ResourceManagerHelper.get(PackType.SERVER_DATA).registerReloadListener(loader);
		ServerPlayConnectionEvents.JOIN.register(loader);
		ServerLifecycleEvents.END_DATA_PACK_RELOAD.register(loader);
	}

	public static class Builder<K, V> {
		private final Registry<K> registry;
		private final ResourceLocation name;
		private final Class<V> vClass;
		private final Codec<V> codec;

		public Builder(Registry<K> registry, ResourceLocation name, Class<V> vClass, Codec<V> codec) {
			this.registry = registry;
			this.name = name;
			this.vClass = vClass;
			this.codec = codec;
		}

		public RegistryEntryAttachment<K, V> build() {
			if (loader == null) {
				throw new IllegalStateException("REAs were not initialized!");
			}
			var rea = new RegistryEntryAttachment<>(this.registry, this.name, this.vClass, this.codec);
			return loader.register(rea);
		}
	}

	private static class Loader implements SimpleResourceReloadListener<Loader.Preparation>, ServerPlayConnectionEvents.Join, ServerLifecycleEvents.EndDataPackReload {
		private static final ResourceLocation ID = id("registry_entry_attachments");
		final Map<Pair<ResourceKey<? extends Registry<?>>, ResourceLocation>, RegistryEntryAttachment<?, ?>> attachments = new HashMap<>();

		public <K, V> RegistryEntryAttachment<K, V> register(RegistryEntryAttachment<K, V> rea) {
			var existing = this.attachments.get(Pair.of(rea.registry.key(), rea.name));
			if (existing != null) {
				if (existing.vClass != rea.vClass || existing.codec != rea.codec) {
					throw new IllegalStateException("Two REAs for registry %s were registered under the same name: %s".formatted(rea.registry.key(), rea.name));
				}

				return (RegistryEntryAttachment<K, V>) existing;
			}

			this.attachments.put(Pair.of(rea.registry.key(), rea.name), rea);
			return rea;
		}

		@Override
		public ResourceLocation getFabricId() {
			return ID;
		}

		@Override
		public CompletableFuture<Preparation> load(ResourceManager manager, ProfilerFiller profiler, Executor executor) {
			Map<Pair<ResourceKey<? extends Registry<?>>, ResourceLocation>, LinkedHashMap<? extends Either<ResourceLocation, ? extends TagKey<?>>, ?>> prepMap = new HashMap<>();
			List<CompletableFuture<Void>> futures = new ArrayList<>();
			for (var rea : this.attachments.values()) {
				futures.add(load(rea, manager, executor).thenAccept(m -> {
					prepMap.put(Pair.of(rea.registry.key(), rea.name), m);
				}));
			}

			return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
				.thenApply($ -> new Preparation(prepMap));
		}

		private static <K, V> CompletableFuture<LinkedHashMap<Either<ResourceLocation, TagKey<K>>, V>> load(RegistryEntryAttachment<K, V> rea, ResourceManager manager, Executor executor) {
			return CompletableFuture.supplyAsync(() -> {
				LinkedHashMap<Either<ResourceLocation, TagKey<K>>, V> map = new LinkedHashMap<>();
				for (var resource : manager.getResourceStack(new ResourceLocation(rea.name.getNamespace(), "attachments/%s/%s/%s.json".formatted(rea.registry.key().location().getNamespace(), rea.registry.key().location().getPath(), rea.name.getPath())))) {
					try(var reader = resource.openAsReader()) {
						var json = JsonParser.parseReader(reader);
						if (!(json instanceof JsonObject jObj)) {
							throw new JsonSyntaxException("Expected a JSON Object");
						}
						if (jObj.has("replace")
							&& jObj.get("replace") instanceof JsonPrimitive prim
							&& prim.isBoolean()
							&& prim.getAsBoolean()) {
							map.clear();
						}
						if (!jObj.has("values")) {
							throw new JsonSyntaxException("Expected a JSON Array or Object for 'values'");
						}
						var values = jObj.get("values");
						if (values instanceof JsonArray valuesArr) {
							for (var value : valuesArr) {
								final Either<ResourceLocation, TagKey<K>> key;
								if (!(value instanceof JsonObject valueObj)) {
									throw new JsonSyntaxException("Expected a JSON Object for element of 'values', got %s instead".formatted(value.toString()));
								}

								try {
									if (valueObj.has("id")) {
										key = Either.left(new ResourceLocation(GsonHelper.getAsString(valueObj, "id")));
									} else if (valueObj.has("tag")) {
										key = Either.right(TagKey.create(rea.registry.key(), new ResourceLocation(GsonHelper.getAsString(valueObj, "tag"))));
									} else {
										throw new JsonSyntaxException("Expected entry to have 'id' or 'tag', had neither");
									}
								} catch (ResourceLocationException e) {
									throw new JsonSyntaxException("Failed to parse resource location", e);
								}

								var valueRes = rea.codec.decode(JsonOps.INSTANCE, valueObj.get("value"));
								addValueOrThrow(key, valueRes, map);
							}
						} else if (values instanceof JsonObject valuesObj) {
							for (var value : valuesObj.entrySet()) {
								final Either<ResourceLocation, TagKey<K>> key;
								try {
									if (value.getKey().startsWith("#")) {
										key = Either.right(TagKey.create(rea.registry.key(), new ResourceLocation(value.getKey().substring(1))));
									} else {
										key = Either.left(new ResourceLocation(value.getKey()));
									}
								} catch (ResourceLocationException e) {
									throw new JsonSyntaxException("Failed to parse resource location", e);
								}

								var valueRes = rea.codec.decode(JsonOps.INSTANCE, value.getValue());
								addValueOrThrow(key, valueRes, map);
							}
						} else {
							throw new JsonSyntaxException("Expected a JSON Object or Array for for 'values', got %s instead".formatted(values.toString()));
						}
					} catch (IOException e) {
						LOGGER.error("Failed to read REA definition {}/{} from pack {}, ignoring:", rea.registry.key().location(), rea.name, resource.sourcePackId(), e);
					} catch (JsonParseException e) {
						LOGGER.error("Failed to parse REA definition {}/{} from pack {}, ignoring:", rea.registry.key().location(), rea.name, resource.sourcePackId(), e);
					}
				}

				return map;
			}, executor);
		}

		private static <K, V> void addValueOrThrow(Either<ResourceLocation, TagKey<K>> key, DataResult<Pair<V, JsonElement>> valueRes, Map<Either<ResourceLocation, TagKey<K>>, V> map) {
			if (valueRes.result().isPresent()) {
				map.put(key, valueRes.result().get().getFirst());
			} else {
				throw new JsonSyntaxException("Failed to parse value: "+valueRes.error().orElseThrow().message());
			}
		}

		@Override
		public CompletableFuture<Void> apply(Preparation data, ResourceManager manager, ProfilerFiller profiler, Executor executor) {
			return CompletableFuture.runAsync(() -> {
				for (var rea : this.attachments.values()) {
					updateRea(rea, data.valuesPerRea);
				}
			}, executor);
		}

		private static <K, V> void updateRea(RegistryEntryAttachment<K,V> rea, Map<Pair<ResourceKey<? extends Registry<?>>, ResourceLocation>, LinkedHashMap<? extends Either<ResourceLocation, ? extends TagKey<?>>,?>> valuesPerRea) {
			LinkedHashMap<Either<ResourceLocation, TagKey<K>>, V> values = (LinkedHashMap<Either<ResourceLocation,net.minecraft.tags.TagKey<K>>, V>) valuesPerRea.get(Pair.of(rea.registry.key(), rea.name));
			rea.map = rea.registry.holders()
				.map(h -> Pair.of(h, getValue(h, values)))
				.filter(p -> Objects.nonNull(p.getSecond()))
				.collect(Collectors.toMap(p -> p.getFirst().value(), Pair::getSecond));
		}

		private static <K, V> V getValue(Holder<K> h, LinkedHashMap<Either<ResourceLocation, TagKey<K>>, V> values) {
			@Nullable V value = null;
			for (var entry : values.entrySet()) {
				if (entry.getKey().map(h::is, h::is)) {
					value = entry.getValue();
				}
			}
			return value;
		}

		@Override
		public void endDataPackReload(MinecraftServer minecraftServer, CloseableResourceManager closeableResourceManager, boolean b) {
			for (var player : minecraftServer.getPlayerList().getPlayers()) {
				this.sendUpdatePacket((ch, buf) -> ServerPlayNetworking.send(player, ch, buf));
			}
		}

		@Override
		public void onPlayReady(ServerGamePacketListenerImpl serverGamePacketListener, PacketSender packetSender, MinecraftServer minecraftServer) {
			this.sendUpdatePacket(packetSender::sendPacket);
		}

		private void sendUpdatePacket(BiConsumer<ResourceLocation, FriendlyByteBuf> packetSender) {
			var buf = PacketByteBufs.create();
			buf.writeVarInt(this.attachments.size());
			for (var entry : this.attachments.entrySet()) {
				buf.writeResourceLocation(entry.getKey().getFirst().location());
				buf.writeResourceLocation(entry.getKey().getSecond());
				buf.markWriterIndex();
				buf.writeInt(0);
				writeAttachmentValues(buf, entry.getValue());
				int idx = buf.writerIndex();
				buf.resetWriterIndex();
				buf.writeInt(idx - (buf.writerIndex() + 4));
				buf.writerIndex(idx);
			}
			packetSender.accept(ID, buf);
		}

		public void loadFromPacket(FriendlyByteBuf buf) {
			int count = buf.readVarInt();
			for (int i = 0; i < count; i++) {
				var key = Pair.of(ResourceKey.createRegistryKey(buf.readResourceLocation()), buf.readResourceLocation());
				int size = buf.readInt();
				var rea = this.attachments.get(key);
				if (rea == null) {
					buf.skipBytes(size);
				} else {
					readAttachmentValues(buf, rea);
				}
			}
		}

		private static <K, V> void writeAttachmentValues(FriendlyByteBuf friendlyByteBuf, RegistryEntryAttachment<K, V> registryEntryAttachment) {
			friendlyByteBuf.writeMap(registryEntryAttachment.map,
				(b, o) -> b.writeVarInt(registryEntryAttachment.registry.getId(o)),
				(b, v) -> b.writeWithCodec(registryEntryAttachment.codec, v));
		}

		private static <K, V> void readAttachmentValues(FriendlyByteBuf friendlyByteBuf, RegistryEntryAttachment<K, V> registryEntryAttachment) {
			registryEntryAttachment.map = friendlyByteBuf.readMap(
				(b) -> registryEntryAttachment.registry.getHolder(b.readVarInt()).orElseThrow().value(),
				(b) -> b.readWithCodec(registryEntryAttachment.codec));
		}

		public record Preparation(
			Map<Pair<ResourceKey<? extends Registry<?>>, ResourceLocation>, LinkedHashMap<? extends Either<ResourceLocation, ? extends TagKey<?>>, ?>> valuesPerRea) {}
	}

	public static class Client {
		public static void clientInit() {
			ClientPlayNetworking.registerGlobalReceiver(Loader.ID, (minecraft, clientPacketListener, friendlyByteBuf, packetSender) -> {
				friendlyByteBuf.retain();
				minecraft.execute(() -> {
					loader.loadFromPacket(friendlyByteBuf);
					friendlyByteBuf.release();
				});
			});
		}
	}
}