path: root/2025/day8.asm

; Advent of Code 2025 - Day 8

bits 64
default rel

global _start

%include "utils/io.asm"
%include "utils/parse.asm"

%define RAW_SIZE   262144
%define MAXN       1024
%define MAXE       (MAXN * (MAXN - 1) / 2)
%define K_LIMIT    1000          ; part 1 connects the K_LIMIT closest pairs

section .data
    input_path db "input/day8.txt", 0

    label1     db "Part 1: ", 0
    label2     db "Part 2: ", 0

section .bss
    raw           resb RAW_SIZE
    x             resq MAXN
    y             resq MAXN
    z             resq MAXN

    edges         resb (MAXE * 16)   ; records: [dist2:8][i:4][j:4]
    edges2        resb (MAXE * 16)   ; radix ping-pong buffer
    counts        resq 256

    parent        resq MAXN
    csize         resq MAXN

    raw_len       resq 1
    n_boxes       resq 1
    m_edges       resq 1
    sorted_ptr    resq 1
    comp_count    resq 1

    part1         resq 1
    part2         resq 1

section .text

_start:
    mov     rdi, input_path
    lea     rsi, [raw]
    mov     rdx, RAW_SIZE
    call    read_file
    mov     [raw_len], rax

    ; parse "X,Y,Z" lines into x/y/z, counting boxes
    lea     rsi, [raw]
    mov     rax, [raw_len]
    lea     r14, [raw + rax]   ; end pointer
    xor     r12, r12           ; box index
parse_line:
    cmp     rsi, r14
    jae     parse_done
    mov     al, [rsi]
    cmp     al, 10
    je      parse_skipnl
    cmp     al, 13
    je      parse_skipnl
    call    parse_u64
    mov     [x + r12*8], rax
    inc     rsi
    call    parse_u64
    mov     [y + r12*8], rax
    inc     rsi
    call    parse_u64
    mov     [z + r12*8], rax
    inc     r12
parse_eol:
    cmp     rsi, r14
    jae     parse_done
    mov     al, [rsi]
    inc     rsi
    cmp     al, 10
    jne     parse_eol
    jmp     parse_line
parse_skipnl:
    inc     rsi
    jmp     parse_line
parse_done:
    mov     [n_boxes], r12

    ; build every pairwise edge with squared distance as the key
    lea     r11, [edges]   ; write pointer
    xor     r12, r12       ; edge count
    xor     r8, r8         ; i
gen_i:
    cmp     r8, [n_boxes]
    jge     gen_done
    mov     r13, [x + r8*8]
    mov     r14, [y + r8*8]
    mov     r15, [z + r8*8]
    lea     r9, [r8 + 1]   ; j = i + 1
gen_j:
    cmp     r9, [n_boxes]
    jge     gen_j_done
    mov     rax, r13
    sub     rax, [x + r9*8]
    imul    rax, rax
    mov     rcx, rax
    mov     rax, r14
    sub     rax, [y + r9*8]
    imul    rax, rax
    add     rcx, rax
    mov     rax, r15
    sub     rax, [z + r9*8]
    imul    rax, rax
    add     rcx, rax        ; key = dx^2 + dy^2 + dz^2
    mov     [r11], rcx
    mov     [r11 + 8], r8d
    mov     [r11 + 12], r9d
    add     r11, 16
    inc     r12
    inc     r9
    jmp     gen_j
gen_j_done:
    inc     r8
    jmp     gen_i
gen_done:
    mov     [m_edges], r12

    ; LSD radix sort by the 5 low bytes of the key (max key < 2^40)
    lea     r11, [edges]    ; src
    lea     r12, [edges2]   ; dst
    xor     r13, r13        ; byte pass 0..4
rsort_pass:
    cmp     r13, 5
    jge     rsort_done
    lea     rdi, [counts]
    mov     rcx, 256
    xor     rax, rax
    rep     stosq

    xor     rcx, rcx        ; k
rs_hist:
    cmp     rcx, [m_edges]
    jge     rs_hist_done
    mov     rax, rcx
    shl     rax, 4
    add     rax, r11
    movzx   edx, byte [rax + r13]
    inc     qword [counts + rdx*8]
    inc     rcx
    jmp     rs_hist
rs_hist_done:
    xor     rcx, rcx        ; t
    xor     r8, r8          ; running offset
rs_prefix:
    cmp     rcx, 256
    jge     rs_prefix_done
    mov     rax, [counts + rcx*8]
    mov     [counts + rcx*8], r8
    add     r8, rax
    inc     rcx
    jmp     rs_prefix
