1 files changed, 157 insertions, 0 deletions
diff --git a/src/main/java/org/chrisoft/trashyaddon/commands/CEntitySelector.java b/src/main/java/org/chrisoft/trashyaddon/commands/CEntitySelector.java
new file mode 100644
index 0000000..8bfa233
--- /dev/null
+++ b/src/main/java/org/chrisoft/trashyaddon/commands/CEntitySelector.java
@@ -0,0 +1,157 @@
+package org.chrisoft.trashyaddon.commands;
+
+import com.google.common.collect.Lists;
+import com.mojang.brigadier.exceptions.CommandSyntaxException;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.function.BiConsumer;
+import java.util.function.Function;
+import java.util.function.Predicate;
+
+import net.minecraft.client.MinecraftClient;
+import net.minecraft.client.world.ClientWorld;
+import net.minecraft.command.argument.EntityArgumentType;
+import net.minecraft.entity.Entity;
+import net.minecraft.entity.EntityType;
+import net.minecraft.predicate.NumberRange;
+import net.minecraft.text.Text;
+import net.minecraft.text.Texts;
+import net.minecraft.util.TypeFilter;
+import net.minecraft.util.function.LazyIterationConsumer;
+import net.minecraft.util.math.Box;
+import net.minecraft.util.math.Vec3d;
+import org.chrisoft.trashyaddon.mixin.ClientWorldMixin;
+import org.jetbrains.annotations.Nullable;
+
+public class CEntitySelector {
+   public static final int MAX_VALUE = Integer.MAX_VALUE;
+   public static final BiConsumer<Vec3d, List<? extends Entity>> ARBITRARY = (pos, entities) -> {
+   };
+   private static final TypeFilter<Entity, ?> PASSTHROUGH_FILTER = new TypeFilter<Entity, Entity>() {
+      public Entity downcast(Entity arg) {
+         return arg;
+      }
+
+      @Override
+      public Class<? extends Entity> getBaseClass() {
+         return Entity.class;
+      }
+   };
+   private final int limit;
+   private final boolean includesNonPlayers;
+   private final boolean localWorldOnly;
+   private final Predicate<Entity> basePredicate;
+   private final NumberRange.DoubleRange distance;
+   private final Function<Vec3d, Vec3d> positionOffset;
+   @Nullable
+   private final Box box;
+   private final BiConsumer<Vec3d, List<? extends Entity>> sorter;
+   private final boolean senderOnly;
+   @Nullable
+   private final String playerName;
+   @Nullable
+   private final UUID uuid;
+   private final TypeFilter<Entity, ?> entityFilter;
+   private final boolean usesAt;
+
+   public CEntitySelector(
+      int count,
+      boolean includesNonPlayers,
+      boolean localWorldOnly,
+      Predicate<Entity> basePredicate,
+      NumberRange.DoubleRange distance,
+      Function<Vec3d, Vec3d> positionOffset,
+      @Nullable Box box,
+      BiConsumer<Vec3d, List<? extends Entity>> sorter,
+      boolean senderOnly,
+      @Nullable String playerName,
+      @Nullable UUID uuid,
+      @Nullable EntityType<?> type,
+      boolean usesAt
+   ) {
+      this.limit = count;
+      this.includesNonPlayers = includesNonPlayers;
+      this.localWorldOnly = localWorldOnly;
+      this.basePredicate = basePredicate;
+      this.distance = distance;
+      this.positionOffset = positionOffset;
+      this.box = box;
+      this.sorter = sorter;
+      this.senderOnly = senderOnly;
+      this.playerName = playerName;
+      this.uuid = uuid;
+      this.entityFilter = (TypeFilter<Entity, ?>)(type == null ? PASSTHROUGH_FILTER : type);
+      this.usesAt = usesAt;
+   }
+
+   public int getLimit() {
+      return this.limit;
+   }
+
+   public boolean includesNonPlayers() {
+      return this.includesNonPlayers;
+   }
+
+   public boolean isSenderOnly() {
+      return this.senderOnly;
+   }
+
+   public boolean isLocalWorldOnly() {
+      return this.localWorldOnly;
+   }
+
+   public boolean usesAt() {
+      return this.usesAt;
+   }
+
+   private int getAppendLimit() {
+      return this.sorter == ARBITRARY ? this.limit : Integer.MAX_VALUE;
+   }
+
+   private Predicate<Entity> getPositionPredicate(Vec3d pos) {
+      Predicate<Entity> predicate = this.basePredicate;
+      if (this.box != null) {
+         Box lv = this.box.offset(pos);
+         predicate = predicate.and(entity -> lv.intersects(entity.getBoundingBox()));
+      }
+
+      if (!this.distance.isDummy()) {
+         predicate = predicate.and(entity -> this.distance.testSqrt(entity.squaredDistanceTo(pos)));
+      }
+
+      return predicate;
+   }
+   public List<? extends Entity> getClientSideEntityMatches(ClientWorld cw, Vec3d sourcePos) throws CommandSyntaxException {
+       Predicate<Entity> p = this.getPositionPredicate(sourcePos);
+       List<Entity> ret = new ArrayList<>();
+       int limit = this.getAppendLimit();
+       if (this.box != null)
+           cw.collectEntitiesByType(this.entityFilter, this.box.offset(sourcePos), p, ret, limit);
+       else {
+           ((ClientWorldMixin)cw).invokeGetEntityLookup().forEach(this.entityFilter, (e) -> {
+               if (p.test(e)) {
+                   ret.add(e);
+                   if (ret.size() >= limit)
+                       return LazyIterationConsumer.NextIteration.ABORT;
+               }
+               return LazyIterationConsumer.NextIteration.CONTINUE;
+           });
+       }
+       return ret;
+   }
+
+   private <T extends Entity> List<T> getEntities(Vec3d pos, List<T> entities) {
+      if (entities.size() > 1) {
+         this.sorter.accept(pos, entities);
+      }
+
+      return entities.subList(0, Math.min(this.limit, entities.size()));
+   }
+
+   public static Text getNames(List<? extends Entity> entities) {
+      return Texts.join(entities, Entity::getDisplayName);
+   }
+}