src/gpu/graphite/StorageContext.h

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

#ifndef skgpu_graphite_StorageContext_DEFINED
#define skgpu_graphite_StorageContext_DEFINED

#include "include/core/SkColorType.h"
#include "include/private/SkTDArray.h"
#include "src/core/SkTHash.h"
#include "src/gpu/graphite/BufferManager.h"
#include "src/gpu/graphite/Caps.h"
#include "src/gpu/graphite/ResourceTypes.h"
#include "src/gpu/graphite/TextureFormat.h"

#include <cstdint>
#include <limits>
#include <optional>
#include <utility>
#include <variant>

class SkGradientBaseShader;

namespace skgpu::graphite {

class DrawContext;
class Recorder;
class TextureProxy;

using StorageContextResult = std::variant<BindBufferInfo, sk_sp<TextureProxy>>;

class StorageContext {
public:
    // TODO (thomsmit): These are currently here for convenience, but in the future these should be
    // caps derived. (e.g. if we include a fallback to RGBA8unorm)
    static constexpr int           kTexelBytes                   = 16;
    static constexpr SkColorType   kColorType                    = kRGBA_F32_SkColorType;
    static constexpr TextureFormat kTextureFormat                = TextureFormat::kRGBA32F;

    StorageContext(const ResourceBindingRequirements&, bool storageBufferSupport);
    StorageContext(int maxFallbackTextureSize,
                   bool storageBufferSupport,
                   int maxStorageBufferBytes =
                           ResourceBindingRequirements::kDefaultMaxStorageBufferBytes);
    ~StorageContext();

    // Resets cached gradient and vertex data. Should only occur at "organic" flush time
    // (e.g. canvas/recorder flush).
    void resetCache();

    // Stub, will call more than resetCache() in the future
    void reset() {
        this->resetCache();
    }

    std::pair<float*, int> allocateGradientData(int numStops, const SkGradientBaseShader* shader);

    // Used to determine the LCM across the stride/alignment requirements of all draws in a
    // drawPass, but does not write any data to the storageContext. Instead, at snapDrawPass() time,
    // the final LCM is used to align both the cached data in finalizePrecachedStorageData() and
    // the incoming appendVertices() stream. Should only be called prior to snapDrawPass().
    void recordAlignment(size_t stride, size_t align);

    // Appends `count * stride` bytes of vertex data into the CPU-side copy, aligned according to
    // `stride` and `align`. Returns the byte offset within the eventual storage buffer.
    uint32_t appendVertices(const void* data, size_t count, size_t stride, size_t align = 1);

    // Called at the beginning of snapping a draw pass. It pads the precached data size to the
    // running LCM so that subsequent appended data begins at a valid aligned offset. Does not
    // perform GPU uploads.
    void finalizePrecachedStorageData();

    // Called after snapping the draw pass. It creates either a single mapped buffer or a fallback
    // texture containing the concatenated cached and append data, then clears the per-pass vertex
    // data.
    std::optional<StorageContextResult> finalize(Recorder*, DrawContext*);

    bool isEmpty() const { return fGradientCache.isEmpty() && fVertexData.empty(); }

    SkDEBUGCODE(int size() const { return fGradientCache.fGradientData.size(); })
    SkDEBUGCODE(int vertexSize() const { return fVertexData.size(); })

#if defined(GPU_TEST_UTILS)
    uint32_t runningLCM() const { return fRunningLCM; }
    int maxStorageBytes() const { return fMaxStorageBytes; }
#endif

private:
    BindBufferInfo finalizeStorageBuffer(Recorder* recorder);
    sk_sp<TextureProxy> finalizeTexture(Recorder* recorder, DrawContext* drawContext);

    struct GradientCache {
        // 1M stops = SkAlign4(1M) + 4M = 5M floats (~20 MB, or 2048x640 RGBA32F texels).
        static constexpr int kMaxGradientStops = 1024 * 1024;

        skia_private::THashMap<const SkGradientBaseShader*, int> fLocalGradientOffsetCache;
        SkTDArray<float> fGradientData;
        size_t fGradientDataSize = 0;

        void reset();

        bool isEmpty() const { return fGradientData.empty(); }
    };

    SkDEBUGCODE(bool fFinalized = false;)

    GradientCache fGradientCache;

    SkTDArray<char> fVertexData;
    uint32_t fRunningLCM;
    int fMaxFallbackTextureSize;
    int fMaxStorageBytes;
    bool fStorageBufferSupport;
};

} // namespace skgpu::graphite

#endif // skgpu_graphite_StorageContext_DEFINED