vide/src/graph.luau

local flags = require "./flags"

export type SourceNode<T> = {
    cache: T,
    [number]: Node<T>
}

export type Node<T> =  {
    cache: T,
    effect:  ((T) -> T) | false,
    cleanups: { () -> () } | false,

    context: { [number]: unknown } | false,

    owned: { Node<T> } | false,
    owner: Node<T> | false,

    parents: { SourceNode<T> },
    [number]: Node<T> -- children
}

local scopes = { n = 0 } :: { [number]: Node<any>, n: number } -- scopes stack

local function efn(err: string)
    local trace = debug.traceback(err, 2)

    if string.find(err, "^effect error stacktrace") then -- if effect error is nested
        trace = string.gsub("    " .. trace, "\n", function() -- indent entire error
            return "\n    "
        end)
    end

    trace ..= "\nsource update stacktrace:"
    return trace
end

local function effect_traceback(err: unknown): string
    local message = tostring(err)
    -- Nested effect error already has full thread stack
    if string.find(message, "^effect error:\n") then return message end
    return `effect error:\n{debug.traceback(message, 2)}`
end

local function ycall<T, U>(fn: (T) -> U, arg: T): (boolean, string|U)

    local thread = coroutine.create(xpcall)
    --local function efn(err: string) return debug.traceback(err, 3) end
    local resume_ok, run_ok, result = coroutine.resume(thread, fn, efn, arg)

    assert(resume_ok)
        
    if coroutine.status(thread) ~= "dead" then
        return false, debug.traceback(thread, "attempt to yield in reactive scope")
    end

    return run_ok, result
end

local function get_scope(): Node<unknown>?
    return scopes[scopes.n]
end

local function assert_stable_scope(): Node<unknown>
    local scope = get_scope()

    if not scope then
        local caller_name = debug.info(2, "n")
        return error(`cannot use {caller_name}() outside a stable or reactive scope`, 0)
    elseif scope.effect then
        error("cannot create a new reactive scope inside another reactive scope", 0)
    end

    return scope
end

local function push_child<T>(parent: SourceNode<any>, child: Node<any>)
    table.insert(parent, child)
    table.insert(child.parents, parent)
end

local function push_scope<T>(node: Node<T>)
    local n = scopes.n + 1
    scopes.n = n
    scopes[n] = node
end

local function pop_scope()
    local n = scopes.n
    scopes.n = n - 1
    scopes[n] = nil
end

local function push_cleanup<T>(node: Node<T>, cleanup: () -> ())
    if node.cleanups then
        table.insert(node.cleanups, cleanup)
    else
        node.cleanups = { cleanup }
    end
end

-- runs every cleanup even if one errors, returns the first error
local function flush_cleanups<T>(node: Node<T>): string?
    local first_err: string?

    if node.cleanups then
        for _, fn in node.cleanups do
            local ok, err: string? = xpcall(fn, debug.traceback)
            if not ok and not first_err then first_err = `cleanup error: {err}` end
        end

        table.clear(node.cleanups)
    end

    return first_err
end

local function find_and_swap_pop<T>(t: { T }, v: T)
    local i = table.find(t, v) :: number
    local n = #t
    t[i] = t[n]
    t[n] = nil
end

local function unparent<T>(node: Node<T>)
    local parents = node.parents

    for i, parent in parents do
        find_and_swap_pop(parent, node)
        parents[i] = nil
    end
end

local destroy_owned: <T>(node: Node<T>) -> string?

-- always fully destroys the node, returns the first cleanup error
local function destroy_node<T>(node: Node<T>): string?
    if flags.strict and table.find(scopes, node) then
        error("attempt to destroy an active scope", 0)
    end

    local err = flush_cleanups(node)
    unparent(node)
    
    if node.owner then
        find_and_swap_pop(node.owner.owned :: { Node<T> }, node)
        node.owner = false
    end

    local owned_err = destroy_owned(node)
    return err or owned_err
end

function destroy_owned<T>(node: Node<T>): string?
    local err: string?

    if node.owned then
        local owned = node.owned
        while owned[1] do
            local owned_err = destroy_node(owned[1])
            err = err or owned_err
        end
    end

    return err
end

local function destroy<T>(node: Node<T>)
    local err = destroy_node(node)
    if err then error(err, 0) end
end

