[dawn][native] Spanify CommandAllocator::AllocateData

Also changes AddNullTerminatedString to not have a return value and only
write to `length`. Change the value stored in `length` to include the
required null-terminator.

Bug: 528305452, 524406299
Change-Id: I9b10a247ee11da7a2fb26cf16bcc4ec6d715d2e3
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/323676
Reviewed-by: Loko Kung <lokokung@google.com>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
diff --git a/src/dawn/native/CommandAllocator.h b/src/dawn/native/CommandAllocator.h
index 7c1599b..5ee1212 100644
--- a/src/dawn/native/CommandAllocator.h
+++ b/src/dawn/native/CommandAllocator.h
@@ -196,28 +196,25 @@
                  alignof(T) <= kMaxAllocatedCommandAlignment)
     T* Allocate(E commandId) {
         Span<std::byte> allocation = Allocate(static_cast<uint32_t>(commandId), sizeof(T));
-        if (allocation.empty()) {
-            return nullptr;
-        }
+        DAWN_CHECK(allocation.data() != nullptr);  // Crash on OOM
 
         T* result = reinterpret_cast<T*>(allocation.data());
         new (result) T;
         return result;
     }
 
-    template <typename T>
-        requires(alignof(T) <= kMaxAllocatedCommandAlignment)
-    T* AllocateData(size_t count) {
+    template <typename T, typename Index>
+        requires(alignof(T) <= kMaxAllocatedCommandAlignment &&
+                 (std::is_same_v<Index, size_t> || !std::is_integral_v<Index>))
+    ityp::span<Index, T> AllocateData(Index count) {
         Span<std::byte> allocation = AllocateData(sizeof(T) * count);
-        if (allocation.empty()) {
-            return nullptr;
-        }
+        DAWN_CHECK(allocation.data() != nullptr);  // Crash on OOM
 
-        Span<T> results = ReinterpretSpan<T>(allocation);
+        ityp::span<Index, T> results = ReinterpretSpan<T, Index>(allocation);
         for (auto& value : results) {
             new (&value) T;
         }
-        return results.data();
+        return results;
     }
 
     size_t GetCommandBlocksCount() const;
diff --git a/src/dawn/native/CommandEncoder.cpp b/src/dawn/native/CommandEncoder.cpp
index 27c078e..4fbb7ed 100644
--- a/src/dawn/native/CommandEncoder.cpp
+++ b/src/dawn/native/CommandEncoder.cpp
@@ -2139,10 +2139,10 @@
         [&](CommandAllocator* allocator) -> MaybeError {
             PushDebugGroupCmd* cmd =
                 allocator->Allocate<PushDebugGroupCmd>(Command::PushDebugGroup);
-            const char* label = AddNullTerminatedString(allocator, groupLabel, &cmd->length);
+            AddNullTerminatedString(allocator, groupLabel, &cmd->length);
 
             mDebugGroupStackSize++;
-            mEncodingContext.PushDebugGroupLabel(std::string_view(label, cmd->length));
+            mEncodingContext.PushDebugGroupLabel(groupLabel);
 
             return {};
         },
@@ -2219,8 +2219,8 @@
             cmd->offset = bufferOffset;
             cmd->size = size;
 
-            uint8_t* inlinedData = allocator->AllocateData<uint8_t>(size);
-            DAWN_UNSAFE_TODO(memcpy(inlinedData, data, size));
+            Span<uint8_t> inlinedData = allocator->AllocateData<uint8_t>(cmd->size);
+            DAWN_UNSAFE_TODO(memcpy(inlinedData.data(), data, size));
 
             mTopLevelBuffers.insert(buffer);
 
diff --git a/src/dawn/native/Commands.cpp b/src/dawn/native/Commands.cpp
index e0afbfa..64010c7 100644
--- a/src/dawn/native/Commands.cpp
+++ b/src/dawn/native/Commands.cpp
@@ -434,22 +434,21 @@
     }
 }
 
