src/ui/testing/util/screenshot_helper.h

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

#ifndef SRC_UI_TESTING_UTIL_SCREENSHOT_HELPER_H_
#define SRC_UI_TESTING_UTIL_SCREENSHOT_HELPER_H_

#include <lib/syslog/cpp/macros.h>
#include <lib/zx/vmo.h>
#include <zircon/status.h>

#include <map>
#include <vector>

namespace ui_testing {

static uint8_t linear_to_srgb(const float val) {
  // Function to convert from linear RGB to sRGB.
  // (https://en.wikipedia.org/wiki/SRGB#From_CIE_XYZ_to_sRGB)
  if (0.f <= val && val <= 0.0031308f) {
    return static_cast<uint8_t>(roundf((val * 12.92f) * 255U));
  } else {
    return static_cast<uint8_t>(roundf(((std::powf(val, 1.0f / 2.4f) * 1.055f) - 0.055f) * 255U));
  }
}

// Represents a Pixel in BGRA format.
// Uses the sRGB color space.
struct Pixel {
  uint8_t blue = 0;
  uint8_t green = 0;
  uint8_t red = 0;
  uint8_t alpha = 0;

  Pixel(uint8_t blue, uint8_t green, uint8_t red, uint8_t alpha)
      : blue(blue), green(green), red(red), alpha(alpha) {}

  static Pixel from_unorm_bgra(float blue, float green, float red, float alpha) {
    return Pixel{linear_to_srgb(blue), linear_to_srgb(green), linear_to_srgb(red),
                 static_cast<uint8_t>(roundf(alpha * 255U))};
  }

  bool operator==(const Pixel& rhs) const {
    return blue == rhs.blue && green == rhs.green && red == rhs.red && alpha == rhs.alpha;
  }

  inline bool operator!=(const Pixel& rhs) const { return !(*this == rhs); }

  bool operator<(const Pixel& other) const {
    return std::tie(blue, green, red, alpha) <
           std::tie(other.blue, other.green, other.red, other.alpha);
  }
};

std::ostream& operator<<(std::ostream& stream, const Pixel& pixel);

// Helper class to get information about a screenshot returned by
// |fuchsia.ui.composition.Screenshot| protocol.
class Screenshot {
 public:
  // BGRA format.
  inline static const Pixel kBlack = Pixel(0, 0, 0, 255);
  inline static const Pixel kBlue = Pixel(255, 0, 0, 255);
  inline static const Pixel kRed = Pixel(0, 0, 255, 255);
  inline static const Pixel kMagenta = Pixel(255, 0, 255, 255);
  inline static const Pixel kGreen = Pixel(0, 255, 0, 255);

  // Params:-
  // |screenshot_vmo| - The VMO returned by |fuchsia.ui.composition.Screenshot.Take| representing
  //                    the screenshot data in BGRA.
  // |width|, |height| - Width and height of the physical display in pixels as
  //                     returned by |fuchsia.ui.display.singleton.Info|.
  // |rotation| - The display rotation value in degrees. The width and the height of the screenshot
  //              are flipped if this value is 90 or 270 degrees as the screenshot shows how content
  //              is seen by the user.
  Screenshot(const zx::vmo& screenshot_vmo, uint64_t width, uint64_t height, int rotation);

  // Returns the |Pixel| located at (x,y) coordinates. |x| and |y| should range from [0,width_) and
  // [0,height_) respectively.
  //
  //  (0,0)________________width_____________(w-1,0)
  //      |                       |         |
  //      |                       | y       |h
  //      |          x            |         |e
  //      |-----------------------X         |i
  //      |                                 |g
  //      |                                 |h
  //      |                                 |t
  //      |_________________________________|
  //(0,h-1)           screenshot             (w-1,h-1)
  //
  // Clients should only use this function to get the pixel data.
  Pixel GetPixelAt(uint64_t x, uint64_t y) const;

  // Counts the frequencies of each color in a screenshot.
  std::map<Pixel, uint32_t> Histogram() const;

  // Returns percentage of pixels that match by comparing two screenshots. Returns 0 if the sizes of
  // the screenshots do not match.
  float ComputeSimilarity(const Screenshot& other) const;

  // Returns a 2D vector of size |height_ * width_|. Each value in the vector corresponds to a pixel
  // in the screenshot.
  std::vector<std::vector<Pixel>> screenshot() const { return screenshot_; }

  uint64_t width() const { return width_; }

  uint64_t height() const { return height_; }

 private:
  // Populates |screenshot_| by converting the linear array of bytes in |screenshot_vmo| of size |4
  // * width_ * height_| to a 2D vector of |Pixel|s of size |height_ * width_|.
  // Note: Size of each pixel is 4 bytes.
  void ExtractScreenshotFromVMO(uint8_t* screenshot_vmo);

  // Returns the |Pixel|s in the |row_index| row of the screenshot.
  std::vector<Pixel> GetPixelsInRow(uint8_t* screenshot_vmo, size_t row_index);

  uint64_t width_ = 0;
  uint64_t height_ = 0;
  std::vector<std::vector<Pixel>> screenshot_;
};

}  // namespace ui_testing

#endif  // SRC_UI_TESTING_UTIL_SCREENSHOT_HELPER_H_