block linking

src/core/jit/block_manager.h

@@ -19,9 +19,15 @@ struct CodeBlock {
     std::set<VAddr> dependencies;
     bool is_linked;
 
+    // Control flow targets for linking
+    VAddr fallthrough_target;    // Next sequential address (if block doesn't end with branch)
+    VAddr branch_target;         // Direct branch target (JMP)
+    void* branch_patch_location; // Location in ARM64 code to patch for direct branch
+
     CodeBlock(VAddr addr, void* code, size_t size, size_t count)
         : ps4_address(addr), arm64_code(code), code_size(size), instruction_count(count),
-          is_linked(false) {}
+          is_linked(false), fallthrough_target(0), branch_target(0),
+          branch_patch_location(nullptr) {}
 };
 
 class BlockManager {

src/core/jit/execution_engine.cpp

@@ -112,6 +112,9 @@ CodeBlock* ExecutionEngine::TranslateBasicBlock(VAddr start_address, size_t max_
     VAddr current_address = start_address;
     size_t instruction_count = 0;
     bool block_end = false;
+    VAddr fallthrough_target = 0;
+    VAddr branch_target = 0;
+    void* branch_patch_location = nullptr;
 
     while (instruction_count < max_instructions && !block_end) {
         ZydisDecodedInstruction instruction;
@@ -129,6 +132,14 @@ CodeBlock* ExecutionEngine::TranslateBasicBlock(VAddr start_address, size_t max_
             break;
         }
 
+        // Track branch target before translation
+        if (instruction.mnemonic == ZYDIS_MNEMONIC_JMP &&
+            operands[0].type == ZYDIS_OPERAND_TYPE_IMMEDIATE) {
+            s64 offset = static_cast<s64>(operands[0].imm.value.s);
+            branch_target = current_address + instruction.length + offset;
+            branch_patch_location = code_generator->getCurr();
+        }
+
         bool translated = translator->TranslateInstruction(instruction, operands, current_address);
         if (!translated) {
             LOG_WARNING(Core, "Failed to translate instruction at {:#x}", current_address);
@@ -136,30 +147,73 @@ CodeBlock* ExecutionEngine::TranslateBasicBlock(VAddr start_address, size_t max_
         }
 
         instruction_count++;
-        current_address += instruction.length;
+        VAddr next_address = current_address + instruction.length;
 
         switch (instruction.mnemonic) {
         case ZYDIS_MNEMONIC_RET:
-        case ZYDIS_MNEMONIC_JMP:
         case ZYDIS_MNEMONIC_CALL:
             block_end = true;
             break;
+        case ZYDIS_MNEMONIC_JMP:
+            block_end = true;
+            break;
         default:
+            // Check for conditional branches (they don't end the block, but we track them)
+            if (instruction.mnemonic >= ZYDIS_MNEMONIC_JO &&
+                instruction.mnemonic <= ZYDIS_MNEMONIC_JZ) {
+                // Conditional branch - block continues with fallthrough
+                // TODO: Track conditional branch targets for linking
+            }
             break;
         }
+
+        current_address = next_address;
     }
 
     if (instruction_count == 0) {
         return nullptr;
     }
 
+    // Set fallthrough target if block doesn't end with unconditional branch/ret
+    if (!block_end || branch_target == 0) {
+        fallthrough_target = current_address;
+    }
+
     size_t code_size = code_generator->getSize();
     code_generator->makeExecutable();
     CodeBlock* block =
         block_manager->CreateBlock(start_address, block_start, code_size, instruction_count);
 
-    LOG_DEBUG(Core, "Translated basic block at {:#x}, {} instructions, {} bytes", start_address,
-              instruction_count, code_size);
+    // Store control flow information
+    block->fallthrough_target = fallthrough_target;
+    block->branch_target = branch_target;
+    block->branch_patch_location = branch_patch_location;
+
+    LOG_DEBUG(Core,
+              "Translated basic block at {:#x}, {} instructions, {} bytes, fallthrough: {:#x}, "
+              "branch: {:#x}",
+              start_address, instruction_count, code_size, fallthrough_target, branch_target);
+
+    // Try to link blocks if targets are available
+    if (branch_target != 0) {
+        CodeBlock* target_block = block_manager->GetBlock(branch_target);
+        if (target_block) {
+            LinkBlock(block, branch_target);
+        } else {
+            // Add dependency for later linking
+            block_manager->AddDependency(start_address, branch_target);
+        }
+    }
+
+    if (fallthrough_target != 0 && branch_target == 0) {
+        // Try to link fallthrough
+        CodeBlock* target_block = block_manager->GetBlock(fallthrough_target);
+        if (target_block) {
+            // For fallthrough, we need to append a branch at the end
+            // This will be handled by linking logic
+            block_manager->AddDependency(start_address, fallthrough_target);
+        }
+    }
 
     return block;
 }
@@ -170,17 +224,82 @@ CodeBlock* ExecutionEngine::TranslateBlock(VAddr ps4_address) {
         return existing;
     }
 
-    return TranslateBasicBlock(ps4_address);
+    CodeBlock* new_block = TranslateBasicBlock(ps4_address);
+    if (!new_block) {
+        return nullptr;
+    }
+
+    // After creating a new block, check if any existing blocks can link to it
+    // This handles the case where we translate a target block after the source
+    for (auto& [addr, block] : block_manager->blocks) {
+        if (block->branch_target == ps4_address && !block->is_linked) {
+            LinkBlock(block.get(), ps4_address);
+        }
+        if (block->fallthrough_target == ps4_address && block->branch_target == 0 &&
+            !block->is_linked) {
+            LinkBlock(block.get(), ps4_address);
+        }
+    }
+
+    return new_block;
 }
 
 void ExecutionEngine::LinkBlock(CodeBlock* block, VAddr target_address) {
     CodeBlock* target_block = block_manager->GetBlock(target_address);
-    if (target_block && !block->is_linked) {
-        void* link_location = static_cast<u8*>(block->arm64_code) + block->code_size - 4;
-        code_generator->setSize(reinterpret_cast<u8*>(link_location) -
-                                static_cast<u8*>(code_generator->getCode()));
-        code_generator->b(target_block->arm64_code);
+    if (!target_block) {
+        return;
+    }
+
+    // Patch the branch instruction if we have a patch location
+    if (block->branch_patch_location && block->branch_target == target_address) {
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+        pthread_jit_write_protect_np(0);
+#endif
+        // Calculate offset from patch location to target
+        s64 offset = reinterpret_cast<s64>(target_block->arm64_code) -
+                     reinterpret_cast<s64>(block->branch_patch_location);
+
+        // Check if we can use a relative branch (within ±128MB)
+        if (offset >= -0x8000000 && offset < 0x8000000) {
+            s32 imm26 = static_cast<s32>(offset / 4);
+            u32* patch_ptr = reinterpret_cast<u32*>(block->branch_patch_location);
+            // Patch the branch instruction: 0x14000000 | (imm26 & 0x3FFFFFF)
+            *patch_ptr = 0x14000000 | (imm26 & 0x3FFFFFF);
+        } else {
+            // Far branch - need to use indirect branch
+            // For now, leave as-is (will use the placeholder branch)
+            LOG_DEBUG(Core, "Branch target too far for direct linking: offset={}", offset);
+        }
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+        pthread_jit_write_protect_np(1);
+        __builtin___clear_cache(static_cast<char*>(block->branch_patch_location),
+                                static_cast<char*>(block->branch_patch_location) + 4);
+#endif
         block->is_linked = true;
+        LOG_DEBUG(Core, "Linked block {:#x} to {:#x}", block->ps4_address, target_address);
+    } else if (block->fallthrough_target == target_address && block->branch_target == 0) {
+        // For fallthrough, append a branch at the end of the block
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+        pthread_jit_write_protect_np(0);
+#endif
+        void* link_location = static_cast<u8*>(block->arm64_code) + block->code_size;
+        s64 offset =
+            reinterpret_cast<s64>(target_block->arm64_code) - reinterpret_cast<s64>(link_location);
+
+        if (offset >= -0x8000000 && offset < 0x8000000) {
+            s32 imm26 = static_cast<s32>(offset / 4);
+            u32* patch_ptr = reinterpret_cast<u32*>(link_location);
+            *patch_ptr = 0x14000000 | (imm26 & 0x3FFFFFF);
+            block->code_size += 4; // Update block size
+        }
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+        pthread_jit_write_protect_np(1);
+        __builtin___clear_cache(static_cast<char*>(link_location),
+                                static_cast<char*>(link_location) + 4);
+#endif
+        block->is_linked = true;
+        LOG_DEBUG(Core, "Linked fallthrough from block {:#x} to {:#x}", block->ps4_address,
+                  target_address);
     }
 }
 

src/core/jit/x86_64_translator.cpp

@@ -571,8 +571,47 @@ bool X86_64Translator::TranslateRet(const ZydisDecodedInstruction& instruction,
 
 bool X86_64Translator::TranslateJmp(const ZydisDecodedInstruction& instruction,
                                     const ZydisDecodedOperand* operands, VAddr address) {
-    LOG_WARNING(Core, "JMP instruction translation needs execution engine integration");
-    return false;
+    const auto& target = operands[0];
+    VAddr target_address = 0;
+
+    // Calculate target address based on operand type
+    if (target.type == ZYDIS_OPERAND_TYPE_IMMEDIATE) {
+        // Direct relative jump: JMP rel32
+        // Target = current_address + instruction.length + offset
+        s64 offset = static_cast<s64>(target.imm.value.s);
+        target_address = address + instruction.length + offset;
+    } else if (target.type == ZYDIS_OPERAND_TYPE_MEMORY) {
+        // Indirect jump: JMP [mem]
+        // Load address from memory into scratch register
+        LoadMemoryOperand(RegisterMapper::SCRATCH_REG, target, 8);
+        // TODO: don't use a dispatcher
+        codegen.br(RegisterMapper::SCRATCH_REG);
+        return true;
+    } else if (target.type == ZYDIS_OPERAND_TYPE_REGISTER) {
+        // Indirect jump: JMP reg
+        int reg = GetArm64Register(target);
+        if (reg == -1) {
+            LOG_ERROR(Core, "Invalid register for JMP");
+            return false;
+        }
+        codegen.br(reg);
+        return true;
+    } else {
+        LOG_ERROR(Core, "Unsupported JMP operand type");
+        return false;
+    }
+
+    // For direct jumps, we need to branch to the target address
+    // Since the target block may not be translated yet, we'll generate
+    // a placeholder that can be patched later during block linking
+    // For now, generate a branch to a dispatcher function
+    // TODO: Implement proper block linking to patch this with direct branch
+
+    // Calculate offset from current code position
+    void* placeholder_target = reinterpret_cast<void*>(target_address);
+    codegen.b(placeholder_target);
+
+    return true;
 }
 
 bool X86_64Translator::TranslateCmp(const ZydisDecodedInstruction& instruction,

tests/CMakeLists.txt

@@ -6,6 +6,7 @@ add_executable(jit_tests
     test_register_mapping.cpp
     test_block_manager.cpp
     test_execution_engine.cpp
+    test_block_linking.cpp
     main.cpp
 )
 

tests/test_block_linking.cpp

@@ -0,0 +1,247 @@
+// SPDX-FileCopyrightText: Copyright 2024 shadPS4 Emulator Project
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+#include "common/decoder.h"
+#include "core/jit/arm64_codegen.h"
+#include "core/jit/block_manager.h"
+#include "core/jit/register_mapping.h"
+#include "core/jit/x86_64_translator.h"
+#include <gtest/gtest.h>
+#include <sys/mman.h>
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+#include <pthread.h>
+#endif
+
+using namespace Core::Jit;
+
+class BlockLinkingTest : public ::testing::Test {
+protected:
+  void SetUp() override {
+    // Allocate executable memory for test code
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+    test_code_buffer = mmap(nullptr, 64 * 1024, PROT_READ | PROT_WRITE,
+                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    ASSERT_NE(test_code_buffer, MAP_FAILED)
+        << "Failed to allocate executable memory for test";
+    pthread_jit_write_protect_np(0);
+#else
+    test_code_buffer =
+        mmap(nullptr, 64 * 1024, PROT_READ | PROT_WRITE | PROT_EXEC,
+             MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    ASSERT_NE(test_code_buffer, MAP_FAILED)
+        << "Failed to allocate executable memory for test";
+#endif
+    codegen = std::make_unique<Arm64CodeGenerator>(64 * 1024, test_code_buffer);
+    register_mapper = std::make_unique<RegisterMapper>();
+    translator = std::make_unique<X86_64Translator>(*codegen, *register_mapper);
+    block_manager = std::make_unique<BlockManager>();
+  }
+
+  void TearDown() override {
+    translator.reset();
+    register_mapper.reset();
+    codegen.reset();
+    block_manager.reset();
+    if (test_code_buffer != MAP_FAILED) {
+      munmap(test_code_buffer, 64 * 1024);
+    }
+  }
+
+  void *test_code_buffer = MAP_FAILED;
+  std::unique_ptr<Arm64CodeGenerator> codegen;
+  std::unique_ptr<RegisterMapper> register_mapper;
+  std::unique_ptr<X86_64Translator> translator;
+  std::unique_ptr<BlockManager> block_manager;
+};
+
+// Test that JMP translation can handle direct immediate addresses
+TEST_F(BlockLinkingTest, TranslateDirectJmp) {
+  // Create a simple x86_64 JMP instruction: JMP +0x1000 (relative jump)
+  // x86_64 encoding: E9 <offset> (near relative jump, 32-bit offset)
+  // E9 00 10 00 00 = JMP +0x1000
+  u8 x86_jmp[] = {0xE9, 0x00, 0x10, 0x00, 0x00};
+
+  ZydisDecodedInstruction instruction;
+  ZydisDecodedOperand operands[ZYDIS_MAX_OPERAND_COUNT];
+
+  ZyanStatus status = Common::Decoder::Instance()->decodeInstruction(
+      instruction, operands, x86_jmp, sizeof(x86_jmp));
+
+  if (!ZYAN_SUCCESS(status)) {
+    GTEST_SKIP()
+        << "Failed to decode JMP instruction - Zydis may not be available";
+  }
+
+  // JMP translation should succeed (even if target isn't linked yet)
+  bool result = translator->TranslateJmp(instruction, operands, 0x400000);
+  EXPECT_TRUE(result) << "JMP translation should succeed";
+  EXPECT_GT(codegen->getSize(), 0) << "JMP should generate ARM64 code";
+}
+
+// Test that we can create two blocks and link them
+TEST_F(BlockLinkingTest, CreateAndLinkBlocks) {
+  VAddr block1_addr = 0x400000;
+  VAddr block2_addr = 0x401000;
+
+  // Allocate separate memory for each block to avoid issues
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+  void *block1_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block1_mem, MAP_FAILED);
+  pthread_jit_write_protect_np(0);
+
+  void *block2_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block2_mem, MAP_FAILED);
+#else
+  void *block1_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block1_mem, MAP_FAILED);
+
+  void *block2_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block2_mem, MAP_FAILED);
+#endif
+
+  // Write simple NOP instructions
+  u32 nop = 0xD503201F; // ARM64 NOP
+  *reinterpret_cast<u32 *>(block1_mem) = nop;
+  *reinterpret_cast<u32 *>(block2_mem) = nop;
+
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+  pthread_jit_write_protect_np(1);
+  mprotect(block1_mem, 4096, PROT_READ | PROT_EXEC);
+  mprotect(block2_mem, 4096, PROT_READ | PROT_EXEC);
+#endif
+
+  // Create blocks
+  CodeBlock *block1 = block_manager->CreateBlock(block1_addr, block1_mem, 4, 1);
+  ASSERT_NE(block1, nullptr);
+
+  CodeBlock *block2 = block_manager->CreateBlock(block2_addr, block2_mem, 4, 1);
+  ASSERT_NE(block2, nullptr);
+
+  // Verify blocks exist
+  EXPECT_EQ(block_manager->GetBlockCount(), 2);
+  EXPECT_NE(block_manager->GetBlock(block1_addr), nullptr);
+  EXPECT_NE(block_manager->GetBlock(block2_addr), nullptr);
+
+  // Test that blocks can be retrieved
+  CodeBlock *retrieved_block1 = block_manager->GetBlock(block1_addr);
+  CodeBlock *retrieved_block2 = block_manager->GetBlock(block2_addr);
+  EXPECT_EQ(retrieved_block1, block1);
+  EXPECT_EQ(retrieved_block2, block2);
+
+  // Cleanup
+  munmap(block1_mem, 4096);
+  munmap(block2_mem, 4096);
+}
+
+// Test that block linking tracks dependencies
+TEST_F(BlockLinkingTest, BlockDependencies) {
+  VAddr block1_addr = 0x400000;
+  VAddr block2_addr = 0x401000;
+
+  // Allocate memory for blocks
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+  void *block1_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block1_mem, MAP_FAILED);
+  pthread_jit_write_protect_np(0);
+
+  void *block2_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block2_mem, MAP_FAILED);
+
+  u32 nop = 0xD503201F;
+  *reinterpret_cast<u32 *>(block1_mem) = nop;
+  *reinterpret_cast<u32 *>(block2_mem) = nop;
+
+  pthread_jit_write_protect_np(1);
+  mprotect(block1_mem, 4096, PROT_READ | PROT_EXEC);
+  mprotect(block2_mem, 4096, PROT_READ | PROT_EXEC);
+#else
+  void *block1_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block1_mem, MAP_FAILED);
+
+  void *block2_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block2_mem, MAP_FAILED);
+
+  u32 nop = 0xD503201F;
+  *reinterpret_cast<u32 *>(block1_mem) = nop;
+  *reinterpret_cast<u32 *>(block2_mem) = nop;
+#endif
+
+  // Create blocks
+  CodeBlock *block1 = block_manager->CreateBlock(block1_addr, block1_mem, 4, 1);
+  CodeBlock *block2 = block_manager->CreateBlock(block2_addr, block2_mem, 4, 1);
+
+  // Add dependency: block1 depends on block2
+  block_manager->AddDependency(block1_addr, block2_addr);
+
+  // Verify dependency is tracked
+  EXPECT_EQ(block1->dependencies.count(block2_addr), 1);
+
+  // Cleanup
+  munmap(block1_mem, 4096);
+  munmap(block2_mem, 4096);
+}
+
+// Test that invalidating a block invalidates dependent blocks
+TEST_F(BlockLinkingTest, InvalidateDependentBlocks) {
+  VAddr block1_addr = 0x400000;
+  VAddr block2_addr = 0x401000;
+
+  // Allocate memory for blocks
+#if defined(__APPLE__) && defined(ARCH_ARM64)
+  void *block1_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block1_mem, MAP_FAILED);
+  pthread_jit_write_protect_np(0);
+
+  void *block2_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block2_mem, MAP_FAILED);
+
+  u32 nop = 0xD503201F;
+  *reinterpret_cast<u32 *>(block1_mem) = nop;
+  *reinterpret_cast<u32 *>(block2_mem) = nop;
+
+  pthread_jit_write_protect_np(1);
+  mprotect(block1_mem, 4096, PROT_READ | PROT_EXEC);
+  mprotect(block2_mem, 4096, PROT_READ | PROT_EXEC);
+#else
+  void *block1_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block1_mem, MAP_FAILED);
+
+  void *block2_mem = mmap(nullptr, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
+                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  ASSERT_NE(block2_mem, MAP_FAILED);
+
+  u32 nop = 0xD503201F;
+  *reinterpret_cast<u32 *>(block1_mem) = nop;
+  *reinterpret_cast<u32 *>(block2_mem) = nop;
+#endif
+
+  // Create blocks with dependency
+  CodeBlock *block1 = block_manager->CreateBlock(block1_addr, block1_mem, 4, 1);
+  CodeBlock *block2 = block_manager->CreateBlock(block2_addr, block2_mem, 4, 1);
+
+  block_manager->AddDependency(block1_addr, block2_addr);
+
+  // Invalidate block2
+  block_manager->InvalidateBlock(block2_addr);
+
+  // block2 should be removed
+  EXPECT_EQ(block_manager->GetBlock(block2_addr), nullptr);
+  // block1 should still exist (dependency tracking doesn't auto-invalidate)
+  // But in a real implementation, we might want to invalidate dependents
+  EXPECT_NE(block_manager->GetBlock(block1_addr), nullptr);
+
+  // Cleanup
+  munmap(block1_mem, 4096);
+  munmap(block2_mem, 4096);
+}