UserspaceEmulator: Make big malloc block lookup O(1) as well

By passing the Region& to the auditing functions, we know exactly which
block we are hitting. This allows us to track big mallocations the same
way we already do chunked ones.

This gets rid of the O(n) scan in find_mallocation() for allocations
larger than the maximum malloc chunk size. :^)

4 files changed, 57 insertions, 55 deletions

diff --git a/DevTools/UserspaceEmulator/MallocTracer.cpp b/DevTools/UserspaceEmulator/MallocTracer.cpp
index 7e4676142a..13bce805ee 100644
--- a/DevTools/UserspaceEmulator/MallocTracer.cpp
+++ b/DevTools/UserspaceEmulator/MallocTracer.cpp
@@ -53,10 +53,6 @@ inline void MallocTracer::for_each_mallocation(Callback callback) const
         }
         return IterationDecision::Continue;
     });
-    for (auto& big_mallocation : m_big_mallocations) {
-        if (callback(big_mallocation) == IterationDecision::Break)
-            return;
-    }
 }
 
 void MallocTracer::target_did_malloc(Badge<SoftCPU>, FlatPtr address, size_t size)
@@ -81,22 +77,22 @@ void MallocTracer::target_did_malloc(Badge<SoftCPU>, FlatPtr address, size_t siz
         return;
     }
 
-    bool is_chunked_allocation = size <= size_classes[num_size_classes - 1];
-    if (is_chunked_allocation) {
-        MallocRegionMetadata* malloc_data = static_cast<MmapRegion&>(*region).malloc_metadata();
-        if (!malloc_data) {
-            auto new_malloc_data = make<MallocRegionMetadata>();
-            malloc_data = new_malloc_data.ptr();
-            static_cast<MmapRegion&>(*region).set_malloc_metadata({}, move(new_malloc_data));
-            malloc_data->address = region->base();
-            malloc_data->chunk_size = mmap_region.read32(offsetof(CommonHeader, m_size)).value();
+    MallocRegionMetadata* malloc_data = static_cast<MmapRegion&>(*region).malloc_metadata();
+    if (!malloc_data) {
+        auto new_malloc_data = make<MallocRegionMetadata>();
+        malloc_data = new_malloc_data.ptr();
+        static_cast<MmapRegion&>(*region).set_malloc_metadata({}, move(new_malloc_data));
+        malloc_data->address = region->base();
+        malloc_data->chunk_size = mmap_region.read32(offsetof(CommonHeader, m_size)).value();
+
+        bool is_chunked_block = malloc_data->chunk_size <= size_classes[num_size_classes - 1];
+        if (is_chunked_block)
             malloc_data->mallocations.resize((ChunkedBlock::block_size - sizeof(ChunkedBlock)) / malloc_data->chunk_size);
-            dbgln("Tracking ChunkedBlock @ {:p} with chunk_size={}, chunk_count={}", malloc_data->address, malloc_data->chunk_size, malloc_data->mallocations.size());
-        }
-        malloc_data->mallocation_for_address(address) = { address, size, true, false, Emulator::the().raw_backtrace(), Vector<FlatPtr>() };
-    } else {
-        m_big_mallocations.append({ address, size, true, false, Emulator::the().raw_backtrace(), Vector<FlatPtr>() });
+        else
+            malloc_data->mallocations.resize(1);
+        dbgln("Tracking malloc block @ {:p} with chunk_size={}, chunk_count={}", malloc_data->address, malloc_data->chunk_size, malloc_data->mallocations.size());
     }
+    malloc_data->mallocation_for_address(address) = { address, size, true, false, Emulator::the().raw_backtrace(), Vector<FlatPtr>() };
 }
 
 ALWAYS_INLINE Mallocation& MallocRegionMetadata::mallocation_for_address(FlatPtr address) const
