-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathasm.py
175 lines (162 loc) · 8.12 KB
/
asm.py
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
from abc import ABC
from ast_exprs import Application, Assignment, AstirExpr, Identifier, InlineASM, Lambda, LambdaDefinition, Literal, Reference, ShuntingYardAlgorithmResults, Symbol, SymbolTable # type: ignore
from common import Cursor, PrimitiveTypes
class ASMFunction:
def __init__(
self, param_to_reg: dict[int, int], name_to_param: dict[str, int]
) -> None:
self.param_to_reg = param_to_reg
self.name_to_param = name_to_param
# The next usable register is calculated by getting
# the last inserted item in the self.param_to_reg
# and by adding 1 to get us to the next x register
self.next_usable_reg = (
next(reversed(param_to_reg)) + 1 if len(param_to_reg) > 0 else 0
)
class ASM(Cursor):
def __init__(
self, input: list[AstirExpr], symbol_tables: dict[int, SymbolTable]
) -> None:
super().__init__(input)
self.lines: list[str] = [
".global _start",
".p2align 3",
]
self.symbol_tables: dict[int, SymbolTable] = symbol_tables
# All tables related to registers most likely
# have keys that map to a Symbol's id.
# The inner dictionary holds the parameter index -> register
# map. TODO: introduce a way to recognize different
# register sizes for a64 (rn)
self.fn_register_store: dict[int, ASMFunction] = {}
self.inside_fn: int | None = None
# # Format (ref_id, register)
# self.ref_id_and_register: list[tuple[int, int]] = []
# self.fn_register_man: dict[int, list[tuple[int, int]]] = {}
# self.current_usable_register = 0
# # This should be cleared after every generate call.
# self.registers_in_use: list[Register] = []
def is_register_reserved(self, register: int) -> bool:
return False
def current_fn(self) -> ASMFunction | None:
if self.inside_fn is None:
return None
return self.fn_register_store[self.inside_fn]
def generate_all(self):
while self.current() is not None:
generated = self.generate()
self.lines.extend(generated)
self.advance()
def lookup_symbol(self, symbol_table: int, symbol_id: int) -> Symbol | None:
if symbol_table not in self.symbol_tables:
return None
_symbol_table: SymbolTable | None = self.symbol_tables[symbol_table]
if _symbol_table is None:
return None
symbol = _symbol_table.lookup_by_id(symbol_id)
return symbol
def generate(self, expr: AstirExpr | None = None) -> list[str]:
c_expr = self.current() if expr is None else expr
to_add: list[str] = []
if isinstance(c_expr, InlineASM):
to_add.extend(c_expr.lines)
elif isinstance(c_expr, Assignment):
if isinstance(c_expr.right, Lambda) and isinstance(c_expr.left, Identifier):
symbols = c_expr.right.definition.parameters.symbols
last_used_register = 0
lambda_param_to_register: dict[int, int] = {}
param_name_to_idx: dict[str, int] = {}
for symbol_idx in symbols:
symbol = symbols[symbol_idx]
if symbol.name == "ret":
# We break instead of continue because "ret"
# should always be the last item in the dict
break
lambda_param_to_register[symbol_idx] = last_used_register
param_name_to_idx[symbol.name] = symbol_idx
last_used_register += 1
asm_function = ASMFunction(lambda_param_to_register, param_name_to_idx)
self.fn_register_store[c_expr.right.symbol_id] = asm_function
# This is so we can parse and get the correct arguments
self.inside_fn = c_expr.right.symbol_id
to_add.append(
f"{c_expr.left.value}: // Symbol ID: {c_expr.right.symbol_id}"
)
# to_add.append(f"// {asm_function.next_usable_reg}")
to_add.extend(self.generate(c_expr.right.body))
#to_add.append("ret")
self.inside_fn = None
elif isinstance(c_expr, ShuntingYardAlgorithmResults):
inside_fn = self.current_fn()
if inside_fn is None:
raise Exception("Out of place sya...")
if len(c_expr.oeprators) > 0:
raise Exception("Invalid shunting yard algorithm")
stack: list[str] = []
c_expr.results.reverse()
while len(c_expr.results) > 0 and (term := c_expr.results.pop()):
# TODO: make some like class method or something
# to make this cleaner??
if isinstance(term, Reference):
stack.extend(self.generate(term))
elif isinstance(term, Literal):
if term.ty != PrimitiveTypes.INT:
raise Exception("Unexpected type.")
stack.append(str(term.val))
elif isinstance(term, str):
if term == "+":
stack.reverse()
(item1, item2) = (stack.pop(), stack.pop())
if not item1.startswith("x"):
register = inside_fn.next_usable_reg
to_add.append(f"mov x{register}, {item1}")
item1 = f"x{register}"
inside_fn.next_usable_reg += 1
print(
f"Adding last two items on stack: {item1}, {item2} = {item1 + item2}"
)
to_add.append(f"add x0, {item1}, {item2}")
elif isinstance(c_expr, Reference):
symbol_in_ref: Symbol | None = self.lookup_symbol(
c_expr.belongs_to, c_expr.symbol_id
)
if symbol_in_ref is None:
raise Exception(f'Failed to lookup referenced symbol "{c_expr.name}"')
if (
(c_fn := self.current_fn())
and c_fn is not None
and c_expr.name in c_fn.name_to_param
and c_fn.name_to_param[c_expr.name] in c_fn.param_to_reg
):
register = c_fn.param_to_reg[c_fn.name_to_param[c_expr.name]]
to_add.append(f"x{register}")
elif isinstance(c_expr, Application):
fn_parameters = self.fn_register_store[c_expr.lambda_ref.symbol_id]
if fn_parameters is None:
raise Exception("Failed to get reserved registers for fn")
elif len(list(fn_parameters.param_to_reg.keys())) != len(c_expr.parameters):
raise Exception("More parameters than reserved registers...")
application_symbol: Symbol | None = self.lookup_symbol(
c_expr.lambda_ref.belongs_to, c_expr.lambda_ref.symbol_id
)
if application_symbol is None:
raise Exception(f"failed to find symbol {c_expr.lambda_ref.symbol_id}")
elif (
not isinstance(application_symbol.val, Lambda)
or c_expr.lambda_ref.symbol_id not in self.fn_register_store
):
raise Exception(
"Expected this symbol to come back to a lambda definition"
)
function_parameters = application_symbol.val.definition.parameters.symbols
for idx, param in enumerate(c_expr.parameters):
if not (0 <= idx < len(list(fn_parameters.param_to_reg.keys()))):
raise Exception(
"Invalid Application...More parameters than reserved registers"
)
reserved_register = fn_parameters.param_to_reg[idx]
if not isinstance(param, Literal) or param.ty != PrimitiveTypes.INT:
raise Exception("TODO: HANDLE MORE THAN JUST LITERALS")
to_add.append(f"mov x{reserved_register}, {param.val}")
to_add.append(f"bl {application_symbol.name}")
return to_add