/*
 * Copyright (c) 2020-2024, Andreas Kling <andreas@ladybird.org>
 * Copyright (c) 2025, Aliaksandr Kalenik <kalenik.aliaksandr@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <LibGC/DeferGC.h>
#include <LibJS/Runtime/DescriptorArray.h>
#include <LibJS/Runtime/ExternalMemory.h>
#include <LibJS/Runtime/Realm.h>
#include <LibJS/Runtime/Shape.h>
#include <LibJS/Runtime/VM.h>

namespace JS {

GC_DEFINE_ALLOCATOR(Shape);
GC_DEFINE_ALLOCATOR(PrototypeChainValidity);

Shape::~Shape() = default;

void Shape::finalize()
{
    Base::finalize();
    if (m_dictionary)
        m_property_storage.property_table.~PropertyTablePtr();
}

GC::Ptr<DescriptorArray> Shape::descriptors() const
{
    VERIFY(!m_dictionary);
    return m_property_storage.descriptors;
}

void Shape::set_descriptors(GC::Ptr<DescriptorArray> descriptors)
{
    VERIFY(!m_dictionary);
    m_property_storage.descriptors = descriptors;
}

Shape::PropertyTable& Shape::property_table()
{
    VERIFY(m_dictionary);
    VERIFY(m_property_storage.property_table);
    return *m_property_storage.property_table;
}

Shape::PropertyTable const& Shape::property_table() const
{
    VERIFY(m_dictionary);
    VERIFY(m_property_storage.property_table);
    return *m_property_storage.property_table;
}

void Shape::become_dictionary_shape()
{
    VERIFY(!m_dictionary);
    new (&m_property_storage.property_table) PropertyTablePtr();
    m_property_storage.property_table = make<PropertyTable>();
    m_dictionary = true;
}

size_t Shape::external_memory_size() const
{
    size_t size = 0;
    if (m_dictionary)
        size += hash_map_external_memory_size(property_table());
    if (m_forward_transitions)
        size += hash_map_external_memory_size(*m_forward_transitions);
    if (m_prototype_transitions)
        size += hash_map_external_memory_size(*m_prototype_transitions);
    if (m_delete_transitions)
        size += hash_map_external_memory_size(*m_delete_transitions);
    if (m_child_prototype_shapes)
        size += vector_external_memory_size(*m_child_prototype_shapes);
    return size;
}

GC::Ref<Shape> Shape::create_dictionary_transition()
{
    auto new_shape = heap().allocate<Shape>(m_realm);
    new_shape->become_dictionary_shape();
    new_shape->m_has_parameter_map = m_has_parameter_map;
    new_shape->m_prototype = m_prototype;
    invalidate_prototype_if_needed_for_new_prototype(new_shape);
    copy_properties_to_dictionary_shape(*new_shape);
    return new_shape;
}

GC::Ptr<Shape> Shape::get_or_prune_cached_forward_transition(TransitionKey const& key)
{
    if (is_prototype_shape())
        return nullptr;
    if (!m_forward_transitions)
        return nullptr;
    auto it = m_forward_transitions->find(key);
    if (it == m_forward_transitions->end())
        return nullptr;
    if (!it->value) {
        // The cached forward transition has gone stale (from garbage collection). Prune it.
        m_forward_transitions->remove(it);
        return nullptr;
    }
    return it->value.ptr();
}

GC::Ptr<Shape> Shape::get_or_prune_cached_delete_transition(PropertyKey const& key)
{
    if (is_prototype_shape())
        return nullptr;
    if (!m_delete_transitions)
        return nullptr;
    auto it = m_delete_transitions->find(key);
    if (it == m_delete_transitions->end())
        return nullptr;
    if (!it->value) {
        // The cached delete transition has gone stale (from garbage collection). Prune it.
        m_delete_transitions->remove(it);
        return nullptr;
    }
    return it->value.ptr();
}

GC::Ptr<Shape> Shape::get_or_prune_cached_prototype_transition(Object* prototype)
{
    if (is_prototype_shape())
        return nullptr;
    if (!m_prototype_transitions)
        return nullptr;
    auto it = m_prototype_transitions->find(prototype);
    if (it == m_prototype_transitions->end())
        return nullptr;
    if (!it->value) {
        // The cached prototype transition has gone stale (from garbage collection). Prune it.
        m_prototype_transitions->remove(it);
        return nullptr;
    }
    return it->value.ptr();
}

GC::Ref<Shape> Shape::create_put_transition(PropertyKey const& property_key, PropertyAttributes attributes)
{
    TransitionKey key { property_key, attributes };
    if (auto existing_shape = get_or_prune_cached_forward_transition(key))
        return *existing_shape;
    auto new_shape = heap().allocate<Shape>(*this, PropertyCountChange::Increment);

    if (descriptors() && descriptors()->size() == m_property_count) {
        descriptors()->set(property_key, { m_property_count, attributes }, m_property_count);
        new_shape->set_descriptors(descriptors());
    } else {
        new_shape->set_descriptors(copy_descriptors());
        new_shape->descriptors()->set(property_key, { m_property_count, attributes }, m_property_count);
    }
    invalidate_prototype_if_needed_for_new_prototype(new_shape);
    if (!is_prototype_shape()) {
        if (!m_forward_transitions)
            m_forward_transitions = make<HashMap<TransitionKey, GC::Weak<Shape>>>();
        m_forward_transitions->set(key, new_shape);
    }
    return new_shape;
}

GC::Ref<Shape> Shape::create_configure_transition(PropertyKey const& property_key, PropertyAttributes attributes)
{
    TransitionKey key { property_key, attributes };
    if (auto existing_shape = get_or_prune_cached_forward_transition(key))
        return *existing_shape;
    auto new_shape = heap().allocate<Shape>(*this, PropertyCountChange::Preserve);
    new_shape->set_descriptors(copy_descriptors());
    new_shape->descriptors()->set_attributes(property_key, attributes, m_property_count);
    invalidate_prototype_if_needed_for_new_prototype(new_shape);
    if (!is_prototype_shape()) {
        if (!m_forward_transitions)
            m_forward_transitions = make<HashMap<TransitionKey, GC::Weak<Shape>>>();
        m_forward_transitions->set(key, new_shape.ptr());
    }
    return new_shape;
}

GC::Ref<Shape> Shape::create_prototype_transition(Object* new_prototype)
{
    if (new_prototype)
        new_prototype->convert_to_prototype_if_needed();
    if (auto existing_shape = get_or_prune_cached_prototype_transition(new_prototype))
        return *existing_shape;
    auto new_shape = heap().allocate<Shape>(*this, new_prototype);
    if (m_dictionary && m_property_count > DescriptorArray::max_descriptor_count) {
        new_shape->become_dictionary_shape();
        copy_properties_to_dictionary_shape(*new_shape);
    } else if (m_dictionary) {
        new_shape->set_descriptors(copy_descriptors());
    } else {
        new_shape->set_descriptors(descriptors());
    }
    invalidate_prototype_if_needed_for_new_prototype(new_shape);
    if (!is_prototype_shape()) {
        if (!m_prototype_transitions)
            m_prototype_transitions = make<HashMap<GC::Ptr<Object>, GC::Weak<Shape>>>();
        m_prototype_transitions->set(new_prototype, new_shape.ptr());
    }
    return new_shape;
}

Shape::Shape(Realm& realm)
    : m_realm(realm)
{
}

Shape::Shape(Shape& previous_shape, PropertyCountChange property_count_change)
    : m_has_parameter_map(previous_shape.m_has_parameter_map)
    , m_realm(previous_shape.m_realm)
    , m_prototype(previous_shape.m_prototype)
    , m_property_count(previous_shape.m_property_count)
{
    switch (property_count_change) {
    case PropertyCountChange::Preserve:
        break;
    case PropertyCountChange::Increment:
        VERIFY(m_property_count < NumericLimits<u32>::max());
        ++m_property_count;
        break;
    case PropertyCountChange::Decrement:
        VERIFY(m_property_count > 0);
        --m_property_count;
        break;
    }
}

Shape::Shape(Shape& previous_shape, Object* new_prototype)
    : m_has_parameter_map(previous_shape.m_has_parameter_map)
    , m_realm(previous_shape.m_realm)
    , m_prototype(new_prototype)
    , m_property_count(previous_shape.m_property_count)
{
}

void Shape::visit_edges(Cell::Visitor& visitor)
{
    Base::visit_edges(visitor);
    visitor.visit(m_realm);
    if (!m_dictionary)
        visitor.visit(m_property_storage.descriptors);
    visitor.visit(m_prototype);

    visitor.ignore(m_prototype_transitions);

    // Child prototype-shape weak refs need no marking; pruning is lazy.
    visitor.ignore(m_child_prototype_shapes);

    // FIXME: The forward transition keys should be weak, but we have to mark them for now in case they go stale.
    if (m_forward_transitions) {
        for (auto& it : *m_forward_transitions)
            it.key.property_key.visit_edges(visitor);
    }

    // FIXME: The delete transition keys should be weak, but we have to mark them for now in case they go stale.
    if (m_delete_transitions) {
        for (auto& it : *m_delete_transitions)
            it.key.visit_edges(visitor);
    }

    visitor.visit(m_prototype_chain_validity);

    // Descriptor arrays mark their own keys; dictionary tables are not cells, so Shape marks their keys directly.
    if (m_dictionary) {
        for (auto& it : property_table())
            it.key.visit_edges(visitor);
    }
}

Optional<PropertyMetadata> Shape::lookup(PropertyKey const& property_key) const
{
    if (m_property_count == 0)
        return {};
    if (m_dictionary) {
        auto property = property_table().get(property_key);
        if (!property.has_value())
            return {};
        return property;
    }
    if (!descriptors())
        return {};
    return descriptors()->lookup(property_key, m_property_count);
}

void Shape::for_each_property_in_insertion_order(Function<void(PropertyKey const&, PropertyMetadata const&)> const& callback) const
{
    if (m_dictionary) {
        for (auto const& [property_key, metadata] : property_table())
            callback(property_key, metadata);
        return;
    }
    if (!descriptors())
        return;
    descriptors()->for_each_in_insertion_order(callback, m_property_count);
}

void Shape::ensure_descriptor_array()
{
    VERIFY(!m_dictionary);
    if (descriptors())
        return;
    set_descriptors(heap().allocate<DescriptorArray>());
}

GC::Ref<DescriptorArray> Shape::copy_descriptors() const
{
    VERIFY(m_property_count <= DescriptorArray::max_descriptor_count);
    if (!m_dictionary && descriptors())
        return heap().allocate<DescriptorArray>(*descriptors(), m_property_count);

    auto descriptors = heap().allocate<DescriptorArray>();
    for_each_property_in_insertion_order([&](auto const& property_key, auto const& metadata) {
        descriptors->set(property_key, metadata, descriptors->size());
    });
    return descriptors;
}

void Shape::copy_properties_to_dictionary_shape(Shape& shape) const
{
    VERIFY(shape.m_dictionary);
    for_each_property_in_insertion_order([&](auto const& property_key, auto const& metadata) {
        shape.property_table().set(property_key, metadata);
    });
    shape.m_property_count = shape.property_table().size();
}

GC::Ref<Shape> Shape::create_delete_transition(PropertyKey const& property_key)
{
    if (auto existing_shape = get_or_prune_cached_delete_transition(property_key))
        return *existing_shape;
    auto new_shape = heap().allocate<Shape>(*this, PropertyCountChange::Decrement);
    new_shape->set_descriptors(copy_descriptors());
    new_shape->descriptors()->remove(property_key, m_property_count);
    invalidate_prototype_if_needed_for_new_prototype(new_shape);
    if (!m_delete_transitions)
        m_delete_transitions = make<HashMap<PropertyKey, GC::Weak<Shape>>>();
    m_delete_transitions->set(property_key, new_shape.ptr());
    return new_shape;
}

void Shape::add_property_without_transition(PropertyKey const& property_key, PropertyAttributes attributes)
{
    invalidate_prototype_if_needed_for_change_without_transition();
    if (m_dictionary) {
        if (property_table().set(property_key, { m_property_count, attributes }) == AK::HashSetResult::InsertedNewEntry) {
            VERIFY(m_property_count < NumericLimits<u32>::max());
            ++m_property_count;
            ++m_dictionary_generation;
        }
        return;
    }

    ensure_descriptor_array();
    if (!descriptors()->lookup(property_key, m_property_count).has_value()) {
        VERIFY(m_property_count < NumericLimits<u32>::max());
        descriptors()->set(property_key, { m_property_count, attributes }, m_property_count);
        ++m_property_count;
        ++m_dictionary_generation;
        return;
    }
    descriptors()->set(property_key, { m_property_count, attributes }, m_property_count);
}

void Shape::set_property_attributes_without_transition(PropertyKey const& property_key, PropertyAttributes attributes)
{
    invalidate_prototype_if_needed_for_change_without_transition();
    VERIFY(is_dictionary());
    auto it = property_table().find(property_key);
    VERIFY(it != property_table().end());
    it->value.attributes = attributes;
    property_table().set(property_key, it->value);
    ++m_dictionary_generation;
}

void Shape::remove_property_without_transition(PropertyKey const& property_key, u32 offset)
{
    invalidate_prototype_if_needed_for_change_without_transition();
    VERIFY(is_dictionary());
    if (property_table().remove(property_key))
        --m_property_count;
    for (auto& it : property_table()) {
        VERIFY(it.value.offset != offset);
        if (it.value.offset > offset)
            --it.value.offset;
    }
    ++m_dictionary_generation;
}

GC::Ref<Shape> Shape::clone_for_prototype()
{
    VERIFY(!is_prototype_shape());
    VERIFY(!m_prototype_chain_validity);
    auto new_shape = heap().allocate<Shape>(m_realm);
    new_shape->m_has_parameter_map = m_has_parameter_map;
    new_shape->m_prototype = m_prototype;
    if (m_dictionary && m_property_count > DescriptorArray::max_descriptor_count) {
        new_shape->become_dictionary_shape();
        copy_properties_to_dictionary_shape(*new_shape);
    } else {
        new_shape->set_descriptors(copy_descriptors());
        new_shape->m_property_count = m_property_count;
    }
    new_shape->m_prototype_chain_validity = heap().allocate<PrototypeChainValidity>();
    if (new_shape->m_prototype)
        new_shape->m_prototype->shape().add_child_prototype_shape(*new_shape);
    return new_shape;
}

void Shape::set_prototype_without_transition(Object* new_prototype)
{
    VERIFY(new_prototype);
    new_prototype->convert_to_prototype_if_needed();
    m_prototype = new_prototype;
}

void Shape::set_prototype_shape()
{
    VERIFY(!is_prototype_shape());
    m_prototype_chain_validity = heap().allocate<PrototypeChainValidity>();
    if (m_prototype)
        m_prototype->shape().add_child_prototype_shape(*this);
}

void Shape::add_child_prototype_shape(GC::Ref<Shape> child)
{
    VERIFY(is_prototype_shape());
    VERIFY(child->is_prototype_shape());
    if (!m_child_prototype_shapes)
        m_child_prototype_shapes = make<Vector<GC::Weak<Shape>>>();
    m_child_prototype_shapes->append(GC::Weak<Shape> { *child });
}

void Shape::invalidate_prototype_if_needed_for_new_prototype(GC::Ref<Shape> new_prototype_shape)
{
    if (!is_prototype_shape())
        return;
    new_prototype_shape->set_prototype_shape();
    m_prototype_chain_validity->set_valid(false);

    invalidate_all_prototype_chains_leading_to_this();

    // The owning object is keeping the same [[Prototype]], so its existing
    // children descend from new_prototype_shape going forward.
    new_prototype_shape->m_child_prototype_shapes = move(m_child_prototype_shapes);
}

void Shape::invalidate_prototype_if_needed_for_change_without_transition()
{
    if (!is_prototype_shape())
        return;
    m_prototype_chain_validity->set_valid(false);
    m_prototype_chain_validity = heap().allocate<PrototypeChainValidity>();

    invalidate_all_prototype_chains_leading_to_this();
}

void Shape::invalidate_all_prototype_chains_leading_to_this()
{
    if (!m_child_prototype_shapes || m_child_prototype_shapes->is_empty())
        return;

    HashTable<Shape*> shapes_to_invalidate;
    Vector<Shape*> worklist;
    auto enqueue_children_of = [&](Shape& shape) {
        if (!shape.m_child_prototype_shapes)
            return;
        // Prune dead weak refs and enqueue the live ones in one pass.
        shape.m_child_prototype_shapes->remove_all_matching([&](GC::Weak<Shape> const& weak) {
            auto child = weak.ptr();
            if (!child)
                return true;
            if (shapes_to_invalidate.set(child.ptr()) == HashSetResult::InsertedNewEntry)
                worklist.append(child.ptr());
            return false;
        });
    };
    enqueue_children_of(*this);
    while (!worklist.is_empty())
        enqueue_children_of(*worklist.take_last());

    for (auto* shape : shapes_to_invalidate) {
        shape->m_prototype_chain_validity->set_valid(false);
        shape->m_prototype_chain_validity = heap().allocate<PrototypeChainValidity>();
    }
}

}