From 8b53764e92173704ab6ddc7b17e468837e61d58f Mon Sep 17 00:00:00 2001
From: paulevsGitch <paulevs@yandex.ru>
Date: Sat, 26 Dec 2020 20:50:39 +0300
Subject: [PATCH] Spawn restriction changes

---
 .../ru/betterend/entity/EntityDragonfly.java  |  9 ++++
 .../ru/betterend/registry/EndEntities.java    |  1 +
 .../java/ru/betterend/util/SpawnHelper.java   | 41 +++----------------
 3 files changed, 15 insertions(+), 36 deletions(-)

diff --git a/src/main/java/ru/betterend/entity/EntityDragonfly.java b/src/main/java/ru/betterend/entity/EntityDragonfly.java
index 7ab75ce0..f808e0d1 100644
--- a/src/main/java/ru/betterend/entity/EntityDragonfly.java
+++ b/src/main/java/ru/betterend/entity/EntityDragonfly.java
@@ -1,11 +1,13 @@
 package ru.betterend.entity;
 
 import java.util.EnumSet;
+import java.util.Random;
 
 import net.minecraft.block.BlockState;
 import net.minecraft.entity.EntityType;
 import net.minecraft.entity.Flutterer;
 import net.minecraft.entity.LivingEntity;
+import net.minecraft.entity.SpawnReason;
 import net.minecraft.entity.ai.TargetFinder;
 import net.minecraft.entity.ai.control.FlightMoveControl;
 import net.minecraft.entity.ai.control.LookControl;
@@ -26,6 +28,8 @@ import net.minecraft.server.world.ServerWorld;
 import net.minecraft.sound.SoundEvent;
 import net.minecraft.util.math.BlockPos;
 import net.minecraft.util.math.Vec3d;
+import net.minecraft.world.Heightmap.Type;
+import net.minecraft.world.ServerWorldAccess;
 import net.minecraft.world.World;
 import net.minecraft.world.WorldView;
 import ru.betterend.registry.EndEntities;
@@ -187,4 +191,9 @@ public class EntityDragonfly extends AnimalEntity implements Flutterer {
 	public PassiveEntity createChild(ServerWorld world, PassiveEntity entity) {
 		return EndEntities.DRAGONFLY.create(world);
 	}
+	
+	public static boolean canSpawn(EntityType<EntityDragonfly> type, ServerWorldAccess world, SpawnReason spawnReason, BlockPos pos, Random random) {
+		int y = world.getChunk(pos).sampleHeightmap(Type.WORLD_SURFACE, pos.getX() & 15, pos.getY() & 15);
+		return y > 0 && pos.getY() >= y;
+	}
 }
diff --git a/src/main/java/ru/betterend/registry/EndEntities.java b/src/main/java/ru/betterend/registry/EndEntities.java
index 1f685f43..16924e6d 100644
--- a/src/main/java/ru/betterend/registry/EndEntities.java
+++ b/src/main/java/ru/betterend/registry/EndEntities.java
@@ -27,6 +27,7 @@ public class EndEntities {
 	public static final EntityType<EntityCubozoa> CUBOZOA = register("cubozoa", SpawnGroup.WATER_AMBIENT, 0.6F, 1F, EntityCubozoa::new, EntityCubozoa.createMobAttributes(), true, MHelper.color(151, 77, 181), MHelper.color(93, 176, 238));
 	
 	public static void register() {
+		SpawnHelper.restrictionAir(DRAGONFLY, EntityDragonfly::canSpawn);
 		SpawnHelper.restrictionLand(END_SLIME, EntityEndSlime::canSpawn);
 		SpawnHelper.restrictionWater(END_FISH, EntityEndFish::canSpawn);
 		SpawnHelper.restrictionLand(SHADOW_WALKER, EntityShadowWalker::canSpawn);
diff --git a/src/main/java/ru/betterend/util/SpawnHelper.java b/src/main/java/ru/betterend/util/SpawnHelper.java
index 71356f7f..a3d9ba15 100644
--- a/src/main/java/ru/betterend/util/SpawnHelper.java
+++ b/src/main/java/ru/betterend/util/SpawnHelper.java
@@ -1,53 +1,22 @@
 package ru.betterend.util;
 
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
-
+import net.fabricmc.fabric.mixin.object.builder.SpawnRestrictionAccessor;
 import net.minecraft.entity.EntityType;
-import net.minecraft.entity.SpawnRestriction;
 import net.minecraft.entity.SpawnRestriction.Location;
 import net.minecraft.entity.SpawnRestriction.SpawnPredicate;
 import net.minecraft.entity.mob.MobEntity;
 import net.minecraft.world.Heightmap.Type;
-import ru.betterend.BetterEnd;
 
 public class SpawnHelper {
-	private static Method regRestriction;
-	
-	public static <T extends MobEntity> void restriction(EntityType<T> entity, Location location, Type heughtmapType, SpawnPredicate<T> predicate) {
-		if (regRestriction != null) {
-			try {
-				regRestriction.invoke(null, entity, location, heughtmapType, predicate);
-			}
-			catch (IllegalAccessException | IllegalArgumentException | InvocationTargetException e) {
-				BetterEnd.LOGGER.error(e.getMessage());
-			}
-		}
-		else {
-			BetterEnd.LOGGER.error("Unable to register spawn restriction, variable is not handled");
-		}
+	public static <T extends MobEntity> void restrictionAir(EntityType<T> entity, SpawnPredicate<T> predicate) {
+		SpawnRestrictionAccessor.callRegister(entity, Location.NO_RESTRICTIONS, Type.MOTION_BLOCKING, predicate);
 	}
 	
 	public static <T extends MobEntity> void restrictionLand(EntityType<T> entity, SpawnPredicate<T> predicate) {
-		restriction(entity, Location.ON_GROUND, Type.MOTION_BLOCKING, predicate);
+		SpawnRestrictionAccessor.callRegister(entity, Location.ON_GROUND, Type.MOTION_BLOCKING, predicate);
 	}
 	
 	public static <T extends MobEntity> void restrictionWater(EntityType<T> entity, SpawnPredicate<T> predicate) {
-		restriction(entity, Location.IN_WATER, Type.MOTION_BLOCKING, predicate);
-	}
-	
-	static {
-		try {
-			for (Method method: SpawnRestriction.class.getDeclaredMethods()) {
-				if (method.getParameterCount() == 4) {
-					regRestriction = method;
-					regRestriction.setAccessible(true);
-					break;
-				}
-			}
-		} 
-		catch (SecurityException e) {
-			BetterEnd.LOGGER.error(e.getMessage());
-		}
+		SpawnRestrictionAccessor.callRegister(entity, Location.IN_WATER, Type.MOTION_BLOCKING, predicate);
 	}
 }