// Copyright 2019 The Dawn Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "dawn_native/CommandValidation.h"

#include "common/BitSetIterator.h"
#include "dawn_native/BindGroup.h"
#include "dawn_native/Buffer.h"
#include "dawn_native/CommandBufferStateTracker.h"
#include "dawn_native/Commands.h"
#include "dawn_native/PassResourceUsageTracker.h"
#include "dawn_native/RenderBundle.h"
#include "dawn_native/RenderPipeline.h"

namespace dawn_native {

    namespace {

        void TrackBindGroupResourceUsage(BindGroupBase* group,
                                         PassResourceUsageTracker* usageTracker) {
            const auto& layoutInfo = group->GetLayout()->GetBindingInfo();

            for (uint32_t i : IterateBitSet(layoutInfo.mask)) {
                wgpu::BindingType type = layoutInfo.types[i];

                switch (type) {
                    case wgpu::BindingType::UniformBuffer: {
                        BufferBase* buffer = group->GetBindingAsBufferBinding(i).buffer;
                        usageTracker->BufferUsedAs(buffer, wgpu::BufferUsage::Uniform);
                    } break;

                    case wgpu::BindingType::StorageBuffer: {
                        BufferBase* buffer = group->GetBindingAsBufferBinding(i).buffer;
                        usageTracker->BufferUsedAs(buffer, wgpu::BufferUsage::Storage);
                    } break;

                    case wgpu::BindingType::SampledTexture: {
                        TextureBase* texture = group->GetBindingAsTextureView(i)->GetTexture();
                        usageTracker->TextureUsedAs(texture, wgpu::TextureUsage::Sampled);
                    } break;

                    case wgpu::BindingType::ReadonlyStorageBuffer: {
                        BufferBase* buffer = group->GetBindingAsBufferBinding(i).buffer;
                        usageTracker->BufferUsedAs(buffer, kReadOnlyStorage);
                    } break;

                    case wgpu::BindingType::Sampler:
                        break;

                    case wgpu::BindingType::StorageTexture:
                        UNREACHABLE();
                        break;
                }
            }
        }

        inline MaybeError ValidateRenderBundleCommand(CommandIterator* commands,
                                                      Command type,
                                                      PassResourceUsageTracker* usageTracker,
                                                      CommandBufferStateTracker* commandBufferState,
                                                      const AttachmentState* attachmentState,
                                                      uint64_t* debugGroupStackSize,
                                                      const char* disallowedMessage) {
            switch (type) {
                case Command::Draw: {
                    commands->NextCommand<DrawCmd>();
                    DAWN_TRY(commandBufferState->ValidateCanDraw());
                } break;

                case Command::DrawIndexed: {
                    commands->NextCommand<DrawIndexedCmd>();
                    DAWN_TRY(commandBufferState->ValidateCanDrawIndexed());
                } break;

                case Command::DrawIndirect: {
                    DrawIndirectCmd* cmd = commands->NextCommand<DrawIndirectCmd>();
                    DAWN_TRY(commandBufferState->ValidateCanDraw());
                    usageTracker->BufferUsedAs(cmd->indirectBuffer.Get(),
                                               wgpu::BufferUsage::Indirect);
                } break;

                case Command::DrawIndexedIndirect: {
                    DrawIndexedIndirectCmd* cmd = commands->NextCommand<DrawIndexedIndirectCmd>();
                    DAWN_TRY(commandBufferState->ValidateCanDrawIndexed());
                    usageTracker->BufferUsedAs(cmd->indirectBuffer.Get(),
                                               wgpu::BufferUsage::Indirect);
                } break;

                case Command::InsertDebugMarker: {
                    InsertDebugMarkerCmd* cmd = commands->NextCommand<InsertDebugMarkerCmd>();
                    commands->NextData<char>(cmd->length + 1);
                } break;

                case Command::PopDebugGroup: {
                    commands->NextCommand<PopDebugGroupCmd>();
                    DAWN_TRY(ValidateCanPopDebugGroup(*debugGroupStackSize));
                    *debugGroupStackSize -= 1;
                } break;

                case Command::PushDebugGroup: {
                    PushDebugGroupCmd* cmd = commands->NextCommand<PushDebugGroupCmd>();
                    commands->NextData<char>(cmd->length + 1);
                    *debugGroupStackSize += 1;
                } break;

                case Command::SetRenderPipeline: {
                    SetRenderPipelineCmd* cmd = commands->NextCommand<SetRenderPipelineCmd>();
                    RenderPipelineBase* pipeline = cmd->pipeline.Get();

                    if (DAWN_UNLIKELY(pipeline->GetAttachmentState() != attachmentState)) {
                        return DAWN_VALIDATION_ERROR("Pipeline attachment state is not compatible");
                    }
                    commandBufferState->SetRenderPipeline(pipeline);
                } break;

                case Command::SetBindGroup: {
                    SetBindGroupCmd* cmd = commands->NextCommand<SetBindGroupCmd>();
                    if (cmd->dynamicOffsetCount > 0) {
                        commands->NextData<uint32_t>(cmd->dynamicOffsetCount);
                    }

                    TrackBindGroupResourceUsage(cmd->group.Get(), usageTracker);
                    commandBufferState->SetBindGroup(cmd->index, cmd->group.Get());
                } break;

                case Command::SetIndexBuffer: {
                    SetIndexBufferCmd* cmd = commands->NextCommand<SetIndexBufferCmd>();

                    usageTracker->BufferUsedAs(cmd->buffer.Get(), wgpu::BufferUsage::Index);
                    commandBufferState->SetIndexBuffer();
                } break;

                case Command::SetVertexBuffer: {
                    SetVertexBufferCmd* cmd = commands->NextCommand<SetVertexBufferCmd>();

                    usageTracker->BufferUsedAs(cmd->buffer.Get(), wgpu::BufferUsage::Vertex);
                    commandBufferState->SetVertexBuffer(cmd->slot);
                } break;

                default:
                    return DAWN_VALIDATION_ERROR(disallowedMessage);
            }

            return {};
        }

    }  // namespace

