Skip to content

Commit e7ac4a9

Browse files
CX330BlakeCX330Blake
committed
Optimize the optimization logic
1 parent 28cc72f commit e7ac4a9

File tree

3 files changed

+91
-53
lines changed

3 files changed

+91
-53
lines changed

‎backends/optimize.py‎

Lines changed: 70 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -208,13 +208,71 @@ def _find_reusable_zero_reg(start_idx: int) -> Optional[str]:
208208
return out
209209

210210

211-
def _x86_mov_reg_imm8(lines: List[str]) -> List[str]:
211+
def _x86_mov_reg_imm_narrow(lines: List[str], arch: str = "x86_64") -> List[str]:
212+
"""Narrow MOV destination to the smallest register that can hold the immediate.
213+
214+
Risky: using 8/16-bit forms will not zero upper bits of the 32/64-bit register.
215+
"""
212216
out: List[str] = []
217+
218+
def _arch_is_64(a: str) -> bool:
219+
al = (a or '').lower()
220+
return 'x86_64' in al or 'amd64' in al or 'x64' in al
221+
222+
is64 = _arch_is_64(arch)
223+
213224
# Capture indentation, spacing, register and immediate
214225
mov_re = re.compile(
215226
r"^(?P<indent>\s*)mov(?P<sep>\s+)(?P<reg>r(?:[0-9]{1,2}|[abcd]x|[sb]p|si|di)|e[abcd]x|e[sb]p|e[sd]i)\s*,\s*(?P<imm>0x[0-9a-fA-F]+|\d+)\s*(?P<cmt>;.*)?$",
216227
re.I,
217228
)
229+
230+
def _narrow_reg(reg: str, bits: int) -> Optional[str]:
231+
r = reg.lower()
232+
if is64:
233+
fam_map = {
234+
'rax': {8: 'al', 16: 'ax', 32: 'eax'},
235+
'rbx': {8: 'bl', 16: 'bx', 32: 'ebx'},
236+
'rcx': {8: 'cl', 16: 'cx', 32: 'ecx'},
237+
'rdx': {8: 'dl', 16: 'dx', 32: 'edx'},
238+
'rsi': {8: 'sil', 16: 'si', 32: 'esi'},
239+
'rdi': {8: 'dil', 16: 'di', 32: 'edi'},
240+
'rbp': {8: 'bpl', 16: 'bp', 32: 'ebp'},
241+
'rsp': {8: 'spl', 16: 'sp', 32: 'esp'},
242+
'eax': {8: 'al', 16: 'ax', 32: 'eax'},
243+
'ebx': {8: 'bl', 16: 'bx', 32: 'ebx'},
244+
'ecx': {8: 'cl', 16: 'cx', 32: 'ecx'},
245+
'edx': {8: 'dl', 16: 'dx', 32: 'edx'},
246+
'esi': {8: 'sil', 16: 'si', 32: 'esi'},
247+
'edi': {8: 'dil', 16: 'di', 32: 'edi'},
248+
'ebp': {8: 'bpl', 16: 'bp', 32: 'ebp'},
249+
'esp': {8: 'spl', 16: 'sp', 32: 'esp'},
250+
}
251+
if r in fam_map and bits in fam_map[r]:
252+
return fam_map[r][bits]
253+
if r.startswith('r') and r[1:].isdigit():
254+
try:
255+
n = int(r[1:])
256+
except Exception:
257+
n = -1
258+
if 8 <= n <= 15:
259+
return f"r{n}{ {8:'b',16:'w',32:'d'}[bits] }".replace(' ', '')
260+
return None
261+
else:
262+
fam_map = {
263+
'eax': {8: 'al', 16: 'ax', 32: 'eax'},
264+
'ebx': {8: 'bl', 16: 'bx', 32: 'ebx'},
265+
'ecx': {8: 'cl', 16: 'cx', 32: 'ecx'},
266+
'edx': {8: 'dl', 16: 'dx', 32: 'edx'},
267+
'esi': {16: 'si', 32: 'esi'}, # no sil in 32-bit mode
268+
'edi': {16: 'di', 32: 'edi'}, # no dil
269+
'ebp': {16: 'bp', 32: 'ebp'}, # no bpl
270+
'esp': {16: 'sp', 32: 'esp'}, # no spl
271+
}
272+
if r in fam_map and bits in fam_map[r]:
273+
return fam_map[r][bits]
274+
return None
275+
218276
for ln in lines:
219277
m = mov_re.match(ln)
220278
if not m:
@@ -230,30 +288,14 @@ def _x86_mov_reg_imm8(lines: List[str]) -> List[str]:
230288
except Exception:
231289
out.append(ln)
232290
continue
233-
if 0 <= imm <= 0xFF:
234-
# Map 64-bit regs to 8-bit low regs (x86_64)
235-
reg8_map = {
236-
'rax': 'al', 'rbx': 'bl', 'rcx': 'cl', 'rdx': 'dl',
237-
'rsi': 'sil', 'rdi': 'dil', 'rbp': 'bpl', 'rsp': 'spl',
238-
'eax': 'al', 'ebx': 'bl', 'ecx': 'cl', 'edx': 'dl',
239-
'esi': 'sil', 'edi': 'dil', 'ebp': 'bpl', 'esp': 'spl',
240-
}
241-
reg8 = None
242-
if reg in reg8_map:
243-
reg8 = reg8_map[reg]
244-
elif reg.startswith('r') and reg[1:].isdigit():
245-
# r8-r15 family
246-
try:
247-
num = int(reg[1:])
248-
if 8 <= num <= 15:
249-
reg8 = f"r{num}b"
250-
except Exception:
251-
reg8 = None
252-
# If we found a valid 8-bit register, emit replacement
253-
if reg8:
254-
out.append(f"{indent}mov{sep}{reg8}, {imm_s}{cmt}")
255-
else:
256-
out.append(ln)
291+
# Determine minimal width that fits the immediate
292+
width = 8 if 0 <= imm <= 0xFF else (16 if imm <= 0xFFFF else (32 if imm <= 0xFFFFFFFF else None))
293+
if width is None:
294+
out.append(ln)
295+
continue
296+
narrow = _narrow_reg(reg, int(width))
297+
if narrow:
298+
out.append(f"{indent}mov{sep}{narrow}, {imm_s}{cmt}")
257299
else:
258300
out.append(ln)
259301
return out
@@ -275,10 +317,10 @@ def default_rules_for_arch(arch: str) -> List[TransformRule]:
275317
rules.append(
276318
TransformRule(
277319
name="mov-reg-imm8-to-mov-reg8",
278-
description="Use mov reg8, imm8 when immediate fits in 8 bits (risky)",
279-
archs=["x86_64"],
320+
description="Use minimal-width destination for immediate (risky)",
321+
archs=["x86", "x86_64"],
280322
enabled=False,
281-
apply=lambda asm: _join(_x86_mov_reg_imm8(_normalize_asm(asm))),
323+
apply=lambda asm, arch=a: _join(_x86_mov_reg_imm_narrow(_normalize_asm(asm), arch)),
282324
)
283325
)
284326
return rules

