Simplify generic ammo handling for autoloaders

CombatExtended-Continued · Dec 8, 2024 · a3d5d36 · a3d5d36
1 parent c2af44b
commit a3d5d36
Show file tree

Hide file tree

Showing 4 changed files with 102 additions and 152 deletions.
diff --git a/Source/CombatExtended/CombatExtended/Comps/CompAmmoListUser.cs b/Source/CombatExtended/CombatExtended/Comps/CompAmmoListUser.cs
@@ -21,9 +21,44 @@ public override List<AmmoSetDef> UsableAmmoSets
             {
                 if (usableAmmoSets.NullOrEmpty())
                 {
-                    foreach (var def in Props.additionalAmmoSets)
+                    var allowedAmmoSets = new List<AmmoSetDef> { Props.ammoSet };
+                    allowedAmmoSets.AddRange(Props.additionalAmmoSets);
+
+                    // Generalize allowed ammosets in generic ammo mode, unless they were already a generic ammoset.
+                    if (Controller.settings.GenericAmmo)
+                    {
+                        foreach (var def in allowedAmmoSets.Select(def => def.similarTo ?? def))
+                        {
+                            if (IsUsableAmmoSet(def))
+                            {
+                                usableAmmoSets.Add(def);
+                            }
+                        }
+
+                        return usableAmmoSets;
+                    }
+
+                    foreach (var def in allowedAmmoSets)
                     {
-                        if (!def.ammoTypes.First().ammo.menuHidden && !IsIdenticalToAny(def))
+                        var similarAmmoSets = DefDatabase<AmmoSetDef>.AllDefs
+                            .Where(other => other.similarTo == def)
+                            .ToList();
+
+                        // Allow specifying a generic ammoset directly, and allow its constituent ammosets in this case.
+                        if (similarAmmoSets.Any())
+                        {
+                            foreach (var similarAmmoSet in similarAmmoSets)
+                            {
+                                if (IsUsableAmmoSet(similarAmmoSet))
+                                {
+                                    usableAmmoSets.Add(similarAmmoSet);
+                                }
+                            }
+
+                            continue;
+                        }
+
+                        if (IsUsableAmmoSet(def))
                         {
                             usableAmmoSets.Add(def);
                         }
@@ -32,5 +67,51 @@ public override List<AmmoSetDef> UsableAmmoSets
                 return usableAmmoSets;
             }
         }
+
+
+        //Merge ammosets with identical ammo usage. I'm worried about its performance and I might want to caculate this during game start up.
+        private bool IsUsableAmmoSet(AmmoSetDef def)
+        {
+            if (def.ammoTypes.First().ammo.menuHidden)
+            {
+                return false;
+            }
+
+            if (usableAmmoSets.NullOrEmpty())
+            {
+                return true;
+            }
+
+            foreach (var v in usableAmmoSets)
+            {
+                if (IsIdenticalTo(v, def))
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+        private bool IsIdenticalTo(AmmoSetDef a, AmmoSetDef b)
+        {
+            if (a.ammoTypes.Count != b.ammoTypes.Count)
+            {
+                return false;
+            }
+            HashSet<AmmoDef> list = new HashSet<AmmoDef>();
+            foreach (var v in a.ammoTypes)
+            {
+                list.Add(v.ammo);
+            }
+            foreach (var n in b.ammoTypes)
+            {
+                if (!list.Contains(n.ammo))
+                {
+                    return false;
+                }
+            }
+            return true;
+        }
     }
 }
diff --git a/Source/CombatExtended/CombatExtended/Comps/CompAmmoListUserGeneric.cs b/Source/CombatExtended/CombatExtended/Comps/CompAmmoListUserGeneric.cs
diff --git a/Source/CombatExtended/CombatExtended/Comps/CompAmmoUserGeneric.cs b/Source/CombatExtended/CombatExtended/Comps/CompAmmoUserGeneric.cs
diff --git a/Source/CombatExtended/CombatExtended/Comps/CompVariableAmmoUser.cs b/Source/CombatExtended/CombatExtended/Comps/CompVariableAmmoUser.cs
@@ -28,45 +28,6 @@ public virtual List<AmmoSetDef> UsableAmmoSets
             }
         }
 