rs_prefix_done:
    xor     rcx, rcx        ; k
rs_scatter:
    cmp     rcx, [m_edges]
    jge     rs_scatter_done
    mov     rax, rcx
    shl     rax, 4
    add     rax, r11        ; src record
    movzx   edx, byte [rax + r13]
    mov     r9, [counts + rdx*8]
    inc     qword [counts + rdx*8]
    mov     r10, r9
    shl     r10, 4
    add     r10, r12        ; dst slot
    mov     r8, [rax]
    mov     [r10], r8
    mov     r8, [rax + 8]
    mov     [r10 + 8], r8
    inc     rcx
    jmp     rs_scatter
rs_scatter_done:
    xchg    r11, r12
    inc     r13
    jmp     rsort_pass
rsort_done:
    mov     [sorted_ptr], r11

    ; union-find init
    xor     rcx, rcx
dsu_init:
    cmp     rcx, [n_boxes]
    jge     dsu_done
    mov     [parent + rcx*8], rcx
    mov     qword [csize + rcx*8], 1
    inc     rcx
    jmp     dsu_init
dsu_done:
    mov     rax, [n_boxes]
    mov     [comp_count], rax
    mov     qword [part2], 0

    ; Kruskal: one sorted pass drives both parts
    mov     r11, [sorted_ptr]
    xor     r15, r15        ; edge index
kloop:
    cmp     r15, [m_edges]
    jge     kdone
    mov     r12, r15
    shl     r12, 4
    add     r12, r11        ; record pointer
    mov     r13d, [r12 + 8]  ; i
    mov     r14d, [r12 + 12] ; j

    mov     edi, r13d
    call    find
    mov     rbx, rax        ; root of i
    mov     edi, r14d
    call    find            ; root of j

    cmp     rbx, rax
    je      k_after         ; already connected

    mov     rcx, [csize + rbx*8]
    mov     rdx, [csize + rax*8]
    cmp     rcx, rdx
    jge     k_no_swap
    xchg    rbx, rax        ; keep rbx as the larger root
k_no_swap:
    mov     [parent + rax*8], rbx
    mov     rcx, [csize + rax*8]
    add     [csize + rbx*8], rcx
    dec     qword [comp_count]

    cmp     qword [comp_count], 1
    jne     k_after
    cmp     qword [part2], 0
    jne     k_after
    mov     rax, [x + r13*8]
    imul    rax, [x + r14*8]
    mov     [part2], rax
k_after:
    cmp     r15, (K_LIMIT - 1)
    jne     k_check_break
    call    compute_part1
k_check_break:
    cmp     r15, (K_LIMIT - 1)
    jl      k_next
    cmp     qword [part2], 0
    je      k_next
    jmp     kdone
k_next:
    inc     r15
    jmp     kloop
kdone:

print_results:
    mov     rcx, label1
    call    print_string
    mov     rax, [part1]
    call    print_int64
    call    print_newline

    mov     rcx, label2
    call    print_string
    mov     rax, [part2]
    call    print_int64
    call    print_newline

    mov     rax, SYS_EXIT
    xor     rdi, rdi
    syscall

; find: union-find root of node RDI, with path halving.
;   Output: RAX = root. Clobbers RAX, RDX, RDI.
find:
.loop:
    mov     rax, [parent + rdi*8]
    cmp     rax, rdi
    je      .done
    mov     rdx, [parent + rax*8]
    mov     [parent + rdi*8], rdx
    mov     rdi, rdx
    jmp     .loop
.done:
    ret

; compute_part1: product of the three largest component sizes into part1.
;   Clobbers RAX, RCX, R8, R9, R10.
compute_part1:
    xor     r8, r8    ; largest
    xor     r9, r9    ; 2nd
    xor     r10, r10  ; 3rd
    xor     rcx, rcx
.loop:
    cmp     rcx, [n_boxes]
    jge     .done
    mov     rax, [parent + rcx*8]
    cmp     rax, rcx
    jne     .next     ; not a root
    mov     rax, [csize + rcx*8]
    cmp     rax, r8
    jle     .try2
    mov     r10, r9
    mov     r9, r8
    mov     r8, rax
    jmp     .next
.try2:
    cmp     rax, r9
    jle     .try3
    mov     r10, r9
    mov     r9, rax
    jmp     .next
.try3:
    cmp     rax, r10
    jle     .next
    mov     r10, rax
.next:
    inc     rcx
    jmp     .loop
.done:
    mov     rax, r8
    imul    rax, r9
    imul    rax, r10
    mov     [part1], rax
    ret