Yaroze.NET/tests/Yaroze.Tests/JIT/BranchJumpTests.cs

using Xunit;
using Yaroze.Core.JIT;

namespace Yaroze.Tests.JIT;

public class BranchJumpTests
{
    [Fact]
    public void BEQ_TakenBranch_UpdatesPC()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // addiu $t1, $zero, 5
        // beq $t0, $t1, target  # Should take branch (5 == 5)
        // nop
        // target: addiu $v0, $zero, 42
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x24090005);   // ADDIU $t1, $zero, 5
        WriteInstruction(memory, 8, 0x11090001);   // BEQ $t0, $t1, +1
        WriteInstruction(memory, 12, 0x00000000);  // NOP (delay slot)
        WriteInstruction(memory, 16, 0x2402002A);  // ADDIU $v0, $zero, 42

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BEQ_NotTakenBranch_ContinuesNormally()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // addiu $t1, $zero, 10
        // beq $t0, $t1, target  # Should NOT take branch (5 != 10)
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x2409000A);   // ADDIU $t1, $zero, 10
        WriteInstruction(memory, 8, 0x11090001);   // BEQ $t0, $t1, +1
        WriteInstruction(memory, 12, 0x00000000);  // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BNE_TakenBranch_UpdatesPC()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // addiu $t1, $zero, 10
        // bne $t0, $t1, target  # Should take branch (5 != 10)
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x2409000A);   // ADDIU $t1, $zero, 10
        WriteInstruction(memory, 8, 0x15090001);   // BNE $t0, $t1, +1
        WriteInstruction(memory, 12, 0x00000000);  // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BNE_NotTakenBranch_ContinuesNormally()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // addiu $t1, $zero, 5
        // bne $t0, $t1, target  # Should NOT take branch (5 == 5)
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x24090005);   // ADDIU $t1, $zero, 5
        WriteInstruction(memory, 8, 0x15090001);   // BNE $t0, $t1, +1
        WriteInstruction(memory, 12, 0x00000000);  // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BLEZ_TakenBranch_NegativeValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, -5
        // blez $t0, target  # Should take branch (-5 <= 0)
        // nop
        WriteInstruction(memory, 0, 0x2408FFFB);   // ADDIU $t0, $zero, -5
        WriteInstruction(memory, 4, 0x19000001);   // BLEZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BLEZ_TakenBranch_ZeroValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 0
        // blez $t0, target  # Should take branch (0 <= 0)
        // nop
        WriteInstruction(memory, 0, 0x24080000);   // ADDIU $t0, $zero, 0
        WriteInstruction(memory, 4, 0x19000001);   // BLEZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BLEZ_NotTakenBranch_PositiveValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // blez $t0, target  # Should NOT take branch (5 > 0)
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x19000001);   // BLEZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BGTZ_TakenBranch_PositiveValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // bgtz $t0, target  # Should take branch (5 > 0)
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x1D000001);   // BGTZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BGTZ_NotTakenBranch_ZeroValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 0
        // bgtz $t0, target  # Should NOT take branch (0 <= 0)
        // nop
        WriteInstruction(memory, 0, 0x24080000);   // ADDIU $t0, $zero, 0
        WriteInstruction(memory, 4, 0x1D000001);   // BGTZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BGTZ_NotTakenBranch_NegativeValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, -5
        // bgtz $t0, target  # Should NOT take branch (-5 <= 0)
        // nop
        WriteInstruction(memory, 0, 0x2408FFFB);   // ADDIU $t0, $zero, -5
        WriteInstruction(memory, 4, 0x1D000001);   // BGTZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BLTZ_TakenBranch_NegativeValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, -5
        // bltz $t0, target  # Should take branch (-5 < 0)
        // nop
        WriteInstruction(memory, 0, 0x2408FFFB);   // ADDIU $t0, $zero, -5
        WriteInstruction(memory, 4, 0x05000001);   // BLTZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BLTZ_NotTakenBranch_ZeroValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 0
        // bltz $t0, target  # Should NOT take branch (0 >= 0)
        // nop
        WriteInstruction(memory, 0, 0x24080000);   // ADDIU $t0, $zero, 0
        WriteInstruction(memory, 4, 0x05000001);   // BLTZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BGEZ_TakenBranch_PositiveValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // bgez $t0, target  # Should take branch (5 >= 0)
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x05010001);   // BGEZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BGEZ_TakenBranch_ZeroValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 0
        // bgez $t0, target  # Should take branch (0 >= 0)
        // nop
        WriteInstruction(memory, 0, 0x24080000);   // ADDIU $t0, $zero, 0
        WriteInstruction(memory, 4, 0x05010001);   // BGEZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BGEZ_NotTakenBranch_NegativeValue()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, -5
        // bgez $t0, target  # Should NOT take branch (-5 < 0)
        // nop
        WriteInstruction(memory, 0, 0x2408FFFB);   // ADDIU $t0, $zero, -5
        WriteInstruction(memory, 4, 0x05010001);   // BGEZ $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void J_UnconditionalJump_UpdatesPC()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // j target
        // nop
        // (instruction encoding: 0x08000010 = j 0x80000040)
        WriteInstruction(memory, 0, 0x08000010);   // J 0x80000040
        WriteInstruction(memory, 4, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void JAL_CallFunction_UpdatesRAandPC()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // jal function
        // nop
        WriteInstruction(memory, 0, 0x0C000010);   // JAL 0x80000040
        WriteInstruction(memory, 4, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
        // Should set $ra to 0x80000008 (return address)
    }

    [Fact]
    public void JR_IndirectJump_UpdatesPC()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 0x100
        // jr $t0
        // nop
        WriteInstruction(memory, 0, 0x24080100);   // ADDIU $t0, $zero, 0x100
        WriteInstruction(memory, 4, 0x01000008);   // JR $t0
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void JALR_IndirectCall_UpdatesRDandPC()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 0x100
        // jalr $ra, $t0
        // nop
        WriteInstruction(memory, 0, 0x24080100);   // ADDIU $t0, $zero, 0x100
        WriteInstruction(memory, 4, 0x0100F809);   // JALR $ra, $t0
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void DelaySlot_ExecutesBeforeBranch()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // beq $t0, $t0, target  # Always taken
        // addiu $v0, $zero, 42   # Delay slot - should execute
        // target: addiu $v1, $zero, 1
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x11080001);   // BEQ $t0, $t0, +1
        WriteInstruction(memory, 8, 0x2402002A);   // ADDIU $v0, $zero, 42 (delay slot)
        WriteInstruction(memory, 12, 0x24030001);  // ADDIU $v1, $zero, 1

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
        // $v0 should be 42 even though branch is taken
    }

    [Fact]
    public void ComplexControlFlow_MultipleConditionals()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 10
        // addiu $t1, $zero, 5
        // bne $t0, $t1, skip  # Should take (10 != 5)
        // addiu $v0, $zero, 1  # Delay slot
        // skip: addiu $v1, $zero, 2
        WriteInstruction(memory, 0, 0x2408000A);   // ADDIU $t0, $zero, 10
        WriteInstruction(memory, 4, 0x24090005);   // ADDIU $t1, $zero, 5
        WriteInstruction(memory, 8, 0x15090001);   // BNE $t0, $t1, +1
        WriteInstruction(memory, 12, 0x24020001);  // ADDIU $v0, $zero, 1 (delay)
        WriteInstruction(memory, 16, 0x24030002);  // ADDIU $v1, $zero, 2

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
    }

    [Fact]
    public void BLTZAL_TakenBranch_SetsReturnAddress()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, -5
        // bltzal $t0, target  # Should take branch (-5 < 0) and set $ra
        // nop
        WriteInstruction(memory, 0, 0x2408FFFB);   // ADDIU $t0, $zero, -5
        WriteInstruction(memory, 4, 0x05100001);   // BLTZAL $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
        // $ra should be set to 0x80000008 (return address)
    }

    [Fact]
    public void BLTZAL_NotTakenBranch_StillSetsReturnAddress()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // bltzal $t0, target  # Should NOT take branch (5 >= 0) but still set $ra
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x05100001);   // BLTZAL $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
        // $ra should still be set even though branch not taken
    }

    [Fact]
    public void BGEZAL_TakenBranch_SetsReturnAddress()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, 5
        // bgezal $t0, target  # Should take branch (5 >= 0) and set $ra
        // nop
        WriteInstruction(memory, 0, 0x24080005);   // ADDIU $t0, $zero, 5
        WriteInstruction(memory, 4, 0x05110001);   // BGEZAL $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
        // $ra should be set to 0x80000008 (return address)
    }

    [Fact]
    public void BGEZAL_NotTakenBranch_StillSetsReturnAddress()
    {
        // Arrange
        byte[] memory = new byte[4096];
        uint baseAddress = 0x80000000;

        // addiu $t0, $zero, -5
        // bgezal $t0, target  # Should NOT take branch (-5 < 0) but still set $ra
        // nop
        WriteInstruction(memory, 0, 0x2408FFFB);   // ADDIU $t0, $zero, -5
        WriteInstruction(memory, 4, 0x05110001);   // BGEZAL $t0, +1
        WriteInstruction(memory, 8, 0x00000000);   // NOP (delay slot)

        var verifier = new LockstepVerifier(memory, baseAddress);

        // Act
        var result = verifier.VerifyBlock(baseAddress);

        // Assert
        Assert.True(result.IsMatch);
        // $ra should still be set even though branch not taken
    }

    private void WriteInstruction(byte[] memory, int offset, uint instruction)
    {
        byte[] bytes = BitConverter.GetBytes(instruction);
        Array.Copy(bytes, 0, memory, offset, 4);
    }
}