Add Buffer::Get[|Const]MappedRange

This CL adds GetMappedRange reusing the existing GetMappedPointerImpl
call in dawn_native. In dawn_wire tracking is added to keep a
Buffer::mMappedData around.

Tests are added to test the result of Get[|Const]MappedRange in all
buffer states.

Bug: dawn:445
Change-Id: I3737dc4d36f31d392839952da0b5c0d10c7c8a88
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/23861
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
diff --git a/dawn.json b/dawn.json
index f313dcf..2575d18 100644
--- a/dawn.json
+++ b/dawn.json
@@ -198,6 +198,14 @@
                 ]
             },
             {
+                "name": "get mapped range",
+                "returns": "void *"
+            },
+            {
+                "name": "get const mapped range",
+                "returns": "void const *"
+            },
+            {
                 "name": "unmap"
             },
             {
@@ -1690,6 +1698,12 @@
     "void": {
         "category": "native"
     },
+    "void *": {
+        "category": "native"
+    },
+    "void const *": {
+        "category": "native"
+    },
     "uint32_t": {
         "category": "native"
     },
diff --git a/dawn_wire.json b/dawn_wire.json
index 5291c3e..f3c1c8c 100644
--- a/dawn_wire.json
+++ b/dawn_wire.json
@@ -97,6 +97,8 @@
             "BufferMapReadAsync",
             "BufferMapWriteAsync",
             "BufferSetSubData",
+            "BufferGetConstMappedRange",
+            "BufferGetMappedRange",
             "DevicePopErrorScope",
             "DeviceSetDeviceLostCallback",
             "DeviceSetUncapturedErrorCallback",
diff --git a/src/dawn_native/Buffer.cpp b/src/dawn_native/Buffer.cpp
index 21a5e43..92c2cea 100644
--- a/src/dawn_native/Buffer.cpp
+++ b/src/dawn_native/Buffer.cpp
@@ -305,6 +305,26 @@
         tracker->Track(this, mMapSerial, true);
     }
 
