Optimize page lookups (#3023)

* Optimize page lookups

* Handle case where these are non-aligned pages

* Non-aligned page check

* Use a binary search method instead

* comment

* Remove unused variable

---------

Co-authored-by: Disconnect3d <dominik.b.czarnota@gmail.com>

pwndbg/aglib/vmmap.py

@@ -6,6 +6,7 @@ import pwndbg
 import pwndbg.aglib.vmmap_custom
 import pwndbg.lib.cache
 import pwndbg.lib.memory
+from pwndbg.dbg import MemoryMap
 
 pwndbg.config.add_param(
     "vmmap-prefer-relpaths",
@@ -15,9 +16,14 @@ pwndbg.config.add_param(
 )
 
 
+@pwndbg.lib.cache.cache_until("start", "stop")
+def get_memory_map() -> MemoryMap:
+    return pwndbg.dbg.selected_inferior().vmmap()
+
+
 @pwndbg.lib.cache.cache_until("start", "stop")
 def get() -> Tuple[pwndbg.lib.memory.Page, ...]:
-    return tuple(pwndbg.dbg.selected_inferior().vmmap().ranges())
+    return tuple(get_memory_map().ranges())
 
 
 @pwndbg.lib.cache.cache_until("start", "stop")
@@ -29,8 +35,9 @@ def find(address: int | pwndbg.dbg_mod.Value | None) -> pwndbg.lib.memory.Page |
     if address < 0:
         return None
 
-    for page in get():
-        if address in page:
-            return page
+    page = get_memory_map().lookup_page(address)
+
+    if page is not None:
+        return page
 
     return pwndbg.aglib.vmmap_custom.explore(address)

pwndbg/dbg/__init__.py

@@ -288,6 +288,11 @@ class MemoryMap:
     A wrapper around a sequence of memory ranges
     """
 
+    pages: tuple[pwndbg.lib.memory.Page, ...]
+
+    def __init__(self, pages: Sequence[pwndbg.lib.memory.Page]):
+        self.pages = tuple(pages)
+
     def is_qemu(self) -> bool:
         """
         Returns whether this memory map was generated from a QEMU target.
@@ -298,7 +303,24 @@ class MemoryMap:
         """
         Returns all ranges in this memory map.
         """
-        raise NotImplementedError()
+        return self.pages
+
+    def lookup_page(self, address: int) -> pwndbg.lib.memory.Page | None:
+        # Binary search for the page
+        lo = 0
+        hi = len(self.pages) - 1
+        while lo <= hi:
+            mid = (hi + lo) // 2
+            page = self.pages[mid]
+            if page.start <= address:
+                if address < page.end:
+                    return page
+
+                lo = mid + 1
+            else:
+                hi = mid - 1
+
+        return None
 
 
 class ExecutionController:

pwndbg/dbg/gdb/__init__.py

@@ -233,17 +233,13 @@ class GDBThread(pwndbg.dbg_mod.Thread):
 
 class GDBMemoryMap(pwndbg.dbg_mod.MemoryMap):
     def __init__(self, qemu: bool, pages: Sequence[pwndbg.lib.memory.Page]):
+        super().__init__(pages)
         self.qemu = qemu
-        self.pages = pages
 
     @override
     def is_qemu(self) -> bool:
         return self.qemu
 
-    @override
-    def ranges(self) -> Sequence[pwndbg.lib.memory.Page]:
-        return self.pages
-
 
 # While this implementation allows breakpoints to be deleted, enabled and
 # disabled from inside the code in a stop handler, GDB does not[1]. Aditionally,

pwndbg/dbg/lldb/__init__.py

@@ -634,17 +634,13 @@ class LLDBValue(pwndbg.dbg_mod.Value):
 
 class LLDBMemoryMap(pwndbg.dbg_mod.MemoryMap):
     def __init__(self, pages: List[pwndbg.lib.memory.Page]):
-        self.pages = pages
+        super().__init__(pages)
 
     @override
     def is_qemu(self) -> bool:
         # TODO/FIXME: Figure a way to detect QEMU later.
         return False
 
-    @override
-    def ranges(self) -> List[pwndbg.lib.memory.Page]:
-        return self.pages
-
 
 class LLDBStopPoint(pwndbg.dbg_mod.StopPoint):
     inner: lldb.SBBreakpoint | lldb.SBWatchpoint