/*
 * Copyright (c) 2021, Ali Mohammad Pur <mpfard@serenityos.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <LibJS/Runtime/GeneratorObject.h>
#include <LibJS/Runtime/GeneratorPrototype.h>
#include <LibJS/Runtime/GlobalObject.h>
#include <LibJS/Runtime/Iterator.h>
#include <LibJS/Runtime/NativeJavaScriptBackedFunction.h>
#include <LibJS/Runtime/VM.h>

namespace JS {

GC_DEFINE_ALLOCATOR(GeneratorObject);

GC::Ref<GeneratorObject> GeneratorObject::create(Realm& realm, Variant<GC::Ref<ECMAScriptFunctionObject>, GC::Ref<NativeJavaScriptBackedFunction>> generating_function, NonnullOwnPtr<ExecutionContext> execution_context)
{
    auto& vm = realm.vm();
    // This is "g1.prototype" in figure-2 (https://tc39.es/ecma262/img/figure-2.png)
    Value generating_function_prototype;

    auto kind = generating_function.visit(
        [](auto function) {
            return function->kind();
        });

    if (kind == FunctionKind::Async) {
        // We implement async functions by transforming them to generator function in the bytecode
        // interpreter. However an async function does not have a prototype and should not be
        // changed thus we hardcode the prototype.
        generating_function_prototype = realm.intrinsics().generator_prototype();
    } else {
        static Bytecode::StaticPropertyLookupCache cache;
        generating_function_prototype = MUST(generating_function.visit([&vm](auto function) {
            static Bytecode::StaticPropertyLookupCache cache;
            return function->get(vm.names.prototype, cache);
        }));
    }

    GC::Ptr<Object> generating_function_prototype_object = nullptr;
    if (!generating_function_prototype.is_nullish())
        generating_function_prototype_object = MUST(generating_function_prototype.to_object(vm));

    auto generating_executable = generating_function.visit(
        [](GC::Ref<ECMAScriptFunctionObject> function) -> GC::Ref<Bytecode::Executable> {
            return function->bytecode_executable().as_nonnull();
        },
        [](GC::Ref<NativeJavaScriptBackedFunction> function) -> GC::Ref<Bytecode::Executable> {
            return function->bytecode_executable();
        });

    auto object = realm.create<GeneratorObject>(realm, generating_function_prototype_object, move(execution_context));
    object->m_generating_executable = generating_executable;
    object->m_yield_continuation = object->m_execution_context->yield_continuation;
    return object;
}

GeneratorObject::GeneratorObject(Realm& realm, Object* prototype, NonnullOwnPtr<ExecutionContext> context, Optional<StringView> generator_brand)
    : Object(realm, prototype)
    , m_execution_context(move(context))
    , m_generator_brand(move(generator_brand))
{
}

void GeneratorObject::visit_edges(Cell::Visitor& visitor)
{
    Base::visit_edges(visitor);
    visitor.visit(m_generating_executable);
    m_execution_context->visit_edges(visitor);
    visitor.visit(m_pending_completion_value);
}

// 27.5.3.2 GeneratorValidate ( generator, generatorBrand ), https://tc39.es/ecma262/#sec-generatorvalidate
ThrowCompletionOr<GeneratorObject::GeneratorState> GeneratorObject::validate(VM& vm, Optional<StringView> const& generator_brand)
{
    // 1. Perform ? RequireInternalSlot(generator, [[GeneratorState]]).
    // 2. Perform ? RequireInternalSlot(generator, [[GeneratorBrand]]).
    // NOTE: Already done by the caller of resume or resume_abrupt, as they wouldn't have a GeneratorObject otherwise.

    // 3. If generator.[[GeneratorBrand]] is not the same value as generatorBrand, throw a TypeError exception.
    if (m_generator_brand != generator_brand)
        return vm.throw_completion<TypeError>(ErrorType::GeneratorBrandMismatch, m_generator_brand.value_or("<empty>"sv), generator_brand.value_or("<empty>"sv));

    // 4. Assert: generator also has a [[GeneratorContext]] internal slot.
    // NOTE: Done by already being a GeneratorObject.

    // 5. Let state be generator.[[GeneratorState]].
    auto state = m_generator_state;

    // 6. If state is executing, throw a TypeError exception.
    if (state == GeneratorState::Executing)
        return vm.throw_completion<TypeError>(ErrorType::GeneratorAlreadyExecuting);

    // 7. Return state.
    return state;
}

ThrowCompletionOr<GeneratorObject::IterationResult> GeneratorObject::execute(VM& vm, Completion const& completion)
{
    // Loosely based on step 4 of https://tc39.es/ecma262/#sec-generatorstart mixed with https://tc39.es/ecma262/#sec-generatoryield at the end.

    set_pending_completion(completion);

    // We should never enter `execute` again after the generator is complete.
    VERIFY(m_yield_continuation != ExecutionContext::no_yield_continuation);

    // Clear yield state so that a normal return (no yield) is detected as done.
    m_execution_context->yield_continuation = ExecutionContext::no_yield_continuation;

    auto result_value = vm.run_executable(vm.running_execution_context(), *m_generating_executable, m_yield_continuation, Value(this));
    clear_pending_completion();

    vm.pop_execution_context();

    if (result_value.is_throw_completion()) {
        // Uncaught exceptions disable the generator.
        m_generator_state = GeneratorState::Completed;
        return result_value.throw_completion();
    }

    auto value = result_value.release_value();
    if (value.is_special_empty_value())
        value = js_undefined();

    m_yield_continuation = m_execution_context->yield_continuation;
    bool done = m_yield_continuation == ExecutionContext::no_yield_continuation;

    m_generator_state = done ? GeneratorState::Completed : GeneratorState::SuspendedYield;

    return IterationResult(value, done);
}

// 27.5.3.3 GeneratorResume ( generator, value, generatorBrand ), https://tc39.es/ecma262/#sec-generatorresume
ThrowCompletionOr<GeneratorObject::IterationResult> GeneratorObject::resume(VM& vm, Value value, Optional<StringView> const& generator_brand)
{
    // 1. Let state be ? GeneratorValidate(generator, generatorBrand).
    auto state = TRY(validate(vm, generator_brand));

    // 2. If state is completed, return CreateIterResultObject(undefined, true).
    if (state == GeneratorState::Completed) {
        return IterationResult(js_undefined(), true);
    }

    // 3. Assert: state is either suspendedStart or suspendedYield.
    VERIFY(state == GeneratorState::SuspendedStart || state == GeneratorState::SuspendedYield);

    // 4. Let genContext be generator.[[GeneratorContext]].
    auto& generator_context = m_execution_context;

    // 5. Let methodContext be the running execution context.
    auto const& method_context = vm.running_execution_context();

    // 6. Suspend methodContext.
    // 8. Push genContext onto the execution context stack; genContext is now the running execution context.
    // NOTE: This is done out of order as to not permanently disable the generator if push_execution_context throws,
    //       as `resume` will immediately throw when [[GeneratorState]] is "executing", never allowing the state to change.
    TRY(vm.push_execution_context(*generator_context, {}));

    // 7. Set generator.[[GeneratorState]] to executing.
    m_generator_state = GeneratorState::Executing;

    // 9. Resume the suspended evaluation of genContext using NormalCompletion(value) as the result of the operation that suspended it. Let result be the value returned by the resumed computation.
    auto result = execute(vm, normal_completion(value));

    // 10. Assert: When we return here, genContext has already been removed from the execution context stack and methodContext is the currently running execution context.
    VERIFY(&vm.running_execution_context() == &method_context);

    // 11. Return ? result.
    return result;
}

// 27.5.3.4 GeneratorResumeAbrupt ( generator, abruptCompletion, generatorBrand ), https://tc39.es/ecma262/#sec-generatorresumeabrupt
ThrowCompletionOr<GeneratorObject::IterationResult> GeneratorObject::resume_abrupt(JS::VM& vm, JS::Completion abrupt_completion, Optional<StringView> const& generator_brand)
{
    // 1. Let state be ? GeneratorValidate(generator, generatorBrand).
    auto state = TRY(validate(vm, generator_brand));

    // 2. If state is suspendedStart, then
    if (state == GeneratorState::SuspendedStart) {
        // a. Set generator.[[GeneratorState]] to completed.
        m_generator_state = GeneratorState::Completed;

        // b. Once a generator enters the completed state it never leaves it and its associated execution context is never resumed. Any execution state associated with generator can be discarded at this point.
        // We don't currently discard anything.

        // c. Set state to completed.
        state = GeneratorState::Completed;
    }

    // 3. If state is completed, then
    if (state == GeneratorState::Completed) {
        // a. If abruptCompletion.[[Type]] is return, then
        if (abrupt_completion.type() == Completion::Type::Return) {
            // i. Return CreateIterResultObject(abruptCompletion.[[Value]], true).
            return IterationResult(abrupt_completion.value(), true);
        }

        // b. Return ? abruptCompletion.
        return abrupt_completion;
    }

    // 4. Assert: state is suspendedYield.
    VERIFY(state == GeneratorState::SuspendedYield);

    // 5. Let genContext be generator.[[GeneratorContext]].
    auto& generator_context = m_execution_context;

    // 6. Let methodContext be the running execution context.
    auto const& method_context = vm.running_execution_context();

    // 7. Suspend methodContext.
    // 9. Push genContext onto the execution context stack; genContext is now the running execution context.
    // NOTE: This is done out of order as to not permanently disable the generator if push_execution_context throws,
    //       as `resume_abrupt` will immediately throw when [[GeneratorState]] is "executing", never allowing the state to change.
    TRY(vm.push_execution_context(*generator_context, {}));

    // 8. Set generator.[[GeneratorState]] to executing.
    m_generator_state = GeneratorState::Executing;

    // 10. Resume the suspended evaluation of genContext using abruptCompletion as the result of the operation that suspended it. Let result be the Completion Record returned by the resumed computation.
    auto result = execute(vm, abrupt_completion);

    // 11. Assert: When we return here, genContext has already been removed from the execution context stack and methodContext is the currently running execution context.
    VERIFY(&vm.running_execution_context() == &method_context);

    // 12. Return ? result.
    return result;
}

}