mt-mandlebrot.lua

local threading = require("threading")

local floor, sqrt = math.floor, math.sqrt
local insert, concat = table.insert, table.concat

local RESET = string.char(27) .. "[0m"
local CLEAR_LINE = string.char(27) .. "[2K"
local SHADES = " .'`^\",:;Il!i><~+_-?][}{1)(|\\/tfjrxnuvczXYUJCLQ0OZmwqpdbkhao*#MW&8%B@$"

local PRESETS = {
    quick = { width = 96, height = 30, iter = 160, zoom = 1.00, cx = -0.75, cy = 0.00 },
    showcase = { width = 132, height = 42, iter = 420, zoom = 1.00, cx = -0.75, cy = 0.00 },
    valley = { width = 132, height = 42, iter = 760, zoom = 62.0, cx = -0.7436439, cy = 0.1318259 },
    stress = { width = 180, height = 60, iter = 950, zoom = 1.00, cx = -0.75, cy = 0.00 },
}

local function clamp(x, lo, hi)
    if x < lo then return lo end
    if x > hi then return hi end
    return x
end


local function default_config()
    local preset = PRESETS.showcase
    local cpus = threading.cpucount()
    local cols = tonumber(os.getenv("COLUMNS") or "") or preset.width
    local w = clamp(cols - 2, 80, preset.width)
    return {
        width = w,
        height = floor(w * preset.height / preset.width),
        iter = preset.iter,
        zoom = preset.zoom,
        cx = preset.cx,
        cy = preset.cy,
        workers = clamp(cpus, 1, 8),
        color = os.getenv("NO_COLOR") == nil,
        baseline = true,
        progress = true,
        preset = "showcase",
    }
end

local function usage()
    io.write([[
LuaJIT-MT Mandelbrot showcase

Usage:
  ./src/luajit mt_mandelbrot.lua [options]

Options:
  -j, --workers N       worker OS threads; default min(cpucount, 8)
  --preset NAME         quick | showcase | valley | stress
  --width N             terminal render width in characters
  --height N            terminal render height in rows
  --iter N              max Mandelbrot iterations per point
  --zoom N              view zoom; higher is deeper
  --center X Y          complex-plane center
  --no-color            disable ANSI 256-color output
  --no-baseline         skip single-thread baseline timing
  --no-progress         do not print progress bar
  -h, --help            show this help

Good show-off commands:
  ./src/luajit demos/mt_mandelbrot.lua
  ./src/luajit demos/mt_mandelbrot.lua -j $(nproc) --preset valley
  ./src/luajit demos/mt_mandelbrot.lua -j $(nproc) --preset stress --no-color
]])
end

local function read_number(argv, i, name)
    local v = tonumber(argv[i])
    if not v then error("expected number after " .. name) end
    return v
end

local function apply_preset(cfg, name)
    local p = PRESETS[name]
    if not p then
        error("unknown preset '" .. tostring(name) .. "' expected quick, showcase, valley, or stress")
    end
    cfg.preset = name
    cfg.width = p.width
    cfg.height = p.height
    cfg.iter = p.iter
    cfg.zoom = p.zoom
    cfg.cx = p.cx
    cfg.cy = p.cy
end

