ServerEntityEvents.ALLOW_FRESH_LOAD and AFTER_FRESH_LOAD event

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/api/event/lifecycle/v1/ServerEntityEvents.java

@@ -18,6 +18,7 @@
 
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.entity.Entity;
+import net.minecraft.world.entity.EntitySpawnReason;
 import net.minecraft.world.entity.EquipmentSlot;
 import net.minecraft.world.entity.LivingEntity;
 import net.minecraft.world.item.ItemStack;
@@ -33,13 +34,43 @@ private ServerEntityEvents() {
 	 * Called when an Entity is loaded into a ServerLevel.
 	 *
 	 * <p>When this event is called, the entity is already in the level.
+	 *
+	 * <p>If you want to get load context use ENTITY_LOAD_CONTEXT scoped value</p>
 	 */
 	public static final Event<ServerEntityEvents.Load> ENTITY_LOAD = EventFactory.createArrayBacked(ServerEntityEvents.Load.class, callbacks -> (entity, level) -> {
 		for (Load callback : callbacks) {
 			callback.onLoad(entity, level);
 		}
 	});
 
+	/**
+	 * Called before an Entity is added to a ServerLevel.
+	 *
+	 * <p>If return value is {@code false} entity will not be added to a server.</p>
+	 *
+	 * <p>Should be used when you want to add another entity instead of added or to disallow add specific entities.</p>
+	 *
+	 * {@snippet :
+	 * ServerEntityEvents.ALLOW_ADD.register((entity, level) -> {
+	 * 	// Spawn with 25% chance zombie instead of creeper
+	 * 	if (entity instanceof Creeper && level.getRandom().nextFloat() <= 0.25f) {
+	 * 	 var zombie = EntityType.ZOMBIE.create(level, null, entity.blockPosition(), EntitySpawnReason.EVENT, true, false);
+	 * 	 if (zombie != null) return !level.addFreshEntity(zombie);
+	 * 	}
+	 * 	return true;
+	 * });
+	 *}
+	 */
+	public static final Event<AllowAdd> ALLOW_ADD = EventFactory.createArrayBacked(AllowAdd.class, callbacks -> (entity, level) -> {
+		boolean bl = true;
+
+		for (AllowAdd callback : callbacks) {
+			bl = bl && callback.allowAdd(entity, level);
+		}
+
+		return bl;
+	});
+
 	/**
 	 * Called when an Entity is unloaded from a ServerLevel.
 	 *
@@ -68,6 +99,11 @@ public interface Load {
 		void onLoad(Entity entity, ServerLevel level);
 	}
 
+	@FunctionalInterface
+	public interface AllowAdd {
+		boolean allowAdd(Entity entity, ServerLevel level);
+	}
+
 	@FunctionalInterface
 	public interface Unload {
 		void onUnload(Entity entity, ServerLevel level);
@@ -77,4 +113,6 @@ public interface Unload {
 	public interface EquipmentChange {
 		void onChange(LivingEntity livingEntity, EquipmentSlot equipmentSlot, ItemStack previousStack, ItemStack currentStack);
 	}
+
+	public static ScopedValue<EntitySpawnReason> ENTITY_LOAD_REASON = ScopedValue.newInstance();
 }

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/EntityTypeMixin.java

@@ -0,0 +1,43 @@
+/*
+ * 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.mixin.event.lifecycle;
+
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerEntityEvents;
+
+import net.minecraft.core.BlockPos;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.entity.EntitySpawnReason;
+import net.minecraft.world.entity.EntityType;
+
+import org.jspecify.annotations.Nullable;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+
+import java.util.function.Consumer;
+
+@Mixin(EntityType.class)
+public class EntityTypeMixin<T extends Entity> {
+	@Inject(method = "spawn(Lnet/minecraft/server/level/ServerLevel;Ljava/util/function/Consumer;Lnet/minecraft/core/BlockPos;Lnet/minecraft/world/entity/EntitySpawnReason;ZZ)Lnet/minecraft/world/entity/Entity;", at = @At(value = "RETURN"))
+	private void entityLoadEvent(ServerLevel level, @Nullable Consumer<T> postSpawnConfig, BlockPos spawnPos, EntitySpawnReason spawnReason, boolean tryMoveDown, boolean movedUp, CallbackInfoReturnable<T> cir) {
+		ScopedValue.where(ServerEntityEvents.ENTITY_LOAD_REASON, spawnReason).run(() -> {
+			ServerEntityEvents.ENTITY_LOAD.invoker().onLoad(cir.getReturnValue(), level);
+		});
+	}
+}

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/ServerLevelEntityCallbacksMixin.java

@@ -16,6 +16,8 @@
 
 package net.fabricmc.fabric.mixin.event.lifecycle;
 
+import net.minecraft.world.entity.EntitySpawnReason;
+
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -37,7 +39,9 @@ abstract class ServerLevelEntityCallbacksMixin {
 
 	@Inject(method = "onTrackingStart(Lnet/minecraft/world/entity/Entity;)V", at = @At("TAIL"))
 	private void invokeEntityLoadEvent(Entity entity, CallbackInfo ci) {
-		ServerEntityEvents.ENTITY_LOAD.invoker().onLoad(entity, this.this$0);
+		ScopedValue.where(ServerEntityEvents.ENTITY_LOAD_REASON, EntitySpawnReason.LOAD).run(() -> {
+			ServerEntityEvents.ENTITY_LOAD.invoker().onLoad(entity, this.this$0);
+		});
 	}
 
 	@Inject(method = "onTrackingEnd(Lnet/minecraft/world/entity/Entity;)V", at = @At("HEAD"))

fabric-lifecycle-events-v1/src/main/java/net/fabricmc/fabric/mixin/event/lifecycle/ServerLevelMixin.java

@@ -18,14 +18,18 @@
 
 import java.util.function.BooleanSupplier;
 
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
 import org.objectweb.asm.Opcodes;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.injection.At;
 import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 
 import net.minecraft.server.level.ServerLevel;
+import net.minecraft.world.entity.Entity;
 
+import net.fabricmc.fabric.api.event.lifecycle.v1.ServerEntityEvents;
 import net.fabricmc.fabric.api.event.lifecycle.v1.ServerTickEvents;
 
 @Mixin(ServerLevel.class)
@@ -40,4 +44,15 @@ private void startLevelTick(BooleanSupplier shouldKeepTicking, CallbackInfo ci)
 	private void endLevelTick(BooleanSupplier shouldKeepTicking, CallbackInfo ci) {
 		ServerTickEvents.END_LEVEL_TICK.invoker().onEndTick((ServerLevel) (Object) this);
 	}
+
+	@WrapOperation(method = "addFreshEntity", at = @At(value = "INVOKE", target = "Lnet/minecraft/server/level/ServerLevel;addEntity(Lnet/minecraft/world/entity/Entity;)Z"))
+	private boolean afterEntityFreshLoad(ServerLevel instance, Entity entity, Operation<Boolean> original) {
+		ServerLevel serverLevel = (ServerLevel) (Object) this;
+
+		if (!ServerEntityEvents.ALLOW_ADD.invoker().allowAdd(entity, serverLevel)) {
+			return false;
+		}
+
+		return original.call(instance, entity);
+	}
 }

fabric-lifecycle-events-v1/src/main/resources/fabric-lifecycle-events-v1.mixins.json

@@ -6,6 +6,7 @@
     "ChunkHolderMixin",
     "ChunkMapMixin",
     "ChunkStatusTasksMixin",
+    "EntityTypeMixin",
     "LevelMixin",
     "LivingEntityMixin",
     "MinecraftServerMixin",

fabric-lifecycle-events-v1/src/testmod/java/net/fabricmc/fabric/test/event/lifecycle/ServerEntityLifecycleTests.java

@@ -50,6 +50,20 @@ public void onInitialize() {
 			}
 		});
 
+		ServerEntityEvents.ALLOW_ADD.register((entity, level) -> {
+			if (PRINT_SERVER_ENTITY_MESSAGES) {
+				logger.info("[SERVER] ALLOW FRESH LOAD {} IN {}", entity, level);
+			}
+
+			return true;
+		});
+
+		ServerEntityEvents.AFTER_FRESH_LOAD.register((entity, level) -> {
+			if (PRINT_SERVER_ENTITY_MESSAGES) {
+				logger.info("[SERVER] FRESH LOADED {} IN {}", entity, level);
+			}
+		});
+
 		ServerEntityEvents.ENTITY_UNLOAD.register((entity, level) -> {
 			this.serverEntities.remove(entity);