Repository: pwndbg/pwndbg
Pull request #2979: A couple disassembly performance optimizations
Author: login: OBarronCS
Created: 2025-05-16T01:14:23Z | Merged: 2025-05-22T09:35:23Z

Small tweaks to avoid some unnecessary function calls related to disassembly. See comments for specific details. Changes save ~1.5 ms per context call.


Diff:
diff --git a/pwndbg/aglib/disasm/arch.py b/pwndbg/aglib/disasm/arch.py
index 76a91bf3f71..b58f2a53b97 100644
--- a/pwndbg/aglib/disasm/arch.py
+++ b/pwndbg/aglib/disasm/arch.py
@@ -1049,3 +1049,23 @@ def _common_binary_op_annotator(
             instruction.annotation = memory_or_register_assign(
                 target_operand.str, math_string, memory_assignment
             )
+
+
+def basic_enhance(ins: PwndbgInstruction) -> None:
+    # Apply syntax highlighting and inline symbol replacement
+    # Used in cases were we don't want to do the full enhancement process
+    # for performance reasons.
+    if pwndbg.config.syntax_highlight:
+        ins.asm_string = syntax_highlight(ins.asm_string)
+
+    if pwndbg.config.disasm_inline_symbols:
+        # Make inline replacements, so `jmp 0x400122` becomes `jmp function_name`
+        for op in ins.operands:
+            if op.type is CS_OP_IMM:
+                op.before_value = op.imm
+
+                if op.before_value >= 0:
+                    op.symbol = MemoryColor.attempt_colorized_symbol(op.before_value)
+
+                if op.symbol:
+                    ins.asm_string = ins.asm_string.replace(hex(op.before_value), op.symbol)
diff --git a/pwndbg/aglib/disasm/disassembly.py b/pwndbg/aglib/disasm/disassembly.py
index 49cd18f4e8f..dfce409d0a8 100644
--- a/pwndbg/aglib/disasm/disassembly.py
+++ b/pwndbg/aglib/disasm/disassembly.py
@@ -19,6 +19,7 @@
 import pwndbg
 import pwndbg.aglib.arch
 import pwndbg.aglib.disasm.aarch64
+import pwndbg.aglib.disasm.arch
 import pwndbg.aglib.disasm.arm
 import pwndbg.aglib.disasm.disassembly
 import pwndbg.aglib.disasm.loongarch64
@@ -98,11 +99,8 @@ def clear_on_reg_mem_change() -> None:
 
 
 @pwndbg.lib.cache.cache_until("objfile")
-def get_disassembler(address: int, cs_info: Tuple[int, int] = None):
-    if cs_info is not None:
-        arch, mode = cs_info
-    else:
-        arch, mode = pwndbg.aglib.arch.get_capstone_constants(address)
+def get_disassembler(cs_info: Tuple[int, int]):
+    arch, mode = cs_info
 
     mode |= CapstoneEndian[pwndbg.aglib.arch.endian]
 
@@ -136,9 +134,12 @@ def get_one_instruction(
 
     cs_info = pwndbg.aglib.arch.get_capstone_constants(address)
     if cs_info is None:
-        return ManualPwndbgInstruction(address)
+        instr = ManualPwndbgInstruction(address)
+        if enhance:
+            pwndbg.aglib.disasm.arch.basic_enhance(instr)
+        return instr
 
-    md = get_disassembler(address, cs_info)
+    md = get_disassembler(cs_info)
     data = pwndbg.aglib.memory.read(address, pwndbg.aglib.arch.max_instruction_size, partial=True)
     for ins in md.disasm(bytes(data), address, 1):
         pwn_ins: PwndbgInstruction = PwndbgInstructionImpl(ins)
diff --git a/pwndbg/commands/rop.py b/pwndbg/commands/rop.py
index ab87ea2c33d..eae1aad1a43 100644
--- a/pwndbg/commands/rop.py
+++ b/pwndbg/commands/rop.py
@@ -24,7 +24,7 @@ def __init__(self, options, start_addr: int):
         self.start_addr = start_addr
         self.__fileName = options.binary
         self.__rawBinary = None
-        self.cs = get_disassembler(pwndbg.aglib.regs.pc)
+        self.cs = get_disassembler(pwndbg.aglib.arch.get_capstone_constants(pwndbg.aglib.regs.pc))
 
         with open(self.__fileName, "rb") as fp:
             self.__rawBinary = fp.read()
diff --git a/pwndbg/emu/emulator.py b/pwndbg/emu/emulator.py
index 3d2c5d1d155..6d4aa7b0407 100644
--- a/pwndbg/emu/emulator.py
+++ b/pwndbg/emu/emulator.py
@@ -827,7 +827,7 @@ def until_syscall_hook_code(self, uc, address, size: int, user_data) -> None:
         )
         self.until_syscall_address = address
 
-    def single_step(self, pc=None) -> Tuple[int, int]:
+    def single_step(self, pc=None, check_instruction=False) -> Tuple[int, int]:
         """Steps one instruction.
 
         Yields:
@@ -844,22 +844,23 @@ def single_step(self, pc=None) -> Tuple[int, int]:
 
         pc = pc or self.pc
 
-        insn = pwndbg.aglib.disasm.disassembly.one_raw(pc)
+        if check_instruction or DEBUG & DEBUG_EXECUTING:
+            insn = pwndbg.aglib.disasm.disassembly.one_raw(pc)
 
-        # If we don't know how to disassemble, bail.
-        if insn is None:
-            debug(DEBUG_EXECUTING, "Can't disassemble instruction at %#x", pc)
-            return self.last_single_step_result
+            # If we don't know how to disassemble, bail.
+            if insn is None:
+                debug(DEBUG_EXECUTING, "Can't disassemble instruction at %#x", pc)
+                return self.last_single_step_result
 
-        if insn.id in BANNED_INSTRUCTIONS.get(self.arch, {}):
-            debug(DEBUG_EXECUTING, "Hit illegal instruction at %#x", pc)
-            return self.last_single_step_result
+            if insn.id in BANNED_INSTRUCTIONS.get(self.arch, {}):
+                debug(DEBUG_EXECUTING, "Hit illegal instruction at %#x", pc)
+                return self.last_single_step_result
 
-        debug(
-            DEBUG_EXECUTING,
-            "# Emulator attempting to single-step at %#x: %s %s",
-            (pc, insn.mnemonic, insn.op_str),
-        )
+            debug(
+                DEBUG_EXECUTING,
+                "# Instruction: attempting to single-step at %#x: %s %s",
+                (pc, insn.mnemonic, insn.op_str),
+            )
 
         try:
             self.single_step_hook_hit_count = 0
diff --git a/pwndbg/enhance.py b/pwndbg/enhance.py
index ddc62634956..e8e9f62bcd4 100644
--- a/pwndbg/enhance.py
+++ b/pwndbg/enhance.py
@@ -112,8 +112,9 @@ def enhance(
         rwx = exe = False
 
     if exe:
-        pwndbg_instr = pwndbg.aglib.disasm.disassembly.one(value)
+        pwndbg_instr = pwndbg.aglib.disasm.disassembly.one_raw(value)
         if pwndbg_instr:
+            pwndbg.aglib.disasm.arch.basic_enhance(pwndbg_instr)
             # For telescoping, we don't want the extra spaces between the mnemonic and operands
             # which are baked in during enhancement. This removes those spaces.
             instr = " ".join(pwndbg_instr.asm_string.split())
diff --git a/tests/qemu-tests/tests/user/test_arm.py b/tests/qemu-tests/tests/user/test_arm.py
index 93e57659974..ae8ed44b40d 100644
--- a/tests/qemu-tests/tests/user/test_arm.py
+++ b/tests/qemu-tests/tests/user/test_arm.py
@@ -745,10 +745,10 @@ def test_arm_it_block(qemu_assembly_run):
         "─────────────────[ DISASM / arm / thumb mode / set emulate on ]─────────────────\n"
         " ► 0x200bc <_start+8>     cmp    r0, #0     0x200bd - 0x0     CPSR => 0x20000030 [ n z C v q j T e a i f ]\n"
         "   0x200be <_start+10>    ittte  eq\n"
-        "   0x200c0 <_start+12>    movs   r1, #1     R1 => 1\n"
-        "   0x200c2 <_start+14>    movs   r2, #2     R2 => 2\n"
-        "   0x200c4 <_start+16>    movs   r2, #3     R2 => 3\n"
-        "   0x200c6 <_start+18>    movs   r1, #4     R1 => 4\n"
+        "   0x200c0 <_start+12>    moveq  r1, #1     R1 => 1\n"
+        "   0x200c2 <_start+14>    moveq  r2, #2     R2 => 2\n"
+        "   0x200c4 <_start+16>    moveq  r2, #3     R2 => 3\n"
+        "   0x200c6 <_start+18>    movne  r1, #4     R1 => 4\n"
         "   0x200c8 <_start+20>    nop    \n"
         "   0x200ca <_start+22>    nop    \n"
         "   0x200cc <_start+24>    nop    \n"
@@ -759,6 +759,25 @@ def test_arm_it_block(qemu_assembly_run):
 
     assert dis_1 == expected_1
 
+
+def test_arm_it_block_step_into(qemu_assembly_run):
+    """
+    Tests 2 things:
+    - Stepping into the IT block doesn't break our usage Unicorn (some of the instructions would appear to be jumps if so)
+    - Instructions have IT mode suffixes
+
+    TODO:
+    - Fix suffixes not appearing here due to https://github.com/capstone-engine/capstone/issues/2702
+    - See comments in https://github.com/pwndbg/pwndbg/pull/2979
+    """
+    qemu_assembly_run(ARM_IT_BLOCK, "arm")
+
+    gdb.execute("si")
+    gdb.execute("si")
+
+    # Prime the instruction cache
+    gdb.execute("context disasm", to_string=True)
+
     # Now, ensure that once we step into the block, the disassembly is still correct.
     gdb.execute("si")
     gdb.execute("si")
@@ -772,7 +791,7 @@ def test_arm_it_block(qemu_assembly_run):
         "─────────────────[ DISASM / arm / thumb mode / set emulate on ]─────────────────\n"
         "   0x200bc <_start+8>     cmp    r0, #0     0x200bd - 0x0     CPSR => 0x20000030 [ n z C v q j T e a i f ]\n"
         "   0x200be <_start+10>    ittte  eq\n"
-        "   0x200c0 <_start+12>    movs   r1, #1     R1 => 1\n"
+        "   0x200c0 <_start+12>    moveq  r1, #1     R1 => 1\n"
         " ► 0x200c2 <_start+14>    movs   r2, #2     R2 => 2\n"
         "   0x200c4 <_start+16>    movs   r2, #3     R2 => 3\n"
         "   0x200c6 <_start+18>    movs   r1, #4     R1 => 4\n"
@@ -807,10 +826,10 @@ def test_arm_it_block_cached_thumb_mode(qemu_assembly_run):
         "    ↓\n"
         "   0x200bc <_start+8>     cmp    r0, #0         0x200bd - 0x0     CPSR => 0x20000030 [ n z C v q j T e a i f ]\n"
         "   0x200be <_start+10>    ittte  eq\n"
-        "   0x200c0 <_start+12>    movs   r1, #1         R1 => 1\n"
-        "   0x200c2 <_start+14>    movs   r2, #2         R2 => 2\n"
-        "   0x200c4 <_start+16>    movs   r2, #3         R2 => 3\n"
-        "   0x200c6 <_start+18>    movs   r1, #4         R1 => 4\n"
+        "   0x200c0 <_start+12>    moveq  r1, #1         R1 => 1\n"
+        "   0x200c2 <_start+14>    moveq  r2, #2         R2 => 2\n"
+        "   0x200c4 <_start+16>    moveq  r2, #3         R2 => 3\n"
+        "   0x200c6 <_start+18>    movne  r1, #4         R1 => 4\n"
         "   0x200c8 <_start+20>    nop    \n"
         "   0x200ca <_start+22>    nop    \n"
         "   0x200cc <_start+24>    nop    \n"


Reviews:
user: login: OBarronCS
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T01:17:44Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR introduces performance optimizations in disassembly and related workflows by reducing unnecessary function calls and streamlining instruction enhancement. Key changes include updating the expected disassembly output in tests, switching from one() to one_raw() in the instruction enhancement code, and modifying emulator single‐step behavior with an additional parameter for instruction checking.

### Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

<details>
<summary>Show a summary per file</summary>

| File                                                   | Description                                                      |
| ------------------------------------------------------ | ---------------------------------------------------------------- |
| tests/qemu-tests/tests/user/test_arm.py                | Updated expected disassembly strings to reflect optimized ops      |
| pwndbg/enhance.py                                      | Replaced one() with one_raw() and added an inline basic enhancement |
| pwndbg/emu/emulator.py                                 | Added a check_instruction parameter and updated debug messaging    |
| pwndbg/commands/rop.py                                 | Adjusted disassembler initialization for correct Capstone setup      |
| pwndbg/aglib/disasm/disassembly.py                     | Updated get_disassembler signature and enhancement call flow         |
| pwndbg/aglib/disasm/arch.py                            | Introduced basic_enhance to apply lightweight inline enhancements      |
</details>



<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**pwndbg/emu/emulator.py:830**
* Consider updating the function's docstring to document the new 'check_instruction' parameter and explain its behavior.
```
def single_step(self, pc=None, check_instruction=False) -> Tuple[int, int>:
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T09:31:44Z
user: login: disconnect3d
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T09:32:55Z
user: login: disconnect3d
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T09:33:16Z
user: login: disconnect3d
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T09:34:30Z
user: login: disconnect3d
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-22T09:35:04Z

Review comments:
diff_hunk: @@ -98,11 +99,8 @@ def clear_on_reg_mem_change() -> None:
 
 
 @pwndbg.lib.cache.cache_until("objfile")
-def get_disassembler(address: int, cs_info: Tuple[int, int] = None):
-    if cs_info is not None:
-        arch, mode = cs_info
-    else:
-        arch, mode = pwndbg.aglib.arch.get_capstone_constants(address)
+def get_disassembler(cs_info: Tuple[int, int]):
path: pwndbg/aglib/disasm/disassembly.py
user: login: OBarronCS
body: Change this function so it no longer takes address - this was messing with the cache behavior, resulting in creating the Capstone object multiple times.
created_at: 2025-05-16T01:15:00Z
author_association: MEMBER
diff_hunk: @@ -844,22 +844,23 @@ def single_step(self, pc=None) -> Tuple[int, int]:
 
         pc = pc or self.pc
 
-        insn = pwndbg.aglib.disasm.disassembly.one_raw(pc)
+        if check_instruction or DEBUG & DEBUG_EXECUTING:
path: pwndbg/emu/emulator.py
user: login: OBarronCS
body: No need to check that this instruction is valid by default - in current usage, this function wouldn't even get called on an invalid instruction.
created_at: 2025-05-16T01:16:12Z
author_association: MEMBER
diff_hunk: @@ -112,8 +112,9 @@ def enhance(
         rwx = exe = False
 
     if exe:
-        pwndbg_instr = pwndbg.aglib.disasm.disassembly.one(value)
+        pwndbg_instr = pwndbg.aglib.disasm.disassembly.one_raw(value)
path: pwndbg/enhance.py
user: login: OBarronCS
body: `one()` goes through the entire enhancement process, which does work unrelated to displaying an instruction while telescoping. 
![image](https://github.com/user-attachments/assets/5adf2061-f63a-426a-8c56-6c74e8a0e54f)
In telescoping, we only care about syntax highlighting and replacing constants with their symbols.
created_at: 2025-05-16T01:17:40Z
author_association: MEMBER
diff_hunk: @@ -112,8 +112,9 @@ def enhance(
         rwx = exe = False
 
     if exe:
-        pwndbg_instr = pwndbg.aglib.disasm.disassembly.one(value)
+        pwndbg_instr = pwndbg.aglib.disasm.disassembly.one_raw(value)
path: pwndbg/enhance.py
user: login: Copilot
body: [nitpick] Consider adding a comment explaining the rationale for using one_raw() over one() to clarify any differences in behavior or performance implications.
created_at: 2025-05-22T09:31:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,11 +99,8 @@ def clear_on_reg_mem_change() -> None:
 
 
 @pwndbg.lib.cache.cache_until("objfile")
-def get_disassembler(address: int, cs_info: Tuple[int, int] = None):
-    if cs_info is not None:
-        arch, mode = cs_info
-    else:
-        arch, mode = pwndbg.aglib.arch.get_capstone_constants(address)
+def get_disassembler(cs_info: Tuple[int, int]):
path: pwndbg/aglib/disasm/disassembly.py
user: login: Copilot
body: Consider updating or adding a docstring for get_disassembler to reflect its new signature and clarify its expected parameters.
created_at: 2025-05-22T09:31:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1049,3 +1049,23 @@ def _common_binary_op_annotator(
             instruction.annotation = memory_or_register_assign(
                 target_operand.str, math_string, memory_assignment
             )
+
+
+def basic_enhance(ins: PwndbgInstruction) -> None:
path: pwndbg/aglib/disasm/arch.py
user: login: Copilot
body: Consider adding a docstring for basic_enhance to explain its purpose and usage, especially given its performance-sensitive context.
created_at: 2025-05-22T09:31:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -98,11 +99,8 @@ def clear_on_reg_mem_change() -> None:
 
 
 @pwndbg.lib.cache.cache_until("objfile")
-def get_disassembler(address: int, cs_info: Tuple[int, int] = None):
-    if cs_info is not None:
-        arch, mode = cs_info
-    else:
-        arch, mode = pwndbg.aglib.arch.get_capstone_constants(address)
+def get_disassembler(cs_info: Tuple[int, int]):
+    arch, mode = cs_info
 
     mode |= CapstoneEndian[pwndbg.aglib.arch.endian]
path: pwndbg/aglib/disasm/disassembly.py
user: login: disconnect3d
body: Btw why do we need to do this? Shouldn't the proper endian be already in place? Can we add a comment here?
created_at: 2025-05-22T09:32:54Z
author_association: MEMBER
diff_hunk: @@ -98,11 +99,8 @@ def clear_on_reg_mem_change() -> None:
 
 
 @pwndbg.lib.cache.cache_until("objfile")
-def get_disassembler(address: int, cs_info: Tuple[int, int] = None):
-    if cs_info is not None:
-        arch, mode = cs_info
-    else:
-        arch, mode = pwndbg.aglib.arch.get_capstone_constants(address)
+def get_disassembler(cs_info: Tuple[int, int]):
+    arch, mode = cs_info
path: pwndbg/aglib/disasm/disassembly.py
user: login: disconnect3d
body: Random comment: since we unpack it anyway, maybe its better to just accept/pass `arch, mode`? Idk, just saying, maybe its better as it is.
created_at: 2025-05-22T09:33:16Z
author_association: MEMBER
diff_hunk: @@ -844,22 +844,23 @@ def single_step(self, pc=None) -> Tuple[int, int]:
 
         pc = pc or self.pc
 
-        insn = pwndbg.aglib.disasm.disassembly.one_raw(pc)
+        if check_instruction or DEBUG & DEBUG_EXECUTING:
path: pwndbg/emu/emulator.py
user: login: disconnect3d
body: Are u saying this func is never called? If so, why do we even have it here?
created_at: 2025-05-22T09:34:29Z
author_association: MEMBER

PR comments:
user: login: OBarronCS
created_at: 2025-05-16T01:45:11Z
body: Interesting - this somehow improves the disassembly of ARM IT blocks - the instructions in the IT block now have their condition attached to the instruction (`moveq`, `movne`).
![image](https://github.com/user-attachments/assets/9428413b-94d5-4ab2-aa5d-9329d0c69d88)
Now that the Capstone object is persistent, capstone somehow now relates those instructions to the `IT` instruction before it? The Thumb2 instructions in the IT block, to my understanding, don't have the condition encoded into them - they are reliant on the IT instruction.
author_association: MEMBER
user: login: OBarronCS
created_at: 2025-05-16T02:10:13Z
body: Very strange behavior:

When stepped on the IT block, all future instructions now have the condition codes in the assembly (even the nops, which are not in the IT block)
![image](https://github.com/user-attachments/assets/e4c43a83-c14c-4f95-87a2-4134c551e928)

Stepping one down, however:
![image](https://github.com/user-attachments/assets/4bc4d450-7b2f-4cfc-a749-bd9f0a15f3e7)

But `nearpc` only makes two of the instructions have the condition?? 
![image](https://github.com/user-attachments/assets/e2256be3-7b94-411f-9d01-adcc2e01ef32)

And disabling the regs section changes the output of disassembly (which could disassemble it when telescoping for PC reg)
![image](https://github.com/user-attachments/assets/4b164df2-c2df-40fd-9d56-c00f4c1d1af9)

Some internal state is being tracked in the Capstone disassembler object?
author_association: MEMBER
user: login: OBarronCS
created_at: 2025-05-16T03:58:13Z
body: Made an issue on https://github.com/capstone-engine/capstone/issues/2702 related to this
author_association: MEMBER
user: login: OBarronCS
created_at: 2025-05-22T06:23:59Z
body: Fixed the ARM IT block tests to reflect the new behavior. In short, before this PR the instructions were always missing their conditional suffix. After this PR, the disassembly is not correct in all cases in the context view (conditional suffixes are missing sometimes, sometimes too many instructions have the conditional suffix) but it is correct when disassembling with nearpc or emulate outside of the context view.

Making the disassembly completely correct in the context view would take larger changes that should be done in a different PR, but the change in behavior (accidentally) encountered here is a step in a better direction for IT instruction blocks.
author_association: MEMBER