minecraft-src/net/minecraft/world/entity/Leashable.java

package net.minecraft.world.entity;

import com.mojang.datafixers.util.Either;
import com.mojang.serialization.Codec;
import java.util.Objects;
import java.util.Optional;
import java.util.UUID;
import net.minecraft.core.BlockPos;
import net.minecraft.core.UUIDUtil;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.network.protocol.game.ClientboundSetEntityLinkPacket;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.entity.decoration.LeashFenceKnotEntity;
import net.minecraft.world.item.Items;
import net.minecraft.world.level.GameRules;
import org.jetbrains.annotations.Nullable;

public interface Leashable {
	String LEASH_TAG = "leash";
	double LEASH_TOO_FAR_DIST = 10.0;
	double LEASH_ELASTIC_DIST = 6.0;

	@Nullable
	Leashable.LeashData getLeashData();

	void setLeashData(@Nullable Leashable.LeashData leashData);

	default boolean isLeashed() {
		return this.getLeashData() != null && this.getLeashData().leashHolder != null;
	}

	default boolean mayBeLeashed() {
		return this.getLeashData() != null;
	}

	default boolean canHaveALeashAttachedToIt() {
		return this.canBeLeashed() && !this.isLeashed();
	}

	default boolean canBeLeashed() {
		return true;
	}

	default void setDelayedLeashHolderId(int delayedLeashHolderId) {
		this.setLeashData(new Leashable.LeashData(delayedLeashHolderId));
		dropLeash((Entity & Leashable)this, false, false);
	}

	default void readLeashData(CompoundTag tag) {
		Leashable.LeashData leashData = (Leashable.LeashData)tag.read("leash", Leashable.LeashData.CODEC).orElse(null);
		if (this.getLeashData() != null && leashData == null) {
			this.removeLeash();
		}

		this.setLeashData(leashData);
	}

	default void writeLeashData(CompoundTag tag, @Nullable Leashable.LeashData leashData) {
		tag.storeNullable("leash", Leashable.LeashData.CODEC, leashData);
	}

	private static <E extends Entity & Leashable> void restoreLeashFromSave(E entity, Leashable.LeashData leashData) {
		if (leashData.delayedLeashInfo != null && entity.level() instanceof ServerLevel serverLevel) {
			Optional<UUID> optional = leashData.delayedLeashInfo.left();
			Optional<BlockPos> optional2 = leashData.delayedLeashInfo.right();
			if (optional.isPresent()) {
				Entity entity2 = serverLevel.getEntity((UUID)optional.get());
				if (entity2 != null) {
					setLeashedTo(entity, entity2, true);
					return;
				}
			} else if (optional2.isPresent()) {
				setLeashedTo(entity, LeashFenceKnotEntity.getOrCreateKnot(serverLevel, (BlockPos)optional2.get()), true);
				return;
			}

			if (entity.tickCount > 100) {
				entity.spawnAtLocation(serverLevel, Items.LEAD);
				entity.setLeashData(null);
			}
		}
	}

	default void dropLeash() {
		dropLeash((Entity & Leashable)this, true, true);
	}

	default void removeLeash() {
		dropLeash((Entity & Leashable)this, true, false);
	}

	default void onLeashRemoved() {
	}

	private static <E extends Entity & Leashable> void dropLeash(E entity, boolean broadcastPacket, boolean dropItem) {
		Leashable.LeashData leashData = entity.getLeashData();
		if (leashData != null && leashData.leashHolder != null) {
			entity.setLeashData(null);
			entity.onLeashRemoved();
			if (entity.level() instanceof ServerLevel serverLevel) {
				if (dropItem) {
					entity.spawnAtLocation(serverLevel, Items.LEAD);
				}

				if (broadcastPacket) {
					serverLevel.getChunkSource().broadcast(entity, new ClientboundSetEntityLinkPacket(entity, null));
				}
			}
		}
	}

	static <E extends Entity & Leashable> void tickLeash(ServerLevel level, E entity) {
		Leashable.LeashData leashData = entity.getLeashData();
		if (leashData != null && leashData.delayedLeashInfo != null) {
			restoreLeashFromSave(entity, leashData);
		}

		if (leashData != null && leashData.leashHolder != null) {
			if (!entity.isAlive() || !leashData.leashHolder.isAlive()) {
				if (level.getGameRules().getBoolean(GameRules.RULE_DOENTITYDROPS)) {
					entity.dropLeash();
				} else {
					entity.removeLeash();
				}
			}

			Entity entity2 = entity.getLeashHolder();
			if (entity2 != null && entity2.level() == entity.level()) {
				float f = entity.distanceTo(entity2);
				if (!entity.handleLeashAtDistance(entity2, f)) {
					return;
				}

				if (f > 10.0) {
					entity.leashTooFarBehaviour();
				} else if (f > 6.0) {
					entity.elasticRangeLeashBehaviour(entity2, f);
					entity.checkSlowFallDistance();
				} else {
					entity.closeRangeLeashBehaviour(entity2);
				}
			}
		}
	}

