Fix VFE Ability issues with (primarily) async time (#431)

The issues are the most problematic with async time enabled. However, it is possible that world time or time on specific maps can "desync" from the rest - in which case those issues would appear even when async time is disabled.

Below is a more detailed list of changes. However, to simplify it - this should fix the bug where using a VFE ability would cause an incorrect cooldown to be applied, as well as cooldown changing when travelling between maps.

Patching utilities:
- Running SetupAsyncTime will also register extra fields needed to access async map comp and some of its fields
  - Previously it would only access MultiplayerGameComp to check if async was enabled or not
- Async time section now includes TimeSnapshot struct
  - It is the same struct as in MP with slight modifications, as it was easier to operate on it without reflection
- Added a timestamp fixer section
  - It uses a specific method and a type to register a delegate that MP will use to fix timestamps

Vanilla Expanded Framework:
- Fixed abilities getting incorrect cooldown duration when used
  - VFECore.Abilities.Ability:GetGizmo will use time snapshot to make gizmos use correct time in constructor, as this is where most gizmos determine if the gizmo is enabled/disabled
  - VFECore.Abilities.Command_Ability:GizmoOnGUIInt will use time snapshot to use correct time when drawing, as this is where the tooltip to display remaining cooldown duration is created
- Added a timestamp fixer for every ability, fixing issues where cooldown would change when travelling between maps
  - It currently is not possible to register implicit timestamp fixers, so each non-abstract ability type needs to have it registered

Vanilla Factions Expanded - Pirates:
- VFEPirates.CommandAbilityToggle:GizmoOnGUIInt also receives the time snapshot patch
  - This is needed as this type is a subtype of Verse.Command_Toggle, and not VFECore.Abilities.Command_Ability

Source/Mods/VanillaExpandedFramework.cs

@@ -4,7 +4,6 @@
 using System.Linq;
 using System.Reflection;
 using System.Reflection.Emit;
-using System.Runtime.CompilerServices;
 using HarmonyLib;
 using Multiplayer.API;
 using RimWorld;
@@ -283,10 +282,13 @@ private static void SyncSetIngredientCommand(SyncWorker sync, ref Command comman
         private static FastInvokeHandler abilityInitMethod;
         private static AccessTools.FieldRef<object, Thing> abilityHolderField;
         private static AccessTools.FieldRef<object, Pawn> abilityPawnField;
+        private static AccessTools.FieldRef<IExposable, int> abilityCooldownField;
         private static ISyncField abilityAutoCastField;
 
         private static void PatchAbilities()
         {
+            PatchingUtilities.SetupAsyncTime();
+
             // Comp holding ability
             // CompAbility
             compAbilitiesType = AccessTools.TypeByName("VFECore.Abilities.CompAbilities");
@@ -305,17 +307,54 @@ private static void PatchAbilities()
             abilityInitMethod = MethodInvoker.GetHandler(AccessTools.Method(type, "Init"));
             abilityHolderField = AccessTools.FieldRefAccess<Thing>(type, "holder");
             abilityPawnField = AccessTools.FieldRefAccess<Pawn>(type, "pawn");
+            abilityCooldownField = AccessTools.FieldRefAccess<int>(type, "cooldown");
             // There's another method taking LocalTargetInfo. Harmony grabs the one we need, but just in case specify the types to avoid ambiguity.
-            MP.RegisterSyncMethod(type, "StartAbilityJob", new SyncType[] { typeof(GlobalTargetInfo[]) });
+            MP.RegisterSyncMethod(type, "StartAbilityJob", [typeof(GlobalTargetInfo[])]);
             MP.RegisterSyncWorker<ITargetingSource>(SyncVEFAbility, type, true);
             abilityAutoCastField = MP.RegisterSyncField(type, "autoCast");
             MpCompat.harmony.Patch(AccessTools.DeclaredMethod(type, "DoAction"),
                 prefix: new HarmonyMethod(typeof(VanillaExpandedFramework), nameof(PreAbilityDoAction)),
                 postfix: new HarmonyMethod(typeof(VanillaExpandedFramework), nameof(PostAbilityDoAction)));
 
+            foreach (var target in type.AllSubclasses().Concat(type))
+            {
+                // Fix timestamp.
+                // We really could use implicit fixers, so we don't have to register one fixer per ability type.
+                if (!target.IsAbstract)
+                    PatchingUtilities.RegisterTimestampFixer(target, MpMethodUtil.MethodOf(FixAbilityTimestamp));
+
+                // We need to set this up in all subtypes that override GetGizmo, as we need to make sure
+                // that the call to Command_Ability constructor has the proper time. We could patch the
+                // constructor of all subtypes of Command_Ability, but the issue with that is that the
+                // gizmo are not guaranteed to be of that specific type. On top of that, their arguments
+                // (which we'd need to use to setup the time snapshot) may have different names, and may
+                // be in a different order. It's just simpler to patch this specific method for simplicity sake,
+                // instead of trying to patch every relevant constructor.
+                var method = AccessTools.DeclaredMethod(target, "GetGizmo");
+                if (method != null)
+                {
+                    MpCompat.harmony.Patch(method,
+                        prefix: new HarmonyMethod(MpMethodUtil.MethodOf(PreGetGizmo)),
+                        finalizer: new HarmonyMethod(MpMethodUtil.MethodOf(RestoreProperTimeSnapshot)));
+                }
+            }
+
             type = AccessTools.TypeByName("VFECore.CompShieldField");
             MpCompat.RegisterLambdaMethod(type, nameof(ThingComp.CompGetWornGizmosExtra), 0);
             MpCompat.RegisterLambdaMethod(type, "GetGizmos", 0, 2);
+
+            // Time snapshot fix for gizmo itself
+            type = AccessTools.TypeByName("VFECore.Abilities.Command_Ability");
+            foreach (var targetType in type.AllSubclasses().Concat(type))
+            {
+                var method = AccessTools.DeclaredMethod(targetType, nameof(Command.GizmoOnGUIInt));
+                if (method != null)
+                {
+                    MpCompat.harmony.Patch(method,
+                        prefix: new HarmonyMethod(MpMethodUtil.MethodOf(PreAbilityGizmoGui)),
+                        finalizer: new HarmonyMethod(MpMethodUtil.MethodOf(RestoreProperTimeSnapshot)));
+                }
+            }
         }
 
         private static void SyncVEFAbility(SyncWorker sync, ref ITargetingSource source)
@@ -394,6 +433,35 @@ private static void PostAbilityDoAction()
             MP.WatchEnd();
         }
 
+        // We need to set the time snapshot when constructing the gizmo since it's disabled in
+        // the constructor, meaning that it will be incorrectly disabled if we don't do it.
+        private static void PreGetGizmo(object __instance, out PatchingUtilities.TimeSnapshot? __state)
+            => __state = SetTemporaryTimeSnapshot(__instance);
+
+        // We need to set the time snapshot when drawing the gizmo GUI, as otherwise it'll
+        // display completely incorrect values for cooldown or will allow for the ability
+        // usage while it should still be on cooldown.
+        public static void PreAbilityGizmoGui(object ___ability, out PatchingUtilities.TimeSnapshot? __state)
+            => __state = SetTemporaryTimeSnapshot(___ability);
+
+        private static PatchingUtilities.TimeSnapshot? SetTemporaryTimeSnapshot(object ability)
+        {
+            if (!MP.IsInMultiplayer)
+                return null;
+
+            var target = abilityPawnField(ability) ?? abilityHolderField(ability);
+            if (target?.Map == null)
+                return null;
+
+            return PatchingUtilities.TimeSnapshot.GetAndSetFromMap(target.Map);
+        }
+
+        public static void RestoreProperTimeSnapshot(PatchingUtilities.TimeSnapshot? __state)
+            => __state?.Set();
+
+        private static ref int FixAbilityTimestamp(IExposable ability)
+            => ref abilityCooldownField(ability);
+
         #endregion
 
         #region Hireable Factions

Source/Mods/VanillaFactionsPirates.cs

@@ -110,12 +110,27 @@ public VanillaFactionsPirates(ModContentPack mod)
                 curseWorkerDisactivateMethod = AccessTools.Method(type, "Disactivate");
                 curseWorkerStartMethod = AccessTools.Method(type, "Start");
             }
-            
+
             // Flecks
             {
                 // Uses GenView.ShouldSpawnMotesAt, which is based on camera position
                 PatchingUtilities.PatchPushPopRand("VFEPirates.IncomingSmoker:ThrowBlackSmoke");
             }
+
+            // Ability cooldown
+            {
+                LongEventHandler.ExecuteWhenFinished(() =>
+                {
+                    // Setup snapshots just in case, but they should be setup by VFE compat.
+                    PatchingUtilities.SetupAsyncTime();
+                    // Just re-use the patch from VFE. It's not a subtype of Command_Ability
+                    // but Verse.Command_Toggle. However it still has the same fields and works
+                    // the same way with cooldowns as the VFE Ability class.
+                    MpCompat.harmony.Patch(AccessTools.DeclaredMethod("VFEPirates.CommandAbilityToggle:GizmoOnGUIInt"),
+                        prefix: new HarmonyMethod(MpMethodUtil.MethodOf(VanillaExpandedFramework.PreAbilityGizmoGui)),
+                        finalizer: new HarmonyMethod(MpMethodUtil.MethodOf(VanillaExpandedFramework.RestoreProperTimeSnapshot)));
+                });
+            }
         }
 
         private static void PreDoWindowContents(Window __instance, ref Color[] __state)

Source/PatchingUtilities.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
@@ -151,8 +152,8 @@ public static void PatchUnityRand(MethodBase method, bool patchPushPop = true)
         #endregion
 
         #region System RNG transpiler
