From f1d7560b3c8fe6676a7bd7ec61b88924cd593e8d Mon Sep 17 00:00:00 2001
From: Raph Levien <raph.levien@gmail.com>
Date: Thu, 25 Nov 2021 22:02:04 -0800
Subject: [PATCH] Tweak extend implementation

The one that takes T is more useful than the one that takes references
to T. When specialization lands, we will be able to have both under the
`extend` name.
---
 piet-gpu-hal/src/bufwrite.rs | 35 ++++++++++++++++++++++++++++++++++-
 tests/src/prefix.rs          |  6 +-----
 tests/src/prefix_tree.rs     |  6 +-----
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/piet-gpu-hal/src/bufwrite.rs b/piet-gpu-hal/src/bufwrite.rs
index eefe66f..980b4b2 100644
--- a/piet-gpu-hal/src/bufwrite.rs
+++ b/piet-gpu-hal/src/bufwrite.rs
@@ -73,7 +73,9 @@ impl BufWrite {
     /// Extend with an iterator over plain data objects.
     ///
     /// Currently, this doesn't panic, just truncates. That may change.
-    pub fn extend<'a, I, T: Pod + 'a>(&mut self, iter: I)
+    // Note: when specialization lands, this can be another impl of
+    // `Extend`.
+    pub fn extend_ref_iter<'a, I, T: Pod + 'a>(&mut self, iter: I)
     where
         I: IntoIterator<Item = &'a T>,
     {
@@ -115,3 +117,34 @@ impl std::ops::DerefMut for BufWrite {
         unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
     }
 }
+
+impl<T: Pod> std::iter::Extend<T> for BufWrite {
+    fn extend<I>(&mut self, iter: I)
+    where
+        I: IntoIterator<Item = T>,
+    {
+        let item_size = std::mem::size_of::<T>();
+        if item_size == 0 {
+            return;
+        }
+        let mut iter = iter.into_iter();
+        let n_remaining = (self.capacity - self.len) / item_size;
+        unsafe {
+            let mut dst = self.ptr.add(self.len);
+            for _ in 0..n_remaining {
+                if let Some(item) = iter.next() {
+                    std::ptr::copy_nonoverlapping(
+                        bytemuck::bytes_of(&item).as_ptr(),
+                        dst,
+                        item_size,
+                    );
+                    self.len += item_size;
+                    dst = dst.add(item_size);
+                } else {
+                    break;
+                }
+            }
+        }
+        // TODO: should we test the iter and panic on overflow?
+    }
+}
diff --git a/tests/src/prefix.rs b/tests/src/prefix.rs
index eb50761..82f471f 100644
--- a/tests/src/prefix.rs
+++ b/tests/src/prefix.rs
@@ -73,11 +73,7 @@ pub unsafe fn run_prefix_test(
         .session
         .create_buffer_with(
             n_elements * 4,
-            |b| {
-                for i in 0..n_elements as u32 {
-                    b.push(&i);
-                }
-            },
+            |b| b.extend(0..n_elements as u32),
             BufferUsage::STORAGE,
         )
         .unwrap();
diff --git a/tests/src/prefix_tree.rs b/tests/src/prefix_tree.rs
index e2a94e0..4fb3423 100644
--- a/tests/src/prefix_tree.rs
+++ b/tests/src/prefix_tree.rs
@@ -51,11 +51,7 @@ pub unsafe fn run_prefix_test(runner: &mut Runner, config: &Config) -> TestResul
         .session
         .create_buffer_with(
             n_elements * 4,
-            |b| {
-                for i in 0..n_elements as u32 {
-                    b.push(&i);
-                }
-            },
+            |b| b.extend(0..n_elements as u32),
             BufferUsage::STORAGE,
         )
         .unwrap();