// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef GPU_COMMAND_BUFFER_SERVICE_SHARED_IMAGE_COMPOUND_IMAGE_BACKING_H_
#define GPU_COMMAND_BUFFER_SERVICE_SHARED_IMAGE_COMPOUND_IMAGE_BACKING_H_

#include <vector>

#include "base/containers/enum_set.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/scoped_refptr.h"
#include "gpu/command_buffer/common/mailbox.h"
#include "gpu/command_buffer/common/shared_image_info.h"
#include "gpu/command_buffer/common/shared_image_usage.h"
#include "gpu/command_buffer/service/memory_tracking.h"
#include "gpu/command_buffer/service/shared_image/shared_image_backing.h"
#include "gpu/command_buffer/service/shared_image/shared_image_factory.h"
#include "gpu/command_buffer/service/shared_image/shared_image_manager.h"
#include "gpu/command_buffer/service/shared_image/shared_image_representation.h"
#include "gpu/command_buffer/service/shared_image/shared_memory_image_backing.h"
#include "gpu/command_buffer/service/shared_memory_region_wrapper.h"
#include "gpu/gpu_gles2_export.h"
#include "gpu/ipc/common/surface_handle.h"
#include "ui/gfx/color_space.h"
#include "ui/gfx/geometry/size.h"