-        private static readonly ConstructorInfo SystemRandConstructor = typeof(System.Random).GetConstructor(Type.EmptyTypes);
-        private static readonly ConstructorInfo SystemRandSeededConstructor = typeof(System.Random).GetConstructor(new[] { typeof(int) });
+        private static readonly ConstructorInfo SystemRandConstructor = typeof(Random).GetConstructor(Type.EmptyTypes);
+        private static readonly ConstructorInfo SystemRandSeededConstructor = typeof(Random).GetConstructor(new[] { typeof(int) });
         private static readonly ConstructorInfo RandRedirectorConstructor = typeof(RandRedirector).GetConstructor(Type.EmptyTypes);
         private static readonly ConstructorInfo RandRedirectorSeededConstructor = typeof(RandRedirector).GetConstructor(new[] { typeof(int) });
 
@@ -697,13 +698,23 @@ private static void PostTryGainMemory(Thought_Memory newThought)
         #region Async Time
 
         private static bool isAsyncTimeSetup = false;
-        private static bool isAsyncTimeSuccessful = false;
+        private static bool isAsyncTimeGameCompSuccessful = false;
+        private static bool isAsyncTimeMapCompSuccessful = false;
+        // Multiplayer
         private static AccessTools.FieldRef<object> multiplayerGameField;
+        // MultiplayerGame
         private static AccessTools.FieldRef<object, object> gameGameCompField;
+        // MultiplayerGameComp
         private static AccessTools.FieldRef<object, bool> gameCompAsyncTimeField;
+        // Extensions
+        private static FastInvokeHandler getAsyncTimeCompForMapMethod;
+        // AsyncTimeComp
+        private static AccessTools.FieldRef<object, int> asyncTimeMapTicksField;
+        private static AccessTools.FieldRef<object, TimeSlower> asyncTimeSlowerField;
+        private static AccessTools.FieldRef<object, TimeSpeed> asyncTimeTimeSpeedIntField;
 
         public static bool IsAsyncTime
-            => isAsyncTimeSuccessful &&
+            => isAsyncTimeGameCompSuccessful &&
                gameCompAsyncTimeField(gameGameCompField(multiplayerGameField()));
 
         public static void SetupAsyncTime()
