grana_gpu/
readback.rs

1use std::collections::VecDeque;
2use std::sync::mpsc::{self, Receiver, TryRecvError};
3
4use wgpu::{
5    Buffer, BufferAsyncError, BufferDescriptor, BufferUsages, CommandEncoder, Device, MapMode,
6    PollType,
7};
8
9const COPY_ALIGNMENT: u64 = 4;
10
11struct Slot {
12    buffer: Buffer,
13    mapping: Option<Receiver<Result<(), BufferAsyncError>>>,
14    reserved: bool,
15}
16
17pub struct Readback {
18    slots: Vec<Slot>,
19    pending: VecDeque<usize>,
20}
21
22impl Readback {
23    pub fn new(device: &Device, label: &'static str, row: u64, slots: usize) -> Self {
24        assert!(
25            row > 0 && row.is_multiple_of(COPY_ALIGNMENT),
26            "a readback row is a whole number of aligned bytes"
27        );
28        assert!(slots > 0, "a readback ring holds at least one slot");
29        let slots = (0..slots)
30            .map(|_| Slot {
31                buffer: device.create_buffer(&BufferDescriptor {
32                    label: Some(label),
33                    size: row,
34                    usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
35                    mapped_at_creation: false,
36                }),
37                mapping: None,
38                reserved: false,
39            })
40            .collect();
41        Self {
42            slots,
43            pending: VecDeque::new(),
44        }
45    }
46
47    pub fn pending(&self) -> bool {
48        !self.pending.is_empty()
49    }
50
51    pub fn reserve(&mut self) -> Option<usize> {
52        let slot = self.slots.iter().position(|slot| !slot.reserved)?;
53        self.slots[slot].reserved = true;
54        self.pending.push_back(slot);
55        Some(slot)
56    }
57
58    pub fn stage(
59        &self,
60        encoder: &mut CommandEncoder,
61        slot: usize,
62        offset: u64,
63        source: &Buffer,
64        size: u64,
65    ) {
66        encoder.copy_buffer_to_buffer(source, 0, &self.slots[slot].buffer, offset, size);
67    }
68
69    pub fn register(&mut self) {
70        for slot in &mut self.slots {
71            if slot.reserved && slot.mapping.is_none() {
72                let (sender, receiver) = mpsc::channel();
73                slot.buffer
74                    .slice(..)
75                    .map_async(MapMode::Read, move |result| {
76                        let _ = sender.send(result);
77                    });
78                slot.mapping = Some(receiver);
79            }
80        }
81    }
82
83    pub fn take(&mut self, device: &Device) -> Option<(usize, Vec<u8>)> {
84        let slot = *self.pending.front()?;
85        device
86            .poll(PollType::Poll)
87            .expect("the graphics device failed while reading back a frame");
88        let receiver = self.slots[slot]
89            .mapping
90            .as_ref()
91            .expect("a pending readback slot is registered before it is taken");
92        match receiver.try_recv() {
93            Ok(result) => result.expect("the readback staging buffer is mappable"),
94            Err(TryRecvError::Empty) => return None,
95            Err(TryRecvError::Disconnected) => panic!("the readback callback fires once"),
96        }
97        let slice = self.slots[slot].buffer.slice(..);
98        let mapped = slice
99            .get_mapped_range()
100            .expect("the mapped readback staging buffer stays mapped");
101        let bytes = mapped.to_vec();
102        drop(mapped);
103        self.slots[slot].buffer.unmap();
104        self.slots[slot].mapping = None;
105        self.slots[slot].reserved = false;
106        self.pending.pop_front();
107        Some((slot, bytes))
108    }
109}
grana_gpu/readback.rs

grana_gpu/
readback.rs
