add option to fix counting of items carried by non-humans

mod-structure/Languages/English/Keyed/IW_LanguageData.xml

@@ -97,4 +97,7 @@
 
   <IW.CountOutsideStockpiles>Count resources outside stockpiles</IW.CountOutsideStockpiles>
   <IW.CountOutsideStockpilesDesc>A bug in the vanilla game prevents "resource items" (e.g. meals and medicine) from being counted if they are not in stockpiles. Enable this option to fix that, or disable if you are using another mod to do that.</IW.CountOutsideStockpilesDesc>
+
+  <IW.CountCarriedByNonHumans>Count resources carried by non-humans</IW.CountCarriedByNonHumans>
+  <IW.CountCarriedByNonHumansDesc>A bug in the vanilla game prevents "resource items" (e.g. meals and medicine) from being counted if they are carried by non-humans, such as hauling animals or mechs, or animals leaving in a caravan. Enable this option to fix that.</IW.CountCarriedByNonHumansDesc>
 </LanguageData>

src/ImprovedWorkbenches/Detours/RecipeWorkerCounter_CountProducts_Detour.cs

@@ -41,6 +41,9 @@ public static void Postfix(ref RecipeWorkerCounter __instance, ref int __result,
                 __result += GetMatchingItemCountOutsideStockpiles(bill, productThingDef);
             }
 
+            if (Main.Instance.ShouldCountCarriedByNonHumans())
+                __result += CountInHonHumanInventory(bill.Map, __instance, bill, productThingDef);
+
             if (!bill.includeEquipped)
                 return;
 
@@ -69,6 +72,28 @@ private static int GetMatchingItemCountOutsideStockpiles(Bill bill, ThingDef pro
             return result;
         }
 