-        //Merge ammosets with identical ammo usage. I'm worried about its performance and I might want to caculate this during game start up.
-        protected bool IsIdenticalToAny(AmmoSetDef def)
-        {
-            if (usableAmmoSets.NullOrEmpty())
-            {
-                return false;
-            }
-            foreach (var v in usableAmmoSets)
-            {
-                if (IsIdenticalTo(v, def))
-                {
-                    return true;
-                }
-            }
-            return false;
-        }
-
-        protected bool IsIdenticalTo(AmmoSetDef a, AmmoSetDef b)
-        {
-            if (a.ammoTypes.Count != b.ammoTypes.Count)
-            {
-                return false;
-            }
-            HashSet<AmmoDef> list = new HashSet<AmmoDef>();
-            foreach (var v in a.ammoTypes)
-            {
-                list.Add(v.ammo);
-            }
-            foreach (var n in b.ammoTypes)
-            {
-                if (!list.Contains(n.ammo))
-                {
-                    return false;
-                }
-            }
-            return true;
-        }
-
-
         public override AmmoSetDef CurAmmoSet => SelectedAmmoSet;
 
         public override void PostExposeData()
@@ -119,29 +80,27 @@ public override IEnumerable<Gizmo> CompGetGizmosExtra()
             {
                 yield return gizmo;
             }
-            if (!Controller.settings.GenericAmmo)
+
+            Command_Action command_Action = new Command_Action();
+            command_Action.defaultLabel = "CE_SelectAmmoSet".Translate();
+            command_Action.defaultDesc = "CE_SelectAmmoSetDesc".Translate();
+            command_Action.icon = ContentFinder<Texture2D>.Get("UI/Buttons/Reload", reportFailure: false);
+            command_Action.action = delegate
             {
-                Command_Action command_Action = new Command_Action();
-                command_Action.defaultLabel = "CE_SelectAmmoSet".Translate();
-                command_Action.defaultDesc = "CE_SelectAmmoSetDesc".Translate();
-                command_Action.icon = ContentFinder<Texture2D>.Get("UI/Buttons/Reload", reportFailure: false);
-                command_Action.action = delegate
+                List<FloatMenuOption> list = new List<FloatMenuOption>();
+                foreach (AmmoSetDef caliber in UsableAmmoSets)
                 {
-                    List<FloatMenuOption> list = new List<FloatMenuOption>();
-                    foreach (AmmoSetDef caliber in UsableAmmoSets)
-                    {
-                        FloatMenuOption item = new FloatMenuOption(
-                           caliber.LabelCap,
-                           delegate { SyncedSelectAmmoSet(caliber); },
-                           caliber.ammoTypes.First().ammo,
-                           priority: MenuOptionPriority.Default,
-                           mouseoverGuiAction: delegate (Rect rect) { ContainedAmmoPopOut(rect, caliber); });
-                        list.Add(item);
-                    }
-                    Find.WindowStack.Add(new FloatMenu(list));
-                };
-                yield return command_Action;
-            }
+                    FloatMenuOption item = new FloatMenuOption(
+                        caliber.LabelCap,
+                        delegate { SyncedSelectAmmoSet(caliber); },
+                        caliber.ammoTypes.First().ammo,
+                        priority: MenuOptionPriority.Default,
+                        mouseoverGuiAction: delegate (Rect rect) { ContainedAmmoPopOut(rect, caliber); });
+                    list.Add(item);
+                }
+                Find.WindowStack.Add(new FloatMenu(list));
+            };
+            yield return command_Action;
         }
 
         public void ContainedAmmoPopOut(Rect rect, AmmoSetDef ammoSet)