local function parse_args(argv)
    local cfg = default_config()
    local i = 1
    while i <= #argv do
        local a = argv[i]
        local eq = a:match("^%-%-workers=(.+)$")
        if a == "-j" or a == "--workers" then
            i = i + 1; cfg.workers = read_number(argv, i, a)
        elseif eq then
            cfg.workers = assert(tonumber(eq), "bad --workers value")
        elseif a == "--preset" then
            i = i + 1; apply_preset(cfg, argv[i])
        elseif a:match("^%-%-preset=") then
            apply_preset(cfg, a:match("^%-%-preset=(.+)$"))
        elseif a == "--width" then
            i = i + 1; cfg.width = read_number(argv, i, a)
        elseif a == "--height" then
            i = i + 1; cfg.height = read_number(argv, i, a)
        elseif a == "--iter" then
            i = i + 1; cfg.iter = read_number(argv, i, a)
        elseif a == "--zoom" then
            i = i + 1; cfg.zoom = read_number(argv, i, a)
        elseif a == "--center" then
            i = i + 1; cfg.cx = read_number(argv, i, a)
            i = i + 1; cfg.cy = read_number(argv, i, a)
        elseif a == "--no-color" then
            cfg.color = false
        elseif a == "--no-baseline" then
            cfg.baseline = false
        elseif a == "--no-progress" then
            cfg.progress = false
        elseif a == "-h" or a == "--help" then
            usage(); os.exit(0)
        else
            error("unknown argument: " .. tostring(a) .. "\ntry --help")
        end
        i = i + 1
    end
    cfg.width = floor(clamp(cfg.width, 20, 400))
    cfg.height = floor(clamp(cfg.height, 8, 200))
    cfg.iter = floor(clamp(cfg.iter, 16, 10000))
    cfg.workers = floor(clamp(cfg.workers, 1, 256))
    if cfg.zoom <= 0 then error("--zoom must be positive") end
    return cfg
end

local function mandel(cx, cy, max_iter)
    local x, y = 0.0, 0.0
    local xx, yy = 0.0, 0.0
    local i = 0
    while xx + yy <= 4.0 and i < max_iter do
        y = 2.0 * x * y + cy
        x = xx - yy + cx
        xx = x * x
        yy = y * y
        i = i + 1
    end
    return i
end

local function color_code(iter, max_iter)
    -- A deterministic 256-color gradient with enough banding to be visible on video.
    local t = iter / max_iter
    local band = floor((sqrt(t) * 180.0 + iter * 13.0) % 216)
    return 16 + band
end

