// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <cstdint>
#include <cstring>
#include <vector>

#include "iree/hal/cts/util/test_base.h"

namespace iree::hal::cts {

using ::testing::ContainerEq;

namespace {
constexpr iree_device_size_t kDefaultAllocationSize = 1024;
}  // namespace

// Tests for buffer mapping (IREE_HAL_BUFFER_USAGE_MAPPING) support and
// for `iree_hal_buffer_*` functions which require buffer mapping.
//
// Most tests write into a buffer then read the (possibly partial) contents
// back via iree_hal_buffer_map_read. Test failures may indicate issues in
// either the writing or the reading code path.
//
// Where applicable, tests are organized in increasing complexity:
//   - write to full buffer
//   - write with an offset and length
//   - write into a subspan of a buffer
class BufferMappingTest : public CtsTestBase<> {
 protected:
  // Allocates a buffer with HOST_VISIBLE + MAPPING usage for mapping tests.
  // This differs from the base class helpers which use DISPATCH_STORAGE.
  void AllocateUninitializedBuffer(iree_device_size_t buffer_size,
                                   iree_hal_buffer_t** out_buffer) {
    iree_hal_buffer_params_t params = {0};
    params.type =
        IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE;
    params.usage =
        IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING;
    iree_hal_buffer_t* buffer = NULL;
    IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer(device_allocator_, params,
                                                      buffer_size, &buffer));
    *out_buffer = buffer;
  }
};

