From 251c5750ee363128f8de9477bdb32853cf9fda7e Mon Sep 17 00:00:00 2001
From: paulevsGitch <paulevs@yandex.ru>
Date: Wed, 13 Jan 2021 14:54:35 +0300
Subject: [PATCH] Entity cooldown map

---
 .../betterend/mixin/common/EntityMixin.java   | 20 ++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/main/java/ru/betterend/mixin/common/EntityMixin.java b/src/main/java/ru/betterend/mixin/common/EntityMixin.java
index 59027c53..fb278a65 100644
--- a/src/main/java/ru/betterend/mixin/common/EntityMixin.java
+++ b/src/main/java/ru/betterend/mixin/common/EntityMixin.java
@@ -1,5 +1,7 @@
 package ru.betterend.mixin.common;
 
+import java.util.Map;
+
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -8,6 +10,8 @@ import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
+import com.google.common.collect.Maps;
+
 import net.minecraft.entity.Entity;
 import net.minecraft.entity.EntityType;
 import net.minecraft.server.world.ServerWorld;
@@ -19,9 +23,8 @@ import ru.betterend.interfaces.TeleportingEntity;
 
 @Mixin(Entity.class)
 public abstract class EntityMixin implements TeleportingEntity {
-	
+	private static final Map<Entity, Long> COOLDOWN = Maps.newHashMap();
 	private BlockPos beExitPos;
-	private long beCooldown;
 	
 	@Shadow
 	public float yaw;
@@ -81,6 +84,7 @@ public abstract class EntityMixin implements TeleportingEntity {
 	@Inject(method = "getTeleportTarget", at = @At("HEAD"), cancellable = true)
 	protected void be_getTeleportTarget(ServerWorld destination, CallbackInfoReturnable<TeleportTarget> info) {
 		if (beExitPos != null) {
+			COOLDOWN.remove(beGetSelf());
 			info.setReturnValue(new TeleportTarget(new Vec3d(beExitPos.getX() + 0.5D, beExitPos.getY(), beExitPos.getZ() + 0.5D), getVelocity(), yaw, pitch));
 		}
 	}
@@ -88,18 +92,20 @@ public abstract class EntityMixin implements TeleportingEntity {
 	@Inject(method = "baseTick", at = @At("TAIL"))
 	public void be_baseTick(CallbackInfo info) {
 		if (hasCooldown()) {
-			this.beCooldown--;
+			Entity key = beGetSelf();
+			long value = COOLDOWN.getOrDefault(key, 0L) - 1;
+			COOLDOWN.put(key, value);
 		}
 	}
 	
 	@Override
 	public long beGetCooldown() {
-		return this.beCooldown;
+		return COOLDOWN.getOrDefault(beGetSelf(), 0L);
 	}
 
 	@Override
 	public void beSetCooldown(long time) {
-		this.beCooldown = time;
+		COOLDOWN.put(beGetSelf(), time);
 	}
 
 	@Override
@@ -111,4 +117,8 @@ public abstract class EntityMixin implements TeleportingEntity {
 	public BlockPos beGetExitPos() {
 		return this.beExitPos;
 	}
+	
+	private Entity beGetSelf() {
+		return (Entity) (Object) this;
+	}
 }