[1.21.3] RegisterRenderStateModifiersEvent for appending custom data to render state objects

patches/net/minecraft/client/renderer/MapRenderer.java.patch

@@ -11,7 +11,16 @@
                  p_362483_.pushPose();
                  p_362483_.translate(
                      (float)maprenderstate$mapdecorationrenderstate.x / 2.0F + 64.0F, (float)maprenderstate$mapdecorationrenderstate.y / 2.0F + 64.0F, -0.02F
-@@ -116,6 +_,7 @@
+@@ -109,13 +_,15 @@
+         p_364922_.texture = this.mapTextureManager.prepareMapTexture(p_361383_, p_363500_);
+         p_364922_.decorations.clear();
+
++        net.neoforged.neoforge.client.renderstate.RenderStateExtensions.onUpdateMapRenderState(p_363500_, p_364922_);
+         for (MapDecoration mapdecoration : p_363500_.getDecorations()) {
+-            p_364922_.decorations.add(this.extractDecorationRenderState(mapdecoration));
++            p_364922_.decorations.add(net.neoforged.neoforge.client.renderstate.RenderStateExtensions.onUpdateMapDecorationRenderState(mapdecoration.type(), p_363500_, p_364922_, this.extractDecorationRenderState(mapdecoration)));
+         }
+     }
 
      private MapRenderState.MapDecorationRenderState extractDecorationRenderState(MapDecoration p_364175_) {
          MapRenderState.MapDecorationRenderState maprenderstate$mapdecorationrenderstate = new MapRenderState.MapDecorationRenderState();

patches/net/minecraft/client/renderer/entity/EntityRenderer.java.patch

@@ -9,6 +9,14 @@
              this.renderNameTag(p_364816_, p_364816_.nameTag, p_114488_, p_114489_, p_114490_);
          }
      }
+@@ -245,6 +_,7 @@
+     public final S createRenderState(T p_361382_, float p_360885_) {
+         S s = this.reusedState;
+         this.extractRenderState(p_361382_, s, p_360885_);
++        net.neoforged.neoforge.client.renderstate.RenderStateExtensions.onUpdateEntityRenderState(this, p_361382_, s);
+         return s;
+     }
+
 @@ -270,7 +_,12 @@
          }
 

patches/net/minecraft/client/renderer/entity/state/EntityRenderState.java.patch

@@ -1,5 +1,14 @@
 --- a/net/minecraft/client/renderer/entity/state/EntityRenderState.java
 +++ b/net/minecraft/client/renderer/entity/state/EntityRenderState.java
