Show shooters who are blocked and friendlies that are blocking them i…

…n cyan and green respectively on the map and colonist bar.

mod-structure/Defs/UpdateFeatureDefs/UpdateFeatureDefs.xml

@@ -35,4 +35,10 @@
         <assemblyVersion>1.0.3</assemblyVersion>
         <content>In Rimworld 1.x, low skill shooters have a larger miss radius than higher skilled ones. This mod now accounts for that. This can be turned off in mod settings to revert to the fixed "one-tile miss radius" assumption, but that may result in friendly fire hits.</content>
     </HugsLib.UpdateFeatureDef>
+
+  <HugsLib.UpdateFeatureDef ParentName="UpdateFeatureBase">
+    <defName>AvoidFriendlyFire_1_0_4</defName>
+    <assemblyVersion>1.0.4</assemblyVersion>
+    <content>Shooters who are being blocked by friendlies (but would otherwise have a target) will now have their name on the map and in the Colonist Bar be written in cyan. Likewise, the first friendly blocking any shooter will show up in green.</content>
+  </HugsLib.UpdateFeatureDef>
 </Defs>

src/AvoidFriendlyFire/AvoidFriendlyFire.csproj

@@ -54,11 +54,15 @@
     <Compile Include="FireProperties.cs" />
     <Compile Include="Main.cs" />
     <Compile Include="Patches\AttackTargetFinder_BestAttackTarget_Patch.cs" />
+    <Compile Include="Patches\TooltipUtility_ShotCalculationTipString_Patch.cs" />
+    <Compile Include="Patches\PawnNameColorUtility_PawnNameColorOf_Patch.cs" />
     <Compile Include="Patches\Pawn_DraftController_GetGizmos_Patch.cs" />
     <Compile Include="Patches\Pawn_DraftController_set_Drafted_Patch.cs" />
     <Compile Include="Patches\Pawn_Kill_Patch.cs" />
     <Compile Include="Patches\Targeter_TargeterUpdate_Patch.cs" />
     <Compile Include="Patches\Verb_CanHitTargetFrom_Patch.cs" />
+    <Compile Include="PawnStatus.cs" />
+    <Compile Include="PawnStatusTracker.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
     <Compile Include="Resources.cs" />
   </ItemGroup>

src/AvoidFriendlyFire/FireManager.cs

