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

#include "cc/tiles/software_image_decode_cache_utils.h"

#include <algorithm>
#include <sstream>
#include <utility>

#include "base/atomic_sequence_num.h"
#include "base/functional/callback_helpers.h"
#include "base/hash/hash.h"
#include "base/memory/discardable_memory_allocator.h"
#include "base/metrics/histogram_macros.h"
#include "base/process/memory.h"
#include "base/trace_event/trace_event.h"
#include "cc/paint/paint_flags.h"
#include "cc/paint/tone_map_util.h"
#include "cc/tiles/mipmap_util.h"
#include "skia/ext/geometry.h"
#include "third_party/skia/include/core/SkColorSpace.h"
#include "third_party/skia/include/core/SkImage.h"
#include "ui/gfx/geometry/skia_conversions.h"

namespace cc {
namespace {
// If the size of the original sized image breaches kMemoryRatioToSubrect but we
// don't need to scale the image, consider caching only the needed subrect.
// The second part that much be true is that we cache only the needed subrect if
// the total size needed for the subrect is at most kMemoryRatioToSubrect *
// (size needed for the full original image).
// Note that at least one of the dimensions has to be at least
// kMinDimensionToSubrect before an image can breach the threshold.
const size_t kMemoryThresholdToSubrect = 64 * 1024 * 1024;
const int kMinDimensionToSubrect = 4 * 1024;
const float kMemoryRatioToSubrect = 0.5f;

// Tracing ID sequence for use in CacheEntry.
base::AtomicSequenceNumber g_next_tracing_id_;

gfx::Rect GetSrcRect(const DrawImage& image) {
  const SkIRect& src_rect = image.src_rect();
  int x = std::max(0, src_rect.x());
  int y = std::max(0, src_rect.y());
  int right = std::min(image.paint_image().width(), src_rect.right());
  int bottom = std::min(image.paint_image().height(), src_rect.bottom());
  if (x >= right || y >= bottom)
    return gfx::Rect();
  return gfx::Rect(x, y, right - x, bottom - y);
}

// Given `base_rect` in the base image with the specified dimensions, return the
// corresponding rectangle in the gainmap image.
SkRect ComputeGainmapRect(SkISize base_image_dimensions,
                          SkISize gain_image_dimensions,
                          SkRect base_rect) {
  SkRect base_image_rect =
      SkRect::MakeSize(SkSize::Make(base_image_dimensions));
  SkRect gain_image_rect =
      SkRect::MakeSize(SkSize::Make(gain_image_dimensions));
  return skia::ScaleSkRectProportional(gain_image_rect, base_image_rect,
                                       base_rect);
}

// Allocate `memory` as discardable, and back `image` with that memory.
// On failure, sets `image` to nullptr.
void AllocateDiscardableSkImage(
    const SkImageInfo& info,
    const SkImageInfo& gainmap_info,
    base::OnceClosure on_no_memory,
    std::unique_ptr<base::DiscardableMemory>& memory,
    sk_sp<SkImage>& image,
    sk_sp<SkImage>& gainmap_image) {
  // Initialize the output images to be empty.
  image = nullptr;
  gainmap_image = nullptr;

  TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("cc.debug"), "AllocateDiscardable");
  size_t size = info.minRowBytes() * info.height() +
                gainmap_info.minRowBytes() * gainmap_info.height();
  memory = base::DiscardableMemoryAllocator::GetInstance()
               ->AllocateLockedDiscardableMemoryWithRetryOrDie(
                   size, std::move(on_no_memory));
  if (!memory->data()) {
    return;
  }
  // SAFETY: The base::DiscardableMemory functionality should provide a safe
  // interface. Fixing it is tracked in https://crbug.com/40586428. This use is
  // safe because we are using the same size parameter here as in the allocation
  // a few lines earlier.
  UNSAFE_BUFFERS(
      base::span<uint8_t> memory_as_span(memory->data_as<uint8_t>(), size));
  auto gainmap_image_span = memory_as_span.subspan(
      info.minRowBytes() * info.height(),
      gainmap_info.minRowBytes() * gainmap_info.height());

  // A more robust scheme would be to unlock `memory` when `image` falls out of
  // scope.
  SkPixmap pixmap(info, memory_as_span.data(), info.minRowBytes());
  image = SkImages::RasterFromPixmap(
      pixmap, [](const void* pixels, void* context) {}, nullptr);
  if (!gainmap_image_span.empty()) {
    SkPixmap gainmap_pixmap(gainmap_info, gainmap_image_span.data(),
                            gainmap_info.minRowBytes());
    gainmap_image = SkImages::RasterFromPixmap(
        gainmap_pixmap, [](const void* pixels, void* context) {}, nullptr);
  }
}

}  // namespace