+        private static int CountInHonHumanInventory(Map billMap, RecipeWorkerCounter counter, Bill_Production bill, ThingDef productThingDef)
+        {
+            int count = 0;
+            foreach( var pawn in billMap.mapPawns.SpawnedColonyAnimals)
+                count += CountPawnThings(pawn, counter, bill, productThingDef);
+            foreach( var pawn in billMap.mapPawns.SpawnedColonyMechs)
+                count += CountPawnThings(pawn, counter, bill, productThingDef);
+            foreach( var pawn in billMap.mapPawns.SpawnedColonyMutantsPlayerControlled)
+                count += CountPawnThings(pawn, counter, bill, productThingDef);
+            return count;
+        }
+
+        private static IEnumerable<Pawn> AllPawnsToCount(MapPawns mapPawns)
+        {
+            if (!Main.Instance.ShouldCountCarriedByNonHumans())
+                return mapPawns.FreeColonistsSpawned;
+            return mapPawns.FreeColonistsSpawned
+                .Concat(mapPawns.SpawnedColonyAnimals)
+                .Concat(mapPawns.SpawnedColonyMechs)
+                .Concat(mapPawns.SpawnedColonyMutantsPlayerControlled);
+        }
+
         private static int CountAway(Map billMap, RecipeWorkerCounter counter, Bill_Production bill, ThingDef productThingDef)
         {
             int count = 0;
@@ -144,7 +169,7 @@ public static int CountAdditionalProducts(RecipeWorkerCounter counter, Bill_Prod
                     && !bill.limitToAllowedStuff)
                 {
                     count += map.resourceCounter.GetCount(def);
-                    foreach (Pawn pawn in map.mapPawns.FreeColonistsSpawned)
+                    foreach (Pawn pawn in AllPawnsToCount(map.mapPawns))
                     {
                         count += CountPawnThings(pawn, counter, bill, def, true);
                     }
@@ -173,7 +198,7 @@ public static int CountAdditionalProducts(RecipeWorkerCounter counter, Bill_Prod
                     if (!bill.includeEquipped)
                     {
                         //Still count Carried Things
-                        foreach (Pawn pawn in map.mapPawns.FreeColonistsSpawned)
+                        foreach (Pawn pawn in AllPawnsToCount(map.mapPawns))
                         {
                             count += CountPawnThings(pawn, counter, bill, def, true);
                         }
@@ -193,7 +218,7 @@ public static int CountAdditionalProducts(RecipeWorkerCounter counter, Bill_Prod
 
                 if (bill.includeEquipped)
                 {
-                    foreach (Pawn pawn in map.mapPawns.FreeColonistsSpawned)
+                    foreach (Pawn pawn in AllPawnsToCount(map.mapPawns))
                     {
                         count += CountPawnThings(pawn, counter, bill, def);
                     }

src/ImprovedWorkbenches/Detours/RecipeWorkerCounter_GetCarriedCount_Patch.cs

@@ -0,0 +1,48 @@
+using System.Collections.Generic;
+using System.Linq;
+using System.Net;
+using System.Reflection.Emit;
+using HarmonyLib;
+using RimWorld;
+using Verse;
+
+namespace ImprovedWorkbenches
+{
+    [HarmonyPatch(typeof(RecipeWorkerCounter))]
+    public static class RecipeWorkerCounter_GetCarriedCount_Patch
+    {
+        [HarmonyPatch(nameof(GetCarriedCount))]
+        public static IEnumerable<CodeInstruction> GetCarriedCount(IEnumerable<CodeInstruction> instructions)
+        {
+            bool found = false;
+            foreach( var instruction in instructions )
+            {
+                if( instruction.opcode == OpCodes.Callvirt
+                    && instruction.operand.ToString().EndsWith( "Verse.MapPawns::get_FreeColonistsSpawned()" ))
+                {
+                    found = true;
+                    yield return new CodeInstruction(OpCodes.Call,
+                        typeof(RecipeWorkerCounter_GetCarriedCount_Patch).GetMethod(nameof(GetCarriedCount_Hook)));
+                }
+                else
+                    yield return instruction;
+            }
+            if( !found )
+                Log.Error( "ImprovedWorkbenches: GetCarriedCount transpiller failed." );
+        }
+
+        private static List<Pawn> cachedResult = new List<Pawn>();
+
+        private static List<Pawn> GetCarriedCount_Hook(MapPawns mapPawns)
+        {
+            if (!Main.Instance.ShouldCountCarriedByNonHumans())
+                return mapPawns.FreeColonistsSpawned;
+            cachedResult.Clear();
+            cachedResult.AddRange(mapPawns.FreeColonistsSpawned);
+            cachedResult.AddRange(mapPawns.SpawnedColonyAnimals);
+            cachedResult.AddRange(mapPawns.SpawnedColonyMechs);
+            cachedResult.AddRange(mapPawns.SpawnedColonyMutantsPlayerControlled);
+            return cachedResult;
+        }
+    }
+}

src/ImprovedWorkbenches/Main.cs

@@ -52,6 +52,10 @@ public override void DefsLoaded()
                 "countOutsideStockpiles", "IW.CountOutsideStockpiles".Translate(),
                 "IW.CountOutsideStockpilesDesc".Translate(), true);
 
+            _countCarriedByNonHumans = Settings.GetHandle(
+                "countCarriedByNonHumans", "IW.CountCarriedByNonHumans".Translate(),
+                "IW.CountCarriedByNonHumansDesc".Translate(), true);
+
             // Integration with other mods
 
             IntegrateWithOutfitter();
@@ -160,6 +164,11 @@ public bool ShouldCountOutsideStockpiles()
             return _countOutsideStockpiles;
         }
 
+        public bool ShouldCountCarriedByNonHumans()
+        {
+            return _countCarriedByNonHumans;
+        }
+
         public ExtendedBillDataStorage GetExtendedBillDataStorage()
         {
             return _extendedBillDataStorage;
@@ -194,6 +203,8 @@ public int GetMaxBills()
 
         private SettingHandle<bool> _countOutsideStockpiles;
 
+        private SettingHandle<bool> _countCarriedByNonHumans;
+
         private ExtendedBillDataStorage _extendedBillDataStorage;
 
         // RImFactory support