forked from CX330Blake/Shellcode-IDE
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathoptimize_panel.py
More file actions
249 lines (231 loc) · 9.98 KB
/
Copy pathoptimize_panel.py
File metadata and controls
249 lines (231 loc) · 9.98 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
from __future__ import annotations
from typing import List, Optional, Callable
try:
from PySide6.QtCore import Qt
from PySide6.QtGui import QTextCharFormat, QColor, QTextCursor, QTextFormat
from PySide6.QtWidgets import (
QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
QCheckBox, QListWidget, QListWidgetItem, QPlainTextEdit, QTextEdit
)
except Exception:
from PySide2.QtCore import Qt # type: ignore
from PySide2.QtGui import QTextCharFormat, QColor, QTextCursor, QTextFormat # type: ignore
from PySide2.QtWidgets import ( # type: ignore
QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
QCheckBox, QListWidget, QListWidgetItem, QPlainTextEdit, QTextEdit
)
from ..backends.optimize import default_rules_for_arch, propose, apply_all, TransformRule, Proposal
class OptimizePanel(QWidget):
def __init__(
self,
get_asm: Callable[[], str],
set_asm: Callable[[str], None],
get_arch: Callable[[], str],
assemble_cb: Optional[Callable[[], None]] = None,
parent=None,
):
super().__init__(parent)
self.get_asm = get_asm
self.set_asm = set_asm
self.get_arch = get_arch
# Optional callback to assemble after applying optimizations (Dev mode)
self.assemble_cb = assemble_cb
layout = QVBoxLayout(self)
# Header controls
top = QHBoxLayout()
self.chk_rule1 = QCheckBox("Push 0 -> xor; push reg")
self.chk_rule2 = QCheckBox("mov reg, imm -> mov reg8/16/32, imm (risky)")
self.btn_apply = QPushButton("Apply All")
top.addWidget(self.chk_rule1)
top.addWidget(self.chk_rule2)
top.addStretch(1)
top.addWidget(self.btn_apply)
layout.addLayout(top)
# Before/After panes
panes = QHBoxLayout()
# Left: Before
self.before_edit = QPlainTextEdit()
self.before_edit.setPlaceholderText("Before (original assembly)")
# Right: After
self.after_edit = QPlainTextEdit()
self.after_edit.setPlaceholderText("After (with selected optimizations)")
try:
for ed in (self.before_edit, self.after_edit):
ed.setReadOnly(True)
ed.setLineWrapMode(QPlainTextEdit.NoWrap) # type: ignore
except Exception:
pass
panes.addWidget(self.before_edit, 1)
panes.addWidget(self.after_edit, 1)
layout.addLayout(panes)
# Note: unified diff view removed per UX request (two panes only)
# Defaults
self.chk_rule1.setChecked(True)
self.chk_rule2.setChecked(True)
# Wire (live preview)
try:
self.chk_rule1.toggled.connect(self.on_preview)
self.chk_rule2.toggled.connect(self.on_preview)
except Exception:
pass
self.btn_apply.clicked.connect(self.on_apply)
# Initial preview
self.on_preview()
def _set_line_highlights(self, edit: QPlainTextEdit, add_lines: set, del_lines: set) -> None:
"""Apply GitHub-like line background highlights to a QPlainTextEdit.
- add_lines: indices to highlight green (entire line width)
- del_lines: indices to highlight red (entire line width)
"""
sels = []
try:
from .highlighters import _good_bad_colors, _tint # reuse theme-aware colors
doc = edit.document()
# Formats using theme colors
good, bad = _good_bad_colors()
# Use stronger tint and mark full-width selection for entire line coloring
add_bg = _tint(good, alpha=72)
del_bg = _tint(bad, alpha=72)
f_add = QTextCharFormat(); f_add.setBackground(add_bg); f_add.setProperty(QTextFormat.FullWidthSelection, True)
f_del = QTextCharFormat(); f_del.setBackground(del_bg); f_del.setProperty(QTextFormat.FullWidthSelection, True)
max_lines = doc.blockCount()
for i in range(max_lines):
block = doc.findBlockByNumber(i)
if not block.isValid():
continue
fmt = f_add if i in add_lines else (f_del if i in del_lines else None)
if fmt is None:
continue
es = QTextEdit.ExtraSelection()
# GitHub-style: mark the whole visual line width, not just text
# Use FullWidthSelection with a zero-length selection at StartOfBlock
cur = QTextCursor(block)
cur.setPosition(block.position()) # no anchor/selection
es.cursor = cur
es.format = fmt
sels.append(es)
edit.setExtraSelections(sels)
except Exception:
try:
edit.setExtraSelections([])
except Exception:
pass
def _set_intraline_highlights(self, edit: QPlainTextEdit, spans_by_line: dict, kind: str) -> None:
"""Highlight only differing parts within lines.
spans_by_line: { line_index: [(start_col, length), ...], ... }
kind: 'add' or 'del' to choose colors.
"""
sels = []
try:
from .highlighters import _good_bad_colors, _tint
doc = edit.document()
fmt = QTextCharFormat()
good, bad = _good_bad_colors()
if kind == 'add':
fmt.setBackground(_tint(good))
fmt.setForeground(good)
else:
fmt.setBackground(_tint(bad))
fmt.setForeground(bad)
for line_no, spans in spans_by_line.items():
block = doc.findBlockByNumber(int(line_no))
if not block.isValid():
continue
base = block.position()
# block.length includes newline; clamp within visible text
max_len = max(0, block.length() - 1)
for (start, length) in spans:
if length <= 0:
continue
s = max(0, min(int(start), max_len))
e = max(0, min(int(start + length), max_len))
if e <= s:
continue
cur = QTextCursor(block)
cur.setPosition(base + s)
cur.setPosition(base + e, QTextCursor.KeepAnchor)
es = QTextEdit.ExtraSelection()
es.cursor = cur
es.format = fmt
sels.append(es)
edit.setExtraSelections(sels)
except Exception:
try:
edit.setExtraSelections([])
except Exception:
pass
def _rules(self) -> List[TransformRule]:
arch = self.get_arch()
rules = default_rules_for_arch(arch)
# Map to checkboxes by name
for r in rules:
if r.name == "push-zero-to-xor-push":
r.enabled = self.chk_rule1.isChecked()
elif r.name == "mov-reg-imm8-to-mov-reg8":
r.enabled = self.chk_rule2.isChecked()
return rules
def on_preview(self):
asm = self.get_asm()
arch = self.get_arch()
rules = self._rules()
props = propose(asm, arch, rules)
# Populate before/after boxes
self.before_edit.setPlainText(asm)
from ..backends.optimize import align_assembly
after = apply_all(asm, arch, rules)
after = align_assembly(after)
self.after_edit.setPlainText(after if after != asm else "; No change with current rules\n" + asm)
# Compute line-level diff highlights (line is the minimum unit)
try:
import difflib
import re
a = (asm or "").splitlines()
b = (after or "").splitlines()
# Normalize whitespace for diffing so formatting-only changes are ignored
def _norm(s: str) -> str:
return re.sub(r"\s+", " ", s).strip()
a_norm = [_norm(x) for x in a]
b_norm = [_norm(x) for x in b]
sm = difflib.SequenceMatcher(None, a_norm, b_norm)
add_after = set()
del_before = set()
for tag, i1, i2, j1, j2 in sm.get_opcodes():
if tag == 'equal':
continue
if tag in ('replace', 'delete'):
del_before.update(range(i1, i2))
if tag in ('replace', 'insert'):
add_after.update(range(j1, j2))
# Ignore blank lines in diffs (do not mark empty-only rows)
try:
del_before = {i for i in del_before if (a[i].strip() != '')}
except Exception:
del_before = del_before
try:
add_after = {j for j in add_after if (b[j].strip() != '')}
except Exception:
add_after = add_after
self._set_line_highlights(self.before_edit, add_lines=set(), del_lines=del_before)
self._set_line_highlights(self.after_edit, add_lines=add_after, del_lines=set())
except Exception:
# Clear highlights on error
self._set_line_highlights(self.before_edit, set(), set())
self._set_line_highlights(self.after_edit, set(), set())
# No unified diff box; rely on line highlights in Before/After panes
def on_apply(self):
asm = self.get_asm()
arch = self.get_arch()
rules = self._rules()
from ..backends.optimize import align_assembly
new_asm = apply_all(asm, arch, rules)
new_asm = align_assembly(new_asm)
if new_asm != asm:
self.set_asm(new_asm)
# Refresh comparison against current editor contents
self.on_preview()
# In Dev mode, after applying all transforms, assemble again to refresh outputs
try:
if callable(self.assemble_cb):
self.assemble_cb()
except Exception:
# Ignore assembly errors here; on_assemble will surface them if needed
pass