Skip to content

Commit 105c5a7

Browse files
r41k0ur41k0u
committed
Cleanup handle_assign
1 parent 933d2a5 commit 105c5a7

File tree

1 file changed

+33
-179
lines changed

1 file changed

+33
-179
lines changed

‎pythonbpf/functions/functions_pass.py‎

Lines changed: 33 additions & 179 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,9 @@
66

77
from pythonbpf.helper import (
88
HelperHandlerRegistry,
9-
handle_helper_call,
109
reset_scratch_pool,
1110
)
1211
from pythonbpf.type_deducer import ctypes_to_ir
13-
from pythonbpf.binary_ops import handle_binary_op
1412
from pythonbpf.expr import eval_expr, handle_expr, convert_to_bool
1513
from pythonbpf.assign_pass import handle_variable_assignment
1614

@@ -80,191 +78,47 @@ def handle_assign(
8078
logger.error(f"Failed to handle assignment to {var_name}")
8179
return
8280

83-
num_types = ("c_int32", "c_int64", "c_uint32", "c_uint64")
84-
8581
logger.info(f"Handling assignment to {ast.dump(target)}")
8682
if not isinstance(target, ast.Name) and not isinstance(target, ast.Attribute):
8783
logger.info("Unsupported assignment target")
8884
return
8985
var_name = target.id if isinstance(target, ast.Name) else target.value.id
9086
rval = stmt.value
91-
if isinstance(target, ast.Attribute):
92-
# struct field assignment
93-
field_name = target.attr
94-
if var_name in local_sym_tab:
95-
struct_type = local_sym_tab[var_name].metadata
96-
struct_info = structs_sym_tab[struct_type]
97-
if field_name in struct_info.fields:
98-
field_ptr = struct_info.gep(
99-
builder, local_sym_tab[var_name].var, field_name
100-
)
101-
val = eval_expr(
102-
func,
103-
module,
104-
builder,
105-
rval,
106-
local_sym_tab,
107-
map_sym_tab,
108-
structs_sym_tab,
109-
)
110-
if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[
111-
1
112-
] == ir.PointerType(ir.IntType(8)):
113-
# TODO: Figure it out, not a priority rn
114-
# Special case for string assignment to char array
115-
# str_len = struct_info["field_types"][field_idx].count
116-
# assign_string_to_array(builder, field_ptr, val[0], str_len)
117-
# print(f"Assigned to struct field {var_name}.{field_name}")
118-
pass
119-
if val is None:
120-
logger.info("Failed to evaluate struct field assignment")
121-
return
122-
logger.info(field_ptr)
123-
builder.store(val[0], field_ptr)
124-
logger.info(f"Assigned to struct field {var_name}.{field_name}")
125-
return
126-
elif isinstance(rval, ast.Constant):
127-
if isinstance(rval.value, bool):
128-
if rval.value:
129-
builder.store(
130-
ir.Constant(ir.IntType(1), 1), local_sym_tab[var_name].var
131-
)
132-
else:
133-
builder.store(
134-
ir.Constant(ir.IntType(1), 0), local_sym_tab[var_name].var
135-
)
136-
logger.info(f"Assigned constant {rval.value} to {var_name}")
137-
elif isinstance(rval.value, int):
138-
# Assume c_int64 for now
139-
# var = builder.alloca(ir.IntType(64), name=var_name)
140-
# var.align = 8
141-
builder.store(
142-
ir.Constant(ir.IntType(64), rval.value), local_sym_tab[var_name].var
143-
)
144-
logger.info(f"Assigned constant {rval.value} to {var_name}")
145-
elif isinstance(rval.value, str):
146-
str_val = rval.value.encode("utf-8") + b"\x00"
147-
str_const = ir.Constant(
148-
ir.ArrayType(ir.IntType(8), len(str_val)), bytearray(str_val)
87+
88+
# struct field assignment
89+
field_name = target.attr
90+
if var_name in local_sym_tab:
91+
struct_type = local_sym_tab[var_name].metadata
92+
struct_info = structs_sym_tab[struct_type]
93+
if field_name in struct_info.fields:
94+
field_ptr = struct_info.gep(
95+
builder, local_sym_tab[var_name].var, field_name
14996
)
150-
global_str = ir.GlobalVariable(
151-
module, str_const.type, name=f"{var_name}_str"
97+
val = eval_expr(
98+
func,
99+
module,
100+
builder,
101+
rval,
102+
local_sym_tab,
103+
map_sym_tab,
104+
structs_sym_tab,
152105
)
153-
global_str.linkage = "internal"
154-
global_str.global_constant = True
155-
global_str.initializer = str_const
156-
str_ptr = builder.bitcast(global_str, ir.PointerType(ir.IntType(8)))
157-
builder.store(str_ptr, local_sym_tab[var_name].var)
158-
logger.info(f"Assigned string constant '{rval.value}' to {var_name}")
159-
else:
160-
logger.info("Unsupported constant type")
161-
elif isinstance(rval, ast.Call):
162-
if isinstance(rval.func, ast.Name):
163-
call_type = rval.func.id
164-
logger.info(f"Assignment call type: {call_type}")
165-
if (
166-
call_type in num_types
167-
and len(rval.args) == 1
168-
and isinstance(rval.args[0], ast.Constant)
169-
and isinstance(rval.args[0].value, int)
170-
):
171-
ir_type = ctypes_to_ir(call_type)
172-
# var = builder.alloca(ir_type, name=var_name)
173-
# var.align = ir_type.width // 8
174-
builder.store(
175-
ir.Constant(ir_type, rval.args[0].value),
176-
local_sym_tab[var_name].var,
177-
)
178-
logger.info(
179-
f"Assigned {call_type} constant {rval.args[0].value} to {var_name}"
180-
)
181-
elif HelperHandlerRegistry.has_handler(call_type):
182-
# var = builder.alloca(ir.IntType(64), name=var_name)
183-
# var.align = 8
184-
val = handle_helper_call(
185-
rval,
186-
module,
187-
builder,
188-
func,
189-
local_sym_tab,
190-
map_sym_tab,
191-
structs_sym_tab,
192-
)
193-
builder.store(val[0], local_sym_tab[var_name].var)
194-
logger.info(f"Assigned constant {rval.func.id} to {var_name}")
195-
elif call_type == "deref" and len(rval.args) == 1:
196-
logger.info(f"Handling deref assignment {ast.dump(rval)}")
197-
val = eval_expr(
198-
func,
199-
module,
200-
builder,
201-
rval,
202-
local_sym_tab,
203-
map_sym_tab,
204-
structs_sym_tab,
205-
)
206-
if val is None:
207-
logger.info("Failed to evaluate deref argument")
208-
return
209-
logger.info(f"Dereferenced value: {val}, storing in {var_name}")
210-
builder.store(val[0], local_sym_tab[var_name].var)
211-
logger.info(f"Dereferenced and assigned to {var_name}")
212-
elif call_type in structs_sym_tab and len(rval.args) == 0:
213-
struct_info = structs_sym_tab[call_type]
214-
ir_type = struct_info.ir_type
215-
# var = builder.alloca(ir_type, name=var_name)
216-
# Null init
217-
builder.store(ir.Constant(ir_type, None), local_sym_tab[var_name].var)
218-
logger.info(f"Assigned struct {call_type} to {var_name}")
219-
else:
220-
logger.info(f"Unsupported assignment call type: {call_type}")
221-
elif isinstance(rval.func, ast.Attribute):
222-
logger.info(f"Assignment call attribute: {ast.dump(rval.func)}")
223-
if isinstance(rval.func.value, ast.Name):
224-
if rval.func.value.id in map_sym_tab:
225-
map_name = rval.func.value.id
226-
method_name = rval.func.attr
227-
if HelperHandlerRegistry.has_handler(method_name):
228-
val = handle_helper_call(
229-
rval,
230-
module,
231-
builder,
232-
func,
233-
local_sym_tab,
234-
map_sym_tab,
235-
structs_sym_tab,
236-
)
237-
builder.store(val[0], local_sym_tab[var_name].var)
238-
else:
239-
# TODO: probably a struct access
240-
logger.info(f"TODO STRUCT ACCESS {ast.dump(rval)}")
241-
elif isinstance(rval.func.value, ast.Call) and isinstance(
242-
rval.func.value.func, ast.Name
243-
):
244-
map_name = rval.func.value.func.id
245-
method_name = rval.func.attr
246-
if map_name in map_sym_tab:
247-
if HelperHandlerRegistry.has_handler(method_name):
248-
val = handle_helper_call(
249-
rval,
250-
module,
251-
builder,
252-
func,
253-
local_sym_tab,
254-
map_sym_tab,
255-
structs_sym_tab,
256-
)
257-
# var = builder.alloca(ir.IntType(64), name=var_name)
258-
# var.align = 8
259-
builder.store(val[0], local_sym_tab[var_name].var)
260-
else:
261-
logger.info("Unsupported assignment call structure")
262-
else:
263-
logger.info("Unsupported assignment call function type")
264-
elif isinstance(rval, ast.BinOp):
265-
handle_binary_op(rval, builder, var_name, local_sym_tab)
266-
else:
267-
logger.info("Unsupported assignment value type")
106+
if isinstance(struct_info.field_type(field_name), ir.ArrayType) and val[
107+
1
108+
] == ir.PointerType(ir.IntType(8)):
109+
# TODO: Figure it out, not a priority rn
110+
# Special case for string assignment to char array
111+
# str_len = struct_info["field_types"][field_idx].count
112+
# assign_string_to_array(builder, field_ptr, val[0], str_len)
113+
# print(f"Assigned to struct field {var_name}.{field_name}")
114+
pass
115+
if val is None:
116+
logger.info("Failed to evaluate struct field assignment")
117+
return
118+
logger.info(field_ptr)
119+
builder.store(val[0], field_ptr)
120+
logger.info(f"Assigned to struct field {var_name}.{field_name}")
121+
return
268122

269123

270124
def handle_cond(

0 commit comments

Comments
 (0)