@@ -106,6 +102,11 @@ ALWAYS_INLINE Mallocation& MallocRegionMetadata::mallocation_for_address(FlatPtr
 
 ALWAYS_INLINE size_t MallocRegionMetadata::chunk_index_for_address(FlatPtr address) const
 {
+    bool is_chunked_block = chunk_size <= size_classes[num_size_classes - 1];
+    if (!is_chunked_block) {
+        // This is a BigAllocationBlock
+        return 0;
+    }
     auto chunk_offset = address - (this->address + sizeof(ChunkedBlock));
     return chunk_offset / this->chunk_size;
 }
@@ -160,26 +161,28 @@ void MallocTracer::target_did_realloc(Badge<SoftCPU>, FlatPtr address, size_t si
     existing_mallocation->malloc_backtrace = Emulator::the().raw_backtrace();
 }
 
-Mallocation* MallocTracer::find_mallocation(FlatPtr address)
+Mallocation* MallocTracer::find_mallocation(const Region& region, FlatPtr address)
 {
-    if (auto* region = Emulator::the().mmu().find_region({ 0x23, address })) {
-        if (region->is_mmap() && static_cast<MmapRegion&>(*region).malloc_metadata()) {
-            auto& malloc_data = *static_cast<MmapRegion&>(*region).malloc_metadata();
-            auto& mallocation = malloc_data.mallocation_for_address(address);
-            if (mallocation.used) {
-                ASSERT(mallocation.contains(address));
-                return &mallocation;
-            }
-            return nullptr;
-        }
-    }
-
-    for (auto& mallocation : m_big_mallocations) {
-        if (mallocation.contains(address))
-            return &mallocation;
-    }
+    if (!region.is_mmap())
+        return nullptr;
+    if (!static_cast<const MmapRegion&>(region).is_malloc_block())
+        return nullptr;
+    auto* malloc_data = static_cast<MmapRegion&>(const_cast<Region&>(region)).malloc_metadata();
+    if (!malloc_data)
+        return nullptr;
+    auto& mallocation = malloc_data->mallocation_for_address(address);
+    if (!mallocation.used)
+        return nullptr;
+    ASSERT(mallocation.contains(address));
+    return &mallocation;
+}
 
-    return nullptr;
+Mallocation* MallocTracer::find_mallocation(FlatPtr address)
+{
+    auto* region = Emulator::the().mmu().find_region({ 0x23, address });
+    if (!region)
+        return nullptr;
+    return find_mallocation(*region, address);
 }
 
 Mallocation* MallocTracer::find_mallocation_before(FlatPtr address)
@@ -208,7 +211,7 @@ Mallocation* MallocTracer::find_mallocation_after(FlatPtr address)
     return found_mallocation;
 }
 
-void MallocTracer::audit_read(FlatPtr address, size_t size)
+void MallocTracer::audit_read(const Region& region, FlatPtr address, size_t size)
 {
     if (!m_auditing_enabled)
         return;
@@ -216,7 +219,7 @@ void MallocTracer::audit_read(FlatPtr address, size_t size)
     if (Emulator::the().is_in_malloc_or_free())
         return;
 
-    auto* mallocation = find_mallocation(address);
+    auto* mallocation = find_mallocation(region, address);
 
     if (!mallocation) {
         reportln("\n=={}==  \033[31;1mHeap buffer overflow\033[0m, invalid {}-byte read at address {:p}", getpid(), size, address);
@@ -250,7 +253,7 @@ void MallocTracer::audit_read(FlatPtr address, size_t size)
     }
 }
 
-void MallocTracer::audit_write(FlatPtr address, size_t size)
+void MallocTracer::audit_write(const Region& region, FlatPtr address, size_t size)
 {
     if (!m_auditing_enabled)
         return;
@@ -258,7 +261,7 @@ void MallocTracer::audit_write(FlatPtr address, size_t size)
     if (Emulator::the().is_in_malloc_or_free())
         return;
 
-    auto* mallocation = find_mallocation(address);
+    auto* mallocation = find_mallocation(region, address);
     if (!mallocation) {
         reportln("\n=={}==  \033[31;1mHeap buffer overflow\033[0m, invalid {}-byte write at address {:p}", getpid(), size, address);
         Emulator::the().dump_backtrace();
@@ -370,5 +373,4 @@ void MallocTracer::dump_leak_report()
     else
         reportln("\n=={}==  \033[31;1m{} leak(s) found: {} byte(s) leaked\033[0m", getpid(), leaks_found, bytes_leaked);
 }
-
 }
diff --git a/DevTools/UserspaceEmulator/MallocTracer.h b/DevTools/UserspaceEmulator/MallocTracer.h
index 01a0f7f700..c04ee43bf4 100644
--- a/DevTools/UserspaceEmulator/MallocTracer.h
+++ b/DevTools/UserspaceEmulator/MallocTracer.h
@@ -26,6 +26,7 @@
 
 #pragma once
 
+#include "SoftMMU.h"
 #include <AK/Badge.h>
 #include <AK/HashMap.h>
 #include <AK/OwnPtr.h>
@@ -71,8 +72,8 @@ public:
     void target_did_free(Badge<SoftCPU>, FlatPtr address);
     void target_did_realloc(Badge<SoftCPU>, FlatPtr address, size_t);
 
-    void audit_read(FlatPtr address, size_t);
-    void audit_write(FlatPtr address, size_t);
+    void audit_read(const Region&, FlatPtr address, size_t);
+    void audit_write(const Region&, FlatPtr address, size_t);
 
     void dump_leak_report();
 
