/*
 * Copyright 2023 Google LLC
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "src/gpu/graphite/compute/VelloComputeSteps.h"

namespace skgpu::graphite {

std::string_view VelloStageName(vello_cpp::ShaderStage stage) {
    auto name = vello_cpp::shader(stage).name();
    return {name.data(), name.length()};
}

WorkgroupSize VelloStageLocalSize(vello_cpp::ShaderStage stage) {
    auto wgSize = vello_cpp::shader(stage).workgroup_size();
    return WorkgroupSize(wgSize.x, wgSize.y, wgSize.z);
}

skia_private::TArray<ComputeStep::WorkgroupBufferDesc> VelloWorkgroupBuffers(
        vello_cpp::ShaderStage stage) {
    auto wgBuffers = vello_cpp::shader(stage).workgroup_buffers();
    skia_private::TArray<ComputeStep::WorkgroupBufferDesc> result;
    if (!wgBuffers.empty()) {
        result.reserve(wgBuffers.size());
        for (const auto& desc : wgBuffers) {
            result.push_back({desc.size_in_bytes, desc.index});
        }
    }
    return result;
}

ComputeStep::NativeShaderSource VelloNativeShaderSource(vello_cpp::ShaderStage stage,
                                                        ComputeStep::NativeShaderFormat format) {
    using NativeFormat = ComputeStep::NativeShaderFormat;

    const auto& shader = vello_cpp::shader(stage);
    ::rust::Str source;
    std::string entryPoint;
    switch (format) {
#ifdef SK_DAWN
        case NativeFormat::kWGSL:
            source = shader.wgsl();
            entryPoint = "main";
            break;
#endif
#ifdef SK_METAL
        case NativeFormat::kMSL:
            source = shader.msl();
            entryPoint = "main_";
            break;
#endif
        default:
            return {std::string_view(), ""};
    }

    return {{source.data(), source.length()}, std::move(entryPoint)};
}

// PathtagReduce
VelloPathtagReduceStep::VelloPathtagReduceStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kMapped,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kMapped,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathtagReduceOutput,
            },
        }) {}

// PathtagScanSmall
VelloPathtagScanSmallStep::VelloPathtagScanSmallStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathtagReduceOutput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_TagMonoid,
            },
        }) {}


// PathtagReduce2
VelloPathtagReduce2Step::VelloPathtagReduce2Step() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_LargePathtagReduceFirstPassOutput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_LargePathtagReduceSecondPassOutput,
            },
        }) {}

// PathtagScan1
VelloPathtagScan1Step::VelloPathtagScan1Step() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_LargePathtagReduceFirstPassOutput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_LargePathtagReduceSecondPassOutput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_LargePathtagScanFirstPassOutput,
            },
        }) {}

// PathtagScanLarge
VelloPathtagScanLargeStep::VelloPathtagScanLargeStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_LargePathtagScanFirstPassOutput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_TagMonoid,
            },
        }) {}

// BboxClear
VelloBboxClearStep::VelloBboxClearStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathBBoxes,
            },
        }) {}

// Pathseg
VelloPathsegStep::VelloPathsegStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_TagMonoid,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Cubics,
            },
        }) {}

// DrawReduce
VelloDrawReduceStep::VelloDrawReduceStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawReduceOutput,
            },
        }) {}

// DrawLeaf
VelloDrawLeafStep::VelloDrawLeafStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawReduceOutput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawMonoid,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_InfoBinData,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipInput,
            },
        }) {}

// ClipReduce
VelloClipReduceStep::VelloClipReduceStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipInput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipBicyclic,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipElement,
            },
        }) {}

// ClipLeaf
VelloClipLeafStep::VelloClipLeafStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipInput,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipBicyclic,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipElement,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawMonoid,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipBBoxes,
            },
        }) {}

// Binning
VelloBinningStep::VelloBinningStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawMonoid,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PathBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ClipBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kClear,
                /*slot=*/kVelloSlot_BumpAlloc,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_InfoBinData,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_BinHeader,
            },
        }) {}

// TileAlloc
VelloTileAllocStep::VelloTileAllocStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawBBoxes,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_BumpAlloc,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Path,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Tile,
            },
        }) {}

// PathCoarseFull
VelloPathCoarseFullStep::VelloPathCoarseFullStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Cubics,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Path,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_BumpAlloc,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Tile,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Segments,
            },
        }) {}

// Backdrop
VelloBackdropStep::VelloBackdropStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Tile,
            },
        }) {}

// BackdropDyn
VelloBackdropDynStep::VelloBackdropDynStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Path,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Tile,
            },
        }) {}

// Coarse
VelloCoarseStep::VelloCoarseStep() : VelloStep(
        /*resources=*/{
            {
                /*type=*/ResourceType::kUniformBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_ConfigUniform,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Scene,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_DrawMonoid,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_BinHeader,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_InfoBinData,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Path,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_Tile,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_BumpAlloc,
            },
            {
                /*type=*/ResourceType::kStorageBuffer,
                /*flow=*/DataFlow::kShared,
                /*policy=*/ResourcePolicy::kNone,
                /*slot=*/kVelloSlot_PTCL,
            },
        }) {}

// Fine
VelloFineStep::VelloFineStep(SkColorType targetFormat)
        : VelloStep(
            /*resources=*/{
                {
                    /*type=*/ResourceType::kUniformBuffer,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_ConfigUniform,
                },
                {
                    /*type=*/ResourceType::kStorageBuffer,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_Segments,
                },
                {
                    /*type=*/ResourceType::kStorageBuffer,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_PTCL,
                },
                {
                    /*type=*/ResourceType::kStorageBuffer,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_InfoBinData,
                },
                {
                    /*type=*/ResourceType::kWriteOnlyStorageTexture,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_OutputImage,
                },
                {
                    /*type=*/ResourceType::kReadOnlyTexture,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_GradientImage,
                },
                {
                    /*type=*/ResourceType::kReadOnlyTexture,
                    /*flow=*/DataFlow::kShared,
                    /*policy=*/ResourcePolicy::kNone,
                    /*slot=*/kVelloSlot_ImageAtlas,
                },
            })
        , fTargetFormat(targetFormat) {}

std::tuple<SkISize, SkColorType> VelloFineStep::calculateTextureParameters(int index, const ResourceDesc&) const {
    return {{}, index == 4 ? fTargetFormat : kRGBA_8888_SkColorType};
}

}  // namespace skgpu::graphite