Skip to content

Commit efadd45

Browse files
willieyzmkannwischer
willieyz
authored and
mkannwischer
committed
Add cfify scripts, enable the --cfify option in autogen
Signed-off-by: willieyz <[email protected]>
1 parent 03e125e commit efadd45

File tree

4 files changed

+282
-4
lines changed

4 files changed

+282
-4
lines changed

‎scripts/autogen‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1289,11 +1289,10 @@ def update_via_simpasm(
12891289
# Determine architecture from filename
12901290
arch = "aarch64" if "aarch64" in infile_full else "x86_64"
12911291

1292-
# TODO: Temporary remvoe the "--cfify", add back when CFI script added.
12931292
cmd = [
12941293
"./scripts/simpasm",
12951294
"--objdump=llvm-objdump",
1296-
# "--cfify",
1295+
"--cfify",
12971296
"--arch=" + arch,
12981297
"-i",
12991298
infile_full,

‎scripts/cfify‎

Lines changed: 279 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
#!/usr/bin/env python3
2+
# Copyright (c) The mlkem-native project authors
3+
# SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT
4+
5+
import sys
6+
import re
7+
import argparse
8+
9+
10+
def add_cfi_directives(text, arch):
11+
lines = text.split("\n")
12+
result = []
13+
i = 0
14+
15+
while i < len(lines):
16+
line = lines[i].rstrip()
17+
18+
if arch == "aarch64":
19+
# Check for SIMD save pattern: stp d8,d9; stp d10,d11; stp d12,d13; stp d14,d15
20+
if i + 3 < len(lines):
21+
pattern_text = "\n".join(lines[i : i + 4])
22+
simd_save_pattern = (
23+
r"(\s*)stp\s+d8,\s*d9,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
24+
r"\s*stp\s+d10,\s*d11,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
25+
r"\s*stp\s+d12,\s*d13,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
26+
r"\s*stp\s+d14,\s*d15,\s*\[sp(?:,\s*#([^]]+))?\]"
27+
)
28+
match = re.match(simd_save_pattern, pattern_text, re.IGNORECASE)
29+
if match:
30+
indent = match.group(1)
31+
offsets = [match.group(j + 2) or "0" for j in range(4)]
32+
for j, reg_pair in enumerate(
33+
[(8, 9), (10, 11), (12, 13), (14, 15)]
34+
):
35+
result.append(lines[i + j].rstrip())
36+
try:
37+
offset_val = int(offsets[j], 0)
38+
result.append(
39+
f"{indent}.cfi_rel_offset d{reg_pair[0]}, 0x{offset_val:x}"
40+
)
41+
result.append(
42+
f"{indent}.cfi_rel_offset d{reg_pair[1]}, 0x{offset_val+8:x}"
43+
)
44+
except:
45+
result.append(
46+
f"{indent}.cfi_rel_offset d{reg_pair[0]}, {offsets[j]}"
47+
)
48+
result.append(
49+
f"{indent}.cfi_rel_offset d{reg_pair[1]}, ({offsets[j]}+8)"
50+
)
51+
i += 4
52+
continue
53+
54+
# Check for SIMD restore pattern: ldp d8,d9; ldp d10,d11; ldp d12,d13; ldp d14,d15
55+
if i + 3 < len(lines):
56+
pattern_text = "\n".join(lines[i : i + 4])
57+
simd_restore_pattern = (
58+
r"(\s*)ldp\s+d8,\s*d9,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
59+
r"\s*ldp\s+d10,\s*d11,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
60+
r"\s*ldp\s+d12,\s*d13,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
61+
r"\s*ldp\s+d14,\s*d15,\s*\[sp(?:,\s*#[^]]+)?\]"
62+
)
63+
match = re.match(simd_restore_pattern, pattern_text, re.IGNORECASE)
64+
if match:
65+
indent = match.group(1)
66+
for j, reg_pair in enumerate(
67+
[(8, 9), (10, 11), (12, 13), (14, 15)]
68+
):
69+
result.append(lines[i + j].rstrip())
70+
result.append(f"{indent}.cfi_restore d{reg_pair[0]}")
71+
result.append(f"{indent}.cfi_restore d{reg_pair[1]}")
72+
i += 4
73+
continue
74+
75+
# Check for GPR save pattern: stp x19,x20 through stp x29,x30
76+
if i + 5 < len(lines):
77+
pattern_text = "\n".join(lines[i : i + 6])
78+
gpr_save_pattern = (
79+
r"(\s*)stp\s+x19,\s*x20,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
80+
r"\s*stp\s+x21,\s*x22,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
81+
r"\s*stp\s+x23,\s*x24,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
82+
r"\s*stp\s+x25,\s*x26,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
83+
r"\s*stp\s+x27,\s*x28,\s*\[sp(?:,\s*#([^]]+))?\]\s*\n"
84+
r"\s*stp\s+x29,\s*x30,\s*\[sp(?:,\s*#([^]]+))?\]"
85+
)
86+
match = re.match(gpr_save_pattern, pattern_text, re.IGNORECASE)
87+
if match:
88+
indent = match.group(1)
89+
offsets = [match.group(j + 2) or "0" for j in range(6)]
90+
for j, reg_pair in enumerate(
91+
[(19, 20), (21, 22), (23, 24), (25, 26), (27, 28), (29, 30)]
92+
):
93+
result.append(lines[i + j].rstrip())
94+
try:
95+
offset_val = int(offsets[j], 0)
96+
result.append(
97+
f"{indent}.cfi_rel_offset x{reg_pair[0]}, 0x{offset_val:x}"
98+
)
99+
result.append(
100+
f"{indent}.cfi_rel_offset x{reg_pair[1]}, 0x{offset_val+8:x}"
101+
)
102+
except:
103+
result.append(
104+
f"{indent}.cfi_rel_offset x{reg_pair[0]}, {offsets[j]}"
105+
)
106+
result.append(
107+
f"{indent}.cfi_rel_offset x{reg_pair[1]}, ({offsets[j]}+8)"
108+
)
109+
i += 6
110+
continue
111+
112+
# Check for GPR restore pattern: ldp x19,x20 through ldp x29,x30
113+
if i + 5 < len(lines):
114+
pattern_text = "\n".join(lines[i : i + 6])
115+
gpr_restore_pattern = (
116+
r"(\s*)ldp\s+x19,\s*x20,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
117+
r"\s*ldp\s+x21,\s*x22,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
118+
r"\s*ldp\s+x23,\s*x24,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
119+
r"\s*ldp\s+x25,\s*x26,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
120+
r"\s*ldp\s+x27,\s*x28,\s*\[sp(?:,\s*#[^]]+)?\]\s*\n"
121+
r"\s*ldp\s+x29,\s*x30,\s*\[sp(?:,\s*#[^]]+)?\]"
122+
)
123+
match = re.match(gpr_restore_pattern, pattern_text, re.IGNORECASE)
124+
if match:
125+
indent = match.group(1)
126+
for j, reg_pair in enumerate(
127+
[(19, 20), (21, 22), (23, 24), (25, 26), (27, 28), (29, 30)]
128+
):
129+
result.append(lines[i + j].rstrip())
130+
result.append(f"{indent}.cfi_restore x{reg_pair[0]}")
131+
result.append(f"{indent}.cfi_restore x{reg_pair[1]}")
132+
i += 6
133+
continue
134+
135+
# Rule 7: add sp, sp, #offset -> .cfi_adjust_cfa_offset (-(offset))
136+
match = re.match(
137+
r"(\s*)add\s+sp,\s*sp,\s*#(0x[0-9a-fA-F]+|\d+)", line, re.IGNORECASE
138+
)
139+
if match:
140+
indent, offset_str = match.groups()
141+
offset = (
142+
int(offset_str, 16)
143+
if offset_str.lower().startswith("0x")
144+
else int(offset_str)
145+
)
146+
result.append(line)
147+
result.append(f"{indent}.cfi_adjust_cfa_offset -{offset:#x}")
148+
i += 1
149+
continue
150+
151+
# Rule 8: sub sp, sp, #offset -> .cfi_adjust_cfa_offset (offset)
152+
match = re.match(
153+
r"(\s*)sub\s+sp,\s*sp,\s*#(0x[0-9a-fA-F]+|\d+)", line, re.IGNORECASE
154+
)
155+
if match:
156+
indent, offset_str = match.groups()
157+
offset = (
158+
int(offset_str, 16)
159+
if offset_str.lower().startswith("0x")
160+
else int(offset_str)
161+
)
162+
result.append(line)
163+
result.append(f"{indent}.cfi_adjust_cfa_offset {offset:#x}")
164+
i += 1
165+
continue
166+
167+
# Rule 2: ret -> .cfi_endproc after ret
168+
match = re.match(r"(\s*)ret\s*$", line, re.IGNORECASE)
169+
if match:
170+
indent = match.group(1)
171+
result.append(line)
172+
result.append(f"{indent}.cfi_endproc")
173+
i += 1
174+
continue
175+
176+
elif arch == "x86_64":
177+
# Check for labels and see if there's a corresponding callq
178+
label_match = re.match(r"^([a-zA-Z_][a-zA-Z0-9_]*):$", line)
179+
if label_match:
180+
label = label_match.group(1)
181+
# Check if this label is called anywhere in the text
182+
if re.search(rf"\s*callq\s+{re.escape(label)}\b", text, re.IGNORECASE):
183+
result.append(line)
184+
result.append(" .cfi_startproc")
185+
i += 1
186+
continue
187+
188+
# x86_64: subq $OFFSET, %rsp -> .cfi_adjust_cfa_offset OFFSET (stack alloc)
189+
match = re.match(
190+
r"(\s*)subq\s+\$(0x[0-9a-fA-F]+|\d+),\s*%rsp", line, re.IGNORECASE
191+
)
192+
if match:
193+
indent, offset_str = match.groups()
194+
offset = (
195+
int(offset_str, 16)
196+
if offset_str.lower().startswith("0x")
197+
else int(offset_str)
198+
)
199+
result.append(line)
200+
result.append(f"{indent}.cfi_adjust_cfa_offset {offset:#x}")
201+
i += 1
202+
continue
203+
204+
# x86_64: addq $OFFSET, %rsp -> .cfi_adjust_cfa_offset -OFFSET (stack free)
205+
match = re.match(
206+
r"(\s*)addq\s+\$(0x[0-9a-fA-F]+|\d+),\s*%rsp", line, re.IGNORECASE
207+
)
208+
if match:
209+
indent, offset_str = match.groups()
210+
offset = (
211+
int(offset_str, 16)
212+
if offset_str.lower().startswith("0x")
213+
else int(offset_str)
214+
)
215+
result.append(line)
216+
result.append(f"{indent}.cfi_adjust_cfa_offset -{offset:#x}")
217+
i += 1
218+
continue
219+
220+
# x86_64: retq -> .cfi_endproc after retq
221+
match = re.match(r"(\s*)retq\s*$", line, re.IGNORECASE)
222+
if match:
223+
indent = match.group(1)
224+
result.append(line)
225+
result.append(f"{indent}.cfi_endproc")
226+
i += 1
227+
continue
228+
229+
result.append(line)
230+
i += 1
231+
232+
return "\n".join(result)
233+
234+
235+
def main():
236+
237+
parser = argparse.ArgumentParser(
238+
description="Add CFI directives to AArch64 assembly"
239+
)
240+
parser.add_argument("-i", "--input", help="Input file (default: stdin)")
241+
parser.add_argument("-o", "--output", help="Output file (default: stdout)")
242+
parser.add_argument(
243+
"--emit-cfi-proc-start",
244+
action="store_true",
245+
help="Emit .cfi_proc_start as first line",
246+
)
247+
parser.add_argument(
248+
"--arch",
249+
choices=["aarch64", "x86_64"],
250+
default="aarch64",
251+
help="Target architecture (default: aarch64)",
252+
)
253+
args = parser.parse_args()
254+
255+
input_file = open(args.input, "r") if args.input else sys.stdin
256+
output_file = open(args.output, "w") if args.output else sys.stdout
257+
258+
try:
259+
# Read all input
260+
text = input_file.read()
261+
262+
# Add initial .cfi_startproc if requested
263+
if args.emit_cfi_proc_start:
264+
text = " .cfi_startproc\n" + text
265+
266+
# Process the text
267+
result = add_cfi_directives(text, args.arch)
268+
269+
# Write output
270+
output_file.write(result)
271+
finally:
272+
if args.input:
273+
input_file.close()
274+
if args.output:
275+
output_file.close()
276+
277+
278+
if __name__ == "__main__":
279+
main()

‎scripts/format‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ if ! command -v black 2>&1 >/dev/null; then
4848
error "black not found. Are you running in a nix shell? See BUILDING.md."
4949
exit 1
5050
fi
51-
black --include "(scripts/tests|scripts/simpasm|scripts/autogen|scripts/check-namespace|\.py$)" "$ROOT"
51+
black --include "(scripts/tests|scripts/simpasm|scripts/cfify|scripts/autogen|scripts/check-namespace|\.py$)" "$ROOT"
5252

5353
info "Formatting c files"
5454
if ! command -v clang-format 2>&1 >/dev/null; then

‎scripts/lint‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ gh_group_end
117117
#gh_group_end
118118

119119
gh_group_start "Linting python scripts with black"
120-
if ! diff=$(black --check --diff -q --include "(scripts/tests|scripts/simpasm|scripts/autogen|scripts/check-namespace|\.py$)" "$ROOT"); then
120+
if ! diff=$(black --check --diff -q --include "(scripts/tests|scripts/simpasm|scripts/cfify|scripts/autogen|scripts/check-namespace|\.py$)" "$ROOT"); then
121121
gh_error_simple "Format error" "$diff"
122122
error "Lint python"
123123
SUCCESS=false

0 commit comments

Comments
 (0)