From c2f24e1567a3ac576c31c7e4fab0d56988e52a56 Mon Sep 17 00:00:00 2001
From: Strokkur24 <133226102+Strokkur424@users.noreply.github.com>
Date: Sat, 11 Jan 2025 22:29:16 +0100
Subject: [PATCH] Feat: Add 'with' methods to CommandSourceStack (#11868)

---
 .../command/brigadier/CommandSourceStack.java | 22 +++++++++++++
 ...y-type-tags-suggestions-in-selectors.patch |  4 +--
 .../commands/CommandSourceStack.java.patch    | 31 +++++++++++++++++++
 .../brigadier/PaperCommandSourceStack.java    | 16 +++++++---
 4 files changed, 67 insertions(+), 6 deletions(-)

diff --git a/paper-api/src/main/java/io/papermc/paper/command/brigadier/CommandSourceStack.java b/paper-api/src/main/java/io/papermc/paper/command/brigadier/CommandSourceStack.java
index ac6f5b754a..109c0a89e2 100644
--- a/paper-api/src/main/java/io/papermc/paper/command/brigadier/CommandSourceStack.java
+++ b/paper-api/src/main/java/io/papermc/paper/command/brigadier/CommandSourceStack.java
@@ -1,5 +1,7 @@
 package io.papermc.paper.command.brigadier;
 
+import com.mojang.brigadier.RedirectModifier;
+import com.mojang.brigadier.tree.CommandNode;
 import org.bukkit.Location;
 import org.bukkit.command.CommandSender;
 import org.bukkit.entity.Entity;
@@ -48,4 +50,24 @@ public interface CommandSourceStack {
      * @return entity that executes this command
      */
     @Nullable Entity getExecutor();
+
+    /**
+     * Creates a new CommandSourceStack object with a different location for redirecting commands to other nodes.
+     *
+     * @param location The location to create a new CommandSourceStack object with
+     * @return The newly created CommandSourceStack
+     * @see #getLocation()
+     * @see com.mojang.brigadier.builder.ArgumentBuilder#fork(CommandNode, RedirectModifier)
+     */
+    CommandSourceStack withLocation(Location location);
+
+    /**
+     * Creates a new CommandSourceStack object with a different executor for redirecting commands to other nodes.
+     *
+     * @param executor The executing entity to create a new CommandSourceStack object with
+     * @return The newly created CommandSourceStack
+     * @see #getExecutor()
+     * @see com.mojang.brigadier.builder.ArgumentBuilder#fork(CommandNode, RedirectModifier)
+     */
+    CommandSourceStack withExecutor(Entity executor);
 }
diff --git a/paper-server/patches/features/0009-Fix-entity-type-tags-suggestions-in-selectors.patch b/paper-server/patches/features/0009-Fix-entity-type-tags-suggestions-in-selectors.patch
index cd578f9f57..b587b90edc 100644
--- a/paper-server/patches/features/0009-Fix-entity-type-tags-suggestions-in-selectors.patch
+++ b/paper-server/patches/features/0009-Fix-entity-type-tags-suggestions-in-selectors.patch
@@ -10,10 +10,10 @@ when if this was fixed on the client, that wouldn't be needed.
 Mojira Issue: https://bugs.mojang.com/browse/MC-235045
 
 diff --git a/net/minecraft/commands/CommandSourceStack.java b/net/minecraft/commands/CommandSourceStack.java
-index 704a63890a06d793f8ac3452838917e7c7335232..75262c8c9eaecb4a88a94f4076d67119c67a97da 100644
+index cf923441da598637be74a5ffa4b4f948e01ff532..cbf32be9235921ebcaca88225120c2ca70a72771 100644
 --- a/net/minecraft/commands/CommandSourceStack.java
 +++ b/net/minecraft/commands/CommandSourceStack.java
-@@ -652,4 +652,20 @@ public class CommandSourceStack implements ExecutionCommandSource<CommandSourceS
+@@ -676,4 +676,20 @@ public class CommandSourceStack implements ExecutionCommandSource<CommandSourceS
          return this.source.getBukkitSender(this);
      }
      // CraftBukkit end
