Follow Up Of the previous commit with the update of TLB update

3 files changed, 103 insertions, 13 deletions

diff --git a/src/core/arm/nce/arm_nce.cpp b/src/core/arm/nce/arm_nce.cpp
index 2bb20b1fd..f08c5c352 100644
--- a/src/core/arm/nce/arm_nce.cpp
+++ b/src/core/arm/nce/arm_nce.cpp
@@ -170,29 +170,54 @@ bool ArmNce::HandleGuestAccessFault(GuestContext* guest_ctx, void* raw_info, voi
     // Get the ArmNce instance from the guest context
     ArmNce* nce = guest_ctx->parent;
 
-    // Check TLB first
-    if (TlbEntry* entry = nce->FindTlbEntry(fault_addr)) {
-        if (!entry->writable && info->si_code == SEGV_ACCERR) {
-            LOG_DEBUG(Core_ARM, "Write to read-only memory at {:X}", fault_addr);
+    // Define a maximum retry count to prevent infinite loops
+    constexpr int max_retries = 3;
+    int retry_count = 0;
+
+    while (retry_count < max_retries) {
+        // Check TLB first
+        if (TlbEntry* entry = nce->FindTlbEntry(fault_addr)) {
+            if (!entry->writable && info->si_code == SEGV_ACCERR) {
+                LOG_DEBUG(Core_ARM, "Write to read-only memory at {:X}", fault_addr);
+                return HandleFailedGuestFault(guest_ctx, raw_info, raw_context);
+            }
+            return true;
+        }
+
+        // TLB miss handling with better error checking
+        if (memory.InvalidateNCE(fault_addr, Memory::CITRON_PAGESIZE)) {
+            const u64 host_addr = reinterpret_cast<u64>(memory.GetPointer(fault_addr));
+
+            if (host_addr) {
+                nce->AddTlbEntry(fault_addr, host_addr, Memory::CITRON_PAGESIZE, true);
+                return true;
+            } else {
+                LOG_DEBUG(Core_ARM, "Failed to get host address for guest address {:X}", fault_addr);
+            }
+        } else {
+            LOG_DEBUG(Core_ARM, "Memory invalidation failed for address {:X}", fault_addr);
+        }
+
+        // Trigger an immediate remap if lookup fails
+        if (!memory.Remap(fault_addr, Memory::CITRON_PAGESIZE)) {
+            LOG_ERROR(Core_ARM, "Immediate remap failed for address {:X}", fault_addr);
             return HandleFailedGuestFault(guest_ctx, raw_info, raw_context);
         }
-        return true;
-    }
 
-    // TLB miss handling with better error checking
-    if (memory.InvalidateNCE(fault_addr, Memory::CITRON_PAGESIZE)) {
+        // Retry adding the TLB entry after remap
         const u64 host_addr = reinterpret_cast<u64>(memory.GetPointer(fault_addr));
-
         if (host_addr) {
             nce->AddTlbEntry(fault_addr, host_addr, Memory::CITRON_PAGESIZE, true);
             return true;
         } else {
-            LOG_DEBUG(Core_ARM, "Failed to get host address for guest address {:X}", fault_addr);
+            LOG_ERROR(Core_ARM, "Failed to get host address after remap for guest address {:X}", fault_addr);
         }
-    } else {
-        LOG_DEBUG(Core_ARM, "Memory invalidation failed for address {:X}", fault_addr);
+
+        // Increment the retry count
+        retry_count++;
     }
 
+    // If all retries fail, handle the fault as a failed guest fault
     return HandleFailedGuestFault(guest_ctx, raw_info, raw_context);
 }
 
@@ -418,6 +443,7 @@ TlbEntry* ArmNce::FindTlbEntry(u64 guest_addr) {
             if (entry.access_count < 1000) { // Prevent overflow
                 entry.access_count++;
             }
+            entry.ref_count++; // Increment reference count
             return &entry;
         }
     }
@@ -474,7 +500,8 @@ void ArmNce::AddTlbEntry(u64 guest_addr, u64 host_addr, u32 size, bool writable)
         .valid = true,
         .writable = writable,
         .access_count = 1,
-        .last_access_time = now // Update the access time
+        .last_access_time = now,
+        .ref_count = 1 // Initialize reference count
     };
 }
 
@@ -509,6 +536,7 @@ struct TlbEntry {
     bool writable;
     u32 access_count;
     std::chrono::steady_clock::time_point last_access_time; // Add this line
+    u32 ref_count; // Add this line
 };
 
 } // namespace Core
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index c598edbc2..dd6ffaf6c 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -1151,4 +1151,51 @@ bool Memory::InvalidateSeparateHeap(void* fault_address) {
 #endif
 }
 
+bool Memory::Remap(u64 guest_addr, u32 size) {
+    // Unmap the old address
+    UnmapRegion(*impl->current_page_table, guest_addr, size, false);
+
+    // Reclaim memory from unreferenced pages
+    ReclaimUnusedMemory();
+
+    // Allocate new memory
+    void* new_memory = std::malloc(size);
+    if (!new_memory) {
+        LOG_ERROR(Core_Memory, "Failed to allocate new memory for remapping address {:X}", guest_addr);
+        return false;
+    }
+
+    // Map the new memory to the guest address
+    MapMemoryRegion(*impl->current_page_table, guest_addr, size, reinterpret_cast<u64>(new_memory), Common::MemoryPermission::ReadWrite, false);
+
+    // Verify the mapping
+    if (GetPointer(guest_addr) != nullptr) {
+        LOG_INFO(Core_Memory, "Successfully remapped address {:X}", guest_addr);
+        return true;
+    } else {
+        LOG_ERROR(Core_Memory, "Failed to remap address {:X}", guest_addr);
+        std::free(new_memory);
+        return false;
+    }
+}
+
+void Memory::ReclaimUnusedMemory() {
+    std::lock_guard lock(m_tlb_mutex);
+
+    for (auto& entry : m_tlb) {
+        if (entry.valid && entry.ref_count == 0) {
+            // Unmap the memory region
+            UnmapRegion(*impl->current_page_table, entry.guest_addr, entry.size, false);
+
+            // Free the memory
+            std::free(reinterpret_cast<void*>(entry.host_addr));
+
+            // Invalidate the TLB entry
+            entry.valid = false;
+
+            LOG_INFO(Core_Memory, "Reclaimed memory for address {:X}", entry.guest_addr);
+        }
+    }
+}
+
 } // namespace Core::Memory
diff --git a/src/core/memory.h b/src/core/memory.h
index 3f315ff7a..eeeee6d65 100644
--- a/src/core/memory.h
+++ b/src/core/memory.h
@@ -492,6 +492,21 @@ public:
 
     bool InvalidateSeparateHeap(void* fault_address);
 
+    /**
+     * Remaps a region of the emulated process address space.
+     *
+     * @param guest_addr The address to begin remapping at.
+     * @param size       The amount of bytes to remap.
+     *
+     * @returns True if remapping is successful, false otherwise.
+     */
+    bool Remap(u64 guest_addr, u32 size);
+
+    /**
+     * Reclaims memory from pages that are no longer used.
+     */
+    void ReclaimUnusedMemory();
+
 private:
     Core::System& system;