-const char* AddNullTerminatedString(CommandAllocator* allocator, StringView s, size_t* length) {
-    std::string_view view = s;
-    *length = view.length();
+void AddNullTerminatedString(CommandAllocator* allocator, std::string_view s, size_t* length) {
+    *length = s.length() + 1;
 
     // Include extra null-terminator character. The string_view may not be null-terminated. It also
     // may already have a null-terminator inside of it, in which case adding the null-terminator is
     // unnecessary. However, this is unlikely, so always include the extra character.
-    char* out = allocator->AllocateData<char>(view.length() + 1);
-    DAWN_UNSAFE_TODO(memcpy(out, view.data(), view.length()));
-    DAWN_UNSAFE_TODO(out[view.length()]) = '\0';
+    Span<char> out = allocator->AllocateData<char>(s.length() + 1);
 
-    return out;
+    // TODO(https://crbug.com/524406299): Use Span::CopyFrom.
+    std::ranges::copy(s, out.begin());
+    out[s.length()] = '\0';
 }
 
 std::string_view NextNullTerminatedString(CommandIterator* iterator, size_t length) {
-    Span<const char> data = iterator->NextData<char>(length + 1);
+    Span<const char> data = iterator->NextData<char>(length);
     DAWN_ASSERT(data[data.size() - 1] == '\0');  // The string is null-terminated.
     return {data.begin(), data.end()};
 }
diff --git a/src/dawn/native/Commands.h b/src/dawn/native/Commands.h
index 2867e49..36853f8 100644
--- a/src/dawn/native/Commands.h
+++ b/src/dawn/native/Commands.h
@@ -477,8 +477,8 @@
 // consuming the correct amount of data from the command iterator.
 void SkipCommand(CommandIterator* commands, Command type);
 
-// Helper function to copy a wgpu::StringView into a safely null-terminated C-string in commands.
-const char* AddNullTerminatedString(CommandAllocator* allocator, StringView s, size_t* length);
+// Helper function to copy a wgpu::StringView into the command stream, writing out its size.
+void AddNullTerminatedString(CommandAllocator* allocator, std::string_view s, size_t* length);
 // Mirror function that gets the same string back as a null-terminated string_view.
 std::string_view NextNullTerminatedString(CommandIterator* iterator, size_t length);
 
diff --git a/src/dawn/native/ProgrammableEncoder.cpp b/src/dawn/native/ProgrammableEncoder.cpp
index 4759b13..8d95890 100644
--- a/src/dawn/native/ProgrammableEncoder.cpp
+++ b/src/dawn/native/ProgrammableEncoder.cpp
@@ -115,10 +115,10 @@
         [&](CommandAllocator* allocator) -> MaybeError {
             PushDebugGroupCmd* cmd =
                 allocator->Allocate<PushDebugGroupCmd>(Command::PushDebugGroup);
-            const char* label = AddNullTerminatedString(allocator, groupLabel, &cmd->length);
+            AddNullTerminatedString(allocator, groupLabel, &cmd->length);
 
             mDebugGroupStackSize++;
-            mEncodingContext->PushDebugGroupLabel(std::string_view(label, cmd->length));
+            mEncodingContext->PushDebugGroupLabel(groupLabel);
 
             return {};
         },
@@ -235,14 +235,11 @@
     cmd->group = group;
     // TODO(https://crbug.com/524554511): Propagate the usage of the BindingIndex type to
     // SetBindGroupCmd and through backends as well.
-    // TODO(https://crbug.com/528305452): Don't convert to uint32_t and instead make AllocateData
-    // handle typed indices and return a span.
     cmd->dynamicOffsetCount = uint32_t{dynamicOffsets.size()};
     if (!dynamicOffsets.empty()) {
-        uint32_t* offsets = allocator->AllocateData<uint32_t>(cmd->dynamicOffsetCount);
+        Span<uint32_t> offsets = allocator->AllocateData<uint32_t>(cmd->dynamicOffsetCount);
         // TODO(https://crbug.com/524406299): Use Span::CopyFrom.
-        DAWN_UNSAFE_TODO(
-            memcpy(offsets, dynamicOffsets.data(), cmd->dynamicOffsetCount * sizeof(uint32_t)));
+        std::ranges::copy(dynamicOffsets, offsets.begin());
     }
 }
 
@@ -258,10 +255,9 @@
     SetImmediatesCmd* cmd = allocator->Allocate<SetImmediatesCmd>(Command::SetImmediates);
     cmd->offset = offset;
     cmd->size = uint32_t(data.size());
-    // TODO(https://crbug.com/528305452): Make AllocateData return a span.
-    uint8_t* immediateDatas = allocator->AllocateData<uint8_t>(data.size());
+    Span<std::byte> immediateDatas = allocator->AllocateData<std::byte>(data.size());
     // TODO(https://crbug.com/524406299): Use Span::CopyFrom.
