// Copyright 2024 The Dawn & Tint Authors
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "src/dawn/native/SharedResourceMemory.h"

#include <algorithm>
#include <span>
#include <utility>

#include "src/dawn/common/Enumerator.h"
#include "src/dawn/native/Buffer.h"
#include "src/dawn/native/ChainUtils.h"
#include "src/dawn/native/Device.h"
#include "src/dawn/native/Queue.h"
#include "src/dawn/native/Texture.h"
#include "src/utils/compiler.h"
#include "src/utils/heap_array.h"

namespace dawn::native {

SharedResource::~SharedResource() = default;

SharedResourceMemoryContents* SharedResource::GetSharedResourceMemoryContents() const {
    return mSharedResourceMemoryContents.Get();
}

SharedResourceMemory::SharedResourceMemory(DeviceBase* device,
                                           ObjectBase::ErrorTag tag,
                                           StringView label)
    : ApiObjectBase(device, tag, label),
      mContents(new SharedResourceMemoryContents(GetWeakRef(this))) {}

SharedResourceMemory::~SharedResourceMemory() = default;

void SharedResourceMemory::DestroyImpl(DestroyReason reason) {}

bool SharedResourceMemoryContents::HasWriteAccess() const {
    return mSharedResourceAccessState == SharedResourceAccessState::Write;
}

bool SharedResourceMemoryContents::HasExclusiveReadAccess() const {
    return mSharedResourceAccessState == SharedResourceAccessState::ExclusiveRead;
}

int SharedResourceMemoryContents::GetReadAccessCount() const {
    return mReadAccessCount;
}

bool SharedResourceMemoryContents::HasAccess() const {
    return mSharedResourceAccessState != SharedResourceAccessState::NotAccessed;
}

void SharedResourceMemory::Initialize() {
    DAWN_CHECK(!IsError());
    mContents = CreateContents();
}

Ref<SharedResourceMemoryContents> SharedResourceMemory::CreateContents() {
    return AcquireRef(new SharedResourceMemoryContents(GetWeakRef(this)));
}

SharedResourceMemoryContents* SharedResourceMemory::GetContents() const {
    return mContents.Get();
}

MaybeValError SharedResourceMemory::ValidateResourceCreatedFromSelf(SharedResource* resource) {
    auto* contents = resource->GetSharedResourceMemoryContents();
    DAWN_INVALID_IF(contents == nullptr, "%s was not created from %s.", resource, this);

    auto* sharedResourceMemory =
        resource->GetSharedResourceMemoryContents()->GetSharedResourceMemory().Promote().Get();
    DAWN_INVALID_IF(sharedResourceMemory != this, "%s created from %s cannot be used with %s.",
                    resource, sharedResourceMemory, this);
    return {};
}

wgpu::Status SharedResourceMemory::APIBeginAccess(
    TextureBase* texture,
    const SharedTextureMemoryBeginAccessDescriptor* descriptor) {
    return GetDevice()->ConsumedError(BeginAccess(texture, descriptor),
                                      "calling %s.BeginAccess(%s).", this, texture)
               ? wgpu::Status::Error
               : wgpu::Status::Success;
}

wgpu::Status SharedResourceMemory::APIBeginAccess(
    BufferBase* buffer,
    const SharedBufferMemoryBeginAccessDescriptor* descriptor) {
    return GetDevice()->ConsumedError(BeginAccess(buffer, descriptor),
                                      "calling %s.BeginAccess(%s).", this, buffer)
               ? wgpu::Status::Error
               : wgpu::Status::Success;
}

template <typename Resource, typename BeginAccessDescriptor>
MaybeUnknownError SharedResourceMemory::BeginAccess(Resource* resource,
                                                    const BeginAccessDescriptor* rawDescriptor) {
    DAWN_TRY(GetDevice()->ValidateIsAlive());
    DAWN_TRY(GetDevice()->ValidateObject(resource));

    UnpackedPtr<BeginAccessDescriptor> descriptor;
    DAWN_TRY_ASSIGN(descriptor, ValidateAndUnpack(rawDescriptor));

    // TODO(https://crbug.com/515272361): Introduce a wgpu::SharedFenceAndSignalValue struct instead
    // of having two arrays that need to have the same size.
    DAWN_INVALID_IF(descriptor->fences.size() != descriptor->signaledValues.size(),
                    "fenceCount (%i) doesn't match signaledValueCount (%i).",
                    descriptor->fences.size(), descriptor->signaledValues.size());

    for (SharedFenceBase* fence : descriptor->fences) {
        DAWN_TRY(GetDevice()->ValidateObject(fence));
    }

    DAWN_TRY(ValidateResourceCreatedFromSelf(resource));

    DAWN_INVALID_IF(resource->IsDestroyed(), "%s has been destroyed.", resource);
    DAWN_INVALID_IF(resource->HasAccess(), "%s is already used to access %s.", resource, this);

    if constexpr (std::is_same_v<Resource, TextureBase>) {
        DAWN_INVALID_IF(resource->GetFormat().IsMultiPlanar() && !descriptor->initialized,
                        "%s with multiplanar format (%s) must be initialized.", resource,
                        resource->GetFormat().format);

        DAWN_INVALID_IF(mContents->HasWriteAccess(), "%s is currently accessed for writing.", this);
        DAWN_INVALID_IF(mContents->HasExclusiveReadAccess(),
                        "%s is currently accessed for exclusive reading.", this);

        if (static_cast<TextureBase*>(resource)->IsReadOnly()) {
            if (descriptor->concurrentRead) {
                DAWN_CHECK(!mExclusiveAccess);
                DAWN_INVALID_IF(!descriptor->initialized, "Concurrent reading an uninitialized %s.",
                                resource);
                ++mContents->mReadAccessCount;
                mContents->mSharedResourceAccessState = SharedResourceAccessState::SimultaneousRead;

            } else {
                DAWN_INVALID_IF(
                    mContents->mReadAccessCount != 0,
                    "Exclusive read access used while %s is currently accessed for reading.", this);
                mContents->mSharedResourceAccessState = SharedResourceAccessState::ExclusiveRead;
                mExclusiveAccess = resource;
            }
        } else {
            DAWN_INVALID_IF(descriptor->concurrentRead, "Concurrent reading read-write %s.",
                            resource);
            DAWN_INVALID_IF(mContents->mReadAccessCount != 0,
                            "Read-Write access used while %s is currently accessed for reading.",
                            this);
            mContents->mSharedResourceAccessState = SharedResourceAccessState::Write;
            mExclusiveAccess = resource;
        }
    } else if constexpr (std::is_same_v<Resource, BufferBase>) {
        DAWN_INVALID_IF(mExclusiveAccess != nullptr,
                        "Cannot begin access with %s on %s which is currently accessed by %s.",
                        resource, this, mExclusiveAccess.Get());
        mContents->mSharedResourceAccessState = SharedResourceAccessState::Write;
        mExclusiveAccess = resource;
    }

    DAWN_TRY(BeginAccessImpl(resource, descriptor));

    for (auto [i, fence] : Enumerate(descriptor->fences)) {
        // Add the fences to mPendingFences if they are not already contained in the list.
        // This loop is O(n*m), but there shouldn't be very many fences.
        auto it = std::find_if(
            mContents->mPendingFences.begin(), mContents->mPendingFences.end(),
            [&](const auto& pendingFence) { return pendingFence.object.Get() == fence; });
        if (it != mContents->mPendingFences.end()) {
            it->signaledValue = std::max(it->signaledValue, descriptor->signaledValues[i]);
            continue;
        }
        mContents->mPendingFences.push_back({fence, descriptor->signaledValues[i]});
    }

    DAWN_CHECK(!resource->IsError());
    resource->OnBeginAccess();
    if constexpr (std::is_same_v<Resource, BufferBase>) {
        if (resource->GetState() == BufferBase::BufferState::MappedAtCreation) {
            // MapAtCreation is where zero-initialization of mappedAtCreation buffers is done, but
            // it is never called again after the initial buffer creation, so zero-initialize
            // manually here.
            if (!descriptor->initialized) {
                resource->GetFullMappedAllocatedRange().FillBytes(std::byte{0});
            }
            resource->SetInitialized(true);
        } else {
            resource->SetInitialized(descriptor->initialized);
        }
    } else {
        resource->SetInitialized(descriptor->initialized);
    }
    return {};
}

wgpu::Status SharedResourceMemory::APIEndAccess(TextureBase* texture,
                                                SharedTextureMemoryEndAccessState* state) {
    return GetDevice()->ConsumedError(EndAccess(texture, state), "calling %s.EndAccess(%s).", this,
                                      texture)
               ? wgpu::Status::Error
               : wgpu::Status::Success;
}

wgpu::Status SharedResourceMemory::APIEndAccess(BufferBase* buffer,
                                                SharedBufferMemoryEndAccessState* state) {
    return GetDevice()->ConsumedError(EndAccess(buffer, state), "calling %s.EndAccess(%s).", this,
                                      buffer)
               ? wgpu::Status::Error
               : wgpu::Status::Success;
}

MaybeValError SharedResourceMemory::BeginAccessImpl(
    TextureBase* texture,
    const UnpackedPtr<SharedTextureMemoryBeginAccessDescriptor>& descriptor) {
    DAWN_UNREACHABLE();
}

MaybeValError SharedResourceMemory::BeginAccessImpl(
    BufferBase* buffer,
    const UnpackedPtr<SharedBufferMemoryBeginAccessDescriptor>& descriptor) {
    DAWN_UNREACHABLE();
}

ResultOrUnknownError<FenceAndSignalValue> SharedResourceMemory::EndAccessImpl(
    TextureBase* texture,
    ExecutionSerial lastUsageSerial,
    UnpackedPtr<SharedTextureMemoryEndAccessState>& state) {
    DAWN_UNREACHABLE();
}

ResultOrUnknownError<FenceAndSignalValue> SharedResourceMemory::EndAccessImpl(
    BufferBase* buffer,
    ExecutionSerial lastUsageSerial,
    UnpackedPtr<SharedBufferMemoryEndAccessState>& state) {
    DAWN_UNREACHABLE();
}

bool SharedResourceMemory::APIIsDeviceLost() const {
    return GetDevice()->IsLost();
}

template <typename Resource, typename EndAccessState>
MaybeUnknownError SharedResourceMemory::EndAccess(Resource* resource, EndAccessState* state) {
    DAWN_TRY(GetDevice()->ValidateObject(resource));
    DAWN_TRY(ValidateResourceCreatedFromSelf(resource));

    DAWN_INVALID_IF(!resource->HasAccess(), "%s is not currently being accessed.", resource);
    if constexpr (std::is_same_v<Resource, TextureBase>) {
        if (static_cast<TextureBase*>(resource)->IsReadOnly()) {
            DAWN_CHECK(!mContents->HasWriteAccess());
            if (mContents->HasExclusiveReadAccess()) {
                DAWN_CHECK(mContents->mReadAccessCount == 0);
                mContents->mSharedResourceAccessState = SharedResourceAccessState::NotAccessed;
                mExclusiveAccess = nullptr;
            } else {
                DAWN_CHECK(mContents->mSharedResourceAccessState ==
                           SharedResourceAccessState::SimultaneousRead);
                DAWN_CHECK(mExclusiveAccess == nullptr);
                --mContents->mReadAccessCount;
                if (mContents->mReadAccessCount == 0) {
                    mContents->mSharedResourceAccessState = SharedResourceAccessState::NotAccessed;
                }
            }
        } else {
            DAWN_CHECK(mContents->mSharedResourceAccessState == SharedResourceAccessState::Write);
            DAWN_CHECK(mContents->mReadAccessCount == 0);
            mContents->mSharedResourceAccessState = SharedResourceAccessState::NotAccessed;
            mExclusiveAccess = nullptr;
        }
    } else if constexpr (std::is_same_v<Resource, BufferBase>) {
        DAWN_INVALID_IF(mExclusiveAccess != resource,
                        "Cannot end access with %s on %s which is currently accessed by %s.",
                        resource, this, mExclusiveAccess.Get());
        auto* buffer = static_cast<BufferBase*>(resource);
        if (buffer->APIGetMapState() != wgpu::BufferMapState::Unmapped) {
            // Unmap the buffer on the caller's behalf instead of failing so that EndAccess always
            // succeeds, regardless of the buffer's map state.
            DAWN_TRY(buffer->Unmap());
        }
        mContents->mSharedResourceAccessState = SharedResourceAccessState::NotAccessed;
        mExclusiveAccess = nullptr;
    }

    PendingFenceList fenceList;

    DAWN_CHECK(!resource->IsError());
    ExecutionSerial lastUsageSerial = resource->OnEndAccess();
    // If the last usage serial is non-zero, the resource was used by GPU commands during this
    // access.
    const bool usedDuringAccess = lastUsageSerial != kBeginningOfGPUTime;

    // If the resource was used, call the error-generating part of the EndAccess implementation to
    // export a fence. This is separated out because writing the output state must happen
    // regardless of whether or not EndAccessInternal succeeds.
    MaybeUnknownError err;
    if (usedDuringAccess) {
        ResultOrUnknownError<FenceAndSignalValue> result =
            EndAccessInternal(lastUsageSerial, resource, state);
        if (result.IsSuccess()) {
            FenceAndSignalValue fence = result.AcquireSuccess();
            // Some backends might not support fence, in those case, a null object might be
            // returned. So skip it.
            if (fence.object) {
                fenceList.push_back(fence);
            }
        } else {
            err = result.AcquireError();
        }
    }

    // The state transitions to NotAccessed if the exclusive access ends, or the last concurrent
    // read ends. When this occurs, acquire any pending fences that may remain. This occurs if
    // the accesses never acquired them. Note that this must happen after EndAccessInternal, which
    // flushes the commands using the resource, since backends may only acquire and wait on the
    // pending fences when those commands are submitted. Otherwise, the fences would be returned
    // here without the submitted commands waiting on them.
    if (mContents->mSharedResourceAccessState == SharedResourceAccessState::NotAccessed) {
        PendingFenceList pendingFences;
        mContents->AcquirePendingFences(&pendingFences);
        fenceList.insert(fenceList.end(), pendingFences.begin(), pendingFences.end());
    }

    // Copy the fences to the output state.
    if (size_t fenceCount = fenceList.size()) {
        auto fences = HeapArray<SharedFenceBase*>(fenceCount);
        auto signaledValues = HeapArray<uint64_t>(fenceCount);

        for (size_t i = 0; i < fenceCount; ++i) {
            fences[i] = ReturnToAPI(std::move(fenceList[i].object));
            signaledValues[i] = fenceList[i].signaledValue;
        }

        state->fences = std::move(fences).MoveToSpan();
        state->signaledValues = std::move(signaledValues).MoveToSpan();
    } else {
        state->fences = {};
        state->signaledValues = {};
    }
    // If EndAccessInternal failed, including because the device is lost (or destroyed), commands
    // using the resource during this access may not have been flushed, may have been dropped
    // without ever executing (e.g. a queue.writeTexture recorded in a passive pending command
    // context, or a submit that failed), or may have been in the middle of executing when the
    // device was lost, even though the resource was already marked as initialized on the CPU side.
    // In that case, conservatively report the resource as uninitialized so that the next user
    // lazily clears it instead of reading back undefined contents. If the resource wasn't used
    // during the access, its initialized state is still the one set on BeginAccess.
    bool dirtyDueToErrors = !err.IsSuccess() && usedDuringAccess;
    state->initialized = resource->IsResourceInitialized() && !dirtyDueToErrors;
    return err;
}

template <typename Resource, typename EndAccessState>
ResultOrUnknownError<FenceAndSignalValue> SharedResourceMemory::EndAccessInternal(
    ExecutionSerial lastUsageSerial,
    Resource* resource,
    EndAccessState* rawState) {
    // If the device is lost (including when it is destroyed), the backend can't record or submit
    // any more commands, and the commands using the resource may have been dropped, so there is
    // nothing to export for them. Return an error so that EndAccess reports the resource as
    // uninitialized, and so that the backend part of EndAccess is skipped entirely.
    DAWN_TRY(GetDevice()->ValidateIsAlive());
    UnpackedPtr<EndAccessState> state;
    DAWN_TRY_ASSIGN(state, ValidateAndUnpack(rawState));
    // Ensure that commands are submitted before exporting fences with the last usage serial.
    DAWN_TRY(GetDevice()->GetQueue()->EnsureCommandsFlushed(lastUsageSerial));
    return EndAccessImpl(resource, lastUsageSerial, state);
}

// SharedResourceMemoryContents

SharedResourceMemoryContents::SharedResourceMemoryContents(
    WeakRef<SharedResourceMemory> sharedResourceMemory)
    : mSharedResourceMemory(std::move(sharedResourceMemory)) {}

const WeakRef<SharedResourceMemory>& SharedResourceMemoryContents::GetSharedResourceMemory() const {
    return mSharedResourceMemory;
}

void SharedResourceMemoryContents::AcquirePendingFences(PendingFenceList* fences) {
    *fences = mPendingFences;
    mPendingFences.clear();
}

}  // namespace dawn::native