/*

 * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#include <AK/Enumerate.h>

#include <LibWasm/AbstractMachine/AbstractMachine.h>

#include <LibWasm/AbstractMachine/BytecodeInterpreter.h>

#include <LibWasm/AbstractMachine/Configuration.h>

#include <LibWasm/AbstractMachine/Interpreter.h>

#include <LibWasm/AbstractMachine/Validator.h>

#include <LibWasm/Types.h>

namespace Wasm {

Optional<FunctionAddress> Store::allocate(ModuleInstance& instance, Module const& module, CodeSection::Code const& code, TypeIndex type_index)

{

    FunctionAddress address { m_functions.size() };

    if (type_index.value() >= instance.types().size())

        return {};

    auto& type = instance.types()[type_index.value()];

    m_functions.empend(WasmFunction { type, instance, module, code });

    return address;

}

Optional<FunctionAddress> Store::allocate(HostFunction&& function)

{

    FunctionAddress address { m_functions.size() };

    m_functions.empend(HostFunction { move(function) });

    return address;

}

Optional<TableAddress> Store::allocate(TableType const& type)

{

    TableAddress address { m_tables.size() };

    Vector<Reference> elements;

    elements.resize(type.limits().min());

    m_tables.empend(TableInstance { type, move(elements) });

    return address;

}

Optional<MemoryAddress> Store::allocate(MemoryType const& type)

{

    MemoryAddress address { m_memories.size() };

    auto instance = MemoryInstance::create(type);

    if (instance.is_error())

        return {};

    m_memories.append(instance.release_value());

    return address;

}

Optional<GlobalAddress> Store::allocate(GlobalType const& type, Value value)

{

    GlobalAddress address { m_globals.size() };

    m_globals.append(GlobalInstance { value, type.is_mutable(), type.type() });

    return address;

}

Optional<DataAddress> Store::allocate_data(Vector<u8> initializer)

{

    DataAddress address { m_datas.size() };

    m_datas.append(DataInstance { move(initializer) });

    return address;

}

Optional<ElementAddress> Store::allocate(ValueType const& type, Vector<Reference> references)

{

    ElementAddress address { m_elements.size() };

    m_elements.append(ElementInstance { type, move(references) });

    return address;

}

FunctionInstance* Store::get(FunctionAddress address)

{

    auto value = address.value();

    if (m_functions.size() <= value)

        return nullptr;

    return &m_functions[value];

}

Module const* Store::get_module_for(Wasm::FunctionAddress address)

{

    auto* function = get(address);

    if (!function || function->has<HostFunction>())

        return nullptr;

    return function->get<WasmFunction>().module_ref().ptr();

}

TableInstance* Store::get(TableAddress address)

{

    auto value = address.value();

    if (m_tables.size() <= value)

        return nullptr;

    return &m_tables[value];

}

MemoryInstance* Store::get(MemoryAddress address)

{

    auto value = address.value();

    if (m_memories.size() <= value)

        return nullptr;

    return &m_memories[value];

}

GlobalInstance* Store::get(GlobalAddress address)

{

    auto value = address.value();

    if (m_globals.size() <= value)

        return nullptr;

    return &m_globals[value];

}

ElementInstance* Store::get(ElementAddress address)

{

    auto value = address.value();

    if (m_elements.size() <= value)

        return nullptr;

    return &m_elements[value];

}

DataInstance* Store::get(DataAddress address)

{

    auto value = address.value();

    if (m_datas.size() <= value)

        return nullptr;

    return &m_datas[value];

}

ErrorOr<void, ValidationError> AbstractMachine::validate(Module& module)

{

    if (module.validation_status() != Module::ValidationStatus::Unchecked) {

        if (module.validation_status() == Module::ValidationStatus::Valid)

            return {};

        return ValidationError { module.validation_error() };

    }

    auto result = Validator {}.validate(module);

    if (result.is_error()) {

        module.set_validation_error(result.error().error_string);

        return result.release_error();

    }

    return {};

}

InstantiationResult AbstractMachine::instantiate(Module const& module, Vector<ExternValue> externs)

{

    if (auto result = validate(const_cast<Module&>(module)); result.is_error())

        return InstantiationError { ByteString::formatted("Validation failed: {}", result.error()) };

    auto main_module_instance_pointer = make<ModuleInstance>();

    auto& main_module_instance = *main_module_instance_pointer;

    main_module_instance.types() = module.type_section().types();

    Vector<Value> global_values;

    Vector<Vector<Reference>> elements;

    ModuleInstance auxiliary_instance;

    for (auto [i, import_] : enumerate(module.import_section().imports())) {

        auto extern_ = externs.at(i);

        auto invalid = import_.description().visit(

            [&](MemoryType const& mem_type) -> Optional<ByteString> {

                if (!extern_.has<MemoryAddress>())

                    return "Expected memory import"sv;

                auto other_mem_type = m_store.get(extern_.get<MemoryAddress>())->type();

                if (other_mem_type.limits().is_subset_of(mem_type.limits()))

                    return {};

                return ByteString::formatted("Memory import and extern do not match: {}-{} vs {}-{}", mem_type.limits().min(), mem_type.limits().max(), other_mem_type.limits().min(), other_mem_type.limits().max());

            },

            [&](TableType const& table_type) -> Optional<ByteString> {

                if (!extern_.has<TableAddress>())

                    return "Expected table import"sv;

                auto other_table_type = m_store.get(extern_.get<TableAddress>())->type();

                if (table_type.element_type() == other_table_type.element_type()

                    && other_table_type.limits().is_subset_of(table_type.limits()))

                    return {};

                return ByteString::formatted("Table import and extern do not match: {}-{} vs {}-{}", table_type.limits().min(), table_type.limits().max(), other_table_type.limits().min(), other_table_type.limits().max());

            },

            [&](GlobalType const& global_type) -> Optional<ByteString> {

                if (!extern_.has<GlobalAddress>())

                    return "Expected global import"sv;

                auto other_global_type = m_store.get(extern_.get<GlobalAddress>())->type();

                if (global_type.type() == other_global_type.type()

                    && global_type.is_mutable() == other_global_type.is_mutable())

                    return {};

                return "Global import and extern do not match"sv;

            },

            [&](FunctionType const& type) -> Optional<ByteString> {

                if (!extern_.has<FunctionAddress>())

                    return "Expected function import"sv;

                auto other_type = m_store.get(extern_.get<FunctionAddress>())->visit([&](WasmFunction const& wasm_func) { return wasm_func.type(); }, [&](HostFunction const& host_func) { return host_func.type(); });

                if (type.results() != other_type.results())

                    return ByteString::formatted("Function import and extern do not match, results: {} vs {}", type.results(), other_type.results());

                if (type.parameters() != other_type.parameters())

                    return ByteString::formatted("Function import and extern do not match, parameters: {} vs {}", type.parameters(), other_type.parameters());

                return {};

            },

            [&](TypeIndex type_index) -> Optional<ByteString> {

                if (!extern_.has<FunctionAddress>())

                    return "Expected function import"sv;

                auto other_type = m_store.get(extern_.get<FunctionAddress>())->visit([&](WasmFunction const& wasm_func) { return wasm_func.type(); }, [&](HostFunction const& host_func) { return host_func.type(); });

                auto& type = module.type_section().types()[type_index.value()];

                if (type.results() != other_type.results())

                    return ByteString::formatted("Function import and extern do not match, results: {} vs {}", type.results(), other_type.results());

                if (type.parameters() != other_type.parameters())

                    return ByteString::formatted("Function import and extern do not match, parameters: {} vs {}", type.parameters(), other_type.parameters());

                return {};

            });

        if (invalid.has_value())

            return InstantiationError { ByteString::formatted("{}::{}: {}", import_.module(), import_.name(), invalid.release_value()) };

    }

    for (auto& entry : externs) {

        if (auto* ptr = entry.get_pointer<GlobalAddress>())

            auxiliary_instance.globals().append(*ptr);

        else if (auto* ptr = entry.get_pointer<FunctionAddress>())

            auxiliary_instance.functions().append(*ptr);

    }

    Vector<FunctionAddress> module_functions;

    module_functions.ensure_capacity(module.function_section().types().size());

    size_t i = 0;

    for (auto& code : module.code_section().functions()) {

        auto type_index = module.function_section().types()[i];

        auto address = m_store.allocate(main_module_instance, module, code, type_index);

        VERIFY(address.has_value());

        auxiliary_instance.functions().append(*address);

        module_functions.append(*address);

        ++i;

    }

    BytecodeInterpreter interpreter(m_stack_info);

    for (auto& entry : module.global_section().entries()) {

        Configuration config { m_store };

        if (m_should_limit_instruction_count)

            config.enable_instruction_count_limit();

        config.set_frame(Frame {

            auxiliary_instance,

            Vector<Value> {},

            entry.expression(),

            1,

        });

        auto result = config.execute(interpreter).assert_wasm_result();

        if (result.is_trap())

            return InstantiationError { ByteString::formatted("Global value construction trapped: {}", result.trap().reason) };

        global_values.append(result.values().first());

    }

    if (auto result = allocate_all_initial_phase(module, main_module_instance, externs, global_values, module_functions); result.has_value())

        return result.release_value();

    for (auto& segment : module.element_section().segments()) {

        Vector<Reference> references;

        for (auto& entry : segment.init) {

            Configuration config { m_store };

            if (m_should_limit_instruction_count)

                config.enable_instruction_count_limit();

            config.set_frame(Frame {

                auxiliary_instance,

                Vector<Value> {},

                entry,

                entry.instructions().size(),

            });

            auto result = config.execute(interpreter).assert_wasm_result();

            if (result.is_trap())

                return InstantiationError { ByteString::formatted("Element construction trapped: {}", result.trap().reason) };

            for (auto& value : result.values()) {

                auto reference = value.to<Reference>();

                references.append(reference);

            }

        }

        elements.append(move(references));

    }

    if (auto result = allocate_all_final_phase(module, main_module_instance, elements); result.has_value())

        return result.release_value();

    size_t index = 0;

    for (auto& segment : module.element_section().segments()) {

        auto current_index = index;

        ++index;

        auto active_ptr = segment.mode.get_pointer<ElementSection::Active>();

        auto elem_instance = m_store.get(main_module_instance.elements()[current_index]);

        if (!active_ptr) {

            if (segment.mode.has<ElementSection::Declarative>())

                *elem_instance = ElementInstance(elem_instance->type(), {});

            continue;

        }

        Configuration config { m_store };

        if (m_should_limit_instruction_count)

            config.enable_instruction_count_limit();

        config.set_frame(Frame {

            auxiliary_instance,

            Vector<Value> {},

            active_ptr->expression,

            1,

        });

        auto result = config.execute(interpreter).assert_wasm_result();

        if (result.is_trap())

            return InstantiationError { ByteString::formatted("Element section initialisation trapped: {}", result.trap().reason) };

        auto d = result.values().first().to<i32>();

        auto table_instance = m_store.get(main_module_instance.tables()[active_ptr->index.value()]);

        if (current_index >= main_module_instance.elements().size())

            return InstantiationError { "Invalid element referenced by active element segment" };

        if (!table_instance || !elem_instance)

            return InstantiationError { "Invalid element referenced by active element segment" };

        Checked<size_t> total_size = elem_instance->references().size();

        total_size.saturating_add(d);

        if (total_size.value() > table_instance->elements().size())

            return InstantiationError { "Table instantiation out of bounds" };

        size_t i = 0;

        for (auto it = elem_instance->references().begin(); it < elem_instance->references().end(); ++i, ++it)

            table_instance->elements()[i + d] = *it;

        // Drop element

        *m_store.get(main_module_instance.elements()[current_index]) = ElementInstance(elem_instance->type(), {});

    }

    for (auto& segment : module.data_section().data()) {

        Optional<InstantiationError> result = segment.value().visit(

            [&](DataSection::Data::Active const& data) -> Optional<InstantiationError> {

                Configuration config { m_store };

                if (m_should_limit_instruction_count)

                    config.enable_instruction_count_limit();

                config.set_frame(Frame {

                    auxiliary_instance,

                    Vector<Value> {},

                    data.offset,

                    1,

                });

                auto result = config.execute(interpreter).assert_wasm_result();

                if (result.is_trap())

                    return InstantiationError { ByteString::formatted("Data section initialisation trapped: {}", result.trap().reason) };

                size_t offset = result.values().first().to<u64>();

                if (main_module_instance.memories().size() <= data.index.value()) {

                    return InstantiationError {

                        ByteString::formatted("Data segment referenced out-of-bounds memory ({}) of max {} entries",

                            data.index.value(), main_module_instance.memories().size())

                    };

                }

                auto maybe_data_address = m_store.allocate_data(data.init);

                if (!maybe_data_address.has_value()) {

                    return InstantiationError { "Failed to allocate a data instance for an active data segment"sv };

                }

                main_module_instance.datas().append(*maybe_data_address);

                auto address = main_module_instance.memories()[data.index.value()];

                auto instance = m_store.get(address);

                Checked<size_t> checked_offset = data.init.size();

                checked_offset += offset;

                if (checked_offset.has_overflow() || checked_offset > instance->size()) {

                    return InstantiationError {

                        ByteString::formatted("Data segment attempted to write to out-of-bounds memory ({}) in memory of size {}",

                            offset, instance->size())

                    };

                }

                if (!data.init.is_empty())

                    instance->data().overwrite(offset, data.init.data(), data.init.size());

                return {};

            },

            [&](DataSection::Data::Passive const& passive) -> Optional<InstantiationError> {

                auto maybe_data_address = m_store.allocate_data(passive.init);

                if (!maybe_data_address.has_value()) {

                    return InstantiationError { "Failed to allocate a data instance for a passive data segment"sv };

                }

                main_module_instance.datas().append(*maybe_data_address);

                return {};

            });

        if (result.has_value())

            return result.release_value();

    }

    if (module.start_section().function().has_value()) {

        auto& functions = main_module_instance.functions();

        auto index = module.start_section().function()->index();

        if (functions.size() <= index.value()) {

            return InstantiationError { ByteString::formatted("Start section function referenced invalid index {} of max {} entries", index.value(), functions.size()) };

        }

        auto result = invoke(functions[index.value()], {});

        if (result.is_trap())

            return InstantiationError { ByteString::formatted("Start function trapped: {}", result.trap().reason) };

    }

    return InstantiationResult { move(main_module_instance_pointer) };

}

Optional<InstantiationError> AbstractMachine::allocate_all_initial_phase(Module const& module, ModuleInstance& module_instance, Vector<ExternValue>& externs, Vector<Value>& global_values, Vector<FunctionAddress>& own_functions)

{

    Optional<InstantiationError> result;

    for (auto& entry : externs) {

        entry.visit(

            [&](FunctionAddress const& address) { module_instance.functions().append(address); },

            [&](TableAddress const& address) { module_instance.tables().append(address); },

            [&](MemoryAddress const& address) { module_instance.memories().append(address); },

            [&](GlobalAddress const& address) { module_instance.globals().append(address); });

    }

    module_instance.functions().extend(own_functions);

    // FIXME: What if this fails?

    for (auto& table : module.table_section().tables()) {

        auto table_address = m_store.allocate(table.type());

        VERIFY(table_address.has_value());

        module_instance.tables().append(*table_address);

    }

    for (auto& memory : module.memory_section().memories()) {

        auto memory_address = m_store.allocate(memory.type());

        VERIFY(memory_address.has_value());

        module_instance.memories().append(*memory_address);

    }

    size_t index = 0;

    for (auto& entry : module.global_section().entries()) {

        auto address = m_store.allocate(entry.type(), move(global_values[index]));

        VERIFY(address.has_value());

        module_instance.globals().append(*address);

        index++;

    }

    for (auto& entry : module.export_section().entries()) {

        Variant<FunctionAddress, TableAddress, MemoryAddress, GlobalAddress, Empty> address {};

        entry.description().visit(

            [&](FunctionIndex const& index) {

                if (module_instance.functions().size() > index.value())

                    address = FunctionAddress { module_instance.functions()[index.value()] };

                else

                    dbgln("Failed to export '{}', the exported address ({}) was out of bounds (min: 0, max: {})", entry.name(), index.value(), module_instance.functions().size());

            },

            [&](TableIndex const& index) {

                if (module_instance.tables().size() > index.value())

                    address = TableAddress { module_instance.tables()[index.value()] };

                else

                    dbgln("Failed to export '{}', the exported address ({}) was out of bounds (min: 0, max: {})", entry.name(), index.value(), module_instance.tables().size());

            },

            [&](MemoryIndex const& index) {

                if (module_instance.memories().size() > index.value())

                    address = MemoryAddress { module_instance.memories()[index.value()] };

                else

                    dbgln("Failed to export '{}', the exported address ({}) was out of bounds (min: 0, max: {})", entry.name(), index.value(), module_instance.memories().size());

            },

            [&](GlobalIndex const& index) {

                if (module_instance.globals().size() > index.value())

                    address = GlobalAddress { module_instance.globals()[index.value()] };

                else

                    dbgln("Failed to export '{}', the exported address ({}) was out of bounds (min: 0, max: {})", entry.name(), index.value(), module_instance.globals().size());

            });

        if (address.has<Empty>()) {

            result = InstantiationError { "An export could not be resolved" };

            continue;

        }

        module_instance.exports().append(ExportInstance {

            entry.name(),

            move(address).downcast<FunctionAddress, TableAddress, MemoryAddress, GlobalAddress>(),

        });

    }

    return result;

}

Optional<InstantiationError> AbstractMachine::allocate_all_final_phase(Module const& module, ModuleInstance& module_instance, Vector<Vector<Reference>>& elements)

{

    size_t index = 0;

    for (auto& segment : module.element_section().segments()) {

        auto address = m_store.allocate(segment.type, move(elements[index]));

        VERIFY(address.has_value());

        module_instance.elements().append(*address);

        index++;

    }

    return {};

}

Result AbstractMachine::invoke(FunctionAddress address, Vector<Value> arguments)

{

    BytecodeInterpreter interpreter(m_stack_info);

    return invoke(interpreter, address, move(arguments));

}

Result AbstractMachine::invoke(Interpreter& interpreter, FunctionAddress address, Vector<Value> arguments)

{

    Configuration configuration { m_store };

    if (m_should_limit_instruction_count)

        configuration.enable_instruction_count_limit();

    return configuration.call(interpreter, address, move(arguments));

}

void Linker::link(ModuleInstance const& instance)

{

    populate();

    if (m_unresolved_imports.is_empty())

        return;

    HashTable<Name> resolved_imports;

    for (auto& import_ : m_unresolved_imports) {

        auto it = instance.exports().find_if([&](auto& export_) { return export_.name() == import_.name; });

        if (!it.is_end()) {

            resolved_imports.set(import_);

            m_resolved_imports.set(import_, it->value());

        }

    }

    for (auto& entry : resolved_imports)

        m_unresolved_imports.remove(entry);

}

void Linker::link(HashMap<Linker::Name, ExternValue> const& exports)

{

    populate();

    if (m_unresolved_imports.is_empty())

        return;

    if (exports.is_empty())

        return;

    HashTable<Name> resolved_imports;

    for (auto& import_ : m_unresolved_imports) {

        auto export_ = exports.get(import_);

        if (export_.has_value()) {

            resolved_imports.set(import_);

            m_resolved_imports.set(import_, export_.value());

        }

    }

    for (auto& entry : resolved_imports)

        m_unresolved_imports.remove(entry);

}

AK::ErrorOr<Vector<ExternValue>, LinkError> Linker::finish()

{

    populate();

    if (!m_unresolved_imports.is_empty()) {

        if (!m_error.has_value())

            m_error = LinkError {};

        for (auto& entry : m_unresolved_imports)

            m_error->missing_imports.append(entry.name);

        return *m_error;

    }

    if (m_error.has_value())

        return *m_error;

    // Result must be in the same order as the module imports

    Vector<ExternValue> exports;

    exports.ensure_capacity(m_ordered_imports.size());

    for (auto& import_ : m_ordered_imports)

        exports.unchecked_append(*m_resolved_imports.get(import_));

    return exports;

}

void Linker::populate()

{

    if (!m_ordered_imports.is_empty())

        return;

    for (auto& import_ : m_module.import_section().imports()) {

        m_ordered_imports.append({ import_.module(), import_.name(), import_.description() });

        m_unresolved_imports.set(m_ordered_imports.last());

    }

}

}