From b372ae961241a79e46179a6c50985b858d4c5bc7 Mon Sep 17 00:00:00 2001
From: Light_Quanta <zch22248@foxmail.com>
Date: Tue, 15 Apr 2025 22:06:10 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E5=8F=96=E6=89=80=E6=9C=89=E5=BC=B9?=
 =?UTF-8?q?=E8=8D=AF=E7=B1=BB=E5=9E=8B=E5=B1=9E=E6=80=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../capability/player/PlayerVariable.java     | 29 +++++++++----------
 .../event/LivingEventHandler.java             |  5 +++-
 .../superbwarfare/tools/AmmoType.java         | 20 ++-----------
 3 files changed, 20 insertions(+), 34 deletions(-)

diff --git a/src/main/java/com/atsuishio/superbwarfare/capability/player/PlayerVariable.java b/src/main/java/com/atsuishio/superbwarfare/capability/player/PlayerVariable.java
index 58cc1b5c3..9a7b3b033 100644
--- a/src/main/java/com/atsuishio/superbwarfare/capability/player/PlayerVariable.java
+++ b/src/main/java/com/atsuishio/superbwarfare/capability/player/PlayerVariable.java
@@ -16,16 +16,15 @@ import net.neoforged.neoforge.network.PacketDistributor;
 import org.jetbrains.annotations.NotNull;
 
 import javax.annotation.ParametersAreNonnullByDefault;
+import java.util.HashMap;
+import java.util.Map;
 
 
 @EventBusSubscriber(modid = Mod.MODID)
 public class PlayerVariable implements INBTSerializable<CompoundTag> {
     private PlayerVariable old = null;
-    public int rifleAmmo = 0;
-    public int handgunAmmo = 0;
-    public int shotgunAmmo = 0;
-    public int sniperAmmo = 0;
-    public int heavyAmmo = 0;
+
+    public Map<AmmoType, Integer> ammo = new HashMap<>();
     public boolean tacticalSprint = false;
     public boolean edit = false;
 
@@ -79,15 +78,13 @@ public class PlayerVariable implements INBTSerializable<CompoundTag> {
     public PlayerVariable copy() {
         var clone = new PlayerVariable();
 
-        clone.rifleAmmo = this.rifleAmmo;
-        clone.handgunAmmo = this.handgunAmmo;
-        clone.shotgunAmmo = this.shotgunAmmo;
-        clone.sniperAmmo = this.sniperAmmo;
-        clone.heavyAmmo = this.heavyAmmo;
+        for (var type : AmmoType.values()) {
+            type.set(clone, type.get(this));
+        }
+
         clone.edit = this.edit;
         clone.tacticalSprint = this.tacticalSprint;
 
-
         return clone;
     }
 
@@ -95,11 +92,11 @@ public class PlayerVariable implements INBTSerializable<CompoundTag> {
     public boolean equals(Object obj) {
         if (!(obj instanceof PlayerVariable other)) return false;
 
-        return rifleAmmo == other.rifleAmmo
-                && handgunAmmo == other.handgunAmmo
-                && shotgunAmmo == other.shotgunAmmo
-                && sniperAmmo == other.sniperAmmo
-                && heavyAmmo == other.heavyAmmo
+        for (var type : AmmoType.values()) {
+            if (type.get(this) != type.get(other)) return false;
+        }
+
+        return tacticalSprint == other.tacticalSprint
                 && edit == other.edit;
     }
 
diff --git a/src/main/java/com/atsuishio/superbwarfare/event/LivingEventHandler.java b/src/main/java/com/atsuishio/superbwarfare/event/LivingEventHandler.java
index ad1bbd240..83eeef41c 100644
--- a/src/main/java/com/atsuishio/superbwarfare/event/LivingEventHandler.java
+++ b/src/main/java/com/atsuishio/superbwarfare/event/LivingEventHandler.java
@@ -52,6 +52,7 @@ import net.neoforged.neoforge.network.PacketDistributor;
 
 import java.util.ArrayList;
 import java.util.Objects;
+import java.util.stream.Stream;
 
 @EventBusSubscriber
 public class LivingEventHandler {
@@ -731,7 +732,9 @@ public class LivingEventHandler {
         if (event.getEntity() instanceof Player player && !player.level().getLevelData().getGameRules().getBoolean(GameRules.RULE_KEEPINVENTORY)) {
             var cap = player.getData(ModAttachments.PLAYER_VARIABLE).watch();
 
-            boolean drop = cap.rifleAmmo + cap.handgunAmmo + cap.shotgunAmmo + cap.sniperAmmo + cap.heavyAmmo > 0;
+            boolean drop = Stream.of(AmmoType.values())
+                    .mapToInt(type -> type.get(cap))
+                    .sum() > 0;
 
             if (drop) {
                 var stack = new ItemStack(ModItems.AMMO_BOX.get());
diff --git a/src/main/java/com/atsuishio/superbwarfare/tools/AmmoType.java b/src/main/java/com/atsuishio/superbwarfare/tools/AmmoType.java
index 2d3ce73a1..a73b2f271 100644
--- a/src/main/java/com/atsuishio/superbwarfare/tools/AmmoType.java
+++ b/src/main/java/com/atsuishio/superbwarfare/tools/AmmoType.java
@@ -62,25 +62,13 @@ public enum AmmoType {
 
     // PlayerVariables
     public int get(PlayerVariable variable) {
-        return switch (this) {
-            case HANDGUN -> variable.handgunAmmo;
-            case RIFLE -> variable.rifleAmmo;
-            case SHOTGUN -> variable.shotgunAmmo;
-            case SNIPER -> variable.sniperAmmo;
-            case HEAVY -> variable.heavyAmmo;
-        };
+        return variable.ammo.get(this);
     }
 
     public void set(PlayerVariable variable, int count) {
         if (count < 0) count = 0;
 
-        switch (this) {
-            case HANDGUN -> variable.handgunAmmo = count;
-            case RIFLE -> variable.rifleAmmo = count;
-            case SHOTGUN -> variable.shotgunAmmo = count;
-            case SNIPER -> variable.sniperAmmo = count;
-            case HEAVY -> variable.heavyAmmo = count;
-        }
+        variable.ammo.put(this, count);
     }
 
     public void add(PlayerVariable variable, int count) {
@@ -90,9 +78,7 @@ public enum AmmoType {
 
     // Entity
     public int get(Entity entity) {
-        var cap = entity.getData(ModAttachments.PLAYER_VARIABLE);
-
-        return get(cap);
+        return get(entity.getData(ModAttachments.PLAYER_VARIABLE));
     }
 
     public void set(Entity entity, int count) {