@@ -714,18 +725,180 @@ public static void SetupAsyncTime()
 
             try
             {
+                // Multiplayer
                 multiplayerGameField = AccessTools.StaticFieldRefAccess<object>(
                     AccessTools.DeclaredField("Multiplayer.Client.Multiplayer:game"));
+            }
+            catch (Exception e)
+            {
+                Log.Error($"Encountered an exception while settings up core async time functionality:\n{e}");
+                // Nothing else will work here without this, just return early.
+                return;
+            }
+
+            try
+            {
+                // MultiplayerGame
                 gameGameCompField = AccessTools.FieldRefAccess<object>(
                     "Multiplayer.Client.MultiplayerGame:gameComp");
+                // MultiplayerGameComp
                 gameCompAsyncTimeField = AccessTools.FieldRefAccess<bool>(
                     "Multiplayer.Client.Comp.MultiplayerGameComp:asyncTime");
 
-                isAsyncTimeSuccessful = true;
+                isAsyncTimeGameCompSuccessful = true;
             }
             catch (Exception e)
             {
-                Log.Error($"Encountered an exception while settings up async time:\n{e}");
+                Log.Error($"Encountered an exception while settings up game async time:\n{e}");
+            }
+
+            try
+            {
+                getAsyncTimeCompForMapMethod = MethodInvoker.GetHandler(
+                    AccessTools.DeclaredMethod("Multiplayer.Client.Extensions:AsyncTime"));
+
+                var type = AccessTools.TypeByName("Multiplayer.Client.AsyncTimeComp");
+                asyncTimeMapTicksField = AccessTools.FieldRefAccess<int>(type, "mapTicks");
+                asyncTimeSlowerField = AccessTools.FieldRefAccess<TimeSlower>(type, "slower");
+                asyncTimeTimeSpeedIntField = AccessTools.FieldRefAccess<TimeSpeed>(type, "timeSpeedInt");
+
+                isAsyncTimeMapCompSuccessful = true;
+            }
+            catch (Exception e)
+            {
+                Log.Error($"Encountered an exception while settings up map async time:\n{e}");
+            }
+        }
+
+        // Taken from MP, GetAndSetFromMap was slightly modified.
+        // https://github.com/rwmt/Multiplayer/blob/master/Source/Client/AsyncTime/SetMapTime.cs#L166-L204
+        // We could access MP's struct and methods using reflection, but there were some issues
+        // in that approach - mainly performance wasn't perfect, as MethodInfo.Invoke call was required.
+        // Having an identical struct in MP Compat won't cause issues, as there's nothing to conflict with MP.
+        // On top of that - if the struct ever gets renamed or moved to a different namespace in MP, it won't
+        // affect MP Compat. However, in case there's any logic changes in MP - they won't be reflected here.
+        // Ideally, we'll make something like this in the MP API.
+        public struct TimeSnapshot
+        {
+            public int ticks;
+            public TimeSpeed speed;
+            public TimeSlower slower;
+
+            public void Set()
+            {
+                Find.TickManager.ticksGameInt = ticks;
+                Find.TickManager.slower = slower;
+                Find.TickManager.curTimeSpeed = speed;
+            }
+
+            public static TimeSnapshot Current()
+            {
+                return new TimeSnapshot
+                {
+                    ticks = Find.TickManager.ticksGameInt,
+                    speed = Find.TickManager.curTimeSpeed,
+                    slower = Find.TickManager.slower
+                };
+            }
+
+            public static TimeSnapshot? GetAndSetFromMap(Map map)
+            {
+                if (map == null) return null;
+                if (!isAsyncTimeMapCompSuccessful) return null;
+
+                var prev = Current();
+
+                var tickManager = Find.TickManager;
+                var mapComp = getAsyncTimeCompForMapMethod(null, map);
+
+                tickManager.ticksGameInt = asyncTimeMapTicksField(mapComp);
+                tickManager.slower = asyncTimeSlowerField(mapComp);
+                tickManager.CurTimeSpeed = asyncTimeTimeSpeedIntField(mapComp);
+
+                return prev;
+            }
+        }
+
+        #endregion
+
+        #region Timestamp Fixer
+
+        private static bool isTimestampFixerInitialized = false;
+        private static Type timestampFixerDelegateType;
+        private static Type timestampFixerListType;
+        private static AccessTools.FieldRef<IDictionary> timestampFieldsDictionaryField;
+
+        public static void RegisterTimestampFixer(Type type, MethodInfo timestampFixerMethod)
+        {
+            if (type == null || timestampFixerMethod == null)
+            {
+                Log.Error($"Trying to register timestamp fixer failed - value null. Type={type.ToStringSafe()}, Method={timestampFixerMethod.ToStringSafe()}");
+                return;
+            }
+
+            InitializeTimestampFixer();
+
+            // Initialize call will display proper errors if needed
+            if (timestampFixerDelegateType == null || timestampFixerListType == null || timestampFieldsDictionaryField == null)
+                return;
+
+            try
+            {
+                var dict = timestampFieldsDictionaryField();
+                IList list;
+                // If the dictionary already contains list of timestamp fixers for
+                // a given type, use that list and add another one to it.
+                if (dict.Contains(type))
+                    list = (IList)dict[type];
+                // If needed, create a new list of timestamp fixers for a given type.
+                else
+                    dict[type] = list = (IList)Activator.CreateInstance(timestampFixerListType);
+
+                // Create a FieldGetter<IExposable> delegate using the provided method
+                list.Add(Delegate.CreateDelegate(timestampFixerDelegateType, timestampFixerMethod));
+            }
+            catch (Exception e)
+            {
+                Log.Error($"Trying to initialize timestamp fixer failed, exception caught:\n{e}");
+            }
+        }
+
+        public static void InitializeTimestampFixer()
+        {
+            if (isTimestampFixerInitialized)
+                return;
+            isTimestampFixerInitialized = true;
+
+            try
+            {
+                var type = AccessTools.TypeByName("Multiplayer.Client.Patches.TimestampFixer");
+                if (type == null)
+                {
+                    Log.Error("Trying to initialize timestamp fixer failed, could not find TimestampFixer type.");
+                    return;
+                }
+
+                // Get the type of the delegate. We need to specify `1 as it's a generic delegate.
+                var delType = AccessTools.Inner(type, "FieldGetter`1");
+                if (delType == null)
+                {
+                    Log.Error("Trying to initialize timestamp fixer failed, could not find FieldGetter inner type.");
+                    return;
+                }
+
+                timestampFieldsDictionaryField = AccessTools.StaticFieldRefAccess<IDictionary>(
+                    AccessTools.DeclaredField(type, "timestampFields"));
+                // The list only accepts FieldGetter<IExposable>
+                timestampFixerDelegateType = delType.MakeGenericType(typeof(IExposable));
+                timestampFixerListType = typeof(List<>).MakeGenericType(timestampFixerDelegateType);
+            }
+            catch (Exception e)
+            {
+                Log.Error($"Trying to initialize timestamp fixer failed, exception caught:\n{e}");
+
+                timestampFixerDelegateType = null;
+                timestampFixerListType = null;
+                timestampFieldsDictionaryField = null;
             }
         }