// static
std::unique_ptr<SoftwareImageDecodeCacheUtils::CacheEntry>
SoftwareImageDecodeCacheUtils::DoDecodeImage(
    const CacheKey& key,
    const PaintImage& paint_image,
    SkColorType color_type,
    PaintImage::GeneratorClientId client_id,
    base::OnceClosure on_no_memory) {
  const SkISize target_size =
      SkISize::Make(key.target_size().width(), key.target_size().height());
  DCHECK(target_size == paint_image.GetSupportedDecodeSize(target_size));

  SkImageInfo target_info =
      SkImageInfo::Make(target_size, color_type, kPremul_SkAlphaType,
                        key.target_color_params().color_space.ToSkColorSpace());

  SkImageInfo target_gainmap_info;
  if (paint_image.HasGainmapInfo()) {
    target_gainmap_info = SkImageInfo::Make(
        paint_image.GetSupportedDecodeSize(target_size, AuxImage::kGainmap),
        color_type, kPremul_SkAlphaType);
  }

  sk_sp<SkImage> target_image;
  sk_sp<SkImage> target_gainmap_image;
  std::unique_ptr<base::DiscardableMemory> target_pixels;
  AllocateDiscardableSkImage(target_info, target_gainmap_info,
                             std::move(on_no_memory), target_pixels,
                             target_image, target_gainmap_image);
  if (!target_image) {
    return nullptr;
  }

  SkPixmap target_pixmap;
  target_image->peekPixels(&target_pixmap);
  // Temporary workaround for migrating HLG and PQ color spaces. The round-trip
  // through gfx::ColorSpace destroys the distinction between HLG and HLGish,
  // and PQ and PQish. Ensure that the decode of these spaces does no
  // conversion. https://issues.skia.org/issues/420956739
  if (key.target_color_params().color_space.GetTransferID() ==
          gfx::ColorSpace::TransferID::PQ ||
      key.target_color_params().color_space.GetTransferID() ==
          gfx::ColorSpace::TransferID::HLG) {
    target_pixmap.setColorSpace(paint_image.GetSkImageInfo().refColorSpace());
  }

  TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("cc.debug"),
               "SoftwareImageDecodeCacheUtils::DoDecodeImage - "
               "decode");
  bool result = paint_image.Decode(target_pixmap, key.frame_key().frame_index(),
                                   AuxImage::kDefault, client_id);
  if (!result) {
    target_pixels->Unlock();
    return nullptr;
  }

  if (target_gainmap_image) {
    SkPixmap target_gainmap_pixmap;
    target_gainmap_image->peekPixels(&target_gainmap_pixmap);
    bool gainmap_result =
        paint_image.Decode(target_gainmap_pixmap, key.frame_key().frame_index(),
                           AuxImage::kGainmap, client_id);
    // If the gainmap fails to decode, just pretend that it wasn't there. Do
    // not fail the base image decode because of problems in the gainmap.
    if (!gainmap_result) {
      target_gainmap_image = nullptr;
    }
  }

  return std::make_unique<CacheEntry>(
      target_image, target_gainmap_image, paint_image.GetHDRMetadata(),
      std::move(target_pixels), SkSize::Make(0, 0));
}