namespace gpu {

class D3DImageBackingFactoryTest;
class SharedImageBackingFactory;
class SharedImageCopyManager;
class SharedImageFactory;
class SharedImageFactoryRef;

// TODO(kylechar): Merge with OzoneImageBacking::AccessStream enum.
//
// LINT.IfChange(SharedImageAccessStream)
enum class SharedImageAccessStream {
  kSkia,
  kOverlay,
  kGL,
  kDawn,
  kDawnBuffer,
  kMemory,
  kVaapi,
  kWebNNTensor,
  kVulkan,
  kMaxValue = kVulkan
};
// LINT.ThenChange(//tools/metrics/histograms/metadata/gpu/enums.xml:SharedImageAccessStream)

GPU_GLES2_EXPORT std::ostream& operator<<(
    std::ostream& os,
    SharedImageAccessStream access_stream);

// Used to represent what access streams a backing can be used for.
using AccessStreamSet = base::EnumSet<SharedImageAccessStream,
                                      SharedImageAccessStream::kSkia,
                                      SharedImageAccessStream::kVulkan>;

// A CompoundImageBacking is a specialized container that manages one or more
// underlying SharedImageBacking instances of different types. It serves as a
// bridge to allow a single SharedImage to be backed by multiple memory types
// (e.g., CPU memory and GPU memory OR multiple GPU memory) and provides the
// necessary interoperability (interop) to synchronize data between them as
// usage requirements change.
//
// ----------------------------------
// Core Architecture and Interop
// ----------------------------------
//
// Initial Creation: It creates one or more backings during initial setup based
// on the provided SharedImageUsageSet.
//
// Dynamic GPU Allocation & Data Sync: If a client requests a new usage that
// the current backings cannot satisfy, CompoundImageBacking can create a new
// GPU backing at runtime. Upon creation, the latest data from the existing
// backings is efficiently and automatically copied to this new backing to
// ensure continuity.
//
// Automated Interop: The container manages the lifecycle and data consistency
// between its members. If a client writes to one backing (e.g., CPU) and later
// requires a different representation (e.g., GPU), the CompoundImageBacking
// handles the synchronization logic internally.
//
// Dynamic Management: To optimize memory, backings can be deleted dynamically
// based on usage and memory pressure, provided that at least one backing
// remains active at all times.
//
// (Note: Dynamic allocation/de-allocations are currently disabled by default
// but is a core architectural feature).
//
// ---------------------------------------
// Critical Constraints and Assumptions
// ---------------------------------------
//
// Shared Memory Limit: A CompoundImageBacking can never have more than one
// SharedMemoryImageBacking.
//
// Mappable Backing Placement: Any type of mappable backing (including
// SharedMemoryImageBacking) must be created during the initial setup and is
// never allocated dynamically. These are strictly stored as the first element
// (elements_[0]).
// TODO(crbug.com/471036798): Add CHECK to ensure that mappable backings are
// never created dynamically.
//
// Persistence: The container must always maintain at least one backing to
// ensure the SharedImage remains valid during dynamic memory adjustments.
//
// Memory Upload Requirements: When combining a shared memory backing with a
// hardware-based GPU backing:
// 1.The GPU backing must implement UploadFromMemory() and ReadbackToMemory() to
// copy to/from shared memory backing.
// 2.The GPU backing must not have its own separate shared memory segment, as it
// relies on the primary shared memory backing for data transfers.
//
// ---------------------------------------
// Thread Safety and Multi-threading
// ---------------------------------------
//
// CompoundImageBacking serves as a thread-safe container when multi-threaded
// access is required (e.g., in DrDC or WebView environments).
//
// 1. Requirement Detection: Thread safety is no longer solely determined based
// on initial shared image usage. Evolution of CSI to a dynamic-CSI have relaxed
// the requirements for clients to specify all shared image usages during
// creation time. Hence capturing the thread safety requirements based on
// initial usage set is no longer enough. The backing automatically determines
// its thread safety requirement during creation by consulting both client usage
// flags and global environmental flags (DrDC/WebView) provided by the factory.
//
// 2. Container Protection: If created as thread-safe, it enables an internal
// lock (via the base class) that protects its metadata, such as the
// |elements_| vector, from concurrent access during Produce*() calls.
//
// 3. Underlying Backings thread safety: The underlying backings/elements
// themselves may or may not be thread-safe. When the container is thread-safe
// (e.g. for DrDC or WebView), any dynamically allocated backing that is not
// thread-safe is treated as "transient." It is owned by the representation and
// destroyed after use, rather than being added to the container's permanent
// elements. This ensures that non-thread-safe backings are never accessed
// concurrently across threads. Upon completion of access, any writes to a
// transient backing are synchronized back to the container's permanent
// backings. Note that indiscriminately making all backings thread-safe
// (especially for internal images like Render Passes) could lead to performance
// regressions, so this transient approach provides safety without unnecessary
// overhead.
class GPU_GLES2_EXPORT CompoundImageBacking
    : public ClearTrackingSharedImageBacking {
 public:
  using CreateBackingCallback =
      base::OnceCallback<void(std::unique_ptr<SharedImageBacking>&)>;

  static bool IsValidSharedMemoryFormat(const gfx::Size& size,
                                        viz::SharedImageFormat format);

  // Remove the SCANOUT flag if |kAllowShmOverlays|.
  static SharedImageUsageSet GetGpuSharedImageUsage(SharedImageUsageSet usage);

  // Creates a backing that contains a shared memory backing and GPU backing
  // provided by `shared_image_factory` based on `usage`. Eventually, instead of
  // creating a shm+gpu backing, this method will have various strategy to
  // allocate different combination of backings based on the `usage`.
  static std::unique_ptr<SharedImageBacking> Create(
      SharedImageFactory* shared_image_factory,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      const Mailbox& mailbox,
      gfx::GpuMemoryBufferHandle handle,
      const SharedImageInfo& si_info);

  // Creates a backing that contains a shared memory backing and GPU backing
  // provided by `shared_image_factory` based on `usage`. Eventually, instead of
  // creating a shm+gpu backing, this method will have various strategy to
  // allocate different combination of backings based on the `usage`.
  // We additionally pass a |buffer_usage| parameter here in order to create a
  // CPU mappable by creating a shared memory handle.
  // TODO(crbug.com/40276878): Remove this method once we figure out the mapping
  // between SharedImageUsage and BufferUsage and no longer need to use
  // BufferUsage.
  static std::unique_ptr<SharedImageBacking> Create(
      SharedImageFactory* shared_image_factory,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      const Mailbox& mailbox,
      const SharedImageInfo& si_info,
      gfx::BufferUsage buffer_usage);

  // Wraps a backing in a CompoundImageBacking. This is used to enable
  // CompoundImageBacking as the default, where it serves as the sole backing.
  // To achieve this, SharedImageFactory creates the standard backing and then
  // wraps it using this method.
  // TODO(crbug.com/448962784): Once CompoundImageBacking is fully enabled by
  // default, it will be directly creating underlying backing itself and
  // SharedImageFactory will be refactored to move most of backing creation
  // logic inside CompoundImageBacking.
  static std::unique_ptr<SharedImageBacking> WrapExternalBacking(
      SharedImageFactory* shared_image_factory,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      std::unique_ptr<SharedImageBacking> backing);

  ~CompoundImageBacking() override;

  // Called by wrapped representations before access. This will update
  // the backing that is going to be accessed if most recent pixels are in
  // a different backing. Returns false if a required content sync failed,
  // in which case the caller must not proceed with the access.
  [[nodiscard]] bool NotifyBeginAccess(SharedImageBacking* backing,
                                       RepresentationAccessMode mode,
                                       SharedImageAccessStream stream);

  // Called by wrapped representations during EndAccess(). This will update the
  // CompoundImageBacking's clear rect with the accessed backing's clear rect it
  // the access was a write access.
  void NotifyEndAccess(SharedImageBacking* backing,
                       RepresentationAccessMode mode);

  // SharedImageBacking implementation.
  void OnContextLost() override;
  void SetPurgeable(bool purgeable) override;
  bool IsPurgeable() const override;
  SharedImageBackingType GetType() const override;
  void Update(std::unique_ptr<gfx::GpuFence> in_fence) override;
  bool CopyToGpuMemoryBuffer() override;
  void CopyToGpuMemoryBufferAsync(
      base::OnceCallback<void(bool)> callback) override;
  gfx::Rect ClearedRect() const override;

  // CompoundImageBacking now supports partial clear for upcoming use
  // cases as it evolves. The cleared rect is now tracked on the compound
  // backing as well as its underlying backings.
  // Some important things to note that is :
  // 1. When a CompoundImageBacking is backed by a single gpu backing, clear
  // rect of CompoundImageBacking will track and reflect clear rect of the
  // underlying backing.
  // 2. When CompoundImageBacking contains more than 1 gpu backing, clear rect
  // of the CompoundImageBacking will track and reflect clear rect of the
  // most recently written backing. Note that when a read is performed from a
  // stale backing, the latest backing content as well as its clear rect will
  // be copied into it.
  // 3. Anytime a copy is performend between backings, the src backing's cleared
  // rect will be xfered to the dst backing.
  // 4. If there is a shm backing, entire CompoundImageBacking as well all the
  // created gpu backings will be marked as cleared always.
  void SetClearedRect(const gfx::Rect& cleared_rect) override;

  void OnAddSecondaryReference() override;

  // CompoundImageBacking is registered as the primary backing while creating a
  // SharedImageRepresentationFactoryRef whereas the underlying
  // elements/backings it holds are not. Since the MarkForDestruction() method
  // in SharedImageRepresentationFactoryRef only runs for primary backing,
  // CompoundImageBacking needs to propagate this call to all its elements.
  void MarkForDestruction() override;
  gfx::GpuMemoryBufferHandle GetGpuMemoryBufferHandle() override;
  scoped_refptr<gfx::NativePixmap> GetNativePixmap() override;

#if BUILDFLAG(IS_ANDROID)
  std::optional<VulkanYCbCrInfo> GetVkCbCrInfo(
      SharedContextState* context_state) override;
#endif

 protected:
  // SharedImageBacking implementation.
  std::unique_ptr<DawnImageRepresentation> ProduceDawn(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker,
      const wgpu::Device& device,
      wgpu::BackendType backend_type,
      std::vector<wgpu::TextureFormat> view_formats,
      scoped_refptr<SharedContextState> context_state) override;
  std::unique_ptr<DawnBufferRepresentation> ProduceDawnBuffer(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker,
      const wgpu::Device& device,
      wgpu::BackendType backend_type,
      scoped_refptr<SharedContextState> context_state) override;
  std::unique_ptr<GLTextureImageRepresentation> ProduceGLTexture(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker) override;
  std::unique_ptr<GLTexturePassthroughImageRepresentation>
  ProduceGLTexturePassthrough(SharedImageManager* manager,
                              MemoryTypeTracker* tracker) override;
  std::unique_ptr<SkiaGaneshImageRepresentation> ProduceSkiaGanesh(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker,
      scoped_refptr<SharedContextState> context_state) override;
  std::unique_ptr<SkiaGraphiteImageRepresentation> ProduceSkiaGraphite(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker,
      scoped_refptr<SharedContextState> context_state) override;
  std::unique_ptr<OverlayImageRepresentation> ProduceOverlay(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker) override;
  std::unique_ptr<WebNNTensorRepresentation> ProduceWebNNTensor(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker) override;
  std::unique_ptr<MemoryImageRepresentation> ProduceMemory(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker) override;
  std::unique_ptr<VideoImageRepresentation> ProduceVideo(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker,
      VideoDevice device) override;

#if BUILDFLAG(ENABLE_VULKAN)
  std::unique_ptr<VulkanImageRepresentation> ProduceVulkan(
      SharedImageManager* manager,
      MemoryTypeTracker* tracker,
      gpu::VulkanDeviceQueue* vulkan_device_queue,
      gpu::VulkanImplementation& vulkan_impl,
      bool needs_detiling) override;
#endif

 private:
  friend class CompoundImageBackingTest;
  friend class D3DImageBackingFactoryTest;

  // Holds one element, aka SharedImageBacking and related information, that
  // makes up the compound.
  struct ElementHolder {
   public:
    ElementHolder();
    ElementHolder(const ElementHolder& other) = delete;
    ElementHolder& operator=(const ElementHolder& other) = delete;
    ElementHolder(ElementHolder&& other);
    ElementHolder& operator=(ElementHolder&& other);
    ~ElementHolder();

    // Will invoke `create_callback` to create backing if
    // required.
    void CreateBackingIfNecessary();

    // Returns the backing. Will call `CreateBackingIfNecessary()`.
    SharedImageBacking* GetBacking();

    AccessStreamSet access_streams;
    uint64_t content_id_ = 0;

    CreateBackingCallback create_callback;
    std::unique_ptr<SharedImageBacking> backing;
  };

  // Creates a backing that contains a shared memory backing and GPU backing
  // provided by `gpu_backing_factory`.
  static std::unique_ptr<SharedImageBacking> CreateSharedMemoryForTesting(
      SharedImageBackingFactory* gpu_backing_factory,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      const Mailbox& mailbox,
      gfx::GpuMemoryBufferHandle handle,
      const SharedImageInfo& si_info);

  // Creates a backing that contains a shared memory backing and GPU backing
  // provided by `gpu_backing_factory`. We additionally pass a |buffer_usage|
  // parameter here in order to create a CPU mappable by creating a shared
  // memory handle.
  // TODO(crbug.com/40276878): Remove this method once we figure out the mapping
  // between SharedImageUsage and BufferUsage and no longer need to use
  // BufferUsage.
  static std::unique_ptr<SharedImageBacking> CreateSharedMemoryForTesting(
      SharedImageBackingFactory* gpu_backing_factory,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      const Mailbox& mailbox,
      const SharedImageInfo& si_info,
      gfx::BufferUsage buffer_usage);

  CompoundImageBacking(
      const Mailbox& mailbox,
      const SharedImageInfo& si_info,
      std::unique_ptr<SharedImageBacking> shm_backing,
      scoped_refptr<SharedImageFactoryRef> shared_image_factory,
      base::WeakPtr<SharedImageBackingFactory> gpu_backing_factory,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      std::optional<gfx::BufferUsage> buffer_usage = std::nullopt);

  CompoundImageBacking(
      std::optional<gfx::BufferUsage> buffer_usage,
      std::unique_ptr<SharedImageBacking> backing,
      scoped_refptr<SharedImageCopyManager> copy_manager,
      scoped_refptr<SharedImageFactoryRef> shared_image_factory);

  base::trace_event::MemoryAllocatorDump* OnMemoryDump(
      const std::string& dump_name,
      base::trace_event::MemoryAllocatorDumpGuid client_guid,
      base::trace_event::ProcessMemoryDump* pmd,
      uint64_t client_tracing_id) override;

  // Returns a SkPixmap for shared memory backing.
  const std::vector<SkPixmap>& GetSharedMemoryPixmaps();

  // Returns the shared memory element used for access stream
  // SharedImageAccessStream::kMemory. There can be only 1 shared memory element
  // at most.
  ElementHolder& GetShmElement();

  // Gets the element corresponding to the backing.
  ElementHolder* GetElement(const SharedImageBacking* backing);

  // Finds the element which has the most recent data/content irrespective of
  // the stream. There could be multiple elements which has the most recent
  // data. This method finds the first element which has most recent data.
  ElementHolder* GetElementWithLatestContent();

  // Gets or allocates a backing for a given `stream` and `params`.
  // It finds a backing that supports the given `stream` and is compatible
  // with the context information in `params` by calling `SupportsAccess()`.
  // If a compatible backing is present, it will either return the
  // backing with the latest content OR will return any supported backing (the
  // first one it finds).
  // If no backing is found, then it will allocate an appropriate backing which
  // can support the `stream`.
  // If CSI is thread-safe and the dynamically allocated backing is not thread
  // safe, then the dynamically created backing will be moved to
  // `out_transient_backing` and not added to `elements_`.
  SharedImageBacking* GetOrAllocateBacking(
      SharedImageAccessStream stream,
      const AccessParams& params,
      std::unique_ptr<SharedImageBacking>& out_transient_backing);

  // Returns the gpu backing from the list of |element_| which has a shm and a
  // gpu backing.
  SharedImageBacking* GetGpuBacking();

  bool HasLatestContent(ElementHolder& element);

  // Sets the element used for `stream` as having the latest content. If
  // `write_access` is true then only that element has the latest content.
  void SetLatestContent(SharedImageAccessStream stream, bool write_access);

  // Runs CreateSharedImage() on `factory` and stores the result in `backing`.
  // If successful this will update the estimated size of compound backing.
  // In multi-threading environment, caller should ensure that the
  // SharedImageBackingFactory is alive while this method is being executed.
  void CreateBackingFromBackingFactory(
      SharedImageBackingFactory* backing_factory,
      std::string debug_label,
      SharedImageUsageSet usage,
      viz::SharedImageFormat backing_format,
      std::unique_ptr<SharedImageBacking>& backing);

  // Method used for lazy backing creation. It will use
  // `shared_image_factory_` to safely find the correct factory and create
  // the backing while holding the factory lock.
  void LazyCreateBacking(SharedImageBackingType factory_type,
                         base::WeakPtr<SharedImageBackingFactory> test_factory,
                         SharedImageUsageSet usage,
                         std::string debug_label,
                         std::unique_ptr<SharedImageBacking>& backing);

  void OnCopyToGpuMemoryBufferComplete(bool success);

  // Computes the thread safety requirement for the backing based on the
  // usage and environment flags from the factory.
  static bool ComputeIsThreadSafe(SharedImageFactoryRef* factory_ref,
                                  SharedImageUsageSet usage);

  // Helper to record GPU.SharedImage.BackingType UMA histogram.
  static void RecordBackingTypeUMA(SharedImageBackingType backing_type);

  // This is required for CompoundImageBacking to be able to query an
  // appropriate SharedImageBackingFactory dynamically based on clients
  // required usage(Produce*) which typically happens after the backing
  // creation time. This uses a thread-safe ref-holder to safely access the
  // factory from any thread.
  scoped_refptr<SharedImageFactoryRef> shared_image_factory_;

  // 64-bit so it never wraps back to a stale element's content id in practice.
  uint64_t latest_content_id_ GUARDED_BY(lock_) = 1;

  // Holds all of the "element" backings that make up this compound backing. For
  // each there is a backing, set of streams and tracking for latest content.
  //
  // It's expected that for each access stream there is exactly one element used
  // to access it. Note that it's possible the backing for a given access stream
  // can't actually support that type of usage, in which case the backing will
  // be null or the ProduceX() call will just fail.
  // As of now, CompoundImageBacking only has 2 backings,i.e., 1 shm and 1 gpu
  // backing. In future, it will evolve into a dynamic CompoundImageBacking
  // where it can have any number of gpu backings and at most 1 cpu backing.
  std::vector<ElementHolder> elements_ GUARDED_BY(lock_);

  base::OnceCallback<void(bool)> pending_copy_to_gmb_callback_;
  scoped_refptr<SharedImageCopyManager> copy_manager_;
  bool has_shm_backing_ = false;
  // Tracks the maximum number of SharedImageBacking elements that were
  // allocated within this CompoundImageBacking instance during its lifetime.
  // This value is recorded in a UMA histogram in the destructor.
  size_t max_elements_allocated_ = 0;
};

}  // namespace gpu

#endif  // GPU_COMMAND_BUFFER_SERVICE_SHARED_IMAGE_COMPOUND_IMAGE_BACKING_H_