chromium/third_party/mediapipe/src/mediapipe/gpu/webgpu/webgpu_texture_buffer.h

#ifndef MEDIAPIPE_GPU_WEBGPU_WEBGPU_TEXTURE_BUFFER_H_
#define MEDIAPIPE_GPU_WEBGPU_WEBGPU_TEXTURE_BUFFER_H_

#include <webgpu/webgpu_cpp.h>

#include <cstdint>
#include <memory>
#include <utility>

#include "absl/status/statusor.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/gpu_buffer_storage.h"
#include "mediapipe/gpu/multi_pool.h"
#include "mediapipe/gpu/reusable_pool.h"
#include "mediapipe/gpu/webgpu/webgpu_service.h"
#include "mediapipe/gpu/webgpu/webgpu_texture_view.h"

namespace mediapipe {

class WebGpuTextureBuffer
    : public internal::GpuBufferStorageImpl<
          WebGpuTextureBuffer, internal::ViewProvider<WebGpuTextureView>> {
 public:
  static std::unique_ptr<WebGpuTextureBuffer> Create(
      const wgpu::Device& device, uint32_t width, uint32_t height,
      GpuBufferFormat format = GpuBufferFormat::kRGBA32);

  static std::unique_ptr<WebGpuTextureBuffer> Create(uint32_t width,
                                                     uint32_t height,
                                                     GpuBufferFormat format);

  WebGpuTextureBuffer(wgpu::Texture texture, uint32_t width, uint32_t height,
                      GpuBufferFormat format)
      : texture_(std::move(texture)),
        width_(width),
        height_(height),
        format_(format) {}

  ~WebGpuTextureBuffer() override { texture_.Destroy(); }

  int width() const override { return width_; }
  int height() const override { return height_; }

  WebGpuTextureView GetReadView(
      internal::types<WebGpuTextureView>) const override;

  WebGpuTextureView GetWriteView(internal::types<WebGpuTextureView>) override;

  GpuBufferFormat format() const override { return format_; }

  // TODO: implement.
  void Reuse() {}

 private:
  wgpu::Texture texture_;
  const uint32_t width_ = 0;
  const uint32_t height_ = 0;
  const GpuBufferFormat format_;
};

class Canvas;
Canvas& RenderToWebGpuCanvas(const std::shared_ptr<WebGpuTextureBuffer>& input);

class WebGpuTextureBufferPool : public ReusablePool<WebGpuTextureBuffer> {
 public:
  static std::shared_ptr<WebGpuTextureBufferPool> Create(
      const wgpu::Device& device, const internal::GpuBufferSpec& spec,
      const MultiPoolOptions& options) {
    return std::make_shared<WebGpuTextureBufferPool>(device, spec, options);
  }

  static absl::StatusOr<std::shared_ptr<WebGpuTextureBuffer>>
  CreateBufferWithoutPool(const internal::GpuBufferSpec& spec);

  WebGpuTextureBufferPool(const wgpu::Device& device,
                          const internal::GpuBufferSpec& spec,
                          const MultiPoolOptions& options)
      : ReusablePool<WebGpuTextureBuffer>(
            [this] {
              return WebGpuTextureBuffer::Create(device_, spec_.width,
                                                 spec_.height, spec_.format);
            },
            options),
        device_(device),
        spec_(spec) {}

 private:
  const wgpu::Device device_;
  const internal::GpuBufferSpec& spec_;
};

class WebGpuTextureBufferMultiPool
    : public MultiPool<WebGpuTextureBufferPool, internal::GpuBufferSpec,
                       std::shared_ptr<WebGpuTextureBuffer>> {
 public:
  explicit WebGpuTextureBufferMultiPool(
      const wgpu::Device& device,
      const MultiPoolOptions& options = kDefaultMultiPoolOptions)
      : MultiPool(
            [this](const internal::GpuBufferSpec& spec,
                   const MultiPoolOptions& options) {
              return WebGpuTextureBufferPool::Create(device_, spec, options);
            },
            options),
        device_(device) {}

  // TODO: This and RenderToWebGpuCanvas use kBGRA32 as defaults, while
  // the rest of our code defaults to kRGBA32 instead. Investigate whether this
  // is intentional or typos leading to accidental over-sharding of pools.
  absl::StatusOr<std::shared_ptr<WebGpuTextureBuffer>> GetBuffer(
      int width, int height,
      GpuBufferFormat format = GpuBufferFormat::kBGRA32) {
    return Get(internal::GpuBufferSpec(width, height, format));
  }

 private:
  const wgpu::Device device_;
};

inline const WebGpuDeviceAttachment<WebGpuTextureBufferMultiPool>
    kWebGpuTexturePool{[](wgpu::Device& device) {
      return WebGpuDeviceAttachment<WebGpuTextureBufferMultiPool>::MakePtr(
          device);
    }};

}  // namespace mediapipe

#endif  // MEDIAPIPE_GPU_WEBGPU_WEBGPU_TEXTURE_BUFFER_H_