+    void* BufferBase::GetMappedRange() {
+        if (GetDevice()->ConsumedError(ValidateGetMappedRange(true))) {
+            return nullptr;
+        }
+        if (mStagingBuffer != nullptr) {
+            return mStagingBuffer->GetMappedPointer();
+        }
+        return GetMappedPointerImpl();
+    }
+
+    const void* BufferBase::GetConstMappedRange() {
+        if (GetDevice()->ConsumedError(ValidateGetMappedRange(false))) {
+            return nullptr;
+        }
+        if (mStagingBuffer != nullptr) {
+            return mStagingBuffer->GetMappedPointer();
+        }
+        return GetMappedPointerImpl();
+    }
+
     void BufferBase::Destroy() {
         if (IsError()) {
             // It is an error to call Destroy() on an ErrorBuffer, but we still need to reclaim the
@@ -406,6 +426,29 @@
         return {};
     }
 
+    MaybeError BufferBase::ValidateGetMappedRange(bool writable) const {
+        DAWN_TRY(GetDevice()->ValidateIsAlive());
+        DAWN_TRY(GetDevice()->ValidateObject(this));
+
+        switch (mState) {
+            // Writeable Buffer::GetMappedRange is always allowed when mapped at creation.
+            case BufferState::MappedAtCreation:
+                return {};
+
+            case BufferState::Mapped:
+                if (writable && !(mUsage & wgpu::BufferUsage::MapWrite)) {
+                    return DAWN_VALIDATION_ERROR("GetMappedRange requires the MapWrite usage");
+                }
+                return {};
+
+            case BufferState::Unmapped:
+            case BufferState::Destroyed:
+                return DAWN_VALIDATION_ERROR("Buffer is not mapped");
+            default:
+                UNREACHABLE();
+        }
+    }
+
     MaybeError BufferBase::ValidateUnmap() const {
         DAWN_TRY(GetDevice()->ValidateIsAlive());
         DAWN_TRY(GetDevice()->ValidateObject(this));
diff --git a/src/dawn_native/Buffer.h b/src/dawn_native/Buffer.h
index 8a1feb7..69e1d72 100644
--- a/src/dawn_native/Buffer.h
+++ b/src/dawn_native/Buffer.h
@@ -61,6 +61,8 @@
         void SetSubData(uint32_t start, uint32_t count, const void* data);
         void MapReadAsync(WGPUBufferMapReadCallback callback, void* userdata);
         void MapWriteAsync(WGPUBufferMapWriteCallback callback, void* userdata);
+        void* GetMappedRange();
+        const void* GetConstMappedRange();
         void Unmap();
         void Destroy();
 
@@ -94,6 +96,7 @@
 
         MaybeError ValidateMap(wgpu::BufferUsage requiredUsage,
                                WGPUBufferMapAsyncStatus* status) const;
+        MaybeError ValidateGetMappedRange(bool writable) const;
         MaybeError ValidateUnmap() const;
         MaybeError ValidateDestroy() const;
 
diff --git a/src/dawn_native/metal/BufferMTL.mm b/src/dawn_native/metal/BufferMTL.mm
index ccfd3b3..46724bf 100644
--- a/src/dawn_native/metal/BufferMTL.mm
+++ b/src/dawn_native/metal/BufferMTL.mm
@@ -120,7 +120,7 @@
     }
 
     void* Buffer::GetMappedPointerImpl() {
-        return reinterpret_cast<uint8_t*>([mMtlBuffer contents]);
+        return [mMtlBuffer contents];
     }
 
     void Buffer::UnmapImpl() {
diff --git a/src/dawn_wire/client/ApiProcs.cpp b/src/dawn_wire/client/ApiProcs.cpp
index da12171..0475284 100644
--- a/src/dawn_wire/client/ApiProcs.cpp
+++ b/src/dawn_wire/client/ApiProcs.cpp
@@ -40,6 +40,16 @@
         buffer->SetSubData(start, count, data);
     }
 
+    void* ClientHandwrittenBufferGetMappedRange(WGPUBuffer cBuffer) {
+        Buffer* buffer = reinterpret_cast<Buffer*>(cBuffer);
+        return buffer->GetMappedRange();
+    }
+
+    const void* ClientHandwrittenBufferGetConstMappedRange(WGPUBuffer cBuffer) {
+        Buffer* buffer = reinterpret_cast<Buffer*>(cBuffer);
+        return buffer->GetConstMappedRange();
+    }
+
     void ClientHandwrittenBufferUnmap(WGPUBuffer cBuffer) {
         Buffer* buffer = reinterpret_cast<Buffer*>(cBuffer);
         buffer->Unmap();
diff --git a/src/dawn_wire/client/Buffer.cpp b/src/dawn_wire/client/Buffer.cpp
index a4f5fad..910ffed 100644
--- a/src/dawn_wire/client/Buffer.cpp
+++ b/src/dawn_wire/client/Buffer.cpp
@@ -122,6 +122,7 @@
         buffer->mSize = descriptor->size;
         // Successfully created staging memory. The buffer now owns the WriteHandle.
         buffer->mWriteHandle = std::move(writeHandle);
+        buffer->mMappedData = result.data;
 
         result.buffer = reinterpret_cast<WGPUBuffer>(buffer);
 
@@ -252,8 +253,8 @@
         // second time. If, for example, buffer.Unmap() is called inside the callback.
         mRequests.erase(requestIt);
 
-        const void* mappedData = nullptr;
         size_t mappedDataLength = 0;
+        const void* mappedData = nullptr;
 
         auto GetMappedData = [&]() -> bool {
             // It is an error for the server to call the read callback when we asked for a map write
@@ -298,7 +299,8 @@
             request.readCallback(WGPUBufferMapAsyncStatus_DeviceLost, nullptr, 0, request.userdata);
             return false;
         } else {
-            request.readCallback(static_cast<WGPUBufferMapAsyncStatus>(status), mappedData,
+            mMappedData = const_cast<void*>(mappedData);
+            request.readCallback(static_cast<WGPUBufferMapAsyncStatus>(status), mMappedData,
                                  static_cast<uint64_t>(mappedDataLength), request.userdata);
             return true;
         }
@@ -316,8 +318,8 @@
         // second time. If, for example, buffer.Unmap() is called inside the callback.
         mRequests.erase(requestIt);
 
-        void* mappedData = nullptr;
         size_t mappedDataLength = 0;
+        void* mappedData = nullptr;
 
         auto GetMappedData = [&]() -> bool {
             // It is an error for the server to call the write callback when we asked for a map read
@@ -355,12 +357,31 @@
                                   request.userdata);
             return false;
         } else {
+            mMappedData = mappedData;
             request.writeCallback(static_cast<WGPUBufferMapAsyncStatus>(status), mappedData,
                                   static_cast<uint64_t>(mappedDataLength), request.userdata);
             return true;
         }
     }
 
+    void* Buffer::GetMappedRange() {
+        if (!IsMappedForWriting()) {
+            ClientDeviceInjectError(reinterpret_cast<WGPUDevice>(device), WGPUErrorType_Validation,
+                                    "Buffer needs to be mapped for writing");
+            return nullptr;
+        }
+        return mMappedData;
+    }
+
+    const void* Buffer::GetConstMappedRange() {
+        if (!IsMappedForWriting() && !IsMappedForReading()) {
+            ClientDeviceInjectError(reinterpret_cast<WGPUDevice>(device), WGPUErrorType_Validation,
+                                    "Buffer needs to be mapped");
+            return nullptr;
+        }
+        return mMappedData;
+    }
+
     void Buffer::Unmap() {
         // Invalidate the local pointer, and cancel all other in-flight requests that would
         // turn into errors anyway (you can't double map). This prevents race when the following
@@ -394,6 +415,7 @@
         } else if (mReadHandle) {
             mReadHandle = nullptr;
         }
+        mMappedData = nullptr;
         ClearMapRequests(WGPUBufferMapAsyncStatus_Unknown);
 
         BufferUnmapCmd cmd;
@@ -405,6 +427,7 @@
         // Cancel or remove all mappings
         mWriteHandle = nullptr;
         mReadHandle = nullptr;
+        mMappedData = nullptr;
         ClearMapRequests(WGPUBufferMapAsyncStatus_Unknown);
 
         BufferDestroyCmd cmd;
@@ -422,4 +445,12 @@
         device->GetClient()->SerializeCommand(cmd);
     }
 
