16.lua

local Vec, allvecs = {}, {}
setmetatable(allvecs, { __mode = "kv" })
local function vec(x, y)
    local key = x .. "," .. y
    local found = allvecs[key]
    if found then return found end
    local v = { x = x, y = y }
    setmetatable(v, Vec)
    allvecs[key] = v
    return v
end
Vec.__add = function(a, b) return vec(a.x + b.x, a.y + b.y) end

local N, E, S, W = vec(0, -1), vec(1, 0), vec(0, 1), vec(-1, 0)
N.l = W ; N.r = E ; E.l = N ; E.r = S ; S.l = E ; S.r = W ; W.l = S ; W.r = N
N.name = "^" ; S.name = "V" ; E.name = ">" ; W.name = "<"

local transformations = {
    ["."] = function(dir) return { dir } end,
    ["-"] = function(dir) return dir.y == 0 and { dir } or { W, E } end,
    ["|"] = function(dir) return dir.y == 0 and { N, S } or { dir } end,
    ["\\"] = function(dir) return dir.y == 0 and { dir.r } or { dir.l } end,
    ["/"] = function(dir) return dir.y == 0 and { dir.l } or { dir.r } end
}

local function readmap(iter)
    local map = {}
    local y = 0
    for line in iter do
        y = y + 1
        local x = 0
        for c in line:gmatch(".") do
            x = x + 1
            map[vec(x, y)] = { tranform = transformations[c], char = c }
        end
    end
    return map
end

local function step(map, state)
    local newbeams = {}
    local visited, energized = state.visited, state.energized
    for _, beam in ipairs(state.beams) do
        local newpos = beam.pos + beam.dir
        local location = map[newpos]
        if location then
            if not energized[newpos] then
                energized[newpos] = true
                energized.total = energized.total + 1
            end
            if not visited[beam.dir][newpos] then
                visited[beam.dir][newpos] = true
                for _, newdir in ipairs(location.tranform(beam.dir)) do
                    table.insert(newbeams, { pos = newpos, dir = newdir })
                end
            end
        end
    end
    state.beams = newbeams
end

local function start(pos, dir)
    return {
        beams = { { pos = pos, dir = dir } },
        visited = { [N] = {}, [S] = {}, [E] = {}, [W] = {} },
        energized = { total = 0 }
    }
end

return {
    readmap = readmap,
    step = step,
    start = start,
    vec = vec,
    N = N,
    S = S,
    W = W,
    E = E
}