    MaybeError ValidateCanPopDebugGroup(uint64_t debugGroupStackSize) {
        if (debugGroupStackSize == 0) {
            return DAWN_VALIDATION_ERROR("Pop must be balanced by a corresponding Push.");
        }
        return {};
    }

    MaybeError ValidateFinalDebugGroupStackSize(uint64_t debugGroupStackSize) {
        if (debugGroupStackSize != 0) {
            return DAWN_VALIDATION_ERROR("Each Push must be balanced by a corresponding Pop.");
        }
        return {};
    }

    MaybeError ValidateRenderBundle(CommandIterator* commands,
                                    const AttachmentState* attachmentState,
                                    PassResourceUsage* resourceUsage) {
        PassResourceUsageTracker usageTracker;
        CommandBufferStateTracker commandBufferState;
        uint64_t debugGroupStackSize = 0;

        Command type;
        while (commands->NextCommandId(&type)) {
            DAWN_TRY(ValidateRenderBundleCommand(commands, type, &usageTracker, &commandBufferState,
                                                 attachmentState, &debugGroupStackSize,
                                                 "Command disallowed inside a render bundle"));
        }

        DAWN_TRY(ValidateFinalDebugGroupStackSize(debugGroupStackSize));
        DAWN_TRY(usageTracker.ValidateRenderPassUsages());
        ASSERT(resourceUsage != nullptr);
        *resourceUsage = usageTracker.AcquireResourceUsage();

        return {};
    }

    MaybeError ValidateRenderPass(CommandIterator* commands,
                                  BeginRenderPassCmd* renderPass,
                                  std::vector<PassResourceUsage>* perPassResourceUsages) {
        PassResourceUsageTracker usageTracker;
        CommandBufferStateTracker commandBufferState;
        uint64_t debugGroupStackSize = 0;

        // Track usage of the render pass attachments
        for (uint32_t i : IterateBitSet(renderPass->attachmentState->GetColorAttachmentsMask())) {
            RenderPassColorAttachmentInfo* colorAttachment = &renderPass->colorAttachments[i];
            TextureBase* texture = colorAttachment->view->GetTexture();
            usageTracker.TextureUsedAs(texture, wgpu::TextureUsage::OutputAttachment);

            TextureViewBase* resolveTarget = colorAttachment->resolveTarget.Get();
            if (resolveTarget != nullptr) {
                usageTracker.TextureUsedAs(resolveTarget->GetTexture(),
                                           wgpu::TextureUsage::OutputAttachment);
            }
        }

        if (renderPass->attachmentState->HasDepthStencilAttachment()) {
            TextureBase* texture = renderPass->depthStencilAttachment.view->GetTexture();
            usageTracker.TextureUsedAs(texture, wgpu::TextureUsage::OutputAttachment);
        }

        Command type;
        while (commands->NextCommandId(&type)) {
            switch (type) {
                case Command::EndRenderPass: {
                    commands->NextCommand<EndRenderPassCmd>();

                    DAWN_TRY(ValidateFinalDebugGroupStackSize(debugGroupStackSize));
                    DAWN_TRY(usageTracker.ValidateRenderPassUsages());
                    ASSERT(perPassResourceUsages != nullptr);
                    perPassResourceUsages->push_back(usageTracker.AcquireResourceUsage());

                    return {};
                } break;

                case Command::ExecuteBundles: {
                    ExecuteBundlesCmd* cmd = commands->NextCommand<ExecuteBundlesCmd>();
                    auto bundles = commands->NextData<Ref<RenderBundleBase>>(cmd->count);
                    for (uint32_t i = 0; i < cmd->count; ++i) {
                        if (DAWN_UNLIKELY(renderPass->attachmentState.Get() !=
                                          bundles[i]->GetAttachmentState())) {
                            return DAWN_VALIDATION_ERROR(
                                "Render bundle is not compatible with render pass");
                        }

                        const PassResourceUsage& usages = bundles[i]->GetResourceUsage();
                        for (uint32_t i = 0; i < usages.buffers.size(); ++i) {
                            usageTracker.BufferUsedAs(usages.buffers[i], usages.bufferUsages[i]);
                        }

                        for (uint32_t i = 0; i < usages.textures.size(); ++i) {
                            usageTracker.TextureUsedAs(usages.textures[i], usages.textureUsages[i]);
                        }
                    }

                    if (cmd->count > 0) {
                        // Reset state. It is invalidated after render bundle execution.
                        commandBufferState = CommandBufferStateTracker{};
                    }

                } break;

                case Command::SetStencilReference: {
                    commands->NextCommand<SetStencilReferenceCmd>();
                } break;

                case Command::SetBlendColor: {
                    commands->NextCommand<SetBlendColorCmd>();
                } break;

                case Command::SetViewport: {
                    commands->NextCommand<SetViewportCmd>();
                } break;

                case Command::SetScissorRect: {
                    commands->NextCommand<SetScissorRectCmd>();
                } break;

                default:
                    DAWN_TRY(ValidateRenderBundleCommand(
                        commands, type, &usageTracker, &commandBufferState,
                        renderPass->attachmentState.Get(), &debugGroupStackSize,
                        "Command disallowed inside a render pass"));
            }
        }

        UNREACHABLE();
        return DAWN_VALIDATION_ERROR("Unfinished render pass");
    }

