From 97dc9187928c5af271d39ea3cc8e12706c18abd1 Mon Sep 17 00:00:00 2001
From: Simon <juliez@outlook.de>
Date: Sat, 18 Nov 2023 09:46:12 +0100
Subject: [PATCH] misc UnitRole changes

---
 megamek/src/megamek/common/UnitRole.java | 87 ++++++++++++++++--------
 1 file changed, 59 insertions(+), 28 deletions(-)

diff --git a/megamek/src/megamek/common/UnitRole.java b/megamek/src/megamek/common/UnitRole.java
index 11f3c99de73..9a9b6759319 100644
--- a/megamek/src/megamek/common/UnitRole.java
+++ b/megamek/src/megamek/common/UnitRole.java
@@ -24,6 +24,9 @@
 import org.apache.logging.log4j.LogManager;
 
 import java.util.Arrays;
+import java.util.function.Predicate;
+
+import static megamek.common.UnitRole.Availability.*;
 
 /**
  * Unit roles as defined by Alpha Strike Companion, used in formation building rules
@@ -32,49 +35,50 @@
  * @author Neoancient
  */
 public enum UnitRole {
-    /** This is the default role; given to units where the role definition is missing from their file. */
-    UNDETERMINED (false),
-
-    AMBUSHER (true),
-    BRAWLER (true),
-    JUGGERNAUT (true),
-    MISSILE_BOAT (true),
-    SCOUT (true),
-    SKIRMISHER (true),
-    SNIPER (true),
-    STRIKER (true),
-    ATTACK_FIGHTER (false),
-    DOGFIGHTER (false),
-    FAST_DOGFIGHTER (false),
-    FIRE_SUPPORT (false),
-    INTERCEPTOR (false),
-    TRANSPORT (false),
 
-    /** This role is used for some large aerospace units that intentionally have none of the combat roles. */
-	NONE (false);
+    /** This is the default role; given to units where the role definition is missing. */
+    UNDETERMINED (ALL),
 
-    private final boolean ground;
+    AMBUSHER (GROUND),
+    BRAWLER (GROUND),
+    JUGGERNAUT (GROUND),
+    MISSILE_BOAT (GROUND),
+    SCOUT (GROUND),
+    SKIRMISHER (GROUND),
+    SNIPER (GROUND),
+    STRIKER (GROUND),
 
-    UnitRole(boolean ground) {
-        this.ground = ground;
-    }
+    ATTACK_FIGHTER (AERO),
+    DOGFIGHTER (AERO),
+    FAST_DOGFIGHTER (AERO),
+    FIRE_SUPPORT (AERO),
+    INTERCEPTOR (AERO),
+    TRANSPORT (AERO),
 
-    public boolean isGroundRole() {
-        return ground;
-    }
+    /** This role is used for some large aerospace units that intentionally have none of the combat roles. */
+	NONE (AERO);
 
+    /** @return True when the given unit may use this role. Used in MML. */
+    @SuppressWarnings("unused")
     public boolean isAvailableTo(BTObject unit) {
-        return (ground == unit.isGround()) && !unit.isUnitGroup();
+        return availableTo.fits(unit) && !unit.isUnitGroup();
     }
 
     public boolean isAnyOf(UnitRole role, UnitRole... otherRoles) {
         return (this == role) || Arrays.stream(otherRoles).anyMatch(otherRole -> this == otherRole);
     }
 
+    /** @return True when this role is not UNDETERMINED. (Returns true for NONE.) */
     public boolean hasRole() {
         return this != UNDETERMINED;
     }
 
+    /**
+     * Parses the given string into the UnitRole if possible and returns it. If it can't parse the string,
+     * logs an error and returns UNDETERMINED. Does not return null.
+     *
+     * @return The UnitRole given as a string or UNDETERMINED.
+     */
     public static UnitRole parseRole(String role) {
         switch (role.toLowerCase()) {
             case "ambusher":
@@ -114,7 +118,7 @@ public static UnitRole parseRole(String role) {
             case "none":
             	return NONE;
             default:
-                LogManager.getLogger().error("Could not parse AS Role " + role);
+                LogManager.getLogger().error("Could not parse role " + role);
                 return UNDETERMINED;
         }
     }
@@ -182,6 +186,9 @@ public boolean qualifiesForRole(AlphaStrikeElement unit) {
      * @return			Boolean value indicating whether the unit meets the qualifications for this role.
      */
     public boolean qualifiesForRole(AlphaStrikeElement unit, double tolerance) {
+        if (!isAvailableTo(unit)) {
+            return false;
+        }
         double score = 0;
         int speed = unit.getPrimaryMovementValue();
         switch (this) {
@@ -408,4 +415,28 @@ public String toString() {
         }
         return sb.toString().trim();
     }
+
+    // PRIVATE
+
+    private final Availability availableTo;
+
+    UnitRole(Availability availableTo) {
+        this.availableTo = availableTo;
+    }
+
+    enum Availability {
+        GROUND(BTObject::isGround),
+        AERO(BTObject::isAerospace),
+        ALL(unit -> true);
+
+        private final Predicate<BTObject> fits;
+
+        Availability(Predicate<BTObject> fits) {
+            this.fits = fits;
+        }
+
+        boolean fits(BTObject unit) {
+            return fits.test(unit);
+        }
+    }
 }