local function render_row(y, width, height, max_iter, zoom, center_x, center_y, use_color)
    local row = {}
    local shade_count = #SHADES

    -- Terminal characters are taller than they are wide, so compress Y a bit.
    local span_x = 3.45 / zoom
    local span_y = 1.75 / zoom
    local cy = center_y + ((y / (height - 1)) - 0.5) * span_y

    for px = 0, width - 1 do
        local cx = center_x + ((px / (width - 1)) - 0.5) * span_x
        local it = mandel(cx, cy, max_iter)
        if it >= max_iter then
            row[#row + 1] = use_color and (string.char(27) .. "[38;5;16m ") or " "
        else
            local idx = 1 + floor((it / max_iter) ^ 0.42 * (shade_count - 1))
            local ch = SHADES:sub(idx, idx)
            if use_color then
                row[#row + 1] = string.char(27) .. "[38;5;" .. color_code(it, max_iter) .. "m" .. ch
            else
                row[#row + 1] = ch
            end
        end
    end

    if use_color then row[#row + 1] = RESET end
    return concat(row)
end

local function draw_progress(done, total, start_t, label)
    local width = 30
    local filled = floor(width * done / total)
    local bar = string.rep("#", filled) .. string.rep("-", width - filled)
    local elapsed = threading.now() - start_t
    io.stderr:write(("\r%s%s [%s] %3d/%-3d %6.2fs"):format(CLEAR_LINE, label, bar, done, total, elapsed))
    io.stderr:flush()
end

local function render_serial(cfg)
    local rows = {}
    local t0 = threading.now()
    for y = 0, cfg.height - 1 do
        rows[y] = render_row(y, cfg.width, cfg.height, cfg.iter, cfg.zoom, cfg.cx, cfg.cy, cfg.color)
    end
    return rows, threading.now() - t0
end

local function worker_main(worker_id, jobs, results, width, height, max_iter, zoom, center_x, center_y, use_color)
    local me = threading.current()
    local runtime_id = me:id()
    local rendered = 0

    while true do
        local y, state = jobs:recv()
        if y == nil then
            -- Closed channel: no more work.
            break
        end

        local row = render_row(y, width, height, max_iter, zoom, center_x, center_y, use_color)
        local ok, err = results:send({ y = y, text = row, worker = worker_id })
        if not ok then
            error("result send failed: " .. tostring(err))
        end
        rendered = rendered + 1
    end

    return worker_id, rendered, runtime_id
end

local function render_parallel(cfg)
    local jobs = threading.channel(cfg.height)
    local results = threading.channel(cfg.height)
    local threads = {}
    local counts = {}
    local runtime_ids = {}
    local rows = {}
    local t0 = threading.now()

    for wid = 1, cfg.workers do
        threads[wid] = threading.spawn(worker_main, wid, jobs, results,
            cfg.width, cfg.height, cfg.iter, cfg.zoom, cfg.cx, cfg.cy, cfg.color)
    end

    for y = 0, cfg.height - 1 do
        local sent, err = jobs:send(y)
        if not sent then error("job send failed: " .. tostring(err)) end
    end
    jobs:close()

    local done = 0
    local last_progress = -1.0
    while done < cfg.height do
        local msg, state = results:recv(0.05)
        if msg ~= nil then
            rows[msg.y] = msg.text
            counts[msg.worker] = (counts[msg.worker] or 0) + 1
            done = done + 1
            local t = threading.now()
            if cfg.progress and (t - last_progress > 0.08 or done == cfg.height) then
                draw_progress(done, cfg.height, t0, "parallel")
                last_progress = t
            end
        elseif state == "timeout" then
            if cfg.progress then draw_progress(done, cfg.height, t0, "parallel") end
        else
            error("results channel closed before all rows arrived")
        end
    end
    if cfg.progress then io.stderr:write("\n") end

    for i = 1, #threads do
        local ok_join, wid_or_err, rendered, runtime_id = threads[i]:join()
        if not ok_join then
            error("worker failed: " .. tostring(wid_or_err))
        end
        runtime_ids[wid_or_err] = runtime_id
        if not counts[wid_or_err] then counts[wid_or_err] = rendered or 0 end
    end

    return rows, threading.now() - t0, counts, runtime_ids
end

local function print_header(cfg)
    io.write("LuaJIT-MT Mandelbrot showcase\n")
    io.write(("  preset=%s  image=%dx%d  iter=%d  zoom=%.3g\n"):format(cfg.preset, cfg.width, cfg.height, cfg.iter,
        cfg.zoom))
    io.write(("  cpus=%d  workers=%d  channels=bounded  ffi=no\n"):format(threading.cpucount(), cfg.workers))
    io.write(("  center=(%.9f, %.9f)\n\n"):format(cfg.cx, cfg.cy))
end

local function print_rows(rows, height)
    for y = 0, height - 1 do
        io.write(rows[y] or "", "\n")
    end
end

local function print_worker_stats(counts, runtime_ids, height)
    io.write("\nworker row distribution:\n")
    for wid = 1, #counts do
        local n = counts[wid] or 0
        local bar_len = floor(32 * n / height)
        io.write(("  worker %02d runtime-id=%-6s rows=%-3d %s\n"):format(
            wid, tostring(runtime_ids[wid] or "?"), n, string.rep("#", bar_len)))
    end
end

local function main()
    local cfg = parse_args(arg or {})
    print_header(cfg)

    local serial_time
    if cfg.baseline then
        io.stderr:write("single-thread baseline...\n")
        local _, dt = render_serial(cfg)
        serial_time = dt
        io.stderr:write(("single-thread baseline: %.3fs\n"):format(serial_time))
    end

    local rows, parallel_time, counts, runtime_ids = render_parallel(cfg)
    print_rows(rows, cfg.height)

    io.write(("\nparallel render: %.3fs with %d worker threads\n"):format(parallel_time, cfg.workers))
    if serial_time then
        io.write(("single-thread:   %.3fs\n"):format(serial_time))
        io.write(("speedup:         %.2fx\n"):format(serial_time / parallel_time))
    end
    print_worker_stats(counts, runtime_ids, cfg.height)
end

local ok_main, err = xpcall(main, debug.traceback)
if not ok_main then
    io.stderr:write(tostring(err), "\n")
    os.exit(1)
end