TEST_P(BufferMappingTest, AllocatorSupportsBufferMapping) {
  iree_hal_buffer_params_t params = {0};
  params.type = IREE_HAL_MEMORY_TYPE_HOST_VISIBLE;
  params.usage = IREE_HAL_BUFFER_USAGE_MAPPING;
  iree_device_size_t allocation_size = 0;
  iree_hal_buffer_compatibility_t compatibility =
      iree_hal_allocator_query_buffer_compatibility(device_allocator_, params,
                                                    kDefaultAllocationSize,
                                                    &params, &allocation_size);
  EXPECT_TRUE(iree_all_bits_set(compatibility,
                                IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE));

  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(allocation_size, &buffer);

  EXPECT_TRUE(
      iree_all_bits_set(iree_hal_buffer_memory_type(buffer), params.type));
  EXPECT_TRUE(
      iree_all_bits_set(iree_hal_buffer_allowed_usage(buffer), params.usage));
  EXPECT_GE(iree_hal_buffer_allocation_size(buffer), allocation_size);

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, ZeroWholeBuffer) {
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(kDefaultAllocationSize, &buffer);

  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer, /*byte_offset=*/0,
                                          IREE_HAL_WHOLE_BUFFER));

  std::vector<uint8_t> actual_data(kDefaultAllocationSize);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer(kDefaultAllocationSize, 0);
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, ZeroWithOffset) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  // Fill the entire buffer then zero only a segment of it.
  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(iree_hal_buffer_map_fill(buffer, /*byte_offset=*/0,
                                          IREE_HAL_WHOLE_BUFFER, &fill_value,
                                          sizeof(fill_value)));
  IREE_ASSERT_OK(
      iree_hal_buffer_map_zero(buffer, /*byte_offset=*/4, /*byte_length=*/8));

  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0xFF, 0xFF, 0xFF, 0xFF};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, ZeroSubspan) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(iree_hal_buffer_map_fill(buffer, /*byte_offset=*/0,
                                          IREE_HAL_WHOLE_BUFFER, &fill_value,
                                          sizeof(fill_value)));

  iree_device_size_t subspan_length = 8;
  iree_hal_buffer_t* buffer_subspan = NULL;
  IREE_ASSERT_OK(
      iree_hal_buffer_subspan(buffer, /*byte_offset=*/4, subspan_length,
                              iree_allocator_system(), &buffer_subspan));

  // Zero part of the subspan.
  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer_subspan, /*byte_offset=*/4,
                                          /*byte_length=*/4));

  // Check the full buffer.
  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0xFF, 0xFF, 0xFF, 0xFF};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  // Also check the subspan.
  std::vector<uint8_t> actual_data_subspan(subspan_length);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(buffer_subspan, /*source_offset=*/0,
                                          actual_data_subspan.data(),
                                          actual_data_subspan.size()));
  std::vector<uint8_t> reference_buffer_subspan{0xFF, 0xFF, 0xFF, 0xFF,  //
                                                0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data_subspan, ContainerEq(reference_buffer_subspan));

  iree_hal_buffer_release(buffer_subspan);
  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, FillEmpty) {
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(kDefaultAllocationSize, &buffer);

  // Zero the buffer then "fill" 0 bytes with a different pattern.
  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer, 0, IREE_HAL_WHOLE_BUFFER));
  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer, /*byte_offset=*/0,
                               /*byte_length=*/0,
                               /*pattern=*/&fill_value,
                               /*pattern_length=*/sizeof(fill_value)));

  // Buffer should still be all zeroes.
  std::vector<uint8_t> actual_data(kDefaultAllocationSize);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer(kDefaultAllocationSize, 0);
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, FillWholeBuffer) {
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(kDefaultAllocationSize, &buffer);

  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer, /*byte_offset=*/0,
                               /*byte_length=*/IREE_HAL_WHOLE_BUFFER,
                               /*pattern=*/&fill_value,
                               /*pattern_length=*/sizeof(fill_value)));

  std::vector<uint8_t> actual_data(kDefaultAllocationSize);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer(kDefaultAllocationSize, fill_value);
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, FillWithOffset) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer, 0, IREE_HAL_WHOLE_BUFFER));
  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer, /*byte_offset=*/4,
                               /*byte_length=*/8,
                               /*pattern=*/&fill_value,
                               /*pattern_length=*/sizeof(fill_value)));

  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, FillSubspan) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer, 0, IREE_HAL_WHOLE_BUFFER));

  iree_device_size_t subspan_length = 8;
  iree_hal_buffer_t* buffer_subspan = NULL;
  IREE_ASSERT_OK(
      iree_hal_buffer_subspan(buffer, /*byte_offset=*/4, subspan_length,
                              iree_allocator_system(), &buffer_subspan));

  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer_subspan, /*byte_offset=*/4,
                               /*byte_length=*/4,
                               /*pattern=*/&fill_value,
                               /*pattern_length=*/sizeof(fill_value)));

  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  std::vector<uint8_t> actual_data_subspan(subspan_length);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(buffer_subspan, /*source_offset=*/0,
                                          actual_data_subspan.data(),
                                          actual_data_subspan.size()));
  std::vector<uint8_t> reference_buffer_subspan{0x00, 0x00, 0x00, 0x00,  //
                                                0xFF, 0xFF, 0xFF, 0xFF};
  EXPECT_THAT(actual_data_subspan, ContainerEq(reference_buffer_subspan));

  iree_hal_buffer_release(buffer_subspan);
  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, ReadData) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  IREE_ASSERT_OK(
      iree_hal_buffer_map_zero(buffer, /*byte_offset=*/0, /*byte_length=*/8));
  uint8_t fill_value = 0xFF;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer, /*byte_offset=*/8,
                               /*byte_length=*/8,
                               /*pattern=*/&fill_value,
                               /*pattern_length=*/sizeof(fill_value)));

  // Read the entire buffer.
  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0xFF, 0xFF, 0xFF, 0xFF,  //
                                        0xFF, 0xFF, 0xFF, 0xFF};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  // Read only a segment of the buffer.
  std::vector<uint8_t> actual_data_offset(8);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(buffer, /*source_offset=*/4,
                                          actual_data_offset.data(),
                                          actual_data_offset.size()));
  std::vector<uint8_t> reference_buffer_offset{0x00, 0x00, 0x00, 0x00,  //
                                               0xFF, 0xFF, 0xFF, 0xFF};
  EXPECT_THAT(actual_data_offset, ContainerEq(reference_buffer_offset));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, ReadDataSubspan) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  // Fill segments with distinct values.
  uint8_t value = 0xAA;
  IREE_ASSERT_OK(iree_hal_buffer_map_fill(buffer, 0, 4, &value, sizeof(value)));
  value = 0xBB;
  IREE_ASSERT_OK(iree_hal_buffer_map_fill(buffer, 4, 4, &value, sizeof(value)));
  value = 0xCC;
  IREE_ASSERT_OK(iree_hal_buffer_map_fill(buffer, 8, 4, &value, sizeof(value)));
  value = 0xDD;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer, 12, 4, &value, sizeof(value)));

  iree_device_size_t subspan_length = 8;
  iree_hal_buffer_t* buffer_subspan = NULL;
  IREE_ASSERT_OK(
      iree_hal_buffer_subspan(buffer, /*byte_offset=*/4, subspan_length,
                              iree_allocator_system(), &buffer_subspan));

  // Read the entire subspan.
  std::vector<uint8_t> actual_data(subspan_length);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(buffer_subspan, /*source_offset=*/0,
                                          actual_data.data(),
                                          actual_data.size()));
  std::vector<uint8_t> reference_buffer{0xBB, 0xBB, 0xBB, 0xBB,  //
                                        0xCC, 0xCC, 0xCC, 0xCC};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  // Read only a segment of the subspan.
  std::vector<uint8_t> actual_data_offset(4);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(buffer_subspan, /*source_offset=*/4,
                                          actual_data_offset.data(),
                                          actual_data_offset.size()));
  std::vector<uint8_t> reference_buffer_offset{0xCC, 0xCC, 0xCC, 0xCC};
  EXPECT_THAT(actual_data_offset, ContainerEq(reference_buffer_offset));

  iree_hal_buffer_release(buffer_subspan);
  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, WriteDataWholeBuffer) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  uint8_t fill_value = 0xFF;
  std::vector<uint8_t> reference_buffer(buffer_size, fill_value);
  IREE_ASSERT_OK(iree_hal_buffer_map_write(buffer, /*target_offset=*/0,
                                           reference_buffer.data(),
                                           reference_buffer.size()));

  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, WriteDataWithOffset) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer, 0, IREE_HAL_WHOLE_BUFFER));

  std::vector<uint8_t> fill_buffer{0x11, 0x22, 0x33, 0x44,  //
                                   0x55, 0x66, 0x77, 0x88};
  IREE_ASSERT_OK(iree_hal_buffer_map_write(
      buffer, /*target_offset=*/4, fill_buffer.data(), fill_buffer.size()));

  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0x11, 0x22, 0x33, 0x44,  //
                                        0x55, 0x66, 0x77, 0x88,  //
                                        0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, WriteDataSubspan) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  IREE_ASSERT_OK(iree_hal_buffer_map_zero(buffer, 0, IREE_HAL_WHOLE_BUFFER));

  iree_device_size_t subspan_length = 8;
  iree_hal_buffer_t* buffer_subspan = NULL;
  IREE_ASSERT_OK(
      iree_hal_buffer_subspan(buffer, /*byte_offset=*/4, subspan_length,
                              iree_allocator_system(), &buffer_subspan));

  std::vector<uint8_t> fill_buffer{0x11, 0x22, 0x33, 0x44};
  IREE_ASSERT_OK(iree_hal_buffer_map_write(buffer_subspan, /*target_offset=*/4,
                                           fill_buffer.data(),
                                           fill_buffer.size()));

  // Check the full buffer.
  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0x11, 0x22, 0x33, 0x44,  //
                                        0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  // Also check the subspan.
  std::vector<uint8_t> actual_data_subspan(subspan_length);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(buffer_subspan, /*source_offset=*/0,
                                          actual_data_subspan.data(),
                                          actual_data_subspan.size()));
  std::vector<uint8_t> reference_buffer_subspan{0x00, 0x00, 0x00, 0x00,  //
                                                0x11, 0x22, 0x33, 0x44};
  EXPECT_THAT(actual_data_subspan, ContainerEq(reference_buffer_subspan));

  iree_hal_buffer_release(buffer_subspan);
  iree_hal_buffer_release(buffer);
}