+    bool Buffer::IsMappedForReading() const {
+        return mReadHandle != nullptr;
+    }
+
+    bool Buffer::IsMappedForWriting() const {
+        return mWriteHandle != nullptr;
+    }
+
 }}  // namespace dawn_wire::client
diff --git a/src/dawn_wire/client/Buffer.h b/src/dawn_wire/client/Buffer.h
index 96cabf6..2dd8dd5 100644
--- a/src/dawn_wire/client/Buffer.h
+++ b/src/dawn_wire/client/Buffer.h
@@ -42,12 +42,17 @@
                                     uint64_t initialDataInfoLength,
                                     const uint8_t* initialDataInfo);
         bool OnMapWriteAsyncCallback(uint32_t requestSerial, uint32_t status);
+        void* GetMappedRange();
+        const void* GetConstMappedRange();
         void Unmap();
 
         void Destroy();
 
         void SetSubData(uint64_t start, uint64_t count, const void* data);
 
+        bool IsMappedForReading() const;
+        bool IsMappedForWriting() const;
+
       private:
         // We want to defer all the validation to the server, which means we could have multiple
         // map request in flight at a single time and need to track them separately.
@@ -70,6 +75,7 @@
         // TODO(enga): Use a tagged pointer to save space.
         std::unique_ptr<MemoryTransferService::ReadHandle> mReadHandle = nullptr;
         std::unique_ptr<MemoryTransferService::WriteHandle> mWriteHandle = nullptr;
+        void* mMappedData = nullptr;
     };
 
 }}  // namespace dawn_wire::client
diff --git a/src/tests/end2end/BufferTests.cpp b/src/tests/end2end/BufferTests.cpp
index 989b919..b0ff2b8 100644
--- a/src/tests/end2end/BufferTests.cpp
+++ b/src/tests/end2end/BufferTests.cpp
@@ -119,6 +119,18 @@
     UnmapBuffer(buffer);
 }
 
