/*

 * Copyright 2024 Google LLC

 *

 * Use of this source code is governed by a BSD-style license that can be

 * found in the LICENSE file.

 */

#include "src/core/SkKnownRuntimeEffects.h"

#include "include/core/SkString.h"

#include "include/effects/SkRuntimeEffect.h"

#include "src/core/SkRuntimeEffectPriv.h"

#include "src/effects/imagefilters/SkMatrixConvolutionImageFilter.h"

namespace SkKnownRuntimeEffects {

namespace {

// This must be kept in sync w/ the version in BlurUtils.h

static constexpr int kMaxBlurSamples = 28;

SkRuntimeEffect* make_blur_1D_effect(int kernelWidth, const SkRuntimeEffect::Options& options) {

    SkASSERT(kernelWidth <= kMaxBlurSamples);

    // The SkSL structure performs two kernel taps; if the kernel has an odd width the last

    // sample will be skipped with the current loop limit calculation.

    SkASSERT(kernelWidth % 2 == 0);

    return SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

            SkStringPrintf(

                    // The coefficients are always stored for the max radius to keep the

                    // uniform block consistent across all effects.

                    "const int kMaxUniformKernelSize = %d / 2;"

                    // But we generate an exact loop over the kernel size. Note that this

                    // program can be used for kernels smaller than the constructed max as long

                    // as the kernel weights for excess entries are set to 0.

                    "const int kMaxLoopLimit = %d / 2;"

                    "uniform half4 offsetsAndKernel[kMaxUniformKernelSize];"

                    "uniform half2 dir;"

                    "uniform shader child;"

                    "half4 main(float2 coord) {"

                        "half4 sum = half4(0);"

                        "for (int i = 0; i < kMaxLoopLimit; ++i) {"

                            "half4 s = offsetsAndKernel[i];"

                            "sum += s.y * child.eval(coord + s.x*dir);"

                            "sum += s.w * child.eval(coord + s.z*dir);"

                        "}"

                        "return sum;"

                    "}", kMaxBlurSamples, kernelWidth).c_str(),

                    options);

}

SkRuntimeEffect* make_blur_2D_effect(int maxKernelSize, const SkRuntimeEffect::Options& options) {

    SkASSERT(maxKernelSize % 4 == 0);

    return SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

            SkStringPrintf(

                    // The coefficients are always stored for the max radius to keep the

                    // uniform block consistent across all effects.

                    "const int kMaxUniformKernelSize = %d / 4;"

                    "const int kMaxUniformOffsetsSize = 2*kMaxUniformKernelSize;"

                    // But we generate an exact loop over the kernel size. Note that this

                    // program can be used for kernels smaller than the constructed max as long

                    // as the kernel weights for excess entries are set to 0.

                    "const int kMaxLoopLimit = %d / 4;"

                    // Pack scalar coefficients into half4 for better packing on std140, and

                    // upload offsets to avoid having to transform the 1D index into a 2D coord

                    "uniform half4 kernel[kMaxUniformKernelSize];"

                    "uniform half4 offsets[kMaxUniformOffsetsSize];"

                    "uniform shader child;"

                    "half4 main(float2 coord) {"

                        "half4 sum = half4(0);"

                        "for (int i = 0; i < kMaxLoopLimit; ++i) {"

                            "half4 k = kernel[i];"

                            "half4 o = offsets[2*i];"

                            "sum += k.x * child.eval(coord + o.xy);"

                            "sum += k.y * child.eval(coord + o.zw);"

                            "o = offsets[2*i + 1];"

                            "sum += k.z * child.eval(coord + o.xy);"

                            "sum += k.w * child.eval(coord + o.zw);"

                        "}"

                        "return sum;"

                    "}", kMaxBlurSamples, maxKernelSize).c_str(),

                    options);

}

enum class MatrixConvolutionImpl {

    kUniformBased,

    kTextureBasedSm,

    kTextureBasedLg,

};

// There are three shader variants:

//    a smaller kernel version that stores the matrix in uniforms and iterates in 1D

//    a larger kernel version that stores the matrix in a 1D texture. The texture version has small

//    and large variants w/ the actual kernel size uploaded as a uniform.

SkRuntimeEffect* make_matrix_conv_effect(MatrixConvolutionImpl impl,

                                         const SkRuntimeEffect::Options& options) {

    // While the uniforms and kernel access are different, pieces of the algorithm are common and

    // defined statically for re-use in the two shaders:

    static const char* kHeaderAndBeginLoopSkSL =

        "uniform int2 size;"

        "uniform int2 offset;"

        "uniform half2 gainAndBias;"

        "uniform int convolveAlpha;" // FIXME not a full  int? Put in a half3 w/ gainAndBias?

        "uniform shader child;"

        "half4 main(float2 coord) {"

            "half4 sum = half4(0);"

            "half origAlpha = 0;"

            "int2 kernelPos = int2(0);"

            "for (int i = 0; i < kMaxKernelSize; ++i) {"

                "if (kernelPos.y >= size.y) { break; }";

    // Used in the inner loop to accumulate convolution sum and increment the kernel position

    static const char* kAccumulateAndIncrementSkSL =

                "half4 c = child.eval(coord + half2(kernelPos) - half2(offset));"

                "if (convolveAlpha == 0) {"

                    // When not convolving alpha, remember the original alpha for actual sample

                    // coord, and perform accumulation on unpremul colors.

                    "if (kernelPos == offset) {"

                        "origAlpha = c.a;"

                    "}"

                    "c = unpremul(c);"

                "}"

                "sum += c*k;"

                "kernelPos.x += 1;"

                "if (kernelPos.x >= size.x) {"

                    "kernelPos.x = 0;"

                    "kernelPos.y += 1;"

                "}";

    // Closes the loop and calculates final color

    static const char* kCloseLoopAndFooterSkSL =

            "}"

            "half4 color = sum*gainAndBias.x + gainAndBias.y;"

            "if (convolveAlpha == 0) {"

                // Reset the alpha to the original and convert to premul RGB

                "color = half4(color.rgb*origAlpha, origAlpha);"

            "} else {"

                // Ensure convolved alpha is within [0, 1]

                "color.a = saturate(color.a);"

            "}"

            // Make RGB valid premul w/ respect to the alpha (either original or convolved)

            "color.rgb = clamp(color.rgb, 0, color.a);"

            "return color;"

        "}";

    static const auto makeTextureEffect = [](int maxTextureKernelSize,

                                             const SkRuntimeEffect::Options& options) {

        return SkMakeRuntimeEffect(

                        SkRuntimeEffect::MakeForShader,

                        SkStringPrintf("const int kMaxKernelSize = %d;"

                                       "uniform shader kernel;"

                                       "uniform half2 innerGainAndBias;"

                                       "%s" // kHeaderAndBeginLoopSkSL

                                               "half k = kernel.eval(half2(half(i) + 0.5, 0.5)).a;"

                                               "k = k * innerGainAndBias.x + innerGainAndBias.y;"

                                               "%s" // kAccumulateAndIncrementSkSL

                                       "%s", // kCloseLoopAndFooterSkSL

                                       maxTextureKernelSize,

                                       kHeaderAndBeginLoopSkSL,

                                       kAccumulateAndIncrementSkSL,

                                       kCloseLoopAndFooterSkSL).c_str(),

                        options);

    };

    switch (impl) {

        case MatrixConvolutionImpl::kUniformBased: {

            return SkMakeRuntimeEffect(

                        SkRuntimeEffect::MakeForShader,

                        SkStringPrintf("const int kMaxKernelSize = %d / 4;"

                                       "uniform half4 kernel[kMaxKernelSize];"

                                       "%s" // kHeaderAndBeginLoopSkSL

                                                "half4 k4 = kernel[i];"

                                                "for (int j = 0; j < 4; ++j) {"

                                                    "if (kernelPos.y >= size.y) { break; }"

                                                    "half k = k4[j];"

                                                    "%s" // kAccumulateAndIncrementSkSL

                                                "}"

                                       "%s", // kCloseLoopAndFooterSkSL

                                       MatrixConvolutionImageFilter::kMaxUniformKernelSize,

                                       kHeaderAndBeginLoopSkSL,

                                       kAccumulateAndIncrementSkSL,

                                       kCloseLoopAndFooterSkSL).c_str(),

                        options);

        }

        case MatrixConvolutionImpl::kTextureBasedSm:

            return makeTextureEffect(MatrixConvolutionImageFilter::kSmallKernelSize, options);

        case MatrixConvolutionImpl::kTextureBasedLg:

            return makeTextureEffect(MatrixConvolutionImageFilter::kLargeKernelSize, options);

    }

    SkUNREACHABLE;

}

} // anonymous namespace

const SkRuntimeEffect* GetKnownRuntimeEffect(StableKey stableKey) {

    SkRuntimeEffect::Options options;

    SkRuntimeEffectPriv::SetStableKey(&options, static_cast<uint32_t>(stableKey));

    SkRuntimeEffectPriv::AllowPrivateAccess(&options);

    switch (stableKey) {

        case StableKey::kInvalid:

            return nullptr;

        // Shaders

        case StableKey::k1DBlur4: {

            static SkRuntimeEffect* s1DBlurEffect = make_blur_1D_effect(4, options);

            return s1DBlurEffect;

        }

        case StableKey::k1DBlur8: {

            static SkRuntimeEffect* s1DBlurEffect = make_blur_1D_effect(8, options);

            return s1DBlurEffect;

        }

        case StableKey::k1DBlur12: {

            static SkRuntimeEffect* s1DBlurEffect = make_blur_1D_effect(12, options);

            return s1DBlurEffect;

        }

        case StableKey::k1DBlur16: {

            static SkRuntimeEffect* s1DBlurEffect = make_blur_1D_effect(16, options);

            return s1DBlurEffect;

        }

        case StableKey::k1DBlur20: {

            static SkRuntimeEffect* s1DBlurEffect = make_blur_1D_effect(20, options);

            return s1DBlurEffect;

        }

        case StableKey::k1DBlur28: {

            static SkRuntimeEffect* s1DBlurEffect = make_blur_1D_effect(28, options);

            return s1DBlurEffect;

        }

        case StableKey::k2DBlur4: {

            static SkRuntimeEffect* s2DBlurEffect = make_blur_2D_effect(4, options);

            return s2DBlurEffect;

        }

        case StableKey::k2DBlur8: {

            static SkRuntimeEffect* s2DBlurEffect = make_blur_2D_effect(8, options);

            return s2DBlurEffect;

        }

        case StableKey::k2DBlur12: {

            static SkRuntimeEffect* s2DBlurEffect = make_blur_2D_effect(12, options);

            return s2DBlurEffect;

        }

        case StableKey::k2DBlur16: {

            static SkRuntimeEffect* s2DBlurEffect = make_blur_2D_effect(16, options);

            return s2DBlurEffect;

        }

        case StableKey::k2DBlur20: {

            static SkRuntimeEffect* s2DBlurEffect = make_blur_2D_effect(20, options);

            return s2DBlurEffect;

        }

        case StableKey::k2DBlur28: {

            static SkRuntimeEffect* s2DBlurEffect = make_blur_2D_effect(28, options);

            return s2DBlurEffect;

        }

        case StableKey::kBlend: {

            static constexpr char kBlendShaderCode[] =

                "uniform shader s, d;"

                "uniform blender b;"

                "half4 main(float2 xy) {"

                    "return b.eval(s.eval(xy), d.eval(xy));"

                "}";

            static const SkRuntimeEffect* sBlendEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kBlendShaderCode,

                                        options);

            return sBlendEffect;

        }

        case StableKey::kLerp: {

            static constexpr char kLerpFilterCode[] =

                "uniform colorFilter cf0;"

                "uniform colorFilter cf1;"

                "uniform half weight;"

                "half4 main(half4 color) {"

                    "return mix(cf0.eval(color), cf1.eval(color), weight);"

                "}";

            static const SkRuntimeEffect* sLerpEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForColorFilter,

                                        kLerpFilterCode,

                                        options);

            return sLerpEffect;

        }

        case StableKey::kMatrixConvUniforms: {

            static const SkRuntimeEffect* sMatrixConvUniformsEffect =

                    make_matrix_conv_effect(MatrixConvolutionImpl::kUniformBased, options);

            return sMatrixConvUniformsEffect;

        }

        case StableKey::kMatrixConvTexSm: {

            static const SkRuntimeEffect* sMatrixConvTexSmEffect =

                    make_matrix_conv_effect(MatrixConvolutionImpl::kTextureBasedSm, options);

            return sMatrixConvTexSmEffect;

        }

        case StableKey::kMatrixConvTexLg: {

            static const SkRuntimeEffect* sMatrixConvTexMaxEffect =

                    make_matrix_conv_effect(MatrixConvolutionImpl::kTextureBasedLg, options);

            return sMatrixConvTexMaxEffect;

        }

        case StableKey::kDecal: {

            static constexpr char kDecalShaderCode[] =

                "uniform shader image;"

                "uniform float4 decalBounds;"

                "half4 main(float2 coord) {"

                    "return sk_decal(image, coord, decalBounds);"

                "}";

            static const SkRuntimeEffect* sDecalEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kDecalShaderCode,

                                        options);

            return sDecalEffect;

        }

        case StableKey::kDisplacement: {

            // NOTE: This uses dot product selection to work on all GLES2 hardware (enforced by

            // public runtime effect restrictions). Otherwise, this would use a "uniform ivec2"

            // and component indexing to convert the displacement color into a vector.

            static constexpr char kDisplacementShaderCode[] =

                "uniform shader displMap;"

                "uniform shader colorMap;"

                "uniform half2 scale;"

                "uniform half4 xSelect;" // Only one of RGBA will be 1, the rest are 0

                "uniform half4 ySelect;"

                "half4 main(float2 coord) {"

                    "return sk_displacement(displMap, colorMap, coord, scale, xSelect, ySelect);"

                "}";

            static const SkRuntimeEffect* sDisplacementEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kDisplacementShaderCode,

                                        options);

            return sDisplacementEffect;

        }

        case StableKey::kLighting: {

            static constexpr char kLightingShaderCode[] =

                "uniform shader normalMap;"

                // Packs surface depth, shininess, material type (0 == diffuse) and light type

                // (< 0 = distant, 0 = point, > 0 = spot)

                "uniform half4 materialAndLightType;"

                "uniform half4 lightPosAndSpotFalloff;" // (x,y,z) are lightPos, w is spot falloff

                                                        // exponent

                "uniform half4 lightDirAndSpotCutoff;" // (x,y,z) are lightDir,

                                                       // w is spot cos(cutoffAngle)

                "uniform half3 lightColor;" // Material's k has already been multiplied in

                "half4 main(float2 coord) {"

                    "return sk_lighting(normalMap, coord,"

                                        /*depth=*/"materialAndLightType.x,"

                                        /*shininess=*/"materialAndLightType.y,"

                                        /*materialType=*/"materialAndLightType.z,"

                                        /*lightType=*/"materialAndLightType.w,"

                                        /*lightPos=*/"lightPosAndSpotFalloff.xyz,"

                                        /*spotFalloff=*/"lightPosAndSpotFalloff.w,"

                                        /*lightDir=*/"lightDirAndSpotCutoff.xyz,"

                                        /*cosCutoffAngle=*/"lightDirAndSpotCutoff.w,"

                                        "lightColor);"

                "}";

            static const SkRuntimeEffect* sLightingEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kLightingShaderCode,

                                        options);

            return sLightingEffect;

        }

        case StableKey::kLinearMorphology: {

            static constexpr char kLinearMorphologyShaderCode[] =

                "uniform shader child;"

                "uniform half2 offset;"

                "uniform half flip;" // -1 converts the max() calls to min()

                "uniform int radius;"

                "half4 main(float2 coord) {"

                    "return sk_linear_morphology(child, coord, offset, flip, radius);"

                "}";

            static const SkRuntimeEffect* sLinearMorphologyEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kLinearMorphologyShaderCode,

                                        options);

            return sLinearMorphologyEffect;

        }

        case StableKey::kMagnifier: {

            static constexpr char kMagnifierShaderCode[] =

                "uniform shader src;"

                "uniform float4 lensBounds;"

                "uniform float4 zoomXform;"

                "uniform float2 invInset;"

                "half4 main(float2 coord) {"

                    "return sk_magnifier(src, coord, lensBounds, zoomXform, invInset);"

                "}";

            static const SkRuntimeEffect* sMagnifierEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kMagnifierShaderCode,

                                        options);

            return sMagnifierEffect;

        }

        case StableKey::kNormal: {

            static constexpr char kNormalShaderCode[] =

                "uniform shader alphaMap;"

                "uniform float4 edgeBounds;"

                "uniform half negSurfaceDepth;"

                "half4 main(float2 coord) {"

                   "return sk_normal(alphaMap, coord, edgeBounds, negSurfaceDepth);"

                "}";

            static const SkRuntimeEffect* sNormalEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kNormalShaderCode,

                                        options);

            return sNormalEffect;

        }

        case StableKey::kSparseMorphology: {

            static constexpr char kSparseMorphologyShaderCode[] =

                "uniform shader child;"

                "uniform half2 offset;"

                "uniform half flip;"

                "half4 main(float2 coord) {"

                    "return sk_sparse_morphology(child, coord, offset, flip);"

                "}";

            static const SkRuntimeEffect* sSparseMorphologyEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForShader,

                                        kSparseMorphologyShaderCode,

                                        options);

            return sSparseMorphologyEffect;

        }

        // Blenders

        case StableKey::kArithmetic: {

            static constexpr char kArithmeticBlenderCode[] =

                "uniform half4 k;"

                "uniform half pmClamp;"

                "half4 main(half4 src, half4 dst) {"

                    "return sk_arithmetic_blend(src, dst, k, pmClamp);"

                "}";

            static const SkRuntimeEffect* sArithmeticEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForBlender,

                                        kArithmeticBlenderCode,

                                        options);

            return sArithmeticEffect;

        }

        // Color Filters

        case StableKey::kHighContrast: {

            static constexpr char kHighContrastFilterCode[] =

                "uniform half grayscale, invertStyle, contrast;"

                "half4 main(half4 color) {"

                    "return half4(sk_high_contrast(color.rgb, grayscale, invertStyle, contrast),"

                                 "color.a);"

                "}";

            static const SkRuntimeEffect* sHighContrastEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForColorFilter,

                                        kHighContrastFilterCode,

                                        options);

            return sHighContrastEffect;

        }

        case StableKey::kLuma: {

            static constexpr char kLumaFilterCode[] =

                "half4 main(half4 color) {"

                    "return sk_luma(color.rgb);"

                "}";

            static const SkRuntimeEffect* sLumaEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForColorFilter,

                                        kLumaFilterCode,

                                        options);

            return sLumaEffect;

        }

        case StableKey::kOverdraw: {

            static constexpr char kOverdrawFilterCode[] =

                "uniform half4 color0, color1, color2, color3, color4, color5;"

                "half4 main(half4 color) {"

                    "return sk_overdraw(color.a, color0, color1, color2, color3, color4, color5);"

                "}";

            static const SkRuntimeEffect* sOverdrawEffect =

                    SkMakeRuntimeEffect(SkRuntimeEffect::MakeForColorFilter,

                                        kOverdrawFilterCode,

                                        options);

            return sOverdrawEffect;

        }

    }

    SkUNREACHABLE;

}

} // namespace SkKnownRuntimeEffects