diff --git a/paper-server/patches/sources/net/minecraft/commands/CommandSourceStack.java.patch b/paper-server/patches/sources/net/minecraft/commands/CommandSourceStack.java.patch
index b8593758f7..bac6f4b64d 100644
--- a/paper-server/patches/sources/net/minecraft/commands/CommandSourceStack.java.patch
+++ b/paper-server/patches/sources/net/minecraft/commands/CommandSourceStack.java.patch
@@ -18,6 +18,37 @@
  
      public CommandSourceStack(
          CommandSource source,
+@@ -187,6 +_,30 @@
+                 this.chatMessageChainer
+             );
+     }
++    
++    // Paper start - Expose 'with' functions from the CommandSourceStack
++    @Override
++    public CommandSourceStack withLocation(org.bukkit.Location location) {
++        return this.getLocation().equals(location)
++            ? this
++            : new CommandSourceStack(
++            this.source,
++            new Vec3(location.x(), location.y(), location.z()),
++            new Vec2(location.getPitch(), location.getYaw()),
++            ((org.bukkit.craftbukkit.CraftWorld) location.getWorld()).getHandle(),
++            this.permissionLevel,
++            this.textName,
++            this.displayName,
++            this.server,
++            this.entity,
++            this.silent,
++            this.resultCallback,
++            this.anchor,
++            this.signingContext,
++            this.chatMessageChainer
++        );
++    }
++    // Paper end - Expose 'with' functions from the CommandSourceStack
+ 
+     public CommandSourceStack withRotation(Vec2 rotation) {
+         return this.rotation.equals(rotation)
 @@ -391,9 +_,44 @@
  
      @Override
diff --git a/paper-server/src/main/java/io/papermc/paper/command/brigadier/PaperCommandSourceStack.java b/paper-server/src/main/java/io/papermc/paper/command/brigadier/PaperCommandSourceStack.java
index 1b1642f306..0f7bbc9193 100644
--- a/paper-server/src/main/java/io/papermc/paper/command/brigadier/PaperCommandSourceStack.java
+++ b/paper-server/src/main/java/io/papermc/paper/command/brigadier/PaperCommandSourceStack.java
@@ -1,21 +1,23 @@
 package io.papermc.paper.command.brigadier;
 
 import com.destroystokyo.paper.brigadier.BukkitBrigadierCommandSource;
+import com.google.common.base.Preconditions;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.phys.Vec2;
 import net.minecraft.world.phys.Vec3;
 import org.bukkit.Location;
 import org.bukkit.command.CommandSender;
+import org.bukkit.craftbukkit.entity.CraftEntity;
 import org.bukkit.entity.Entity;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
+import org.jspecify.annotations.NonNull;
+import org.jspecify.annotations.Nullable;
 
 public interface PaperCommandSourceStack extends CommandSourceStack, BukkitBrigadierCommandSource {
 
     net.minecraft.commands.CommandSourceStack getHandle();
 
     @Override
-    default @NotNull Location getLocation() {
+    default @NonNull Location getLocation() {
         Vec2 rot = this.getHandle().getRotation();
         Vec3 pos = this.getHandle().getPosition();
         Level level = this.getHandle().getLevel();
@@ -24,7 +26,7 @@ public interface PaperCommandSourceStack extends CommandSourceStack, BukkitBriga
     }
 
     @Override
-    @NotNull
+    @NonNull
     default CommandSender getSender() {
         return this.getHandle().getBukkitSender();
     }
@@ -40,6 +42,12 @@ public interface PaperCommandSourceStack extends CommandSourceStack, BukkitBriga
         return nmsEntity.getBukkitEntity();
     }
 
+    @Override
+    default CommandSourceStack withExecutor(@NonNull Entity executor) {
+        Preconditions.checkNotNull(executor, "Executor cannot be null.");
+        return this.getHandle().withEntity(((CraftEntity) executor).getHandle());
+    }
+
     // OLD METHODS
     @Override
     default org.bukkit.entity.Entity getBukkitEntity() {