TEST_P(BufferMappingTest, CopyData) {
  iree_hal_buffer_t* buffer_a = NULL;
  iree_hal_buffer_t* buffer_b = NULL;
  AllocateUninitializedBuffer(kDefaultAllocationSize, &buffer_a);
  AllocateUninitializedBuffer(kDefaultAllocationSize, &buffer_b);

  uint8_t fill_value = 0x07;
  IREE_ASSERT_OK(
      iree_hal_buffer_map_fill(buffer_a, /*byte_offset=*/0,
                               /*byte_length=*/kDefaultAllocationSize,
                               /*pattern=*/&fill_value,
                               /*pattern_length=*/sizeof(fill_value)));
  IREE_ASSERT_OK(iree_hal_buffer_map_copy(
      /*source_buffer=*/buffer_a,
      /*source_offset=*/0, /*target_buffer=*/buffer_b, /*target_offset=*/0,
      /*data_length=*/kDefaultAllocationSize));

  std::vector<uint8_t> reference_buffer(kDefaultAllocationSize, fill_value);
  std::vector<uint8_t> actual_data(kDefaultAllocationSize);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer_b, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_release(buffer_a);
  iree_hal_buffer_release(buffer_b);
}

// Maps a buffer range for reading from device -> host.
TEST_P(BufferMappingTest, MapRangeRead) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  uint8_t fill_value = 0xEF;
  IREE_ASSERT_OK(iree_hal_buffer_map_fill(buffer, /*byte_offset=*/0,
                                          IREE_HAL_WHOLE_BUFFER, &fill_value,
                                          sizeof(fill_value)));

  iree_hal_buffer_mapping_t mapping;
  IREE_ASSERT_OK(iree_hal_buffer_map_range(
      buffer, IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ,
      /*byte_offset=*/0, /*byte_length=*/buffer_size, &mapping));
  EXPECT_EQ(buffer, mapping.buffer);
  EXPECT_GE(mapping.contents.data_length, (iree_host_size_t)buffer_size);

  std::vector<uint8_t> reference_buffer(buffer_size, fill_value);
  std::vector<uint8_t> mapping_data(
      mapping.contents.data,
      mapping.contents.data + mapping.contents.data_length);
  EXPECT_THAT(mapping_data, ContainerEq(reference_buffer));

  iree_hal_buffer_unmap_range(&mapping);
  iree_hal_buffer_release(buffer);
}