+// Test the result of GetMappedRange when mapped for reading.
+TEST_P(BufferMapReadTests, GetMappedRange) {
+    wgpu::BufferDescriptor descriptor;
+    descriptor.size = 4;
+    descriptor.usage = wgpu::BufferUsage::MapRead | wgpu::BufferUsage::CopyDst;
+    wgpu::Buffer buffer = device.CreateBuffer(&descriptor);
+
+    const void* mappedData = MapReadAsyncAndWait(buffer);
+    ASSERT_EQ(mappedData, buffer.GetConstMappedRange());
+    UnmapBuffer(buffer);
+}
+
 DAWN_INSTANTIATE_TEST(BufferMapReadTests, D3D12Backend(), MetalBackend(), OpenGLBackend(), VulkanBackend());
 
 class BufferMapWriteTests : public DawnTest {
@@ -253,6 +265,19 @@
     }
 }
 
+// Test the result of GetMappedRange when mapped for writing.
+TEST_P(BufferMapWriteTests, GetMappedRange) {
+    wgpu::BufferDescriptor descriptor;
+    descriptor.size = 4;
+    descriptor.usage = wgpu::BufferUsage::MapWrite | wgpu::BufferUsage::CopySrc;
+    wgpu::Buffer buffer = device.CreateBuffer(&descriptor);
+
+    void* mappedData = MapWriteAsyncAndWait(buffer);
+    ASSERT_EQ(mappedData, buffer.GetMappedRange());
+    ASSERT_EQ(mappedData, buffer.GetConstMappedRange());
+    UnmapBuffer(buffer);
+}
+
 DAWN_INSTANTIATE_TEST(BufferMapWriteTests, D3D12Backend(), MetalBackend(), OpenGLBackend(), VulkanBackend());
 
 // TODO(enga): These tests should use the testing toggle to initialize resources to 1.
@@ -505,6 +530,19 @@
     ASSERT_NE(nullptr, result.data);
 }
 
+// Test the result of GetMappedRange when mapped at creation.
+TEST_P(CreateBufferMappedTests, GetMappedRange) {
+    wgpu::BufferDescriptor descriptor;
+    descriptor.size = 4;
+    descriptor.usage = wgpu::BufferUsage::CopyDst;
+    wgpu::CreateBufferMappedResult result;
+    result = device.CreateBufferMapped(&descriptor);
+
+    ASSERT_EQ(result.data, result.buffer.GetMappedRange());
+    ASSERT_EQ(result.data, result.buffer.GetConstMappedRange());
+    result.buffer.Unmap();
+}
+
 DAWN_INSTANTIATE_TEST(CreateBufferMappedTests,
                       D3D12Backend(),
                       D3D12Backend({}, {"use_d3d12_resource_heap_tier2"}),
diff --git a/src/tests/unittests/validation/BufferValidationTests.cpp b/src/tests/unittests/validation/BufferValidationTests.cpp
index 0f74d66..d1ececf 100644
--- a/src/tests/unittests/validation/BufferValidationTests.cpp
+++ b/src/tests/unittests/validation/BufferValidationTests.cpp
@@ -662,3 +662,166 @@
         buf.Unmap();
     }
 }
