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

@@ -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?
+    }
+}

tests/src/prefix.rs

@@ -73,11 +73,7 @@ pub unsafe fn run_prefix_test(
         .session
         .create_buffer_with(
             n_elements * 4,
-            |b| {
+            |b| b.extend(0..n_elements as u32),
-                for i in 0..n_elements as u32 {
-                    b.push(&i);
-                }
-            },
             BufferUsage::STORAGE,
         )
         .unwrap();

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| {
+            |b| b.extend(0..n_elements as u32),
-                for i in 0..n_elements as u32 {
-                    b.push(&i);
-                }
-            },
             BufferUsage::STORAGE,
         )
         .unwrap();