local update_queue = { n = 0 } :: { n: number, [number]: Node<any> }

-- returns whether the cached value changed, and an error instead of throwing
local function try_evaluate_node<T>(node: Node<T>): (boolean, string?)
    if flags.strict then
        if table.find(scopes, node) then
            return false, "a scope, that should rerun due to the update of a source, is already active"
        end

        local initial_value = node.cache
        local err: string?

        for i = 1, 2 do
            local cur_value = node.cache

            local cleanup_err = flush_cleanups(node)
            local owned_err = destroy_owned(node)
            err = err or cleanup_err or owned_err
    
            push_scope(node)
            local ok, new_value = ycall(node.effect :: (T) -> T, cur_value)
            pop_scope()
            
            if not ok then
                return false, `effect error stacktrace\n{new_value :: string}`
            end

            node.cache = new_value :: T
        end

        return initial_value ~= node.cache, err
    else
        local cur_value = node.cache

        local cleanup_err = flush_cleanups(node)
        local owned_err = destroy_owned(node)

        push_scope(node)
        local ok, new_value = xpcall(node.effect :: (T) -> T, effect_traceback, node.cache)
        pop_scope()

        if not ok then
            return false, new_value
        end
    
        node.cache = new_value
        return cur_value ~= new_value, cleanup_err or owned_err
    end
end

local function evaluate_node<T>(node: Node<T>): boolean
    local changed, err = try_evaluate_node(node)
    if err then error(err, 0) end
    return changed
end

local function queue_children_for_update<T>(node: SourceNode<T>)
    local i = update_queue.n
    while node[1] do
        i += 1
        update_queue[i] = node[1]
        unparent(node[1])
    end
    update_queue.n = i
end

local function get_update_queue_length()
    return update_queue.n
end

-- queued nodes are already unparented, so every one must be evaluated even if
-- another errors; the first error is thrown after the queue is flushed
local function flush_update_queue(from: number)
    local first_err: string?

    local i = from + 1
    while i <= update_queue.n do
        local node = update_queue[i]
        --assert(node.effect)

        -- check if node is still owned in case destroyed after queued
        if node.owner then
            local changed, err = try_evaluate_node(node)
            if changed then queue_children_for_update(node) end
            first_err = first_err or err
        end

        update_queue[i] = false :: any -- false instead of nil to avoid sparse
        i += 1
    end
    
    update_queue.n = from

    if first_err then error(first_err, 0) end
end

local function update_descendants<T>(root: SourceNode<T>)
    local n0 = update_queue.n
    queue_children_for_update(root)

    if flags.batch then return end

    flush_update_queue(n0)
end

local function push_scope_as_child_of<T>(node: SourceNode<T>)
    local scope = get_scope()
    if scope and scope.effect then -- do not track nodes with no effect
        push_child(node, scope)
    end
end

local function create_node<T>(owner: false | Node<any>, effect: false | (T) -> T, value: T): Node<T>
    local node: Node<T> = {
        cache = value,
        effect = effect,
        cleanups = false,

        context = false,

        owner = owner,
        owned = false,

        parents = {},
    }

    if owner then
        if owner.owned then
            table.insert(owner.owned, node)
        else
            owner.owned = { node }
        end
    end

    return node
end

local function create_source_node<T>(value: T): SourceNode<T>
    return { cache = value }
end

local function get_children<T>(node: Node<T>): { Node<unknown> }
    return { unpack(node) } :: { Node<any> }
end

local function set_context<T>(node: Node<T>, key: number, value: unknown)
    if node.context then
        node.context[key] = value
    else
        node.context = { [key] = value }
    end
end

return table.freeze {
    push_scope = push_scope,
    pop_scope = pop_scope,
    evaluate_node = evaluate_node,
    get_scope = get_scope,
    assert_stable_scope = assert_stable_scope,
    push_cleanup = push_cleanup,
    destroy = destroy,
    flush_cleanups = flush_cleanups,
    push_scope_as_child_of = push_scope_as_child_of,
    update_descendants = update_descendants,
    push_child = push_child,
    create_node = create_node,
    create_source_node = create_source_node,
    get_children = get_children,
    flush_update_queue = flush_update_queue,
    get_update_queue_length = get_update_queue_length,
    set_context = set_context,
    scopes = scopes,

    q = update_queue
}