Skip to content

IEnumerator Patching

BlazingTwist edited this page Dec 19, 2022 · 1 revision

Sometimes you come across Methods that return 'IEnumerator', patching these can be tricky without BTHU.

For example, imagine a Game has a Poison mechanic. While you're poisoned, you take damage every 10 seconds.

public class Player : MonoBehaviour {
    public bool Alive { get; private set; }
    public int Health { get; private set; }
    public int PoisonStacks { get; private set; }

    public void TakeDamage(int damage) {
        Health -= damage;
        if (Health < 0) {
            Alive = false;
        }
    }
    
    public IEnumerator PoisonCoroutine() {
        while (Alive) {
            if (PoisonStacks > 0) {
                TakeDamage(PoisonStacks);
                PoisonStacks--;
            }
            yield return new WaitForSeconds(10f);
        }
    }
}

We'd like to change the amount of poison damage the player takes, let's try altering the 'PoisonCoroutine' Method.

Notice that 'PoisonCoroutine' returns an 'IEnumerator'.
Do you notice something else?
It may yield (i.e. "return") more than one value, in fact, we can't even tell how many values will be yielding at compile time.
If you don't find that weird, you already know everything I'm about to tell you.
But if you do find it weird, read on. Let's explore how the compiler does this magic.

While we cannot return an unknown number of values from our Method, we can create a Method that returns something else every time it is called.
The compiler does exactly that. It generates a new class in 'Player' that implements 'IEnumerator'.
The resulting Player class looks like this:

public class Player : MonoBehaviour {

    public bool Alive { get; private set; }
    public int Health { get; private set; }
    public int PoisonStacks { get; private set; }

    public void TakeDamage(int damage) {
        Health -= damage;
        if (Health < 0) {
            Alive = false;
        }
    }

    public IEnumerator PoisonCoroutine() {
        PoisonCoroutine_IEnumerator coroutine = new PoisonCoroutine_IEnumerator(0);
        coroutine.player = this;
        return coroutine;
    }

    private class PoisonCoroutine_IEnumerator : IEnumerator<object> {
        public Player player;
        public object Current { get; private set; }

        private int state;

        public PoisonCoroutine_IEnumerator(int state) {
            this.state = state;
        }

        public bool MoveNext() {
            if (state < 0) {
                return false;
            }

            if (!player.Alive) {
                // if the player died, stop this coroutine.
                state = -1;
                return false;
            }

            if (player.PoisonStacks > 0) {
                player.TakeDamage(player.PoisonStacks);
                player.PoisonStacks--;
            }

            Current = new WaitForSeconds(10f);
            return true;
        }

        public void Reset() {
            throw new NotSupportedException();
        }

        public void Dispose() { }
    }
    
}

Okay, that makes sense.
There's a 'state' variable that indicates when the Enumerator is exhausted (no more values will be returned),
Along with the Player instance and an object 'Current' storing the most recent return value.

But wait, where did the poison damage logic go?
'PoisonCoroutine' now only contains the code required to construct the Enumerator, if we tried to patch it, we won't find the code we're looking for.
Instead, we want to patch the 'MoveNext' Method of the generated class - this is where BTHU comes in.

We can find the Method info of the 'MoveNext' Method with this one-liner.

MethodInfo moveNextMethod = PatcherUtils.FindIEnumeratorMoveNext(AccessTools.DeclaredMethod(typeof(Player), nameof(Player.PoisonCoroutine)))

Now we can apply our patches as usual.


Midfix Patch

This basic example writes "Player is taking posion damage" to the console, just before 'PoisonCoroutine' calls 'TakeDamage'.

[HarmonyPatch]
public class Player_PoisonCoroutine_Patch {
    
    [HarmonyTargetMethod]
    private static MethodBase MoveNext() {
        return PatcherUtils.FindIEnumeratorMoveNext(AccessTools.DeclaredMethod(typeof(Player), nameof(Player.PoisonCoroutine)));
    }

    private static MidFixInstructionMatcher Matcher() {
        MethodInfo takeDamage = AccessTools.DeclaredMethod(typeof(Player), nameof(Player.TakeDamage));
        return new MidFixInstructionMatcher(
                expectedMatches: 1,
                postfixInstructionSequence: new[] {
                        InstructionMask.MatchInstruction(OpCodes.Call, takeDamage),
                }
        );
    }

    [BTHarmonyMidFix(nameof(Matcher))]
    private static void Midfix() {
        Console.WriteLine("Player is taking poison damage");
    }
}

For more information on Midfix Patching, go here


Transpiler Patch

This example makes the Player take twice as much poison damage during each tick.

[HarmonyPatch]
public class Player_PoisonCoroutine_Patch {

    private static readonly ManualLogSource logger = Logger.CreateLogSource(nameof(Player_PoisonCoroutine_Patch));
    
    [HarmonyTargetMethods]
    private static IEnumerable<MethodBase> MoveNext() {
        yield return PatcherUtils.FindIEnumeratorMoveNext(AccessTools.DeclaredMethod(typeof(Player), nameof(Player.PoisonCoroutine)));
    }

    [HarmonyTranspiler]
    private static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> codeInstructions) {
        List<CodeInstruction> instructions = codeInstructions.ToList();
        MethodInfo takeDamage = AccessTools.DeclaredMethod(typeof(Player), nameof(Player.TakeDamage));
        MethodInfo applyPoisonDamageModifier = SymbolExtensions.GetMethodInfo(() => ApplyPoisonDamageModifier(0));

        CodeReplacementPatch patch = new CodeReplacementPatch(
                expectedMatches: 1,
                postfixInstructions: new List<InstructionMask> {
                        InstructionMask.MatchInstruction(OpCodes.Call, takeDamage),
                },
                insertInstructions: new List<CodeInstruction> {
                        new CodeInstruction(OpCodes.Call, applyPoisonDamageModifier),
                }
        );
        patch.ApplySafe(instructions, logger);

        return instructions;
    }

    private static int ApplyPoisonDamageModifier(int poisonDamage) {
        return poisonDamage * 2;
    }

}

For more information on Transpiler Patching, go here


See also: