/*
 * Copyright 2022 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/PaintParamsKey.h"

#include "include/private/SkAttributes.h"
#include "include/private/SkLog.h"
#include "src/core/SkArenaAlloc.h"
#include "src/core/SkAutoMalloc.h"
#include "src/core/SkBase64.h"
#include "src/core/SkStringView.h"
#include "src/gpu/graphite/Caps.h"
#include "src/gpu/graphite/KeyHelpers.h"
#include "src/gpu/graphite/ShaderCodeDictionary.h"

using namespace skia_private;

namespace skgpu::graphite {

//--------------------------------------------------------------------------------------------------
// PaintParamsKeyBuilder

#ifdef SK_DEBUG

void PaintParamsKeyBuilder::checkReset() {
    SkASSERT(!fLocked);
    SkASSERT(fData.empty());
    SkASSERT(fStack.empty());
    SkASSERT(!fHasError);
}

void PaintParamsKeyBuilder::pushStack(int32_t codeSnippetID) {
    SkASSERT(!fLocked);
    SkASSERT(fDict->isValidID(codeSnippetID));
    // If the kError ID is pushed, fHasError must have been set already.
    SkASSERT(codeSnippetID != (int) BuiltInCodeSnippetID::kError || fHasError);

    if (!fStack.empty()) {
        fStack.back().fNumActualChildren++;
        SkASSERT(fStack.back().fNumActualChildren <= fStack.back().fNumExpectedChildren);
    }

    const ShaderSnippet* snippet = fDict->getEntry(codeSnippetID);
    fStack.push_back({codeSnippetID, snippet->fNumChildren});
}

void PaintParamsKeyBuilder::validateData(size_t dataSize) {
    // Check that the size of the embedded data fits our self-imposed limit.
    SkASSERT(dataSize <= PaintParamsKey::kEmbeddedDataSizeLimit);

    SkASSERT(!fLocked);
    SkASSERT(!fStack.empty()); // addData() called within code snippet block
    // Check that addData() is only called for snippets that support it and is only called once
    const ShaderSnippet* snippet = fDict->getEntry(fStack.back().fCodeSnippetID);
    SkASSERT(snippet->storesSamplerDescData());
    SkASSERT(fStack.back().fDataSize < 0);

    fStack.back().fDataSize = SkTo<int>(dataSize);
}

void PaintParamsKeyBuilder::popStack() {
    SkASSERT(!fLocked);
    SkASSERT(!fStack.empty());
    SkASSERT(fStack.back().fNumActualChildren == fStack.back().fNumExpectedChildren);
    const bool expectsData = fDict->getEntry(fStack.back().fCodeSnippetID)->storesSamplerDescData();
    const bool hasData = fStack.back().fDataSize >= 0;
    SkASSERT(expectsData == hasData);
    fStack.pop_back();
}

void PaintParamsKeyBuilder::validateReplacement(int32_t oldCodeSnippetID,
                                                int32_t newCodeSnippetID) {
    // Rules for allowed replacements:
    // 1. BuiltInCodeSnippetIDs only
    // 2. Have the same number of children (and 0 if allowChildren is false)
    // 3. Have the same uniform and texture signature
    // 4. No extra data
    //
    // This ensures that replacing on a block by block basis produces a valid ShaderNode tree
    // after modification, and that extracted uniforms and textures for the original key can be
    // used with the new key.
    SkASSERT(oldCodeSnippetID >= 0 && oldCodeSnippetID < kBuiltInCodeSnippetIDCount);
    SkASSERT(newCodeSnippetID >= 0 && newCodeSnippetID < kBuiltInCodeSnippetIDCount);

    const ShaderSnippet* oldSnippet = fDict->getEntry(oldCodeSnippetID);
    const ShaderSnippet* newSnippet = fDict->getEntry(newCodeSnippetID);

    SkASSERT(oldSnippet->fNumChildren == newSnippet->fNumChildren);

    // Same signature, not caring about the actual variable names.
    SkASSERT(oldSnippet->fUniforms.size() == newSnippet->fUniforms.size());
    for (int u = 0; u < oldSnippet->fUniforms.size(); ++u) {
        const Uniform& oldU = oldSnippet->fUniforms[u];
        const Uniform& newU = newSnippet->fUniforms[u];
        SkASSERT(oldU.type() == newU.type());
        SkASSERT(oldU.count() == newU.count());
        SkASSERT(oldU.isPaintColor() == newU.isPaintColor());
    }

    SkASSERT(oldSnippet->fTexturesAndSamplers.size() == newSnippet->fTexturesAndSamplers.size());

    SkASSERT(!oldSnippet->storesSamplerDescData());
    SkASSERT(!newSnippet->storesSamplerDescData());
}

#endif // SK_DEBUG

//--------------------------------------------------------------------------------------------------
// PaintParamsKey

PaintParamsKey PaintParamsKey::clone(SkArenaAlloc* arena) const {
    int32_t* newData = arena->makeArrayDefault<int32_t>(fData.size());
    SK_UNSAFE_TODO(memcpy(newData, fData.data(), fData.size_bytes()));
    return PaintParamsKey({newData, fData.size()});
}

ShaderNode* PaintParamsKey::createNode(const ShaderCodeDictionary* dict,
                                       int* currentIndex,
                                       SkArenaAlloc* arena) const {
    SkASSERT(*currentIndex < SkTo<int>(fData.size()));
    const int32_t index = (*currentIndex)++;
    const int32_t id = fData[index];
    SkASSERT(id >= 0); // Otherwise we're calling createNode on a data length somehow

    const ShaderSnippet* entry = dict->getEntry(id);
    if (!entry) {
        SKIA_LOG_E("Unknown snippet ID in key: %d", id);
        return nullptr;
    }

    SkSpan<const uint32_t> dataSpan = {};
    if (entry->storesSamplerDescData()) {
        // If a snippet stores data, then the subsequent paint key index signifies the length of
        // its data. Determine this data length and iterate currentIndex past it.
        const int storedDataLengthIdx = (*currentIndex)++;
        SkASSERT(storedDataLengthIdx < SkTo<int>(fData.size()));
        const int dataLength = EncodeDataSize(fData[storedDataLengthIdx]);
        // Building ShaderNodes is assumed to be done for keys built by PaintParamsKeyBuilder or
        // already validated for deserialization, so we consider them trusted.
        SkASSERT(dataLength >= 0 && dataLength <= kEmbeddedDataSizeLimit);
        SkASSERT(storedDataLengthIdx + dataLength < SkTo<int>(fData.size()));

        // Gather the data contents (length can now be inferred by the consumers of the data) to
        // pass into ShaderNode creation. Iterate the paint key index past the data indices.
        // For the embedded data, we re-interpret the values as uint32_t's since that was what was
        // written by addData().
        SkSpan<const int32_t> signedDataSpan = fData.subspan(storedDataLengthIdx + 1, dataLength);
        dataSpan = {reinterpret_cast<const uint32_t*>(signedDataSpan.data()), dataLength};
        *currentIndex += dataLength;
    }

    ShaderNode** childArray = arena->makeArray<ShaderNode*>(entry->fNumChildren);
    for (int i = 0; i < entry->fNumChildren; ++i) {
        ShaderNode* child = this->createNode(dict, currentIndex, arena);
        if (!child) {
            return nullptr;
        }
        SK_UNSAFE_TODO(childArray[i]) = child;
    }

    return arena->make<ShaderNode>(entry,
                                   SkSpan(childArray, entry->fNumChildren),
                                   id,
                                   index,
                                   dataSpan);
}

// Traverse a ShaderNode tree, attempting to lift any coordinate modification expressions.
// Returns whether any of the given nodes need local coordinate inputs after lifting.
bool lift_coord_expressions(SkSpan<ShaderNode*> nodes, int* availableVaryings) {
    bool anyNeedLocalCoords = false;

    for (ShaderNode* node : nodes) {
        bool curNeedsLocalCoords =
                SkToBool(node->requiredFlags() & SnippetRequirementFlags::kLocalCoords);

        // Lift expressions from nodes whose liftable expressions are on coordinate inputs.
        if (*availableVaryings > 0 && curNeedsLocalCoords &&
            node->entry()->fLiftableExpressionType ==
                    ShaderSnippet::LiftableExpressionType::kLocalCoords) {
            --*availableVaryings;

#if !defined(SK_USE_LEGACY_UNIFORM_LIFTING_GRAPHITE)
            // We can potentially lift the nested expressions under here as well.
            const bool childNeedsOurCoords =
                    lift_coord_expressions(node->children(), availableVaryings);
            // If no child needs our lifted coords, we can omit them from the fragment shader
            // entirely, and only use them in the vertex shader for calculating other coords.
            if (!childNeedsOurCoords) {
                node->setOmitExpressionFlag();
            } else {
                node->setLiftExpressionFlag();
            }
#else
            node->setLiftExpressionFlag();
#endif
            // Since we lifted the coordinate expression here, this node no longer needs a local
            // coords argument.
            curNeedsLocalCoords = false;
            node->unsetLocalCoordsFlag();

#if !defined(SK_USE_LEGACY_UNIFORM_LIFTING_GRAPHITE)
        // If the node passes through its local coords to its children, we check if those perform
        // modifications that can be lifted.
        } else if (*availableVaryings > 0 &&
                   node->requiredFlags() & SnippetRequirementFlags::kPassthroughLocalCoords) {
            // Assume that this node doesn't need local coordinates unless its actual shader snippet
            // entry does, or one of its children does even after accounting for lifting.
            const bool entryNeedsLocalCoords = node->entry()->needsLocalCoords();
            const bool childNeedsLocalCoords = lift_coord_expressions(node->children(),
                                                                      availableVaryings);
            curNeedsLocalCoords = entryNeedsLocalCoords || childNeedsLocalCoords;
            if (!curNeedsLocalCoords) {
                node->unsetLocalCoordsFlag();
            }
#endif
        }

        anyNeedLocalCoords |= curNeedsLocalCoords;
    }

    return anyNeedLocalCoords;
}

// Traverse a list of ShaderNodes, attempting to lift any expressions that resolve to a color.
// For now, this does not recurse into ShaderNodes' lists of children. In practice we only lift
// solid color expressions, and we only care to lift such expressions if there is no other fragment
// shader work (i.e., if the solid color expression is a root node in a shader's ShaderNode tree).
// If there is other fragment shader work, we'll likely be accessing other fragment shader uniforms,
// the color value will likely be cached, and lifting may not be worth the extra varying.
void lift_color_expressions(SkSpan<ShaderNode*> nodes, int* availableVaryings) {
#if !defined(SK_USE_LEGACY_UNIFORM_LIFTING_GRAPHITE)
    for (ShaderNode* node : nodes) {
        if (*availableVaryings > 0 &&
            node->entry()->fLiftableExpressionType ==
                    ShaderSnippet::LiftableExpressionType::kPriorStageOutput) {
            --*availableVaryings;
            node->setLiftExpressionFlag();
        }
    }
#endif
}

RootNodesInfo PaintParamsKey::getRootNodes(const Caps* caps,
                                           const ShaderCodeDictionary* dict,
                                           const RuntimeEffectDictionary* rteDict,
                                           SkArenaAlloc* arena,
                                           int availableVaryings,
                                           bool canLiftCoords) const {
    // TODO: Once the PaintParamsKey creation is organized to represent a single tree starting at
    // the final blend, there will only be a single root node and this can be simplified.
    // For now, we don't know how many roots there are, so collect them into a local array before
    // copying into the arena.
    const int keySize = SkTo<int>(fData.size());

    RootNodesInfo rootsInfo;
    // Normal PaintParams creation will have up to 4 roots for the different stages.
    STArray<4, ShaderNode*> roots;
    STArray<2, ShaderNode*> liftableRoots;
    int currentIndex = 0;
    while (currentIndex < keySize) {
        int32_t blockMarker = fData[currentIndex++];
        if (blockMarker >= 0) {
            return {}; // a bad key
        }
        RootBlockType type = static_cast<RootBlockType>(blockMarker);
        ShaderNode* root = this->createNode(dict, &currentIndex, arena);
        if (!root) {
            return {}; // a bad key
        }
        roots.push_back(root);
        switch (type) {
            case RootBlockType::kSrcColor:
                SkASSERT(!rootsInfo.fSrcColor);
                rootsInfo.fSrcColor = root;
                liftableRoots.push_back(root);
                break;
            case RootBlockType::kFinalBlend:
                SkASSERT(!rootsInfo.fFinalBlend);
                rootsInfo.fFinalBlend = root;
                liftableRoots.push_back(root);
                break;
            case RootBlockType::kClip:
                SkASSERT(!rootsInfo.fClip);
                rootsInfo.fClip = root;
                break;
            case RootBlockType::kMeshShader:
                SkASSERT(!rootsInfo.fMeshShader);
                rootsInfo.fMeshShader = root;
                break;
            default:
                SkUNREACHABLE;
        }
    }

    if (rootsInfo.fMeshShader) {
        rootsInfo.fMeshSpec = rteDict->findMeshSpec(rootsInfo.fMeshShader->codeSnippetId());
        if (!rootsInfo.fMeshSpec) {
            // Couldn't find the SkMeshSpecification for the mesh shader snippet so the key is bad.
            return {};
        }
    }

    // See what expressions we can lift to the vertex shader.
    if (canLiftCoords) {
        lift_coord_expressions(liftableRoots, &availableVaryings);
    }
    // Don't lift constant expressions if we're using regular UBOs, since lifting is likely only
    // beneficial if we're avoiding a storage buffer access.
    if (caps->storageBufferSupport()) {
        lift_color_expressions(liftableRoots, &availableVaryings);
    }

    // Copy the accumulated roots into a span stored in the arena
    const ShaderNode** rootSpan = arena->makeArray<const ShaderNode*>(roots.size());
    SK_UNSAFE_TODO(memcpy(rootSpan, roots.data(), roots.size_bytes()));
    rootsInfo.fRoots = SkSpan(rootSpan, roots.size());
    return rootsInfo;
}

static void append_as_base64(SkString* str, SkSpan<const int32_t> data) {
    str->append("(");
    str->appendU32(data.size());
    str->append(": ");
    // Encode data in base64 to shorten it
    const size_t srcDataSize = data.size() * sizeof(uint32_t); // size in bytes
    SkAutoMalloc encodedData{SkBase64::EncodedSize(srcDataSize)};
    char* dst = static_cast<char*>(encodedData.get());
    size_t encodedLen = SkBase64::Encode(data.data(), srcDataSize, dst);
    str->append(dst, encodedLen);
    str->append(")");
}

static int key_to_string(const Caps* caps,
                         SkString* str,
                         const ShaderCodeDictionary* dict,
                         SkSpan<const int32_t> keyData,
                         int currentIndex,
                         int indent) {
    SkASSERT(currentIndex < SkTo<int>(keyData.size()));

    const bool multiline = indent >= 0;
    if (multiline) {
        // Format for multi-line printing
        str->appendf("%*c", 2 * indent, ' ');
    }

    int32_t id = keyData[currentIndex++];
    // Skip the root block headers.
    if (id < 0) {
        if (multiline) {
            switch (static_cast<RootBlockType>(id)) {
                case RootBlockType::kSrcColor:
                    str->append("[RootSrcColor] ");
                    break;
                case RootBlockType::kFinalBlend:
                    str->append("[RootFinalBlend] ");
                    break;
                case RootBlockType::kClip:
                    str->append("[RootClip] ");
                    break;
                case RootBlockType::kMeshShader:
                    str->append("[RootMeshShader] ");
                    break;
                default:
                    SkUNREACHABLE;
            }
        }
        id = keyData[currentIndex++];
    }

    auto entry = dict->getEntry(id);
    if (!entry) {
        str->append("Unknown(");
        str->appendS32(id);
        str->append(")");
        return currentIndex;
    }

    // Single lined Composes get shortened to just a plus between its two children, e.g. Compose [ A
    // B ] becomes A+B. We don't do a similar prettification for Blend [ A B C ] to (A B)+C because
    // that's not quite as readable and they aren't nearly as common within keys. To make sure
    // chains of Composes are not ambiguous, we only consolidate cases where the inner node is not
    // Compose, e.g. Compose [ A Compose [ B C ]] => A+B+C but [ Compose [ Compose [ A B ] C ] does
    // not become A+B+C
    static constexpr int32_t kComposeID = (int32_t) BuiltInCodeSnippetID::kCompose;
    SkASSERT(dict->getEntry(kComposeID)->fNumChildren == 2);
    const bool prettyCompose =
            // single-lined Compose block
            id == kComposeID && !multiline &&
            // that doesn't have an inner Compose child
            currentIndex < SkTo<int>(keyData.size()) && keyData[currentIndex] != kComposeID;

    if (prettyCompose) {
        SkASSERT(entry->fNumChildren == 2); // ordered [inner, outer]
        currentIndex = key_to_string(caps, str, dict, keyData, currentIndex, indent);
        str->append("+");
        return key_to_string(caps, str, dict, keyData, currentIndex, indent);
    }

    str->append(entry->fName);

    if (entry->storesSamplerDescData()) {
        SkASSERT(currentIndex + 1 < SkTo<int>(keyData.size()));

        // If an entry stores data, then the next key value reports the quantity of key indices that
        // are used to house the data for this snippet. This way, we know how many indices to
        // iterate over in order to capture the snippet's data before we may encounter another
        // snippet ID.
        // For example:
        // [snippetId using 2 indices worth of data] [2] [dataValue0] [dataValue1] [next snippet ID]
        const int dataIndexCount = PaintParamsKey::EncodeDataSize(keyData[currentIndex++]);
        // Printing keys is assumed to be done for keys that were built by PaintParamsKeyBuilder or
        // already validated for deserialization, so we consider them trusted.
        SkASSERT(dataIndexCount >= 0 && dataIndexCount <= PaintParamsKey::kEmbeddedDataSizeLimit);
        SkASSERT(currentIndex + dataIndexCount < SkTo<int>(keyData.size()));

        bool descriptiveFormAppended = false;
        if (dataIndexCount == 0) {
            // We shorten the string for the common case of no extra data.
            str->append("(0)");
            descriptiveFormAppended = true;
        } else {
            SkASSERTF(dataIndexCount == 2 || dataIndexCount == 3, "count %d", dataIndexCount);

            // Attempt to append the sampler data as human-readable YCbCr information
            SamplerDesc s(keyData[currentIndex],
                          keyData[currentIndex+1],
                          /* extFormatMSB= */ dataIndexCount == 3 ? keyData[currentIndex+2] : 0);

            if (s.isImmutable()) {
                std::string tmp = caps->toString(s.immutableSamplerInfo());
                if (!tmp.empty()) {
                    str->append("(");
                    str->append(tmp);
                    str->append(")");
                    descriptiveFormAppended = true;
                }
            }
        }

        if (!descriptiveFormAppended) {
            append_as_base64(str, { &keyData[currentIndex], dataIndexCount });
        }

        // Increment current index past the indices which contain data
        currentIndex += dataIndexCount;
    }

    if (entry->fNumChildren > 0) {
        if (multiline) {
            str->append(":\n");
            indent++;
        } else {
            str->append("[");
        }

        for (int i = 0; i < entry->fNumChildren; ++i) {
            if (i > 0) {
                str->append(", ");
            }
            currentIndex = key_to_string(caps, str, dict, keyData, currentIndex, indent);
        }

        if (!multiline) {
            str->append("]");
        }
    }

    if (multiline && entry->fNumChildren == 0) {
        str->append("\n");
    }
    return currentIndex;
}

SkString PaintParamsKey::toString(const Caps* caps,
                                  const ShaderCodeDictionary* dict) const {
    SkString str;
    const int keySize = SkTo<int>(fData.size());
    for (int currentIndex = 0; currentIndex < keySize;) {
        currentIndex = key_to_string(caps, &str, dict, fData, currentIndex, /*indent=*/-1);
        str.append(" ");
    }
    return str.isEmpty() ? SkString("(empty)") : str;
}

#ifdef SK_DEBUG

void PaintParamsKey::dump(const Caps* caps,
                          const ShaderCodeDictionary* dict,
                          UniquePaintParamsID id) const {
    const int keySize = SkTo<int>(fData.size());

    SkDebugf("--------------------------------------\n");
    SkDebugf("PaintParamsKey %u (keySize: %d): ", id.asUInt(), keySize);
    const int32_t* data = fData.data();
    for (int i = 0; i < keySize; ++i) {
        SkDebugf("%x ", SK_UNSAFE_TODO((uint32_t)data[i]));
    }
    SkDebugf("\n");

    int currentIndex = 0;
    while (currentIndex < keySize) {
        SkString nodeStr;
        currentIndex = key_to_string(caps, &nodeStr, dict, fData, currentIndex, /*indent=*/1);
        SkDebugf("%s", nodeStr.c_str());
    }
}

#endif // SK_DEBUG

namespace {

// check a single block and, recursively, all its children
[[nodiscard]] bool is_block_valid(const ShaderCodeDictionary* dict,
                                  SkSpan<const int32_t> keyData,
                                  int* currentIndex) {
    if (*currentIndex >= SkTo<int>(keyData.size())) {
        return false;
    }

    int32_t id = keyData[(*currentIndex)++];
    if (id >= kBuiltInCodeSnippetIDCount &&
        !SkKnownRuntimeEffects::IsSkiaKnownRuntimeEffect(id) &&
        !dict->isUserDefinedKnownRuntimeEffect(id)) {
        return false;
    }

    auto entry = dict->getEntry(id);
    if (!entry) {
        return false;
    }

    if (entry->storesSamplerDescData()) {
        if (*currentIndex >= SkTo<int>(keyData.size())) {
            return false;
        }

        int dataLength = keyData[(*currentIndex)++];
        // This `dataLength` is untrusted, so check that it doesn't overflow EncodeDataSize
        // and that matches expectations of a valid length (i.e. it started out negative and is
        // now positive and less than the key data limit).
        if (dataLength >= 0 ||
            dataLength < PaintParamsKey::EncodeDataSize(PaintParamsKey::kEmbeddedDataSizeLimit)) {
            // Would not produce a valid size after decoding
            return false;
        }
        dataLength = PaintParamsKey::EncodeDataSize(dataLength);
        SkASSERT(dataLength >= 0 && dataLength <= PaintParamsKey::kEmbeddedDataSizeLimit);
        if (*currentIndex + dataLength > SkTo<int>(keyData.size())) {
            return false;
        }

        *currentIndex += dataLength;
    }

    if (entry->fNumChildren > 0) {
        for (int i = 0; i < entry->fNumChildren; ++i) {
            if (!is_block_valid(dict, keyData, currentIndex)) {
                return false;
            }
        }
    }

    return true;
}

} // anonymous namespace


bool PaintParamsKey::isSerializable(const ShaderCodeDictionary* dict) const {
    const int keySize = SkTo<int>(fData.size());

    int currentIndex = 0;
    while (currentIndex < keySize) {
        // Ensure root nodes have their headers set properly, if not the key is malformed.
        if (fData[currentIndex++] >= 0) {
            return false;
        }
        if (!is_block_valid(dict, fData, &currentIndex)) {
            return false;
        }
    }

    return true;
}

} // namespace skgpu::graphite