sketch_2d/buffer/
stream.rs

1/*
2 * Created on Mon Jul 03 2023
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7use std::{borrow::Cow, num::NonZeroU64, ops::Range};
8
9use wgpu::{Buffer, BufferAddress, BufferBinding, BufferSlice, BufferUsages, Device, Queue};
10
11use super::GrowingBuffer;
12
13pub type StreamRange = Range<BufferAddress>;
14
15#[derive(Debug)]
16pub struct BufferStream<'a> {
17    buffer: GrowingBuffer<'a>,
18    data: Vec<u8>,
19}
20
21impl<'a> BufferStream<'a> {
22    pub fn new(label: Option<Cow<'a, str>>, usages: BufferUsages) -> Self {
23        Self {
24            buffer: GrowingBuffer::new(label, usages | BufferUsages::COPY_DST, true),
25            data: Vec::new(),
26        }
27    }
28
29    /// Return next writer starting from end of buffer
30    pub fn next_writer(&mut self) -> StreamWriter {
31        let offset = self.data.len();
32
33        StreamWriter {
34            write_buffer: &mut self.data,
35            offset: offset as BufferAddress,
36        }
37    }
38
39    /// Write slice of data and return written range
40    pub fn write_slice(&mut self, data: &[u8]) -> StreamRange {
41        let mut writer = self.next_writer();
42        writer.write(data);
43        writer.finish()
44    }
45
46    /// Finish streaming and upload memory buffer to gpu
47    pub fn finish(&mut self, device: &Device, queue: &Queue) -> StreamBuffer {
48        let size = self.data.len() as BufferAddress;
49
50        let (buffer, mapped) = self.buffer.get(device, size);
51
52        if size > 0 {
53            if mapped {
54                buffer
55                    .slice(..size)
56                    .get_mapped_range_mut()
57                    .copy_from_slice(&self.data);
58                buffer.unmap();
59            } else {
60                queue.write_buffer(buffer, 0, &self.data);
61            }
62
63            self.data.clear();
64        }
65
66        StreamBuffer { buffer }
67    }
68}
69
70#[derive(Debug)]
71pub struct StreamWriter<'a> {
72    write_buffer: &'a mut Vec<u8>,
73
74    offset: BufferAddress,
75}
76
77impl<'a> StreamWriter<'a> {
78    pub fn write(&mut self, data: &[u8]) {
79        self.write_buffer.extend_from_slice(data);
80    }
81
82    /// Finish writer.
83    /// Returns range of written data.
84    pub fn finish(self) -> StreamRange {
85        self.offset..(self.write_buffer.len() as BufferAddress)
86    }
87}
88
89#[derive(Debug, Clone, Copy)]
90pub struct StreamBuffer<'a> {
91    buffer: &'a Buffer,
92}
93
94impl<'a> StreamBuffer<'a> {
95    pub fn slice(&self, range: StreamRange) -> BufferSlice<'a> {
96        self.buffer.slice(range)
97    }
98
99    pub fn binding(&self, range: StreamRange) -> BufferBinding<'a> {
100        BufferBinding {
101            buffer: self.buffer,
102            offset: range.start,
103            size: NonZeroU64::new(range.end - range.start),
104        }
105    }
106}
107
108#[cfg(test)]
109#[test]
110pub fn write_test() {
111    use std::time::Instant;
112
113    use wgpu::BufferUsages;
114
115    let mut stream = BufferStream::new(None, BufferUsages::COPY_DST | BufferUsages::VERTEX);
116
117    let start = Instant::now();
118
119    for _ in 0..59999 {
120        let mut entry = stream.next_writer();
121        entry.write(&[1]);
122
123        entry.finish();
124    }
125
126    println!("Writing took {} microseconds", start.elapsed().as_micros());
127}