    MaybeError ValidateComputePass(CommandIterator* commands,
                                   std::vector<PassResourceUsage>* perPassResourceUsages) {
        PassResourceUsageTracker usageTracker;
        CommandBufferStateTracker commandBufferState;
        uint64_t debugGroupStackSize = 0;

        Command type;
        while (commands->NextCommandId(&type)) {
            switch (type) {
                case Command::EndComputePass: {
                    commands->NextCommand<EndComputePassCmd>();

                    DAWN_TRY(ValidateFinalDebugGroupStackSize(debugGroupStackSize));
                    DAWN_TRY(usageTracker.ValidateComputePassUsages());
                    ASSERT(perPassResourceUsages != nullptr);
                    perPassResourceUsages->push_back(usageTracker.AcquireResourceUsage());
                    return {};
                } break;

                case Command::Dispatch: {
                    commands->NextCommand<DispatchCmd>();
                    DAWN_TRY(commandBufferState.ValidateCanDispatch());
                } break;

                case Command::DispatchIndirect: {
                    DispatchIndirectCmd* cmd = commands->NextCommand<DispatchIndirectCmd>();
                    DAWN_TRY(commandBufferState.ValidateCanDispatch());
                    usageTracker.BufferUsedAs(cmd->indirectBuffer.Get(),
                                              wgpu::BufferUsage::Indirect);
                } break;

                case Command::InsertDebugMarker: {
                    InsertDebugMarkerCmd* cmd = commands->NextCommand<InsertDebugMarkerCmd>();
                    commands->NextData<char>(cmd->length + 1);
                } break;

                case Command::PopDebugGroup: {
                    commands->NextCommand<PopDebugGroupCmd>();
                    DAWN_TRY(ValidateCanPopDebugGroup(debugGroupStackSize));
                    debugGroupStackSize--;
                } break;

                case Command::PushDebugGroup: {
                    PushDebugGroupCmd* cmd = commands->NextCommand<PushDebugGroupCmd>();
                    commands->NextData<char>(cmd->length + 1);
                    debugGroupStackSize++;
                } break;

                case Command::SetComputePipeline: {
                    SetComputePipelineCmd* cmd = commands->NextCommand<SetComputePipelineCmd>();
                    ComputePipelineBase* pipeline = cmd->pipeline.Get();
                    commandBufferState.SetComputePipeline(pipeline);
                } break;

                case Command::SetBindGroup: {
                    SetBindGroupCmd* cmd = commands->NextCommand<SetBindGroupCmd>();
                    if (cmd->dynamicOffsetCount > 0) {
                        commands->NextData<uint32_t>(cmd->dynamicOffsetCount);
                    }

                    TrackBindGroupResourceUsage(cmd->group.Get(), &usageTracker);
                    commandBufferState.SetBindGroup(cmd->index, cmd->group.Get());
                } break;

                default:
                    return DAWN_VALIDATION_ERROR("Command disallowed inside a compute pass");
            }
        }

        UNREACHABLE();
        return DAWN_VALIDATION_ERROR("Unfinished compute pass");
    }

}  // namespace dawn_native