+
+// Test that it is invalid to call GetMappedRange on an unmapped buffer.
+TEST_F(BufferValidationTest, GetMappedRangeOnUnmappedBuffer) {
+    // Unmapped at creation case.
+    {
+        wgpu::BufferDescriptor desc;
+        desc.size = 4;
+        desc.usage = wgpu::BufferUsage::CopySrc;
+        wgpu::Buffer buf = device.CreateBuffer(&desc);
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+
+    // Unmapped after CreateBufferMapped case.
+    {
+        wgpu::Buffer buf = CreateBufferMapped(4, wgpu::BufferUsage::CopySrc).buffer;
+        buf.Unmap();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+
+    // Unmapped after MapReadAsync case.
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        buf.MapReadAsync(ToMockBufferMapReadCallback, nullptr);
+        EXPECT_CALL(*mockBufferMapReadCallback,
+                    Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+            .Times(1);
+        queue.Submit(0, nullptr);
+        buf.Unmap();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+
+    // Unmapped after MapWriteAsync case.
+    {
+        wgpu::Buffer buf = CreateMapWriteBuffer(4);
+        buf.MapWriteAsync(ToMockBufferMapWriteCallback, nullptr);
+        EXPECT_CALL(*mockBufferMapWriteCallback,
+                    Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+            .Times(1);
+        queue.Submit(0, nullptr);
+        buf.Unmap();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+}
+
+// Test that it is invalid to call GetMappedRange on a destroyed buffer.
+TEST_F(BufferValidationTest, GetMappedRangeOnDestroyedBuffer) {
+    // Destroyed after creation case.
+    {
+        wgpu::BufferDescriptor desc;
+        desc.size = 4;
+        desc.usage = wgpu::BufferUsage::CopySrc;
+        wgpu::Buffer buf = device.CreateBuffer(&desc);
+        buf.Destroy();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+
+    // Destroyed after CreateBufferMapped case.
+    {
+        wgpu::Buffer buf = CreateBufferMapped(4, wgpu::BufferUsage::CopySrc).buffer;
+        buf.Destroy();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+
+    // Destroyed after MapReadAsync case.
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        buf.MapReadAsync(ToMockBufferMapReadCallback, nullptr);
+        EXPECT_CALL(*mockBufferMapReadCallback,
+                    Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+            .Times(1);
+        queue.Submit(0, nullptr);
+        buf.Destroy();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+
+    // Destroyed after MapWriteAsync case.
+    {
+        wgpu::Buffer buf = CreateMapWriteBuffer(4);
+        buf.MapWriteAsync(ToMockBufferMapWriteCallback, nullptr);
+        EXPECT_CALL(*mockBufferMapWriteCallback,
+                    Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+            .Times(1);
+        queue.Submit(0, nullptr);
+        buf.Destroy();
+
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+        ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetConstMappedRange()));
+    }
+}
+
+// Test that it is invalid to call GetMappedRange on a buffer afterMapReadAsync
+TEST_F(BufferValidationTest, GetMappedRangeOnMappedForReading) {
+    wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+    buf.MapReadAsync(ToMockBufferMapReadCallback, nullptr);
+    EXPECT_CALL(*mockBufferMapReadCallback,
+                Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+        .Times(1);
+    queue.Submit(0, nullptr);
+
+    ASSERT_DEVICE_ERROR(ASSERT_EQ(nullptr, buf.GetMappedRange()));
+}
+
+// Test valid cases to call GetMappedRange on a buffer.
+TEST_F(BufferValidationTest, GetMappedRangeValidCases) {
+    // GetMappedRange after CreateBufferMapped case.
+    {
+        wgpu::CreateBufferMappedResult result = CreateBufferMapped(4, wgpu::BufferUsage::CopySrc);
+        ASSERT_NE(result.buffer.GetConstMappedRange(), nullptr);
+        ASSERT_EQ(result.buffer.GetConstMappedRange(), result.buffer.GetMappedRange());
+        ASSERT_EQ(result.buffer.GetConstMappedRange(), result.data);
+    }
+
+    // GetMappedRange after MapReadAsync case.
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        buf.MapReadAsync(ToMockBufferMapReadCallback, nullptr);
+
+        const void* mappedPointer = nullptr;
+        EXPECT_CALL(*mockBufferMapReadCallback,
+                    Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+            .WillOnce(SaveArg<1>(&mappedPointer));
+
+        queue.Submit(0, nullptr);
+
+        ASSERT_NE(buf.GetConstMappedRange(), nullptr);
+        ASSERT_EQ(buf.GetConstMappedRange(), mappedPointer);
+    }
+
+    // GetMappedRange after MapWriteAsync case.
+    {
+        wgpu::Buffer buf = CreateMapWriteBuffer(4);
+        buf.MapWriteAsync(ToMockBufferMapWriteCallback, nullptr);
+
+        const void* mappedPointer = nullptr;
+        EXPECT_CALL(*mockBufferMapWriteCallback,
+                    Call(WGPUBufferMapAsyncStatus_Success, Ne(nullptr), 4u, _))
+            .WillOnce(SaveArg<1>(&mappedPointer));
+
+        queue.Submit(0, nullptr);
+
+        ASSERT_NE(buf.GetConstMappedRange(), nullptr);
+        ASSERT_EQ(buf.GetConstMappedRange(), buf.GetMappedRange());
+        ASSERT_EQ(buf.GetConstMappedRange(), mappedPointer);
+    }
+}