[SharedCache] Have VM track ranges rather than pages

There are typically only a few dozen mappings, while there can be
millions of pages. This reduces the amount of time spent populating the
mapping from region to file accessor along with the memory usage of the
same.

2 files changed, 37 insertions, 37 deletions

diff --git a/view/sharedcache/core/VM.cpp b/view/sharedcache/core/VM.cpp
index c351a146..5aca51d9 100644
--- a/view/sharedcache/core/VM.cpp
+++ b/view/sharedcache/core/VM.cpp
@@ -475,10 +475,6 @@ void MMappedFileAccessor::Read(void* dest, size_t address, size_t length)
 
 VM::VM(size_t pageSize, bool safe) : m_pageSize(pageSize), m_safe(safe)
 {
-	unsigned bits, var = (m_pageSize - 1 < 0) ? -(m_pageSize - 1) : m_pageSize - 1;
-	for (bits = 0; var != 0; ++bits)
-		var >>= 1;
-	m_pageSizeBits = bits;
 }
 
 VM::~VM()
@@ -498,36 +494,25 @@ void VM::MapPages(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, uint64_t se
 		throw MappingPageAlignmentException();
 	}
 
-	size_t pagesRemainingCount = size / m_pageSize;
-	for (size_t i = 0; i < size; i += m_pageSize)
+	auto accessor = MMappedFileAccessor::Open(std::move(dscView), sessionID, filePath, postAllocationRoutine);
+	auto [it, inserted] = m_map.insert_or_assign({vm_address, vm_address + size}, PageMapping(std::move(filePath), std::move(accessor), fileoff));
+	if (m_safe && !inserted)
 	{
-		// Our pages will be delimited by shifting off the page size
-		// So, 0x12345000 will become 0x12345 (assuming m_pageSize is 0x1000)
-		auto page = (vm_address + (i)) >> m_pageSizeBits;
-		if (m_map.count(page) != 0)
-		{
-			if (m_safe)
-			{
-				BNLogWarn("Remapping page 0x%lx (i == 0x%lx) (a: 0x%zx, f: 0x%zx)", page, i, vm_address, fileoff);
-				throw MappingCollisionException();
-			}
-		}
-		m_map.insert_or_assign(page, PageMapping(filePath, MMappedFileAccessor::Open(dscView, sessionID, filePath, postAllocationRoutine), i + fileoff));
+		BNLogWarn("Remapping page 0x%zx (f: 0x%zx)", vm_address, fileoff);
+		throw MappingCollisionException();
 	}
 }
 
 std::pair<PageMapping, size_t> VM::MappingAtAddress(size_t address)
 {
-	// Get the page (e.g. 0x12345678 will become 0x12345 on 0x1000 aligned caches)
-	auto page = address >> m_pageSizeBits;
-	if (auto f = m_map.find(page); f != m_map.end())
+	if (auto it = m_map.find(address); it != m_map.end())
 	{
 		// The PageMapping object returned contains the page, and more importantly, the file pointer (there can be
-		// multiple in newer caches) This is relevant for reading out the data in the rest of this file. The second item
-		// in this pair is created by taking the fileOffset (which will be a page but with the trailing bits (e.g.
-		// 0x12345000)
-		//      and will add the "extra" bits lopped off when determining the page. (e.g. 0x12345678 -> 0x678)
-		return {f->second, f->second.fileOffset + (address & (m_pageSize - 1))};
+		// multiple in newer caches) This is relevant for reading out the data in the rest of this file.
+		// The second item in the returned pair is the offset of `address` within the file.
+		auto& range = it->first;
+		auto& mapping = it->second;
+		return {mapping, mapping.fileOffset + (address - range.start)};
 	}
 
 	throw MappingReadException();
@@ -536,14 +521,8 @@ std::pair<PageMapping, size_t> VM::MappingAtAddress(size_t address)
 
 bool VM::AddressIsMapped(uint64_t address)
 {
-	try
-	{
-		MappingAtAddress(address);
-		return true;
-	}
-	catch (...)
-	{}
-	return false;
+	auto it = m_map.find(address);
+	return it != m_map.end();
 }
 
 
diff --git a/view/sharedcache/core/VM.h b/view/sharedcache/core/VM.h
index b8e5c59b..e47cf15e 100644
--- a/view/sharedcache/core/VM.h
+++ b/view/sharedcache/core/VM.h
@@ -183,7 +183,7 @@ struct PageMapping {
 	std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> fileAccessor;
     size_t fileOffset;
 	PageMapping(std::string filePath, std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> fileAccessor, size_t fileOffset)
-		: filePath(filePath), fileAccessor(fileAccessor), fileOffset(fileOffset) {}
+		: filePath(std::move(filePath)), fileAccessor(std::move(fileAccessor)), fileOffset(fileOffset) {}
 };
 
 
@@ -215,9 +215,30 @@ class VMReader;
 
 
 class VM {
-    std::map<size_t, PageMapping> m_map;
+
+    // Represents a range of addresses [start, end).
+    // Note that `end` is not included within the range.
+    struct AddressRange {
+        size_t start;
+        size_t end;
+
+        bool operator<(const AddressRange& b) const {
+            return start < b.start || (start == b.start && end < b.end);
+        }
+
+        friend bool operator<(const AddressRange& range, size_t address) {
+            return range.end <= address;
+        }
+
+        friend bool operator<(size_t address, const AddressRange& range) {
+            return address < range.start;
+        }
+    };
+
+    // A map keyed by address ranges that can be looked up via any
+    // address within a range thanks to C++14's transparent comparators.
+    std::map<AddressRange, PageMapping, std::less<>> m_map;
     size_t m_pageSize;
-    size_t m_pageSizeBits;
     bool m_safe;
 
     friend VMReader;