‎ui/main_window.py‎

Lines changed: 20 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
from PySide6.QtGui import QFont, QAction, QPalette, QColor, QIcon, QPixmap, QPainter # type: ignore # QAction is in QtGui on Qt6
1111
from PySide6.QtWidgets import ( # type: ignore
1212
QApplication,
13-
QCheckBox,
1413
QComboBox,
1514
QGridLayout,
1615
QGroupBox,
@@ -37,7 +36,6 @@
3736
from PySide2.QtWidgets import ( # type: ignore
3837
QAction, # QAction is in QtWidgets on Qt5
3938
QApplication,
40-
QCheckBox,
4139
QComboBox,
4240
QGridLayout,
4341
QGroupBox,
@@ -186,18 +184,8 @@ def _is_null_pat(p) -> bool:
186184
tb.addWidget(QLabel("Arch:"))
187185
tb.addWidget(self.arch_combo)
188186
tb.addSeparator()
189-
# Block assemble toggle: allow labels/forward refs
190-
self.chk_block_asm = QCheckBox("Allow labels")
191-
try:
192-
self.chk_block_asm.setToolTip("Block assemble: preserve labels and forward references")
193-
except Exception:
194-
pass
195-
try:
196-
# Default off for safety; user can enable when needed
197-
self.chk_block_asm.setChecked(False)
198-
except Exception:
199-
pass
200-
tb.addWidget(self.chk_block_asm)
187+
# Labels are always allowed (block assembly mode enabled)
188+
# Checkbox removed; always assemble with labels preserved.
201189
tb.addSeparator()
202190
tb.addAction(self.act_assemble)
203191
tb.addAction(self.act_disassemble)
@@ -288,9 +276,9 @@ def _mode_provider() -> str:
288276
self.output_text = QPlainTextEdit()
289277
self._apply_mono(self.output_text)
290278
self.output_text.setReadOnly(True)
291-
# Apply comfortable inner padding similar to Syscalls tab visuals
279+
# Use default padding to match Optimize tab appearance (no extra viewport margins)
292280
try:
293-
self._apply_inner_padding(self.output_text, margin_px=10, viewport_pad=(6, 6, 6, 6))
281+
self._apply_inner_padding(self.output_text, margin_px=4, viewport_pad=(0, 0, 0, 0))
294282
except Exception:
295283
pass
296284
# Header row with "Send to Dev mode" action
@@ -1342,22 +1330,31 @@ def _sync_editors_padding_to_syscalls(self, margins, spacing: Optional[int] = No
13421330
doc_pad = max(2, min(16, doc_pad))
13431331
except Exception:
13441332
doc_pad = 6
1333+
# Editors that should follow Syscalls-style padding (inputs, validation)
13451334
for edit in (
13461335
getattr(self, 'asm_edit', None),
13471336
getattr(self, 'hex_edit', None),
1348-
getattr(self, 'output_text', None),
1349-
getattr(self, 'inline_text', None),
1337+
getattr(self, 'validation_text', None),
1338+
):
1339+
if edit is None:
1340+
continue
1341+
try:
1342+
self._apply_inner_padding(edit, margin_px=doc_pad, viewport_pad=pads)
1343+
except Exception:
1344+
continue
1345+
# Editors that should match Optimize tab style (no extra viewport margins, minimal document margin)
1346+
for edit in (
1347+
getattr(self, 'output_text', None), # Disassembly view
1348+
getattr(self, 'inline_text', None), # Shellcode outputs
13501349
getattr(self, 'hex_text', None),
13511350
getattr(self, 'hll_text', None),
1352-
getattr(self, 'opcode_text', None),
1351+
getattr(self, 'opcode_text', None), # Debug tab
13531352
getattr(self, 'debug_asm_text', None),
1354-
getattr(self, 'validation_text', None),
13551353
):
13561354
if edit is None:
13571355
continue
13581356
try:
1359-
# Apply both visible outer padding and an inner document margin
1360-
self._apply_inner_padding(edit, margin_px=doc_pad, viewport_pad=pads)
1357+
self._apply_inner_padding(edit, margin_px=4, viewport_pad=(0, 0, 0, 0))
13611358
except Exception:
13621359
continue
13631360

@@ -1779,8 +1776,7 @@ def on_assemble(self):
17791776
return
17801777
arch = self.arch_combo.currentText() or "x86_64"
17811778
try:
1782-
allow_labels = bool(getattr(self, 'chk_block_asm', None) and self.chk_block_asm.isChecked())
1783-
data = self.adapter.assemble(asm, arch_name=arch, platform_name=None, addr=0, allow_labels=allow_labels)
1779+
data = self.adapter.assemble(asm, arch_name=arch, platform_name=None, addr=0, allow_labels=True)
17841780
except Exception as e:
17851781
QMessageBox.critical(self, "Assemble Error", f"{e}\n\n{traceback.format_exc()}")
17861782
return

‎ui/optimize_panel.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ def __init__(
4040
# Header controls
4141
top = QHBoxLayout()
4242
self.chk_rule1 = QCheckBox("Push 0 -> xor; push reg")
43-
self.chk_rule2 = QCheckBox("mov reg, imm8 -> mov reg8, imm8 (risky)")
43+
self.chk_rule2 = QCheckBox("mov reg, imm -> mov reg8/16/32, imm (risky)")
4444
self.btn_apply = QPushButton("Apply All")
4545
top.addWidget(self.chk_rule1)
4646
top.addWidget(self.chk_rule2)

0 commit comments

Comments
 (0)