Implement support for MEMFindContainHeap when using mapped memory

source/function_replacements.cpp

@@ -1,7 +1,6 @@
 #include "function_replacements.h"
-
 #include "memory_mapping.h"
-
+#include <coreinit/memheap.h>
 
 DECL_FUNCTION(uint32_t, KiEffectiveToPhysical, uint32_t addressSpace, uint32_t virtualAddress) {
     uint32_t result = real_KiEffectiveToPhysical(addressSpace, virtualAddress);
@@ -52,7 +51,6 @@ DECL_FUNCTION(uint32_t, KiIsEffectiveRangeValid, uint32_t addressSpace, uint32_t
     return result;
 }
 
-
 // clang-format off
 #define k_memcpy ((void(*)(void *, void *, uint32_t))(0xfff09e44))
 // clang-format on
@@ -71,6 +69,15 @@ DECL_FUNCTION(uint32_t, KiGetOrPutUserData, void *src, uint32_t size, void *dst,
     return real_KiGetOrPutUserData(src, size, dst, isRead);
 }
 
+DECL_FUNCTION(MEMHeapHandle, MEMFindContainHeap, void *block) {
+    auto result = MemoryMapping_MEMFindContainHeap(block);
+    if (result == nullptr) {
+        return real_MEMFindContainHeap(block);
+    }
+
+    return result;
+}
+
 // clang-format off
 function_replacement_data_t function_replacements[] __attribute__((section(".data"))) = {
         REPLACE_FUNCTION_VIA_ADDRESS(sCheckDataRange,                       0x3200cf60, 0x0100cf60),
@@ -80,6 +87,7 @@ function_replacement_data_t function_replacements[] __attribute__((section(".dat
         REPLACE_FUNCTION_VIA_ADDRESS(KiIsEffectiveRangeValid,               0xffee0d6c, 0xffee0d6c),
         REPLACE_FUNCTION_VIA_ADDRESS(IPCKDriver_ValidatePhysicalAddress,    0xfff0cb5c, 0xfff0cb5c),
         REPLACE_FUNCTION_VIA_ADDRESS(KiGetOrPutUserData,                    0xffee0794, 0xffee0794),
+        REPLACE_FUNCTION(MEMFindContainHeap,                    LIBRARY_COREINIT, MEMFindContainHeap),
 };
 // clang-format on
 

source/memory_mapping.cpp

@@ -1,7 +1,9 @@
 #include "memory_mapping.h"
 #include <coreinit/cache.h>
+#include <coreinit/memblockheap.h>
 #include <coreinit/memdefaultheap.h>
 #include <coreinit/memexpheap.h>
+#include <coreinit/memlist.h>
 #include <coreinit/memorymap.h>
 #include <coreinit/thread.h>
 
@@ -429,6 +431,29 @@ void *MemoryMapping_allocVideoMemory(uint32_t size, int32_t align) {
     return MemoryMapping_allocEx(size, align, true);
 }
 
+// clang-format off
+#define FindHeapContainingBlock ((MEMHeapHandle (*) (MEMMemoryList *, void *) )(0x101C400 + 0x2f2d8))
+// clang-format on
+
+MEMHeapHandle MemoryMapping_MEMFindContainHeap(void *block) {
+    for (int32_t i = 0; /* waiting for a break */; i++) {
+        if (mem_mapping[i].physical_addresses == nullptr) {
+            break;
+        }
+        uint32_t effectiveAddress = mem_mapping[i].effective_start_address;
+        auto heapHandle           = (MEMHeapHandle) effectiveAddress;
+        auto *heap                = (MEMExpHeap *) heapHandle;
+        if (block >= heap->header.dataStart &&
+            block < heap->header.dataEnd) {
+#pragma GCC diagnostic ignored "-Waddress-of-packed-member"
+            auto child = FindHeapContainingBlock(&heap->header.list, block);
+            return child ? child : heapHandle;
+        }
+    }
+
+    return nullptr;
+}
+
 void MemoryMapping_free(void *ptr) {
     if (ptr == nullptr) {
         return;

source/memory_mapping.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <coreinit/memheap.h>
 #include <cstddef>
 
 #ifdef __cplusplus
@@ -206,6 +207,8 @@ void *MemoryMapping_allocVideoMemory(uint32_t size, int32_t align);
 
 void MemoryMapping_free(void *ptr);
 
+MEMHeapHandle MemoryMapping_MEMFindContainHeap(void *block);
+
 uint32_t MemoryMapping_getAreaSizeFromPageTable(uint32_t start, uint32_t maxSize);
 
 uint32_t MemoryMapping_PhysicalToEffective(uint32_t phyiscalAddress);