diff --git a/lang/src/codegen.asm b/lang/src/codegen.asm
index b58e0f1..349abbe 100644
--- a/lang/src/codegen.asm
+++ b/lang/src/codegen.asm
@@ -29,7 +29,7 @@ section .rdata
     EPILOGUE_LEN equ $ - EPILOGUE
     MOV_RAX db "mov rax, "
     MOV_RAX_LEN equ $ - MOV_RAX
-    JMP_EPILOGUE db 10, "jmp .epilogue", 10
+    JMP_EPILOGUE db "jmp .epilogue", 10
     JMP_EPILOGUE_LEN equ $ - JMP_EPILOGUE
     DOT_ARGS db ".args:", 10
     DOT_ARGS_LEN equ $ - DOT_ARGS
@@ -39,11 +39,11 @@ section .rdata
     DOT_PROLOGUE_LEN equ $ - DOT_PROLOGUE
     DOT_EPILOGUE db ".epilogue:", 10
     DOT_EPILOGUE_LEN equ $ - DOT_EPILOGUE
-    JMP_ARGS db 10, "jmp .args", 10
+    JMP_ARGS db  "jmp .args", 10
     JMP_ARGS_LEN equ $ - JMP_ARGS
-    JMP_BODY db 10, "jmp .body", 10
+    JMP_BODY db  "jmp .body", 10
     JMP_BODY_LEN equ $ - JMP_BODY
-    JMP_PROLOGUE db 10, "jmp .prologue", 10
+    JMP_PROLOGUE db "jmp .prologue", 10
     JMP_PROLOGUE_LEN equ $ - JMP_PROLOGUE
     SUB_RSP db "sub rsp, "
     SUB_RSP_LEN equ $ - SUB_RSP
@@ -273,8 +273,8 @@ stackvar_cmp:
 ;; struct FunctionCtx {
 ;;   current_stack_size: u64,
 ;;   stack_vars: Vec<(u64, u64)>,
-;;   register_bitset: u128,
-;;   dirtied_register_bitset: u128,
+;;   register_bitset: u16,
+;;   dirtied_register_bitset: u16,
 ;; }
 ;; end-structs
 
@@ -296,41 +296,27 @@ codegen_allocate_register:
     ; dirtied registers are those that have been used in the function and need
     ; to be saved/restored in the prologue/epilogue
 
-    mov rax, [rdi + 48]         ; register_bitset
+    mov ax, word [rdi + 48]     ; register_bitset
     xor rcx, rcx
     ; flip bits
-    not rax
-    ; find first set bit
-    bsf rcx, rax
-    test rcx, rcx
-    jnz .found
-    mov rax, [rdi + 56]         ; higher 64 bits
-    not rax
-    bsf rcx, rax
-    test rcx, rcx
+    not ax
+    test ax, ax
     jz .panic
-    add rcx, 8
+    ; find first set bit
+    bsf cx, ax
 
 .found:
     mov rbx, rcx
-    cmp rcx, 8
-    jl .set_low
-    sub rcx, 8
-    bts rax, rcx
-    mov [rdi + 56], rax         ; update register_bitset
-    mov rax, [rdi + 72]         ; dirtied_register_bitset
-    bts rax, rcx
-    mov [rdi + 72], rax         ; update dirtied_register_bitset
-    jmp .done
-.set_low:
-    bts rax, rcx
-    mov [rdi + 48], rax         ; update register_bitset
-    mov rax, [rdi + 64]         ; dirtied_register_bitset
-    bts rax, rcx
-    mov [rdi + 64], rax         ; update dirtied_register_bitset
+    mov ax, word [rdi + 48]     ; update register_bitset
+    bts ax, cx
+    mov word [rdi + 48], ax     ; update register_bitset
+    mov ax, word [rdi + 52]     ; dirtied_register_bitset
+    bts ax, cx
+    mov word [rdi + 52], ax     ; update dirtied_register_bitset
 
 .done:
-    pop rax
+    mov rax, rbx
+    pop rbx
     pop rbp
     ret
 .panic:
@@ -349,8 +335,9 @@ codegen_function:
     push r13
 
     ; scratch [104..120]
-    ; dirtied-register-bitset [88..104]  [a,b,c,d,si,di,bp,sp,8,9,10,11,12,13,14,15]
-    ; register-bitset [72..88]  [a,b,c,d,si,di,bp,sp,8,9,10,11,12,13,14,15]
+    ; scratch2 [80..104]
+    ; dirtied-register-bitset [76..80]  [a,b,c,d,si,di,bp,sp,8,9,10,11,12,13,14,15]
+    ; register-bitset [72..76]  [a,b,c,d,si,di,bp,sp,8,9,10,11,12,13,14,15]
     ; stack-vars: Vec<(index, offset)> [32..72]
     ; current_stack_size: [24..32]
     ; func_idx [16..24]