// static
std::unique_ptr<SoftwareImageDecodeCacheUtils::CacheEntry>
SoftwareImageDecodeCacheUtils::GenerateCacheEntryFromCandidate(
    const CacheKey& key,
    const DecodedDrawImage& candidate_image,
    bool needs_extract_subset,
    SkColorType color_type) {
  TRACE_EVENT0(
      TRACE_DISABLED_BY_DEFAULT("cc.debug"),
      "SoftwareImageDecodeCacheUtils::GenerateCacheEntryFromCandidate");

  // Let `decoded_pixmap` be the candidate image's pixels that we will be
  // copying and potentially scaling. Let `decoded_gainmap_image` be
  // the candidate gainmap image (if it exists).
  SkPixmap decoded_pixmap;
  bool result = candidate_image.image()->peekPixels(&decoded_pixmap);
  DCHECK(result) << key.ToString();
  sk_sp<SkImage> decoded_gainmap_image = candidate_image.gainmap_image();

  // Compute the actual source rect of `decoded_pixmap` that will be used,
  // and let `decoded_pixmap_sub_rect` an SkPixmap with just that rect.
  SkIRect src_rect = SkIRect::MakeSize(decoded_pixmap.dimensions());
  if (needs_extract_subset) {
    result = src_rect.intersect(gfx::RectToSkIRect(key.src_rect()));
    DCHECK(result) << key.ToString();
  }
  SkPixmap decoded_pixmap_sub_rect;
  result = decoded_pixmap.extractSubset(&decoded_pixmap_sub_rect, src_rect);
  DCHECK(result) << key.ToString();

  // Compute the corresponding SkRect of the gainmap image. Note that this is
  // not an SkIRect (because it doesn't necessarily at integer bounds).
  SkRect src_gainmap_rect;
  if (decoded_gainmap_image) {
    src_gainmap_rect = ComputeGainmapRect(decoded_pixmap.dimensions(),
                                          decoded_gainmap_image->dimensions(),
                                          SkRect::Make(src_rect));
  }

  SkImageInfo target_info =
      SkImageInfo::Make(gfx::SizeToSkISize(key.target_size()), color_type,
                        kPremul_SkAlphaType, decoded_pixmap.refColorSpace());
  SkImageInfo target_gainmap_info;
  if (decoded_gainmap_image) {
    // Set the target gainmap image size to the target base image's size. But,
    // don't supersample the gainmap, so take the minimum with the gainmap's
    // source rectangle size.
    SkISize target_gainmap_max_size =
        SkSize(src_gainmap_rect.width(), src_gainmap_rect.height()).toCeil();
    SkISize target_gainmap_size = SkISize::Make(
        std::min(key.target_size().width(), target_gainmap_max_size.width()),
        std::min(key.target_size().height(), target_gainmap_max_size.height()));
    target_gainmap_info =
        decoded_gainmap_image->imageInfo().makeDimensions(target_gainmap_size);
  }

  sk_sp<SkImage> target_image;
  sk_sp<SkImage> target_gainmap_image;
  std::unique_ptr<base::DiscardableMemory> target_pixels;
  // TODO(crbug.com/40095682): If this turns into a crasher, pass an actual
  // "free memory" closure.
  AllocateDiscardableSkImage(target_info, target_gainmap_info,
                             base::DoNothing(), target_pixels, target_image,
                             target_gainmap_image);
  if (!target_image) {
    return nullptr;
  }

  // Populate the pixels of `target_pixmap` from `decoded_pixmap_sub_rect`
  // by `scalePixels`. The implementation will optimize this to `readPixels` if
  // possible.
  SkPixmap target_pixmap;
  target_image->peekPixels(&target_pixmap);
  result = decoded_pixmap_sub_rect.scalePixels(
      target_pixmap, PaintFlags::FilterQualityToSkSamplingOptions(
                         PaintFlags::FilterQuality::kMedium));
  DCHECK(result) << key.ToString();

  // Populate the pixels of `target_gainmap_image` using a `drawImage` from
  // `src_gainmap_rect`.
  if (target_gainmap_image) {
    SkPixmap target_gainmap_pixmap;
    target_gainmap_image->peekPixels(&target_gainmap_pixmap);
    auto canvas = SkCanvas::MakeRasterDirect(
        target_gainmap_pixmap.info(), target_gainmap_pixmap.writable_addr(),
        target_gainmap_pixmap.rowBytes());
    canvas->drawImageRect(
        decoded_gainmap_image, src_gainmap_rect,
        SkRect::MakeSize(SkSize::Make(target_gainmap_info.dimensions())),
        SkSamplingOptions(SkFilterMode::kLinear), nullptr,
        SkCanvas::kStrict_SrcRectConstraint);
  }

  return std::make_unique<CacheEntry>(
      target_image, target_gainmap_image, candidate_image.hdr_metadata(),
      std::move(target_pixels),
      SkSize::Make(-key.src_rect().x(), -key.src_rect().y()));
}