	default boolean handleLeashAtDistance(Entity leashHolder, float distance) {
		return true;
	}

	default void leashTooFarBehaviour() {
		this.dropLeash();
	}

	default void closeRangeLeashBehaviour(Entity entity) {
	}

	default void elasticRangeLeashBehaviour(Entity leashHolder, float distance) {
		legacyElasticRangeLeashBehaviour((Entity & Leashable)this, leashHolder, distance);
	}

	private static <E extends Entity & Leashable> void legacyElasticRangeLeashBehaviour(E entity, Entity leashHolder, float distance) {
		double d = (leashHolder.getX() - entity.getX()) / distance;
		double e = (leashHolder.getY() - entity.getY()) / distance;
		double f = (leashHolder.getZ() - entity.getZ()) / distance;
		entity.setDeltaMovement(entity.getDeltaMovement().add(Math.copySign(d * d * 0.4, d), Math.copySign(e * e * 0.4, e), Math.copySign(f * f * 0.4, f)));
	}

	default void setLeashedTo(Entity leashHolder, boolean broadcastPacket) {
		setLeashedTo((Entity & Leashable)this, leashHolder, broadcastPacket);
	}

	private static <E extends Entity & Leashable> void setLeashedTo(E entity, Entity leashHolder, boolean broadcastPacket) {
		Leashable.LeashData leashData = entity.getLeashData();
		if (leashData == null) {
			leashData = new Leashable.LeashData(leashHolder);
			entity.setLeashData(leashData);
		} else {
			leashData.setLeashHolder(leashHolder);
		}

		if (broadcastPacket && entity.level() instanceof ServerLevel serverLevel) {
			serverLevel.getChunkSource().broadcast(entity, new ClientboundSetEntityLinkPacket(entity, leashHolder));
		}

		if (entity.isPassenger()) {
			entity.stopRiding();
		}
	}

	@Nullable
	default Entity getLeashHolder() {
		return getLeashHolder((Entity & Leashable)this);
	}

	@Nullable
	private static <E extends Entity & Leashable> Entity getLeashHolder(E entity) {
		Leashable.LeashData leashData = entity.getLeashData();
		if (leashData == null) {
			return null;
		} else {
			if (leashData.delayedLeashHolderId != 0 && entity.level().isClientSide) {
				Entity var3 = entity.level().getEntity(leashData.delayedLeashHolderId);
				if (var3 instanceof Entity) {
					leashData.setLeashHolder(var3);
				}
			}

			return leashData.leashHolder;
		}
	}

	public static final class LeashData {
		public static final Codec<Leashable.LeashData> CODEC = Codec.xor(UUIDUtil.CODEC.fieldOf("UUID").codec(), BlockPos.CODEC)
			.xmap(
				Leashable.LeashData::new,
				leashData -> {
					if (leashData.leashHolder instanceof LeashFenceKnotEntity leashFenceKnotEntity) {
						return Either.right(leashFenceKnotEntity.getPos());
					} else {
						return leashData.leashHolder != null
							? Either.left(leashData.leashHolder.getUUID())
							: (Either)Objects.requireNonNull(leashData.delayedLeashInfo, "Invalid LeashData had no attachment");
					}
				}
			);
		int delayedLeashHolderId;
		@Nullable
		public Entity leashHolder;
		@Nullable
		public Either<UUID, BlockPos> delayedLeashInfo;

		private LeashData(Either<UUID, BlockPos> delayedLeashInfo) {
			this.delayedLeashInfo = delayedLeashInfo;
		}

		LeashData(Entity leashHolder) {
			this.leashHolder = leashHolder;
		}

		LeashData(int delayedLeashInfoId) {
			this.delayedLeashHolderId = delayedLeashInfoId;
		}

		public void setLeashHolder(Entity leashHolder) {
			this.leashHolder = leashHolder;
			this.delayedLeashInfo = null;
			this.delayedLeashHolderId = 0;
		}
	}
}