@@ -369,8 +356,9 @@ codegen_function:
     mov rcx, 16                 ; initial capacity
     call vec_init_with
 
-    bts qword [rsp + 72], 7     ; mark rsp as used
-    bts qword [rsp + 72], 6     ; mark rbp as used
+    bts word [rsp + 72], 7     ; mark rsp as used
+    bts word [rsp + 72], 6     ; mark rbp as used
+    bts word [rsp + 72], 0     ; mark rax as used
 
     ; push "section .text\n"
     mov rdi, [rsp]              ; ctx
@@ -502,7 +490,7 @@ codegen_function:
     ; "pop {dirtied registers}\n"
     mov rdi, [rsp]              ; ctx
     lea rdi, [rdi + 8]          ; &ctx.text
-    mov rsi, [rsp + 24]         ; &function_ctx
+    lea rsi, [rsp + 24]         ; &function_ctx
     mov rdx, 0                  ; push = false
     call codegen_push_pop_dirtied_registers
 
@@ -513,13 +501,20 @@ codegen_function:
     mov rdx, ADD_RSP_LEN
     call vec_extend
 
-    mov rdi, [rsp + 24]         ; &function_ctx
+    lea rdi, [rsp + 24]         ; &function_ctx
     mov rdi, [rdi + 0]          ; current_stack_size
     lea rsi, [rsp + 104]        ; scratch
     mov rdx, 16                 ; buffer length
     mov rcx, 10                 ; radix
     call int_to_str2
 
+    ; ; push "\n"
+    ; mov rdi, [rsp]              ; ctx
+    ; lea rdi, [rdi + 8]          ; &ctx.text
+    ; lea rsi, [rel COLON_NL]
+    ; inc rsi
+    ; call vec_push
+
     mov rdi, [rsp]              ; ctx
     lea rdi, [rdi + 8]          ; &ctx.text
     lea rsi, [rax + rdx]
@@ -549,13 +544,20 @@ codegen_function:
     mov rdx, SUB_RSP_LEN
     call vec_extend
 
-    mov rdi, [rsp + 24]         ; &function_ctx
+    lea rdi, [rsp + 24]         ; &function_ctx
     mov rdi, [rdi + 0]          ; current_stack_size
     lea rsi, [rsp + 104]        ; scratch
     mov rdx, 16                 ; buffer length
     mov rcx, 10                 ; radix
     call int_to_str2
 
+    ; ; push "\n"
+    ; mov rdi, [rsp]              ; ctx
+    ; lea rdi, [rdi + 8]          ; &ctx.text
+    ; lea rsi, [rel COLON_NL]
+    ; inc rsi
+    ; call vec_push
+
     mov rdi, [rsp]              ; ctx
     lea rdi, [rdi + 8]          ; &ctx.text
     lea rsi, [rax + rdx]
@@ -567,7 +569,7 @@ codegen_function:
     ; "push{dirtied registers}\n"
     mov rdi, [rsp]              ; ctx
     lea rdi, [rdi + 8]          ; &ctx.text
-    mov rsi, [rsp + 24]         ; &function_ctx
+    lea rsi, [rsp + 24]         ; &function_ctx
     mov rdx, 1                  ; push = false
     call codegen_push_pop_dirtied_registers
 
@@ -578,10 +580,11 @@ codegen_function:
     mov rdx, JMP_ARGS_LEN
     call vec_extend
 
+    .epilogue:
     add rsp, 120
-    push r13
-    push r14
-    push r15
+    pop r13
+    pop r14
+    pop r15
     pop rbx
     pop rbp
     ret
@@ -605,11 +608,11 @@ codegen_push_pop_dirtied_registers:
 
     mov byte [rsp + 29 + 3], 10 ; newline
     mov qword [rsp + 16], -1
-    mov rax, ' hsup'
+    mov rax, 'push '
     mov qword [rsp + 24], rax
     test rdx, rdx
     jz .skip_setup_pop
-    mov rax, '  pop'
+    mov rax, 'pop  '
     mov qword [rsp + 24], rax
 
     mov rdi, [rsp]              ; text
@@ -617,53 +620,24 @@ codegen_push_pop_dirtied_registers:
     mov [rsp + 16], rax
 .skip_setup_pop:
 
-    ; low: 0b01000000 high: 0b00001111
+    ; abcdsdpp_89abcdef
+    ; 01000000_00001111
     ; preserved registers:
