AoC 2024 Day 9 Part 2

#! /usr/bin/env lua

function main()
    local file_spans, free_spans = read_input()
    local result = defragment_spans(file_spans, free_spans)
    print(result)
end

function defragment_spans(file_spans, free_spans)
    local checksum = 0

    local impossible_sizes = { false, false, false, false, false, false, false, false, false }
    local impossible_sizes_length = 0

    local last_file_span_index = 0
    local file_spans_length = #file_spans
    for file_span_index = file_spans_length, 1, -1 do
        local file_span = file_spans[file_span_index]
        local file_span_size = file_span.size
        local file_span_start_block_index = 0

        local free_span = nil
        local candidate_free_span = nil

        if (impossible_sizes[file_span_size]) then
            goto continue
        end

        file_span_start_block_index = file_span.start_block_index

        candidate_free_span = free_spans
        while (candidate_free_span ~= nil) do
            if (candidate_free_span.start_block_index >= file_span_start_block_index) then
                break
            end

            if (candidate_free_span.size >= file_span_size) then
                free_span = candidate_free_span
                break
            end

            candidate_free_span = candidate_free_span.next
        end

        if (free_span == nil) then
            impossible_sizes[file_span_size] = true

            impossible_sizes_length = impossible_sizes_length + 1
            if (impossible_sizes_length == 9) then
                last_file_span_index = file_span_index
                break
            end

            goto continue
        end

        file_span.start_block_index = free_span.start_block_index

        new_free_span_size = free_span.size - file_span_size
        if (new_free_span_size ~= 0) then
            free_span.start_block_index = free_span.start_block_index + file_span_size
            free_span.size = new_free_span_size
        else
            if (free_span == free_spans) then
                free_spans = free_span.next
            end

            if (free_span.prev ~= nil) then
                free_span.prev.next = free_span.next
            end

            if (free_span.next ~= nil) then
                free_span.next.prev = free_span.prev
            end

            free_span.prev = nil
            free_span.next = nil
        end

        ::continue::

        checksum = checksum +
            file_span.size * file_span.file_index * (file_span.start_block_index + ((file_span.size - 1) / 2))
    end

    for file_span_index = last_file_span_index, 1, -1 do
        local file_span = file_spans[file_span_index]
        checksum = checksum +
            file_span.size * file_span.file_index * (file_span.start_block_index + ((file_span.size - 1) / 2))
    end

    return math.tointeger(checksum)
end

function read_input()
    local file_spans = {}
    local free_spans = nil
    local prev_free_span = nil

    local block_index = 0
    local is_file_span = true
    local file_index = 0
    while (true) do
        local input = io.read(1)
        if (input == '\n') then
            break
        end

        local span = {
            start_block_index = block_index,
            size = tonumber(input)
        }

        if (is_file_span) then
            span.file_index = file_index
            table.insert(file_spans, span)
        else
            span.next = nil

            if (free_spans == nil) then
                span.prev = nil
                free_spans = span
            else
                span.prev = prev_free_span
                prev_free_span.next = span
            end

            prev_free_span = span
        end

        block_index = block_index + span.size

        if (is_file_span) then
            is_file_span = false
        else
            is_file_span = true
            file_index = file_index + 1
        end
    end

    return file_spans, free_spans
end

main()