nitpick-bootstrap/patch_codegen.py at main · alternative-intelligence-cp/nitpick-bootstrap · GitHub

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
import re

# 1. Update boot/codegen.ll
with open('boot/codegen.ll', 'r') as f:
    cg = f.read()

# Update %CodeGen struct
cg = re.sub(
    r'%CodeGen = type \{ ptr, i64, i64, i32, i32, i32, i32, ptr, ptr, i64 \}',
    r'%CodeGen = type { ptr, i64, i64, i32, i32, i32, i32, ptr, ptr, i64, ptr }',
    cg
)

# Update boot.cg.init signature
cg = re.sub(
    r'define void @boot\.cg\.init\(ptr %cg, ptr %ast_arena, ptr %src, i64 %src_len\) \{',
    r'define void @boot.cg.init(ptr %cg, ptr %ast_arena, ptr %src, i64 %src_len, ptr %ta) {\n  %ta_p = getelementptr %CodeGen, ptr %cg, i32 0, i32 10\n  store ptr %ta, ptr %ta_p, align 8',
    cg
)

# Update c_var_decl to use get_extra instead of token_kind_to_type
cg = re.sub(
    r'%v_type_op = call i32 @boot\.ast\.get_op_kind\(ptr %ast, i32 %v_type_node\)\s+%v_type_idx = call i32 @boot\.types\.token_kind_to_type\(i32 %v_type_op\)',
    r'%v_type_idx = call i32 @boot.ast.get_extra(ptr %ast, i32 %node_idx)',
    cg
)

# Update emit_type_by_idx to look up the kind and handle structs
emit_type_header = """define void @boot.cg.emit_type_by_idx(ptr %cg, i32 %idx) {
entry:
  %ta_p = getelementptr %CodeGen, ptr %cg, i32 0, i32 10
  %ta = load ptr, ptr %ta_p, align 8
  %kind = call i32 @boot.types.get_kind(ptr %ta, i32 %idx)
  %is_struct = icmp eq i32 %kind, 51
  br i1 %is_struct, label %emit_struct, label %check_ptr
check_ptr:
  %is_ptr = icmp eq i32 %kind, 55
  br i1 %is_ptr, label %t_ptr, label %check_arr
check_arr:
  %is_arr = icmp eq i32 %kind, 50
  br i1 %is_arr, label %t_ptr, label %switch_prim
emit_struct:
  %tp = call ptr @boot.types.type_ptr(ptr %ta, i32 %idx)
  %ns_p = getelementptr %NpkType, ptr %tp, i32 0, i32 3
  %ns = load i32, ptr %ns_p, align 4
  %nl_p = getelementptr %NpkType, ptr %tp, i32 0, i32 4
  %nl = load i32, ptr %nl_p, align 4
  %src_p = getelementptr %CodeGen, ptr %cg, i32 0, i32 8
  %src = load ptr, ptr %src_p, align 8
  %pref = alloca [4 x i8], align 1
  store [4 x i8] c"%St.", ptr %pref, align 1
  call void @boot.cg.emit(ptr %cg, ptr %pref, i64 4)
  call void @boot.cg.emit_name(ptr %cg, ptr %src, i32 %ns, i32 %nl)
  ret void
switch_prim:
  switch i32 %kind, label %default ["""

cg = re.sub(
    r'define void @boot\.cg\.emit_type_by_idx\(ptr %cg, i32 %idx\) \{\s*entry:\s*switch i32 %idx, label %default \[',
    emit_type_header,
    cg
)

with open('boot/codegen.ll', 'w') as f:
    f.write(cg)

# 2. Update boot/typechecker.ll
with open('boot/typechecker.ll', 'r') as f:
    tc = f.read()

# Revert res_val_ok
tc = re.sub(
    r'%res_inner_kind = call i32 @boot\.types\.get_kind\(ptr %ta_mem, i32 %res_inner_type\)\s+call void @boot\.ast\.set_extra\(ptr %ast, i32 %node_idx, i32 %res_inner_kind\)',
    r'call void @boot.ast.set_extra(ptr %ast, i32 %node_idx, i32 %res_inner_type)',
    tc
)

# Revert idx_found
tc = re.sub(
    r'%f_kind = call i32 @boot\.types\.get_kind\(ptr %ta_mem, i32 %f_type\)\s+call void @boot\.ast\.set_extra\(ptr %ast, i32 %node_idx, i32 %f_kind\)',
    r'call void @boot.ast.set_extra(ptr %ast, i32 %node_idx, i32 %f_type)',
    tc
)

# Revert register_func
tc = re.sub(
    r'%ret_kind = call i32 @boot\.types\.get_kind\(ptr %ta_f, i32 %ret_type_val\)\s+call void @boot\.ast\.set_flags\(ptr %ast, i32 %node_idx, i32 %ret_kind\)',
    r'call void @boot.ast.set_flags(ptr %ast, i32 %node_idx, i32 %ret_type_val)',
    tc
)

# Add set_extra to check_var_decl
tc = re.sub(
    r'%sym = call i32 @boot\.types\.scope_define\(ptr %sca, ptr %sa, i32 %scope, i32 %ns, i32 %nl, i32 1, i32 %decl_type\) ; SYMBOL_VAR=1',
    r'%sym = call i32 @boot.types.scope_define(ptr %sca, ptr %sa, i32 %scope, i32 %ns, i32 %nl, i32 1, i32 %decl_type)\n  call void @boot.ast.set_extra(ptr %ast, i32 %node_idx, i32 %decl_type)',
    tc
)

with open('boot/typechecker.ll', 'w') as f:
    f.write(tc)

# 3. Update boot/main.ll
with open('boot/main.ll', 'r') as f:
    main = f.read()

main = re.sub(
    r'call void @boot\.cg\.init\(ptr %cg, ptr %ast, ptr %src, i64 %src_len\)',
    r'call void @boot.cg.init(ptr %cg, ptr %ast, ptr %src, i64 %src_len, ptr %ta)',
    main
)

with open('boot/main.ll', 'w') as f:
    f.write(main)

print("Patch applied.")