4.wsf

# Advent of Code 2021
# Day 4: Giant Squid
# https://adventofcode.com/2021/day/4

    "@array"  import
    "@matrix" import
    "@mem"    import
    "@string" import

    0 10 ',' call array.read_base_delim # numbers
    jz .read_numbers_error
    ^ ^ readi retrieve # boards
    2dup
.read_boards:
    1- ^ jn .read_board_break
    ^1 readc # '\n'
    ^1 5 call read_square_board
    swap 25+ swap
    jmp .read_boards
.read_board_break:
    2drop

    2dup call lookup_number_indices
    2drop end

    0 0 # i winners
.loop_numbers: # IVs: i winners
    ^1 retrieve swap
    0 swap # j
.loop_boards: # IVs: j winners
    # numbers boards i nums[i] j winners
    ^1 25* ^6 + ^3 call check_board
    ^ jn .loop_boards_next
    swap 1+ swap
    ^1 1 j= .first_winner
    ^1 ^6 j= .last_winner
.loop_boards_next:
    drop
    swap 1+ swap
    ^1 ^5 j< .loop_boards
    2slide
    swap 1+ swap
    ^1 ^4 j< .loop_numbers
    4drop end
.first_winner:
    ^ printi '\n' printc # Part 1
    jmp .loop_boards_next
.last_winner:
    6slide printi '\n' printc # Part 2
    end

# read_square_board reads a square grid of two-digit numbers into the
# heap.
# (addr n -- )
read_square_board:
    ^ * 1- ^1 + swap
.read_square_board_loop:
    2dup j< .read_square_board_ret
    ^
    ^ ^ readc retrieve
    ^ ' ' j= .read_square_board_space
    '0' -
.read_square_board_next_digit:
    10*
    ^1 ^ readc retrieve '0' - +
    ^1 readc # discard ' ' or '\n'
    store
    1+
    jmp .read_square_board_loop
.read_square_board_ret:
    2drop ret
.read_square_board_space:
    drop 0 jmp .read_square_board_next_digit

# check_board marks the number on the board and checks whether the board
# has won. Returns the board score when it is a winner, that is, the sum
# of the remaining numbers, multiplied by the picked number. Otherwise,
# returns negative.
# (addr num -- sum)
check_board:
    swap
    ^ retrieve -2 j= .check_board_already_won
    ^ 25 ^3 -1 call mem.replace_all
    jz .check_board_not_found
    0 # i
.check_board_loop:
    # num addr i
    ^1 5 5 ^3 -1 call matrix.count_row 5 j= .check_board_winner
    ^1 5 5 ^3 -1 call matrix.count_col 5 j= .check_board_winner
    1+
    ^ 5 j< .check_board_loop
    drop
.check_board_not_found:
    -1 2slide ret
.check_board_already_won:
    -2 2slide ret
.check_board_winner:
    drop
    ^ 25 call array.sum
    ^1 25 -1 call array.count +
    swap -2 store # Mark as won
    * ret

# lookup_number_indices replaces each number on the boards with the
# index that it will be picked.
# (numbers boards -- )
lookup_number_indices:
    25* ^1 + 1- # max
    ^1 # i
.lookup_number_indices_loop:
    ^ 0 ^1 ^ retrieve call array.index_of store
    1+
    ^ ^2 j< .lookup_number_indices_loop
    3drop ret

# (x -- )
.read_numbers_error:
    -1 "Error: invalid numbers format"
    call string.println_neg_stack drop end