src/ui/testing/util/screenshot_helper.cc

// 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.

#include "screenshot_helper.h"

#include <lib/zx/vmar.h>
#include <zircon/status.h>

#include <memory>
#include <unordered_map>

namespace ui_testing {
namespace {

constexpr uint64_t kBytesPerPixel = 4;

}  // namespace

Screenshot::Screenshot(const zx::vmo& screenshot_vmo, uint64_t width, uint64_t height, int rotation)
    : width_(width), height_(height) {
  if (rotation == 90 || rotation == 270) {
    std::swap(width_, height_);
  }

  // Populate |screenshot_| from |screenshot_vmo|.
  uint64_t vmo_size;
  screenshot_vmo.get_prop_content_size(&vmo_size);

  FX_CHECK(vmo_size == kBytesPerPixel * width_ * height_);

  uint8_t* vmo_host = nullptr;
  auto status = zx::vmar::root_self()->map(ZX_VM_PERM_READ, /*vmar_offset*/ 0, screenshot_vmo,
                                           /*vmo_offset*/ 0, vmo_size,
                                           reinterpret_cast<uintptr_t*>(&vmo_host));

  FX_CHECK(status == ZX_OK);

  ExtractScreenshotFromVMO(vmo_host);

  // Unmap the pointer.
  uintptr_t address = reinterpret_cast<uintptr_t>(vmo_host);
  status = zx::vmar::root_self()->unmap(address, vmo_size);
  FX_CHECK(status == ZX_OK);
}

std::ostream& operator<<(std::ostream& stream, const Pixel& pixel) {
  return stream << "{Pixel:"
                << " r:" << static_cast<unsigned int>(pixel.red)
                << " g:" << static_cast<unsigned int>(pixel.green)
                << " b:" << static_cast<unsigned int>(pixel.blue)
                << " a:" << static_cast<unsigned int>(pixel.alpha) << "}";
}

Pixel Screenshot::GetPixelAt(uint64_t x, uint64_t y) const {
  FX_CHECK(x >= 0 && x < width_ && y >= 0 && y < height_) << "Index out of bounds";
  return screenshot_[y][x];
}

std::map<Pixel, uint32_t> Screenshot::Histogram() const {
  std::map<Pixel, uint32_t> histogram;
  FX_CHECK(screenshot_.size() == height_ && screenshot_[0].size() == width_);

  for (size_t i = 0; i < height_; i++) {
    for (size_t j = 0; j < width_; j++) {
      histogram[screenshot_[i][j]]++;
    }
  }

  return histogram;
}

float Screenshot::ComputeSimilarity(const Screenshot& other) const {
  if (width() != other.width() || height() != other.height())
    return 0;

  uint64_t num_matching_pixels = 0;
  for (uint64_t x = 0; x < width(); ++x) {
    for (uint64_t y = 0; y < height(); ++y) {
      if (GetPixelAt(x, y) == other.GetPixelAt(x, y))
        ++num_matching_pixels;
    }
  }
  return 100.f * static_cast<float>(num_matching_pixels) / static_cast<float>((width() * height()));
}

void Screenshot::ExtractScreenshotFromVMO(uint8_t* screenshot_vmo) {
  FX_CHECK(screenshot_vmo);

  for (size_t i = 0; i < height_; i++) {
    // The head index of the ith row in the screenshot is |i* width_* KbytesPerPixel|.
    screenshot_.push_back(GetPixelsInRow(screenshot_vmo, i));
  }
}

std::vector<Pixel> Screenshot::GetPixelsInRow(uint8_t* screenshot_vmo, size_t row_index) {
  std::vector<Pixel> row;

  for (size_t col_idx = 0; col_idx < static_cast<size_t>(width_ * kBytesPerPixel);
       col_idx += kBytesPerPixel) {
    // Each row in the screenshot has |kBytesPerPixel * width_| elements. Therefore in order to
    // reach the first pixel of the |row_index| row, we have to jump |row_index * width_ *
    // kBytesPerPixel| positions.
    auto pixel_start_index = row_index * width_ * kBytesPerPixel;

    // Every |kBytesPerPixel| bytes represents the BGRA values of a pixel. Skip |kBytesPerPixel|
    // bytes
    // to get to the BGRA values of the next pixel. Each row in a screenshot has |kBytesPerPixel *
    // width_| bytes of data.
    // Example:-
    // auto data = TakeScreenshot();
    // data[0-3] -> RGBA of pixel 0.
    // data[4-7] -> RGBA pf pixel 1.
    row.emplace_back(screenshot_vmo[pixel_start_index + col_idx],
                     screenshot_vmo[pixel_start_index + col_idx + 1],
                     screenshot_vmo[pixel_start_index + col_idx + 2],
                     screenshot_vmo[pixel_start_index + col_idx + 3]);
  }

  return row;
}

}  // namespace ui_testing