@@ -8,13 +8,23 @@ namespace AvoidFriendlyFire
 {
     public class FireManager
     {
+        public bool SkipNextCheck;
+
         private readonly Dictionary<int, Dictionary<int, CachedFireCone>> _cachedFireCones
             = new Dictionary<int, Dictionary<int, CachedFireCone>>();
 
         private int _lastCleanupTick;
 
         public bool CanHitTargetSafely(FireProperties fireProperties)
         {
+            if (SkipNextCheck)
+            {
+                SkipNextCheck = false;
+                return true;
+            }
+
+            Main.Instance.PawnStatusTracker.Remove(fireProperties.Caster);
+
             HashSet<int> fireCone = GetOrCreatedCachedFireConeFor(fireProperties);
             if (fireCone == null)
                 return true;
@@ -52,6 +62,8 @@ public bool CanHitTargetSafely(FireProperties fireProperties)
                 if (IsPawnWearingUsefulShield(pawn))
                     continue;
 
+                Main.Instance.PawnStatusTracker.AddBlockedShooter(fireProperties.Caster, pawn);
+
                 return false;
             }
 

src/AvoidFriendlyFire/FireProperties.cs

@@ -20,7 +20,7 @@ public class FireProperties
     {
         public IntVec3 Target;
 
-        public Map CasterMap => _caster.Map;
+        public Map CasterMap => Caster.Map;
 
         public IntVec3 Origin;
 
@@ -30,17 +30,16 @@ public class FireProperties
 
         public int TargetIndex => CasterMap.cellIndices.CellToIndex(Target);
 
-
-        private readonly Thing _caster;
+        public Pawn Caster { get; }
 
         private readonly Verb _weaponVerb;
 
         public FireProperties(Pawn caster, IntVec3 target)
         {
             Target = target;
-            _caster = caster;
+            Caster = caster;
             _weaponVerb = GetEquippedWeaponVerb(caster);
-            Origin = _caster.Position;
+            Origin = Caster.Position;
         }
 
         public bool ArePointsVisibleAndValid()
@@ -78,13 +77,13 @@ public float GetAimOnTargetChance()
         {
             var distance = (Target - Origin).LengthHorizontal;
 
-            var factorFromShooterAndDist = ShotReport.HitFactorFromShooter(_caster, distance);
+            var factorFromShooterAndDist = ShotReport.HitFactorFromShooter(Caster, distance);
 
             var factorFromEquipment = _weaponVerb.verbProps.GetHitChanceFactor(
                 _weaponVerb.EquipmentSource, distance);
 
             var factorFromWeather = 1f;
-            if (!_caster.Position.Roofed(CasterMap) || !Target.Roofed(CasterMap))
+            if (!Caster.Position.Roofed(CasterMap) || !Target.Roofed(CasterMap))
             {
                 factorFromWeather = CasterMap.weatherManager.CurWeatherAccuracyMultiplier;
             }

src/AvoidFriendlyFire/Main.cs

@@ -20,6 +20,8 @@ public class Main : HugsLib.ModBase
 
         private ExtendedDataStorage _extendedDataStorage;
 
+        public PawnStatusTracker PawnStatusTracker { get; } = new PawnStatusTracker();
+
         private FireManager _fireManager;
 
         private SettingHandle<bool> _showOverlay;
@@ -47,6 +49,7 @@ public override void Tick(int currentTick)
             if (!IsModEnabled())
                 return;
             _fireManager?.RemoveExpiredCones(currentTick);
+            PawnStatusTracker.RemoveExpired();
         }
 
         public override void WorldLoaded()
@@ -64,6 +67,7 @@ public override void MapLoaded(Map map)
             base.MapLoaded(map);
             _fireManager = new FireManager();
             _fireConeOverlay = new FireConeOverlay();
+            PawnStatusTracker.Reset();
         }
 
         public override void DefsLoaded()

src/AvoidFriendlyFire/Patches/PawnNameColorUtility_PawnNameColorOf_Patch.cs

@@ -0,0 +1,27 @@
+using Harmony;
+using UnityEngine;
+using Verse;
+
+namespace AvoidFriendlyFire
+{
+    [HarmonyPatch(typeof(PawnNameColorUtility), "PawnNameColorOf")]
+    public class PawnNameColorUtility_PawnNameColorOf_Patch
+    {
+        public static void Postfix(ref Pawn pawn, ref Color __result)
+        {
+            if (!Main.Instance.IsModEnabled())
+                return;
+
+            if (Main.Instance.PawnStatusTracker.IsAShooter(pawn))
+            {
+                __result = Color.cyan;
+                return;
+            }
+
+            if (Main.Instance.PawnStatusTracker.IsABlocker(pawn))
+            {
+                __result = Color.green;
+            }
+        }
+    }
+}

src/AvoidFriendlyFire/Patches/Pawn_DraftController_set_Drafted_Patch.cs

@@ -8,9 +8,6 @@ public class Pawn_DraftController_set_Drafted_Patch
     {
         public static void Postfix(Pawn_DraftController __instance, bool value)
         {
-            if (!Main.Instance.ShouldEnableWhenUndrafted())
-                return;
-
             if (value)
                 return;
 
@@ -22,6 +19,11 @@ public static void Postfix(Pawn_DraftController __instance, bool value)
             if (!extendedDataStore.CanTrackPawn(pawn))
                 return;
 
+            Main.Instance.PawnStatusTracker.Remove(pawn);
+
+            if (!Main.Instance.ShouldEnableWhenUndrafted())
+                return;
+
             var pawnData = extendedDataStore.GetExtendedDataFor(pawn);
             pawnData.AvoidFriendlyFire = true;
         }

src/AvoidFriendlyFire/Patches/Pawn_Kill_Patch.cs

@@ -20,6 +20,8 @@ public static void Postfix(ref Pawn __instance)
                 if (!__instance.Dead)
                     return;
 
+                Main.Instance.PawnStatusTracker.KillOff(__instance);
+
                 Main.Instance.GetExtendedDataStorage().DeleteExtendedDataFor(__instance);
 
             }

src/AvoidFriendlyFire/Patches/TooltipUtility_ShotCalculationTipString_Patch.cs

@@ -0,0 +1,23 @@
+using Harmony;
+using Verse;
+
+namespace AvoidFriendlyFire
+{
+    [HarmonyPatch(typeof(TooltipUtility), "ShotCalculationTipString")]
+    public class TooltipUtility_ShotCalculationTipString_Patch
+    {
+        public static bool Prefix(ref Thing target)
+        {
+            if (Find.Selector.SingleSelectedThing is Pawn pawn)
+            {
+                if (pawn != target && pawn.equipment?.Primary != null
+                    && pawn.equipment.PrimaryEq.PrimaryVerb is Verb_LaunchProjectile)
+                {
+                    Main.Instance.GetFireManager().SkipNextCheck = true;
+                }
+            }
+
+            return true;
+        }
+    }
+}

src/AvoidFriendlyFire/PawnStatus.cs

@@ -0,0 +1,31 @@
+using System.Collections.Generic;
+using Verse;
+
+namespace AvoidFriendlyFire
+{
+    public class PawnStatus
+    {
+        public readonly Pawn Shooter;
+
+        public readonly Pawn Blocker;
+
+        private int _expiryTime;
+
+        public PawnStatus(Pawn shooter, Pawn blocker)
+        {
+            Shooter = shooter;
+            Blocker = blocker;
+            Refresh();
+        }
+
+        public void Refresh()
+        {
+            _expiryTime = Find.TickManager.TicksGame + 20;
+        }
+
+        public bool IsExpired()
+        {
+            return Find.TickManager.TicksGame >= _expiryTime;
+        }
+    }
+}

src/AvoidFriendlyFire/PawnStatusTracker.cs

@@ -0,0 +1,52 @@
+using System.Collections.Generic;
+using System.Linq;
+using Verse;
+
+namespace AvoidFriendlyFire
+{
+    public class PawnStatusTracker
+    {
+        private readonly List<PawnStatus> _shooters = new List<PawnStatus>();
+
+        public void AddBlockedShooter(Pawn shooter, Pawn blocker)
+        {
+            var existingEntry = _shooters.FirstOrDefault(ps => ps.Shooter == shooter);
+            if (existingEntry != null)
+            {
+                existingEntry.Refresh();
+                return;
+            }
+            _shooters.Add(new PawnStatus(shooter, blocker));
+        }
+
+        public bool IsAShooter(Pawn pawn)
+        {
+            return _shooters.Any(ps => ps.Shooter == pawn);
+        }
+
+        public bool IsABlocker(Pawn pawn)
+        {
+            return _shooters.Any(ps => ps.Blocker == pawn);
+        }
+
+        public void KillOff(Pawn pawn)
+        {
+            _shooters.RemoveAll(ps => ps.Shooter == pawn || ps.Blocker == pawn);
+        }
+
+        public void Remove(Pawn pawn)
+        {
+            _shooters.RemoveAll(ps => ps.Shooter == pawn);
+        }
+
+        public void RemoveExpired()
+        {
+            _shooters.RemoveAll(ps => ps.IsExpired());
+        }
+
+        public void Reset()
+        {
+            _shooters.Clear();
+        }
+    }
+}