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

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.ClientHooks.onUpdateRenderState(this, p_361382_, s);
+         return s;
+     }
+
 @@ -270,7 +_,12 @@
          }
 
@@ -23,11 +31,12 @@
          if (flag) {
              p_361028_.nameTag = this.getNameTag(p_362104_);
              p_361028_.nameTagAttachment = p_362104_.getAttachments().getNullable(EntityAttachment.NAME_TAG, 0, p_362104_.getYRot(p_362204_));
-@@ -302,5 +_,7 @@
+@@ -302,5 +_,8 @@
          }
 
          p_361028_.displayFireAnimation = p_362104_.displayFireAnimation();
 +
 +        p_361028_.partialTick = p_362204_;
++        net.neoforged.neoforge.attachment.AttachmentInternals.copyEntityAttachments(p_362104_, p_362104_, false);
      }
  }

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

@@ -1,10 +1,39 @@
 --- a/net/minecraft/client/renderer/entity/state/EntityRenderState.java
 +++ b/net/minecraft/client/renderer/entity/state/EntityRenderState.java
-@@ -27,6 +_,7 @@
+@@ -7,7 +_,7 @@
+ import net.neoforged.api.distmarker.OnlyIn;
+
+ @OnlyIn(Dist.CLIENT)
+-public class EntityRenderState {
++public class EntityRenderState extends net.neoforged.neoforge.attachment.AttachmentHolder implements net.neoforged.neoforge.client.extensions.IRenderStateExtension {
+     public double x;
+     public double y;
+     public double z;
+@@ -27,6 +_,27 @@
      public Vec3 nameTagAttachment;
      @Nullable
      public EntityRenderState.LeashState leashState;
 +    public float partialTick;
++    private java.util.Map<net.minecraft.util.context.ContextKey<?>, Object> extensions = java.util.Map.of(); // NeoForge: allow generic extension of render states
++
++    @SuppressWarnings("unchecked")
++    @Nullable
++    @Override
++    public <T> T getRenderData(net.minecraft.util.context.ContextKey<T> key) {
++        return (T) extensions.getOrDefault(key, null);
++    }
++
++    @Override
++    public <T> void setRenderData(net.minecraft.util.context.ContextKey<T> key, @Nullable T data) {
++        if (extensions == java.util.Map.<net.minecraft.util.context.ContextKey<?>, Object>of()) {
++            extensions = new it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap<>();
++        }
++        if (data == null) {
++            extensions.remove(key);
++        } else {
++            extensions.put(key, data);
++        }
++    }
 
      @OnlyIn(Dist.CLIENT)
      public static class LeashState {

src/main/java/net/neoforged/neoforge/attachment/AttachmentInternals.java

@@ -53,6 +53,10 @@ public static void copyEntityAttachments(Entity from, Entity to, boolean isDeath
         copyAttachments(from.registryAccess(), from, to, isDeath ? type -> type.copyOnDeath : type -> true);
     }
 
+    public static <H extends AttachmentHolder> void copyEntityAttachments(Entity from, H to, boolean isDeath) {
+        copyAttachments(from.registryAccess(), from, to, isDeath ? type -> type.copyOnDeath : type -> true);
+    }
+
     @SubscribeEvent
     public static void onPlayerClone(PlayerEvent.Clone event) {
         event.getEntity().copyAttachmentsFrom(event.getOriginal(), event.isWasDeath());

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

@@ -78,6 +78,8 @@
 import net.minecraft.client.renderer.blockentity.BlockEntityRenderer;
 import net.minecraft.client.renderer.chunk.RenderChunkRegion;
 import net.minecraft.client.renderer.culling.Frustum;
+import net.minecraft.client.renderer.entity.EntityRenderer;
+import net.minecraft.client.renderer.entity.state.EntityRenderState;
 import net.minecraft.client.renderer.entity.state.HumanoidRenderState;
 import net.minecraft.client.renderer.texture.TextureAtlas;
 import net.minecraft.client.renderer.texture.TextureAtlasSprite;
@@ -135,6 +137,7 @@
 import net.neoforged.fml.common.EventBusSubscriber;
 import net.neoforged.fml.common.asm.enumextension.ExtensionInfo;
 import net.neoforged.neoforge.client.entity.animation.json.AnimationTypeManager;
+import net.neoforged.neoforge.client.entity.state.EntityRenderStateModifier;
 import net.neoforged.neoforge.client.event.AddSectionGeometryEvent;
 import net.neoforged.neoforge.client.event.CalculateDetachedCameraDistanceEvent;
 import net.neoforged.neoforge.client.event.CalculatePlayerTurnEvent;
@@ -156,6 +159,7 @@
 import net.neoforged.neoforge.client.event.RegisterKeyMappingsEvent;
 import net.neoforged.neoforge.client.event.RegisterMaterialAtlasesEvent;
 import net.neoforged.neoforge.client.event.RegisterParticleProvidersEvent;
+import net.neoforged.neoforge.client.event.RegisterRenderStateModifiersEvent;
 import net.neoforged.neoforge.client.event.RegisterShadersEvent;
 import net.neoforged.neoforge.client.event.RegisterSpriteSourceTypesEvent;
 import net.neoforged.neoforge.client.event.RenderArmEvent;
@@ -272,6 +276,15 @@ public static boolean onDrawHighlight(LevelRenderer context, Camera camera, Bloc
         return NeoForge.EVENT_BUS.post(new RenderHighlightEvent.Block(context, camera, target, deltaTracker, poseStack, bufferSource, forTranslucentBlocks)).isCanceled();
     }
 
+    public static <E extends Entity, S extends EntityRenderState> void onUpdateRenderState(EntityRenderer<E, S> renderer, E entity, S renderState) {
+        var modifiers = RenderStateExtensions.getCachedEntityModifiers(renderer);
+        if (!modifiers.isEmpty()) {
+            for (EntityRenderStateModifier<E, S> modifier : modifiers) {
+                modifier.accept(entity, renderState);
+            }
+        }
+    }
+
     public static void dispatchRenderStage(RenderLevelStageEvent.Stage stage, LevelRenderer levelRenderer, @Nullable PoseStack poseStack, Matrix4f modelViewMatrix, Matrix4f projectionMatrix, int renderTick, Camera camera, Frustum frustum) {
         var mc = Minecraft.getInstance();
         var profiler = Profiler.get();
@@ -990,6 +1003,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/RenderStateExtensions.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client;
+
+import it.unimi.dsi.fastutil.objects.ObjectArrayList;
+import it.unimi.dsi.fastutil.objects.Reference2ObjectArrayMap;
+import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import net.minecraft.Util;
+import net.minecraft.client.renderer.entity.EntityRenderer;
+import net.minecraft.client.renderer.entity.state.EntityRenderState;
+import net.minecraft.world.entity.Entity;
+import net.neoforged.neoforge.client.entity.state.EntityRenderStateModifier;
+import org.jetbrains.annotations.ApiStatus;
+
+public final class RenderStateExtensions {
+    private RenderStateExtensions() {}
+
+    private static final Map<Class<? extends EntityRenderer<?, ?>>, Collection<EntityRenderStateModifier<?, ?>>> ENTITY_EXTENSIONS = new Reference2ObjectArrayMap<>();
+
+    private static final Map<Class<? extends EntityRenderer<?, ?>>, Collection<EntityRenderStateModifier<?, ?>>> ENTITY_CACHE = Util.make(new Reference2ObjectOpenHashMap<>(), map -> map.defaultReturnValue(List.of()));
+
+    @SuppressWarnings("unchecked")
+    static <E extends Entity, S extends EntityRenderState> Collection<EntityRenderStateModifier<E, S>> getCachedEntityModifiers(EntityRenderer<E, S> renderer) {
+        return (Collection<EntityRenderStateModifier<E, S>>) (Object) ENTITY_CACHE.computeIfAbsent((Class<? extends EntityRenderer<E, S>>) renderer.getClass(), aClass -> {
+            var list = new ObjectArrayList<EntityRenderStateModifier<?, ?>>();
+            for (var entry : ENTITY_EXTENSIONS.entrySet()) {
+                if (aClass.isInstance(entry.getClass())) {
+                    list.addAll(entry.getValue());
+                }
+            }
+            if (list.isEmpty()) {
+                return null;
+            }
+            return list;
+        });
+    }
+
+    @ApiStatus.Internal
+    public static <E extends Entity, S extends EntityRenderState> void registerExtender(Class<? extends EntityRenderer<E, S>> baseRenderer, EntityRenderStateModifier<E, S> modifier) {
+        ENTITY_EXTENSIONS.computeIfAbsent(baseRenderer, aClass -> new ObjectArrayList<>()).add(modifier);
+    }
+}

src/main/java/net/neoforged/neoforge/client/entity/state/EntityRenderStateModifier.java

@@ -0,0 +1,13 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client.entity.state;
+
+import net.minecraft.client.renderer.entity.state.EntityRenderState;
+import net.minecraft.world.entity.Entity;
+
+public interface EntityRenderStateModifier<E extends Entity, S extends EntityRenderState> {
+    void accept(E entity, S renderState);
+}

src/main/java/net/neoforged/neoforge/client/entity/state/package-info.java

@@ -0,0 +1,13 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+@FieldsAreNonnullByDefault
+@MethodsReturnNonnullByDefault
+@ParametersAreNonnullByDefault
+package net.neoforged.neoforge.client.entity.state;
+
+import javax.annotation.ParametersAreNonnullByDefault;
+import net.minecraft.FieldsAreNonnullByDefault;
+import net.minecraft.MethodsReturnNonnullByDefault;

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

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.neoforge.client.event;
+
+import net.minecraft.client.renderer.entity.EntityRenderer;
+import net.minecraft.client.renderer.entity.state.EntityRenderState;
+import net.minecraft.world.entity.Entity;
+import net.neoforged.bus.api.Event;
+import net.neoforged.fml.event.IModBusEvent;
+import net.neoforged.neoforge.client.RenderStateExtensions;
+import net.neoforged.neoforge.client.entity.state.EntityRenderStateModifier;
+import org.jetbrains.annotations.ApiStatus;
+
+public class RegisterRenderStateModifiersEvent extends Event implements IModBusEvent {
+    @ApiStatus.Internal
+    public RegisterRenderStateModifiersEvent() {}
+
+    public <E extends Entity, S extends EntityRenderState> void registerEntityModifier(Class<? extends EntityRenderer<E, S>> baseRenderer, EntityRenderStateModifier<E, S> modifier) {
+        RenderStateExtensions.registerExtender(baseRenderer, modifier);
+    }
+}

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

@@ -0,0 +1,32 @@
+/*
+ * 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 org.jetbrains.annotations.Nullable;
+
+public interface IRenderStateExtension {
+    @Nullable
+    <T> T getRenderData(ContextKey<T> key);
+
+    <T> void setRenderData(ContextKey<T> key, @Nullable T 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;
+    }
+
+    default <T> T getRenderDataOrDefault(ContextKey<T> key, T defaultVal) {
+        T data = getRenderData(key);
+        if (data == null) {
+            return defaultVal;
+        }
+        return data;
+    }
+}

tests/src/main/java/net/neoforged/neoforge/debug/client/ClientEventTests.java

@@ -6,23 +6,28 @@
 package net.neoforged.neoforge.debug.client;
 
 import com.mojang.blaze3d.vertex.BufferBuilder;
+import com.mojang.math.Axis;
 import java.util.Map;
 import net.minecraft.client.Minecraft;
 import net.minecraft.client.renderer.RenderType;
+import net.minecraft.client.renderer.entity.player.PlayerRenderer;
 import net.minecraft.client.renderer.texture.OverlayTexture;
 import net.minecraft.client.resources.model.ModelResourceLocation;
 import net.minecraft.core.SectionPos;
 import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.network.chat.Component;
 import net.minecraft.world.item.ItemDisplayContext;
 import net.minecraft.world.item.Items;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.util.context.ContextKey;
 import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.levelgen.SingleThreadedRandomSource;
 import net.neoforged.api.distmarker.Dist;
 import net.neoforged.neoforge.client.event.AddSectionGeometryEvent;
 import net.neoforged.neoforge.client.event.ClientChatEvent;
 import net.neoforged.neoforge.client.event.ClientPlayerChangeGameTypeEvent;
 import net.neoforged.neoforge.client.event.RegisterRenderBuffersEvent;
+import net.neoforged.neoforge.client.event.RegisterRenderStateModifiersEvent;
 import net.neoforged.neoforge.client.event.RenderLevelStageEvent;
 import net.neoforged.neoforge.client.event.RenderPlayerEvent;
 import net.neoforged.neoforge.client.model.data.ModelData;
@@ -126,4 +131,31 @@ static void renderPlayerEvent(final DynamicTest test) {
             test.requestConfirmation(Minecraft.getInstance().player, Component.literal("Is an iron block rendered above you in third-person?"));
         });
     }
+
+    @TestHolder(description = { "" }, enabledByDefault = true)
+    static void updateRenderState(final DynamicTest test) {
+        var key = new ContextKey<Integer>(ResourceLocation.fromNamespaceAndPath(test.createModId(), "test"));
+        var testAttachment = test.registrationHelper().attachments().registerSimpleAttachment("test", () -> 3);
+        test.whenEnabled(listeners -> {
+            listeners.mod().addListener((RegisterRenderStateModifiersEvent event) -> {
+                event.registerEntityModifier(PlayerRenderer.class, (entity, renderState) -> {
+                    renderState.setRenderData(key, 5);
+                });
+            });
+            listeners.forge().addListener((RenderPlayerEvent.Post event) -> {
+                int numRender = event.getRenderState().getData(testAttachment);
+                int xRotation = event.getRenderState().getRenderDataOrDefault(key, -1);
+                var poseStack = event.getPoseStack();
+                poseStack.pushPose();
+                for (int i = 0; i < numRender; i++) {
+                    poseStack.translate(0, 1, 0);
+                    poseStack.pushPose();
+                    poseStack.mulPose(Axis.XP.rotation(xRotation));
+                    Minecraft.getInstance().getBlockRenderer().renderSingleBlock(Blocks.CALCITE.defaultBlockState(), poseStack, event.getMultiBufferSource(), event.getPackedLight(), OverlayTexture.NO_OVERLAY, ModelData.EMPTY, RenderType.solid());
+                    poseStack.popPose();
+                }
+                poseStack.popPose();
+            });
+        });
+    }
 }