From 506f89584118e194b0d0f542d26aed36b136ae79 Mon Sep 17 00:00:00 2001
From: OBarronCS <55004530+OBarronCS@users.noreply.github.com>
Date: Tue, 27 May 2025 02:56:34 -0700
Subject: [PATCH] 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       | 15 +++++++++++----
 pwndbg/dbg/__init__.py      | 24 +++++++++++++++++++++++-
 pwndbg/dbg/gdb/__init__.py  |  6 +-----
 pwndbg/dbg/lldb/__init__.py |  6 +-----
 4 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/pwndbg/aglib/vmmap.py b/pwndbg/aglib/vmmap.py
index 86c668df3..e607ccac1 100644
--- a/pwndbg/aglib/vmmap.py
+++ b/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)
diff --git a/pwndbg/dbg/__init__.py b/pwndbg/dbg/__init__.py
index 22b01d7d6..13e08da0b 100644
--- a/pwndbg/dbg/__init__.py
+++ b/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:
diff --git a/pwndbg/dbg/gdb/__init__.py b/pwndbg/dbg/gdb/__init__.py
index b20243527..bd757e94f 100644
--- a/pwndbg/dbg/gdb/__init__.py
+++ b/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,
diff --git a/pwndbg/dbg/lldb/__init__.py b/pwndbg/dbg/lldb/__init__.py
index 0f0c6142d..077fcd2f3 100644
--- a/pwndbg/dbg/lldb/__init__.py
+++ b/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