-    DAWN_UNSAFE_TODO(memcpy(immediateDatas, data.data(), data.size()));
+    std::ranges::copy(data, immediateDatas.begin());
 }
 
 MaybeError ProgrammableEncoder::SetResourceTable(ResourceTableBase* table,
diff --git a/src/dawn/native/RenderPassEncoder.cpp b/src/dawn/native/RenderPassEncoder.cpp
index 9e81a2f..1dc1ce8 100644
--- a/src/dawn/native/RenderPassEncoder.cpp
+++ b/src/dawn/native/RenderPassEncoder.cpp
@@ -371,11 +371,8 @@
                 allocator->Allocate<ExecuteBundlesCmd>(Command::ExecuteBundles);
             cmd->count = renderBundles.size();
 
-            // TODO(https://crbug.com/528305452): Make AllocateData handle typed indices and return
-            // a span.
-            Span<Ref<RenderBundleBase>> bundles = DAWN_UNSAFE_TODO(
-                {allocator->AllocateData<Ref<RenderBundleBase>>(renderBundles.size()),
-                 renderBundles.size()});
+            Span<Ref<RenderBundleBase>> bundles =
+                allocator->AllocateData<Ref<RenderBundleBase>>(renderBundles.size());
             for (auto [i, bundle] : Enumerate(bundles)) {
                 // TODO(https://crbug.com/524406299): Use Span::CopyFrom.
                 bundles[i] = renderBundles[i];
diff --git a/src/dawn/tests/unittests/CommandAllocatorTests.cpp b/src/dawn/tests/unittests/CommandAllocatorTests.cpp
index 485451e..f0d3362 100644
--- a/src/dawn/tests/unittests/CommandAllocatorTests.cpp
+++ b/src/dawn/tests/unittests/CommandAllocatorTests.cpp
@@ -146,7 +146,7 @@
         immediates->size = mySize;
         immediates->offset = myOffset;
 
-        uint32_t* values = allocator.AllocateData<uint32_t>(5);
+        Span<uint32_t> values = allocator.AllocateData<uint32_t>(size_t{5u});
         for (size_t i = 0; i < 5; i++) {
             values[i] = myValues[i];
         }
@@ -391,35 +391,31 @@
 };
 
 // Test for overflows in Allocate's computations, size 1 variant
-TEST(CommandAllocator, AllocationOverflow_1) {
+TEST(CommandAllocatorDeathTest, AllocationOverflow_1) {
     CommandAllocator allocator;
-    AlignedStruct<1>* data =
-        allocator.AllocateData<AlignedStruct<1>>(std::numeric_limits<size_t>::max() / 1);
-    ASSERT_EQ(data, nullptr);
+    EXPECT_DEATH_IF_SUPPORTED(
+        allocator.AllocateData<AlignedStruct<1>>(std::numeric_limits<size_t>::max() / 1u), "");
 }
 
 // Test for overflows in Allocate's computations, size 2 variant
-TEST(CommandAllocator, AllocationOverflow_2) {
+TEST(CommandAllocatorDeathTest, AllocationOverflow_2) {
     CommandAllocator allocator;
-    AlignedStruct<2>* data =
-        allocator.AllocateData<AlignedStruct<2>>(std::numeric_limits<size_t>::max() / 2);
-    ASSERT_EQ(data, nullptr);
+    EXPECT_DEATH_IF_SUPPORTED(
+        allocator.AllocateData<AlignedStruct<2>>(std::numeric_limits<size_t>::max() / 2u), "");
 }
 
 // Test for overflows in Allocate's computations, size 4 variant
-TEST(CommandAllocator, AllocationOverflow_4) {
+TEST(CommandAllocatorDeathTest, AllocationOverflow_4) {
     CommandAllocator allocator;
-    AlignedStruct<4>* data =
-        allocator.AllocateData<AlignedStruct<4>>(std::numeric_limits<size_t>::max() / 4);
-    ASSERT_EQ(data, nullptr);
+    EXPECT_DEATH_IF_SUPPORTED(
+        allocator.AllocateData<AlignedStruct<4>>(std::numeric_limits<size_t>::max() / 4u), "");
 }
 
 // Test for overflows in Allocate's computations, size 8 variant
-TEST(CommandAllocator, AllocationOverflow_8) {
+TEST(CommandAllocatorDeathTest, AllocationOverflow_8) {
     CommandAllocator allocator;
-    AlignedStruct<8>* data =
-        allocator.AllocateData<AlignedStruct<8>>(std::numeric_limits<size_t>::max() / 8);
-    ASSERT_EQ(data, nullptr);
+    EXPECT_DEATH_IF_SUPPORTED(
+        allocator.AllocateData<AlignedStruct<8>>(std::numeric_limits<size_t>::max() / 8u), "");
 }
 
 template <int DefaultValue>
@@ -450,14 +446,17 @@
 TEST(CommandAllocator, AllocateDataDefaultInitializes) {
     CommandAllocator allocator;
 
-    IntWithDefault<33>* int33 = allocator.AllocateData<IntWithDefault<33>>(1);
+    Span<IntWithDefault<33>> int33 = allocator.AllocateData<IntWithDefault<33>>(size_t{1});
+    ASSERT_EQ(int33.size(), 1);
     ASSERT_EQ(int33[0].value, 33);
 
-    IntWithDefault<34>* int34 = allocator.AllocateData<IntWithDefault<34>>(2);
+    Span<IntWithDefault<34>> int34 = allocator.AllocateData<IntWithDefault<34>>(size_t{2});
+    ASSERT_EQ(int34.size(), 2);
     ASSERT_EQ(int34[0].value, 34);
     ASSERT_EQ(int34[0].value, 34);
 
-    IntWithDefault<35>* int35 = allocator.AllocateData<IntWithDefault<35>>(3);
+    Span<IntWithDefault<35>> int35 = allocator.AllocateData<IntWithDefault<35>>(size_t{3});
+    ASSERT_EQ(int35.size(), 3);
     ASSERT_EQ(int35[0].value, 35);
     ASSERT_EQ(int35[1].value, 35);
     ASSERT_EQ(int35[2].value, 35);