+@@ -7,7 +_,7 @@
+ import net.neoforged.api.distmarker.OnlyIn;
+
+ @OnlyIn(Dist.CLIENT)
+-public class EntityRenderState {
++public class EntityRenderState extends net.neoforged.neoforge.client.renderstate.BaseRenderState {
+     public double x;
+     public double y;
+     public double z;
 @@ -27,6 +_,7 @@
      public Vec3 nameTagAttachment;
      @Nullable

patches/net/minecraft/client/renderer/state/MapRenderState.java.patch

@@ -1,9 +1,18 @@
 --- a/net/minecraft/client/renderer/state/MapRenderState.java
 +++ b/net/minecraft/client/renderer/state/MapRenderState.java
-@@ -17,6 +_,7 @@
+@@ -10,13 +_,14 @@
+ import net.neoforged.api.distmarker.OnlyIn;
+
+ @OnlyIn(Dist.CLIENT)
+-public class MapRenderState {
++public class MapRenderState extends net.neoforged.neoforge.client.renderstate.BaseRenderState {
+     @Nullable
+     public ResourceLocation texture;
+     public final List<MapRenderState.MapDecorationRenderState> decorations = new ArrayList<>();
 
      @OnlyIn(Dist.CLIENT)
-     public static class MapDecorationRenderState {
+-    public static class MapDecorationRenderState {
++    public static class MapDecorationRenderState extends net.neoforged.neoforge.client.renderstate.BaseRenderState {
 +        public net.minecraft.core.Holder<net.minecraft.world.level.saveddata.maps.MapDecorationType> type;
          @Nullable
          public TextureAtlasSprite atlasSprite;

src/main/java/net/neoforged/neoforge/client/ClientHooks.java

@@ -180,6 +180,7 @@
 import net.neoforged.neoforge.client.gui.GuiLayerManager;
 import net.neoforged.neoforge.client.gui.map.MapDecorationRendererManager;
 import net.neoforged.neoforge.client.model.data.ModelData;
+import net.neoforged.neoforge.client.renderstate.RegisterRenderStateModifiersEvent;
 import net.neoforged.neoforge.common.NeoForge;
 import net.neoforged.neoforge.common.NeoForgeMod;
 import net.neoforged.neoforge.forge.snapshots.ForgeSnapshotsModClient;
@@ -990,6 +991,7 @@ public static void initClientHooks(Minecraft mc, ReloadableResourceManager resou
         ModLoader.postEvent(new RegisterClientReloadListenersEvent(resourceManager));
         ModLoader.postEvent(new EntityRenderersEvent.RegisterLayerDefinitions());
         ModLoader.postEvent(new EntityRenderersEvent.RegisterRenderers());
+        ModLoader.postEvent(new RegisterRenderStateModifiersEvent());
         ClientTooltipComponentManager.init();
         EntitySpectatorShaderManager.init();
         ClientHooks.onRegisterKeyMappings(mc.options);

src/main/java/net/neoforged/neoforge/client/extensions/IRenderStateExtension.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client.extensions;
+
+import net.minecraft.util.context.ContextKey;
+import net.neoforged.neoforge.client.renderstate.BaseRenderState;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Extension class for render state objects. Implemented by {@link BaseRenderState} for
+ * simple class extension.
+ */
+public interface IRenderStateExtension {
+    /**
+     * Gets the object associated with the given key.
+     * 
+     * @param key Static key reference object
+     * @return The object associated with the key or null if the key is not present.
+     * @param <T> Type of render data
+     */
+    @Nullable
+    <T> T getRenderData(ContextKey<T> key);
+
+    /**
+     * Sets the object associated with the given key. Key should be stored statically for later retrieval of the object.
+     * 
+     * @param key  Static key reference object
+     * @param data Object to store for custom rendering
+     * @param <T>  Type of render data
+     */
+    <T> void setRenderData(ContextKey<T> key, @Nullable T data);
+
+    /**
+     * Gets the value or throws an exception. Should be used in cases where the data must be present.
+     * 
+     * @param key Static key reference object
+     * @return The data associate with the key
+     * @param <T> Type of render data
+     */
+    default <T> T getRenderDataOrThrow(ContextKey<T> key) {
+        T data = getRenderData(key);
+        if (data == null) {
+            throw new IllegalStateException("No value associated for key " + key);
+        }
+        return data;
+    }
+
+    /**
+     * Gets the value or returns the default object if an object is not present
+     * 
+     * @param key        Static key reference object
+     * @param defaultVal Default value if an object is not present
+     * @return Value from the render data or the given default value if value is not present
+     * @param <T> Type of render data
+     */
+    default <T> T getRenderDataOrDefault(ContextKey<T> key, T defaultVal) {
+        T data = getRenderData(key);
+        if (data == null) {
+            return defaultVal;
+        }
+        return data;
+    }
+}

src/main/java/net/neoforged/neoforge/client/renderstate/BaseRenderState.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client.renderstate;
+
+import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
+import java.util.Map;
+import net.minecraft.client.renderer.entity.state.EntityRenderState;
+import net.minecraft.util.context.ContextKey;
+import net.neoforged.neoforge.client.extensions.IRenderStateExtension;
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Extension class for RenderState objects (ie {@link EntityRenderState}).
+ * Allows modders to add arbitrary data onto render states for use in custom rendering.
+ */
+public abstract class BaseRenderState implements IRenderStateExtension {
+    protected Map<ContextKey<?>, Object> extensions = new Reference2ObjectOpenHashMap<>();
+
+    @SuppressWarnings("unchecked")
+    @Nullable
+    @Override
+    public <T> T getRenderData(ContextKey<T> key) {
+        return (T) extensions.get(key);
+    }
+
+    @Override
+    public <T> void setRenderData(ContextKey<T> key, @Nullable T data) {
+        if (data != null) {
+            extensions.put(key, data);
+        } else {
+            extensions.remove(key);
+        }
+    }
+
+    @ApiStatus.Internal
+    public void resetRenderData() {
+        extensions.clear();
+    }
+}

src/main/java/net/neoforged/neoforge/client/renderstate/MapDecorationRenderStateModifier.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client.renderstate;
+
+import net.minecraft.client.renderer.state.MapRenderState;
+import net.minecraft.world.level.saveddata.maps.MapDecorationType;
+import net.minecraft.world.level.saveddata.maps.MapItemSavedData;
+import net.neoforged.neoforge.client.gui.map.IMapDecorationRenderer;
+
+/**
+ * Function interface for render state modifiers that target MapDecorations. Useful for adding custom data for rendering
+ * in {@link IMapDecorationRenderer}s.
+ */
+@FunctionalInterface
+public interface MapDecorationRenderStateModifier {
+    /**
+     * Called when the registered {@link MapDecorationType} is added to a {@link MapRenderState}.
+     *
+     * @param mapItemSavedData         The map SavedData.
+     * @param mapRenderState           The render state of the map after the texture has been set and custom data is added.
+     * @param mapDecorationRenderState The decoration render state after vanilla has set it up.
+     */
+    void accept(MapItemSavedData mapItemSavedData, MapRenderState mapRenderState, MapRenderState.MapDecorationRenderState mapDecorationRenderState);
+}

src/main/java/net/neoforged/neoforge/client/renderstate/RegisterRenderStateModifiersEvent.java

@@ -0,0 +1,142 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client.renderstate;
+
+import com.google.common.reflect.TypeParameter;
+import com.google.common.reflect.TypeToken;
+import java.lang.reflect.ParameterizedType;
+import java.util.function.BiConsumer;
+import net.minecraft.client.renderer.entity.EntityRenderer;
+import net.minecraft.client.renderer.entity.state.EntityRenderState;
+import net.minecraft.client.renderer.state.MapRenderState;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.util.context.ContextKey;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.level.saveddata.maps.MapDecorationType;
+import net.minecraft.world.level.saveddata.maps.MapItemSavedData;
+import net.neoforged.bus.api.Event;
+import net.neoforged.fml.LogicalSide;
+import net.neoforged.fml.event.IModBusEvent;
+import net.neoforged.neoforge.client.extensions.IRenderStateExtension;
+import org.jetbrains.annotations.ApiStatus;
+
+/**
+ * Fired for registering modifier functions for various render state objects. Useful for gathering context for
+ * custom rendering with objects that are not your own.
+ *
+ * <p>This event is fired on the mod-specific event bus, only on the {@linkplain LogicalSide#CLIENT logical client}.</p>
+ */
+public class RegisterRenderStateModifiersEvent extends Event implements IModBusEvent {
+    @ApiStatus.Internal
+    public RegisterRenderStateModifiersEvent() {}
+
+    /**
+     * Registers a render state modifier for {@link EntityRenderState}s which are run after all vanilla data is
+     * extracted. Can add custom data to the map using {@link EntityRenderState#setRenderData(ContextKey, Object)}.
+     * Any subclasses of the passed renderer class will also have this modifier applied.
+     *
+     * <pre>
+     * <code>
+     *     event.registerEntityModifier(new TypeToken<LivingEntityRenderer<LivingEntity, LivingEntityRenderState, ?>>() {}, (entity, renderState) -> {
+     *         . . .
+     *     });
+     * </code>
+     * </pre>
+     * 
+     * @param baseRenderer Entity renderer class. Any subclasses will also apply this modifier.
+     * @param modifier     The function for modifying the {@link EntityRenderState} and adding custom render data.
+     * @param <E>          The type of the entity
+     * @param <S>          The specific render state type
+     */
+    public <E extends Entity, S extends EntityRenderState> void registerEntityModifier(TypeToken<? extends EntityRenderer<? extends E, ? extends S>> baseRenderer, BiConsumer<E, S> modifier) {
+        ensureParametersMatchBounds(baseRenderer);
+        RenderStateExtensions.registerEntity(baseRenderer.getRawType(), modifier);
+    }
+
+    /**
+     * Convenience method for cases where generics are not present. Registers a render state modifier for
+     * {@link EntityRenderState}s which are run after all vanilla data is extracted. Can add custom data to the map
+     * using {@link EntityRenderState#setRenderData(ContextKey, Object)}. Any subclasses of the passed renderer class
+     * will also have this modifier applied.
+     *
+     * <pre>
+     * <code>
+     *     event.registerEntityModifier(PlayerRenderer.class, (entity, renderState) -> {
+     *         . . .
+     *     });
+     * </code>
+     * </pre>
+     *
+     * @param baseRenderer Entity renderer class. Any subclasses will also apply this modifier.
+     * @param modifier     The function for modifying the {@link EntityRenderState} and adding custom render data.
+     * @param <E>          The type of the entity
+     * @param <S>          The specific render state type
+     */
+    public <E extends Entity, S extends EntityRenderState> void registerEntityModifier(Class<? extends EntityRenderer<? extends E, ? extends S>> baseRenderer, BiConsumer<E, S> modifier) {
+        ensureParametersMatchBounds(TypeToken.of(baseRenderer));
+        RenderStateExtensions.registerEntity(baseRenderer, modifier);
+    }
+
+    /**
+     * Registers a render state modifier for {@link MapRenderState}s which are run after the texture has been set
+     * and before decorations have been added. Can add custom data to the map using
+     * {@link IRenderStateExtension#setRenderData(ContextKey, Object)}.
+     *
+     * @param modifier The function for modifying the {@link net.minecraft.client.renderer.state.MapRenderState} and adding custom render data.
+     */
+    public void registerMapModifier(BiConsumer<MapItemSavedData, MapRenderState> modifier) {
+        RenderStateExtensions.registerMap(modifier);
+    }
+
+    /**
+     * Registers a render state modifier for {@link MapRenderState.MapDecorationRenderState}s which are run after
+     * vanilla map decoration data has been set. Can add custom data to the map using
+     * {@link IRenderStateExtension#setRenderData(ContextKey, Object)}.
+     *
+     * @param mapDecorationTypeKey Key for the registered {@link MapDecorationType}
+     * @param modifier             The function for modifying the {@link MapRenderState.MapDecorationRenderState} and adding custom render data.
+     */
+    public void registerMapDecorationModifier(ResourceKey<MapDecorationType> mapDecorationTypeKey, MapDecorationRenderStateModifier modifier) {
+        RenderStateExtensions.registerMapDecoration(mapDecorationTypeKey, modifier);
+    }
+
+    private static void ensureParametersMatchBounds(TypeToken<? extends EntityRenderer<? extends Entity, ? extends EntityRenderState>> baseRenderer) {
+        if (baseRenderer.getType() instanceof ParameterizedType parameterizedType) {
+            Class<?> bound = baseRenderer.getRawType();
+            ParameterizedType parameterized = parameterizedType;
+            do {
+                var userArgs = parameterized.getActualTypeArguments();
+                var typeArgs = bound.getTypeParameters();
+
+                for (int i = 0; i < userArgs.length; i++) {
+                    var userArg = userArgs[i];
+                    var userToken = Container.of(TypeToken.of(userArg));
+                    var typeArg = typeArgs[i];
+                    for (var singleBound : typeArg.getBounds()) {
+                        var token = Container.of(TypeToken.of(singleBound));
+                        if (!token.isSubtypeOf(userToken)) {
+                            throw new IllegalArgumentException("%s does not match expected type parameter %s".formatted(userArg, singleBound));
+                        }
+                    }
+                }
+
+                if (!(parameterized.getOwnerType() instanceof ParameterizedType parameterizedOwner)) {
+                    break;
+                }
+                parameterized = parameterizedOwner;
+                bound = bound.getEnclosingClass();
+            } while (bound != null);
+        }
+    }
+
+    @SuppressWarnings("unused")
+    private record Container<X>() {
+        private static <Z> TypeToken<Container<Z>> of(TypeToken<Z> parameter) {
+            return new TypeToken<Container<Z>>() {}
+                    .where(new TypeParameter<>() {}, parameter);
+        }
+    }
+}