@@ -80,13 +81,12 @@ private:
     template<typename Callback>
     void for_each_mallocation(Callback callback) const;
 
+    Mallocation* find_mallocation(const Region&, FlatPtr);
     Mallocation* find_mallocation(FlatPtr);
     Mallocation* find_mallocation_before(FlatPtr);
     Mallocation* find_mallocation_after(FlatPtr);
     bool is_reachable(const Mallocation&) const;
 
-    Vector<Mallocation> m_big_mallocations;
-
     bool m_auditing_enabled { true };
 };
 
diff --git a/DevTools/UserspaceEmulator/MmapRegion.cpp b/DevTools/UserspaceEmulator/MmapRegion.cpp
index 5f5aaf9354..19fbac94ef 100644
--- a/DevTools/UserspaceEmulator/MmapRegion.cpp
+++ b/DevTools/UserspaceEmulator/MmapRegion.cpp
@@ -75,7 +75,7 @@ ValueWithShadow<u8> MmapRegion::read8(FlatPtr offset)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_read(base() + offset, 1);
+            tracer->audit_read(*this, base() + offset, 1);
     }
 
     ASSERT(offset < size());
@@ -92,7 +92,7 @@ ValueWithShadow<u16> MmapRegion::read16(u32 offset)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_read(base() + offset, 2);
+            tracer->audit_read(*this, base() + offset, 2);
     }
 
     ASSERT(offset + 1 < size());
@@ -109,7 +109,7 @@ ValueWithShadow<u32> MmapRegion::read32(u32 offset)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_read(base() + offset, 4);
+            tracer->audit_read(*this, base() + offset, 4);
     }
 
     ASSERT(offset + 3 < size());
@@ -126,7 +126,7 @@ ValueWithShadow<u64> MmapRegion::read64(u32 offset)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_read(base() + offset, 8);
+            tracer->audit_read(*this, base() + offset, 8);
     }
 
     ASSERT(offset + 7 < size());
@@ -143,7 +143,7 @@ void MmapRegion::write8(u32 offset, ValueWithShadow<u8> value)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_write(base() + offset, 1);
+            tracer->audit_write(*this, base() + offset, 1);
     }
 
     ASSERT(offset < size());
@@ -161,7 +161,7 @@ void MmapRegion::write16(u32 offset, ValueWithShadow<u16> value)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_write(base() + offset, 2);
+            tracer->audit_write(*this, base() + offset, 2);
     }
 
     ASSERT(offset + 1 < size());
@@ -179,7 +179,7 @@ void MmapRegion::write32(u32 offset, ValueWithShadow<u32> value)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_write(base() + offset, 4);
+            tracer->audit_write(*this, base() + offset, 4);
     }
 
     ASSERT(offset + 3 < size());
@@ -198,7 +198,7 @@ void MmapRegion::write64(u32 offset, ValueWithShadow<u64> value)
 
     if (is_malloc_block()) {
         if (auto* tracer = Emulator::the().malloc_tracer())
-            tracer->audit_write(base() + offset, 8);
+            tracer->audit_write(*this, base() + offset, 8);
     }
 
     ASSERT(offset + 7 < size());
diff --git a/DevTools/UserspaceEmulator/SoftMMU.cpp b/DevTools/UserspaceEmulator/SoftMMU.cpp
index 4c4e1fad6e..b8c2f3b941 100644
--- a/DevTools/UserspaceEmulator/SoftMMU.cpp
+++ b/DevTools/UserspaceEmulator/SoftMMU.cpp
@@ -260,7 +260,7 @@ bool SoftMMU::fast_fill_memory8(X86::LogicalAddress address, size_t size, ValueW
         if (auto* tracer = Emulator::the().malloc_tracer()) {
             // FIXME: Add a way to audit an entire range of memory instead of looping here!
             for (size_t i = 0; i < size; ++i) {
-                tracer->audit_write(address.offset() + (i * sizeof(u8)), sizeof(u8));
+                tracer->audit_write(*region, address.offset() + (i * sizeof(u8)), sizeof(u8));
             }
         }
     }
@@ -285,7 +285,7 @@ bool SoftMMU::fast_fill_memory32(X86::LogicalAddress address, size_t count, Valu
         if (auto* tracer = Emulator::the().malloc_tracer()) {
             // FIXME: Add a way to audit an entire range of memory instead of looping here!
             for (size_t i = 0; i < count; ++i) {
-                tracer->audit_write(address.offset() + (i * sizeof(u32)), sizeof(u32));
+                tracer->audit_write(*region, address.offset() + (i * sizeof(u32)), sizeof(u32));
             }
         }
     }