ladybird/Libraries/LibWeb/IndexedDB/Internal/Database.cpp

/*
 * Copyright (c) 2024-2025, stelar7 <dudedbz@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/NeverDestroyed.h>
#include <LibGC/Heap.h>
#include <LibWeb/IndexedDB/IDBDatabase.h>
#include <LibWeb/IndexedDB/IDBTransaction.h>
#include <LibWeb/IndexedDB/Internal/Algorithms.h>
#include <LibWeb/IndexedDB/Internal/ConnectionQueueHandler.h>
#include <LibWeb/IndexedDB/Internal/Database.h>

namespace Web::IndexedDB {

using IDBDatabaseMapping = HashMap<StorageAPI::StorageKey, HashMap<String, GC::Root<Database>>>;
static IDBDatabaseMapping& idb_databases()
{
    static NeverDestroyed<IDBDatabaseMapping> databases;
    return *databases;
}

void Database::for_each_database(AK::Function<void(Database&)> const& visitor)
{
    for (auto const& [key, mapping] : idb_databases()) {
        for (auto const& [_, database] : mapping) {
            if (!database)
                continue;
            visitor(*database);
        }
    }
}

GC_DEFINE_ALLOCATOR(Database);

Database::~Database() = default;

GC::Ref<Database> Database::create(GC::Heap& heap, String const& name)
{
    return heap.allocate<Database>(name);
}

void Database::visit_edges(Visitor& visitor)
{
    Base::visit_edges(visitor);
    visitor.visit(m_upgrade_transaction);
    visitor.visit(m_object_stores);

    if (m_pending_connection_wait.has_value())
        visitor.visit(m_pending_connection_wait->callback);
}

GC::Ptr<ObjectStore> Database::object_store_with_name(String const& name) const
{
    for (auto const& object_store : m_object_stores) {
        if (object_store->name() == name)
            return object_store;
    }

    return nullptr;
}

Vector<GC::Weak<Database>> Database::for_key(StorageAPI::StorageKey const& key)
{
    Vector<GC::Weak<Database>> databases;
    for (auto const& database_mapping : idb_databases().get(key).value_or({})) {
        databases.append(*database_mapping.value);
    }

    return databases;
}

RequestList& ConnectionQueueHandler::for_key_and_name(StorageAPI::StorageKey const& key, String const& name)
{
    auto& instance = ConnectionQueueHandler::the();
    auto maybe_connection = instance.m_open_requests.find_if([&key, &name](Connection const& connection) {
        return connection.storage_key == key && connection.name == name;
    });

    if (!maybe_connection.is_end())
        return (*maybe_connection)->request_list;

    auto new_connection = adopt_ref(*new Connection(key, name));
    instance.m_open_requests.append(new_connection);
    return new_connection->request_list;
}

Optional<Database&> Database::for_key_and_name(StorageAPI::StorageKey const& key, String const& name)
{
    auto database_mapping = idb_databases().ensure(key, [] { return HashMap<String, GC::Root<Database>>(); });
    if (auto maybe_database = database_mapping.get(name); maybe_database.has_value())
        return *maybe_database.value();
    return {};
}

ErrorOr<GC::Ref<Database>> Database::create_for_key_and_name(GC::Heap& heap, StorageAPI::StorageKey const& key, String const& name)
{
    auto database_mapping = TRY(idb_databases().try_ensure(key, [] {
        return HashMap<String, GC::Root<Database>>();
    }));

    auto value = Database::create(heap, name);

    database_mapping.set(name, value);
    idb_databases().set(key, database_mapping);

    return value;
}

ErrorOr<void> Database::delete_for_key_and_name(StorageAPI::StorageKey const& key, String const& name)
{
    // FIXME: Is a missing entry a failure?
    auto maybe_database_mapping = idb_databases().get(key);
    if (!maybe_database_mapping.has_value())
        return {};

    auto& database_mapping = maybe_database_mapping.value();
    auto maybe_database = database_mapping.get(name);
    if (!maybe_database.has_value())
        return {};

    auto did_remove = database_mapping.remove(name);
    if (!did_remove)
        return {};

    idb_databases().set(key, database_mapping);

    return {};
}

void Database::associate(GC::Ref<IDBDatabase> connection)
{
    m_associated_connections.append(connection);
}

void Database::dissociate(IDBDatabase& connection)
{
    m_associated_connections.remove_first_matching([&](auto& entry) { return entry == &connection; });
}

GC::Ref<Database::AssociatedConnections> Database::associated_connections_as_heap_vector(GC::Heap& heap)
{
    auto connections = heap.allocate<AssociatedConnections>();
    for (auto& associated_connection : m_associated_connections) {
        if (associated_connection)
            connections->elements().append(*associated_connection);
    }
    return connections;
}

GC::RootVector<GC::Ref<IDBDatabase>> Database::associated_connections_as_root_vector()
{
    GC::RootVector<GC::Ref<IDBDatabase>> connections {};
    for (auto& connection : m_associated_connections) {
        if (connection)
            connections.append(*connection);
    }
    return connections;
}

GC::Ref<Database::AssociatedConnections> Database::associated_connections_as_heap_vector_except(GC::Heap& heap, IDBDatabase& connection)
{
    auto connections = heap.allocate<AssociatedConnections>();
    for (auto& associated_connection : m_associated_connections) {
        if (associated_connection && associated_connection != &connection)
            connections->elements().append(*associated_connection);
    }
    return connections;
}

void Database::wait_for_connections_to_close(ReadonlySpan<GC::Ref<IDBDatabase>> connections, GC::Ref<GC::Function<void()>> after_all)
{
    bool all_closed = true;
    for (auto const& entry : connections) {
        if (entry->state() != ConnectionState::Closed) {
            all_closed = false;
            break;
        }
    }

    if (all_closed) {
        queue_a_database_task(after_all);
        return;
    }

    VERIFY(!m_pending_connection_wait.has_value());
    Vector<GC::Weak<IDBDatabase>> weak_connections;
    weak_connections.ensure_capacity(connections.size());
    for (auto const& connection : connections)
        weak_connections.unchecked_append(connection);
    m_pending_connection_wait = PendingConnectionWait {
        .connections = move(weak_connections),
        .callback = after_all,
    };
}

void Database::check_pending_connection_wait()
{
    if (!m_pending_connection_wait.has_value())
        return;

    auto& wait = m_pending_connection_wait.value();
    for (auto const& connection : wait.connections) {
        if (connection && connection->state() != ConnectionState::Closed)
            return;
    }

    // All connections are closed - invoke the callback and clear the wait
    auto callback = wait.callback;
    m_pending_connection_wait.clear();
    queue_a_database_task(callback);
}

}