// CacheKey --------------------------------------------------------------------
// static
SoftwareImageDecodeCacheUtils::CacheKey
SoftwareImageDecodeCacheUtils::CacheKey::FromDrawImage(const DrawImage& image,
                                                       SkColorType color_type) {
  DCHECK(!image.paint_image().IsTextureBacked());

  const auto& paint_image = image.paint_image();
  const PaintImage::FrameKey frame_key = image.frame_key();
  const PaintImage::Id stable_id = paint_image.stable_id();

  // If the image will be tone mapped or contains a gain map, then do not bake
  // the tone mapping or color conversion in at decode time. This is an explicit
  // trade-off to:
  // * Avoid re-decoding an image every time it is drawn with a different HDR
  //   headroom (improving performance when transitioning to HDR).
  // * Require running tone mapping shader at every draw (degrading performance
  //   when drawing repeatedly at the same target headroom).
  // * Maximize the code sharing between the software and GPU raster path. This
  //   is the honest motivation for making the trade-off in this direction.
  // * A better solution could be to separately cache decode and tone map
  //   results (and only ever cache one tone mapping for a given image).
  TargetColorParams target_color_params = image.target_color_params();
  target_color_params.hdr_headroom = std::nullopt;
  if (paint_image.HasGainmapInfo() ||
      ToneMapUtil::UseGlobalToneMapFilter(paint_image.color_space())) {
    if (paint_image.color_space()) {
      target_color_params.color_space =
          gfx::ColorSpace(*paint_image.color_space());
    }
  }

  const SkSize& scale = image.scale();
  // If the src_rect falls outside of the image, we need to clip it since
  // otherwise we might end up with uninitialized memory in the decode process.
  // Note that the scale is still unchanged and the target size is now a
  // function of the new src_rect.
  const gfx::Rect& src_rect = GetSrcRect(image);

  // Start with the exact target size. However, this will be adjusted below to
  // be either a mip level, the original size, or a subrect size. This is done
  // to keep memory accounting correct.
  gfx::Size target_size(
      SkScalarRoundToInt(std::abs(src_rect.width() * scale.width())),
      SkScalarRoundToInt(std::abs(src_rect.height() * scale.height())));

  // If the target size is empty, then we'll be skipping the decode anyway, so
  // the filter quality doesn't matter. Early out instead.
  if (target_size.IsEmpty()) {
    return CacheKey(frame_key, stable_id, kSubrectAndScale, false,
                    paint_image.may_be_lcp_candidate(), src_rect, target_size,
                    target_color_params);
  }

  ProcessingType type = kOriginal;
  bool is_nearest_neighbor =
      image.filter_quality() == PaintFlags::FilterQuality::kNone;
  int mip_level = MipMapUtil::GetLevelForSize(src_rect.size(), target_size);
  // If any of the following conditions hold, then use at most low filter
  // quality and adjust the target size to match the original image:
  // - Quality is none: We need a pixelated image, so we can't upgrade it.
  // - Mip level is 0: The required mip is the original image, so just use low
  //   filter quality.
  // - Matrix is not decomposable: There's perspective on this image and we
  //   can't determine the size, so use the original.
  if (is_nearest_neighbor || mip_level == 0 ||
      !image.matrix_is_decomposable()) {
    type = kOriginal;
    // Update the size to be the original image size.
    target_size = gfx::Size(paint_image.width(), paint_image.height());
  } else {
    type = kSubrectAndScale;
    // Update the target size to be a mip level size.
    target_size = MipMapUtil::GetSizeForLevel(src_rect.size(), mip_level);
  }

  // If the original image is large, we might want to do a subrect instead if
  // the subrect would be kMemoryRatioToSubrect times smaller.
  if (type == kOriginal && (paint_image.width() >= kMinDimensionToSubrect ||
                            paint_image.height() >= kMinDimensionToSubrect)) {
    base::CheckedNumeric<size_t> checked_original_size = 4u;
    checked_original_size *= paint_image.width();
    checked_original_size *= paint_image.height();
    size_t original_size = checked_original_size.ValueOrDefault(
        std::numeric_limits<size_t>::max());

    base::CheckedNumeric<size_t> checked_src_rect_size = 4u;
    checked_src_rect_size *= src_rect.width();
    checked_src_rect_size *= src_rect.height();
    size_t src_rect_size = checked_src_rect_size.ValueOrDefault(
        std::numeric_limits<size_t>::max());

    // If the sizes are such that we get good savings by subrecting, then do
    // that. Also update the target size to be the src rect size since that's
    // the rect we want to use.
    if (original_size > kMemoryThresholdToSubrect &&
        src_rect_size <= original_size * kMemoryRatioToSubrect) {
      type = kSubrectOriginal;
      target_size = src_rect.size();
    }
  }

  return CacheKey(frame_key, stable_id, type, is_nearest_neighbor,
                  image.paint_image().may_be_lcp_candidate(), src_rect,
                  target_size, target_color_params);
}