// Maps a buffer range for writing from host -> device.
TEST_P(BufferMappingTest, MapRangeWrite) {
  iree_device_size_t buffer_size = 16;
  iree_hal_buffer_t* buffer = NULL;
  AllocateUninitializedBuffer(buffer_size, &buffer);

  iree_hal_buffer_mapping_t mapping;
  IREE_ASSERT_OK(iree_hal_buffer_map_range(
      buffer, IREE_HAL_MAPPING_MODE_SCOPED,
      IREE_HAL_MEMORY_ACCESS_DISCARD_WRITE,
      /*byte_offset=*/0, /*byte_length=*/buffer_size, &mapping));
  EXPECT_EQ(buffer, mapping.buffer);
  EXPECT_GE(mapping.contents.data_length, (iree_host_size_t)buffer_size);

  // Write into the mapped memory, flush for device access, then read back.
  uint8_t fill_value = 0x12;
  std::memset(mapping.contents.data, fill_value, buffer_size);
  IREE_ASSERT_OK(
      iree_hal_buffer_mapping_flush_range(&mapping, /*byte_offset=*/0,
                                          /*byte_length=*/buffer_size));
  std::vector<uint8_t> actual_data(buffer_size);
  IREE_ASSERT_OK(iree_hal_buffer_map_read(
      buffer, /*source_offset=*/0, actual_data.data(), actual_data.size()));
  std::vector<uint8_t> reference_buffer(buffer_size, fill_value);
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_buffer_unmap_range(&mapping);
  iree_hal_buffer_release(buffer);
}

CTS_REGISTER_TEST_SUITE(BufferMappingTest);

}  // namespace iree::hal::cts