-    mov rax, [rsi + 48 + 16]    ; dirtied_register_bitset low
-    mov rbx, 0b01000000
-    and rax, rbx
-    test rax, rax
-    jz .skip_rbp
-
-    mov rdi, 2
-    mov rsi, 8
-    lea rdx, [rsp + 29]
-    call get_register_name
-
-    mov rax, -1
-    cmp [rsp + 16], rax
-    jne .rbx_pop
-
-    mov rdi, [rsp]              ; text
-    lea rsi, [rsp + 24]
-    mov rdx, 9
-    call vec_extend
-    jmp .skip_rbp
-.rbx_pop:
-
-    mov rdi, [rsp]              ; text
-    lea rsi, [rsp + 24]
-    mov rdx, 9
-    mov rcx, [rsp + 16]         ; text.len()
-    call vec_insert_many
-
-.skip_rbp:
-    mov rax, [rsi + 48 + 16 + 8]  ; dirtied_register_bitset high
-    mov rbx, 0b00001111
-    and rax, rbx
-    test rax, rax
+    mov bx, word [rsi + 48 + 4]    ; dirtied_register_bitset
+    mov ax, 0b11110000_00000010
+    and bx, ax
+    test bx, bx
     jz .done
 
-    mov r15, 8
+    mov r15, 16
     xor r14, r14
 .reg_loop:
     cmp r14, r15
     jge .done
-    bt rax, r14
+    bt bx, r14w
     jnc .next_reg
 
-    mov rdi, 8
-    add rdi, r14
+    mov rdi, r14
     mov rsi, 8
     lea rdx, [rsp + 29]
     call get_register_name
@@ -680,11 +654,11 @@ codegen_push_pop_dirtied_registers:
 .reg_pop:
 
     mov rdi, [rsp]              ; text
-    lea rsi, [rsp + 24]
-    mov rdx, 9
-    mov rcx, [rsp + 16]         ; text.len()
+    mov rsi, [rsp + 16]         ; text.len()
+    lea rdx, [rsp + 24]
+    mov rcx, 9
     call vec_insert_many
-
+    nop
 
 .next_reg:
     inc r14
@@ -823,14 +797,14 @@ codegen_expr:
     mov [rsp + 40], rdx
 
     mov rdi, [rsp]              ; ctx
-    mov rdi, [rdi + 8]          ; &ctx.text
+    lea rdi, [rdi + 8]          ; &ctx.text
     lea rsi, [rsp + 32]         ; dst
     lea rdx, [rsp + 16]         ; src
     call codegen_move_dst_src
 
     ; push "jmp .epilogue\n"
     mov rdi, [rsp]              ; ctx
-    mov rdi, [rdi + 8]          ; &ctx.text
+    lea rdi, [rdi + 8]          ; &ctx.text
     lea rsi, [rel JMP_EPILOGUE]
     mov rdx, JMP_EPILOGUE_LEN
     call vec_extend
@@ -871,9 +845,9 @@ codegen_expr:
     mov word [rsp + 34], 0                 ; Operand.len = 0
 
     mov rdi, [rsp]              ; ctx
-    mov rdi, [rdi + 8]          ; &ctx.text
+    lea rdi, [rdi + 8]          ; &ctx.text
     lea rsi, [rsp + 16]         ; dst
-    lea rdx, [rsp + 16]         ; src
+    lea rdx, [rsp + 32]         ; src
     call codegen_move_dst_src
 
     mov rax, qword [rsp + 16]
@@ -1252,6 +1226,12 @@ codegen_move_dst_src:
     call codegen_write_operand
 
 .epilogue:
+    mov rdi, [rsp]              ; *text
+    lea rsi, [rel COLON_NL]
+    inc rsi
+    mov rdx, 1
+    call vec_extend
+
     add rsp, 24
     pop rbx
     pop rbp
diff --git a/lang/tests/codegen.rs b/lang/tests/codegen.rs
index 541166c..0d31fb4 100644
--- a/lang/tests/codegen.rs
+++ b/lang/tests/codegen.rs
@@ -76,14 +76,16 @@ fn main() {
 
             let mut codegen = util::defs::CodegenCtx {
                 ast: &mut ast,
-                text: util::vec::Vec::new(),
+                text: util::vec::Vec::new_with(1024),
             };
 
             util::defs::codegen_function(&mut codegen, expr_id);
 
+            eprintln!("Code generation completed.");
+
             println!(
                 "Generated code:\n{}",
-                core::str::from_utf8(codegen.text.as_slice()).unwrap()
+                core::str::from_utf8_unchecked(codegen.text.as_slice())
             );
         };
     }
diff --git a/lang/tests/shared/defs.rs b/lang/tests/shared/defs.rs
index 5bbdb2e..51f9e0b 100644
--- a/lang/tests/shared/defs.rs
+++ b/lang/tests/shared/defs.rs
@@ -216,8 +216,8 @@ pub struct CodegenCtx {
 pub struct FunctionCtx {
     pub current_stack_size: u64,
     pub stack_vars: Vec<(u64, u64)>,
-    pub register_bitset: u128,
-    pub dirtied_register_bitset: u128,
+    pub register_bitset: u16,
+    pub dirtied_register_bitset: u16,
 }
 
 #[repr(C)]