SoftwareImageDecodeCacheUtils::CacheKey::CacheKey(
    PaintImage::FrameKey frame_key,
    PaintImage::Id stable_id,
    ProcessingType type,
    bool is_nearest_neighbor,
    bool may_be_lcp_candidate,
    const gfx::Rect& src_rect,
    const gfx::Size& target_size,
    const TargetColorParams& target_color_params)
    : frame_key_(frame_key),
      stable_id_(stable_id),
      type_(type),
      is_nearest_neighbor_(is_nearest_neighbor),
      may_be_lcp_candidate_(may_be_lcp_candidate),
      src_rect_(src_rect),
      target_size_(target_size),
      target_color_params_(target_color_params) {
  if (type == kOriginal) {
    hash_ = frame_key_.hash();
  } else {
    // TODO(vmpstr): This is a mess. Maybe it's faster to just search the vector
    // always (forwards or backwards to account for LRU).
    uint64_t src_rect_hash = base::HashInts(
        static_cast<uint64_t>(base::HashInts(src_rect_.x(), src_rect_.y())),
        static_cast<uint64_t>(
            base::HashInts(src_rect_.width(), src_rect_.height())));

    uint64_t target_size_hash =
        base::HashInts(target_size_.width(), target_size_.height());

    hash_ = base::HashInts(base::HashInts(src_rect_hash, target_size_hash),
                           frame_key_.hash());
  }
  // Include the target color space in the hash regardless of scaling.
  hash_ = base::HashInts(hash_, target_color_params.GetHash());
}

SoftwareImageDecodeCacheUtils::CacheKey::CacheKey(const CacheKey& other) =
    default;

SoftwareImageDecodeCacheUtils::CacheKey&
SoftwareImageDecodeCacheUtils::CacheKey::operator=(const CacheKey& other) =
    default;

std::string SoftwareImageDecodeCacheUtils::CacheKey::ToString() const {
  std::ostringstream str;
  str << "frame_key[" << frame_key_.ToString() << "]\ntype[";
  switch (type_) {
    case kOriginal:
      str << "Original";
      break;
    case kSubrectOriginal:
      str << "SubrectOriginal";
      break;
    case kSubrectAndScale:
      str << "SubrectAndScale";
      break;
  }
  str << "]\nis_nearest_neightbor[" << is_nearest_neighbor_ << "]\nsrc_rect["
      << src_rect_.ToString() << "]\ntarget_size[" << target_size_.ToString()
      << "]\ntarget_color_params[" << target_color_params_.ToString()
      << "]\nhash[" << hash_ << "]";
  return str.str();
}

// CacheEntry ------------------------------------------------------------------
SoftwareImageDecodeCacheUtils::CacheEntry::CacheEntry()
    : tracing_id_(g_next_tracing_id_.GetNext()) {}

SoftwareImageDecodeCacheUtils::CacheEntry::CacheEntry(
    sk_sp<SkImage> image,
    sk_sp<SkImage> gainmap_image,
    const std::optional<gfx::HDRMetadata>& hdr_metadata,
    std::unique_ptr<base::DiscardableMemory> in_memory,
    const SkSize& src_rect_offset)
    : is_locked(true),
      memory(std::move(in_memory)),
      image_(std::move(image)),
      gainmap_image_(std::move(gainmap_image)),
      hdr_metadata_(hdr_metadata),
      src_rect_offset_(src_rect_offset),
      tracing_id_(g_next_tracing_id_.GetNext()) {
  DCHECK(memory);
}

SoftwareImageDecodeCacheUtils::CacheEntry::~CacheEntry() {
  DCHECK(!is_locked);
}

void SoftwareImageDecodeCacheUtils::CacheEntry::MoveImageMemoryTo(
    CacheEntry* entry) {
  DCHECK(!is_budgeted);
  DCHECK_EQ(ref_count, 0);

  // Copy/move most things except budgeted and ref counts.
  entry->decode_failed = decode_failed;
  entry->is_locked = is_locked;
  is_locked = false;

  entry->memory = std::move(memory);
  entry->src_rect_offset_ = std::move(src_rect_offset_);
  entry->image_ = std::move(image_);
  entry->gainmap_image_ = std::move(gainmap_image_);
  entry->hdr_metadata_ = hdr_metadata_;
}

bool SoftwareImageDecodeCacheUtils::CacheEntry::Lock() {
  if (!memory)
    return false;

  DCHECK(!is_locked);
  bool success = memory->Lock();
  if (!success) {
    memory = nullptr;
    return false;
  }
  is_locked = true;
  return true;
}

void SoftwareImageDecodeCacheUtils::CacheEntry::Unlock() {
  if (!memory)
    return;

  DCHECK(is_locked);
  memory->Unlock();
  is_locked = false;
}

}  // namespace cc