1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

//! Easy screenshot taking.
//!
//! Executes compositing into a special target if we want to take a screenshot.
//! This is necessary because `composite` is expected to write directly into the final output
//! from which we can't read back and of which we don't control the format.
//!
//! This comes with the perk that we can do extra things here if we want!
//!
//! TODO(andreas): One thing to add would be more anti-aliasing!
//! We could render the same image with subpixel moved camera in order to get super-sampling without hitting texture size limitations.
//! Or alternatively try to render the images in several tiles 🤔. In any case this would greatly improve quality!

use parking_lot::Mutex;

use crate::{
    allocator::GpuReadbackError,
    texture_info::Texture2DBufferInfo,
    wgpu_resources::{GpuTexture, TextureDesc},
    DebugLabel, GpuReadbackBuffer, GpuReadbackIdentifier, RenderContext,
};

/// Type used as user data on the gpu readback belt.
struct ReadbackBeltMetadata<T: 'static + Send + Sync> {
    extent: wgpu::Extent3d,
    user_data: T,
}

pub struct ScreenshotProcessor {
    screenshot_texture: GpuTexture,
    screenshot_readback_buffer: Mutex<GpuReadbackBuffer>,
}

impl ScreenshotProcessor {
    /// The texture format used for screenshots.
    pub const SCREENSHOT_COLOR_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8Unorm;

    pub fn new<T: 'static + Send + Sync>(
        ctx: &RenderContext,
        view_name: &DebugLabel,
        resolution: glam::UVec2,
        readback_identifier: GpuReadbackIdentifier,
        readback_user_data: T,
    ) -> Self {
        let size = wgpu::Extent3d {
            width: resolution.x,
            height: resolution.y,
            depth_or_array_layers: 1,
        };
        let buffer_info = Texture2DBufferInfo::new(Self::SCREENSHOT_COLOR_FORMAT, size);
        let screenshot_readback_buffer = Mutex::new(ctx.gpu_readback_belt.lock().allocate(
            &ctx.device,
            &ctx.gpu_resources.buffers,
            buffer_info.buffer_size_padded,
            readback_identifier,
            Box::new(ReadbackBeltMetadata {
                extent: size,
                user_data: readback_user_data,
            }),
        ));

        let screenshot_texture = ctx.gpu_resources.textures.alloc(
            &ctx.device,
            &TextureDesc {
                label: format!("{view_name} - ScreenshotProcessor").into(),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: Self::SCREENSHOT_COLOR_FORMAT,
                usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
            },
        );

        Self {
            screenshot_texture,
            screenshot_readback_buffer,
        }
    }

    pub fn begin_render_pass<'a>(
        &'a self,
        view_name: &DebugLabel,
        encoder: &'a mut wgpu::CommandEncoder,
    ) -> wgpu::RenderPass<'a> {
        re_tracing::profile_function!();

        let pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: DebugLabel::from(format!("{view_name} - screenshot")).get(),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &self.screenshot_texture.default_view,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
        });

        pass
    }

    pub fn end_render_pass(
        &self,
        encoder: &mut wgpu::CommandEncoder,
    ) -> Result<(), GpuReadbackError> {
        self.screenshot_readback_buffer.lock().read_texture2d(
            encoder,
            wgpu::ImageCopyTexture {
                texture: &self.screenshot_texture.texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            self.screenshot_texture.texture.size(),
        )
    }

    /// Returns the oldest received screenshot results for a given identifier and user data type.
    ///
    /// It is recommended to call this method repeatedly until it returns `None` to ensure that all
    /// pending data is flushed.
    ///
    /// Ready data that hasn't been retrieved for more than a frame will be discarded.
    ///
    /// See also [`crate::view_builder::ViewBuilder::schedule_screenshot`]
    pub fn next_readback_result<T: 'static + Send + Sync>(
        ctx: &RenderContext,
        identifier: GpuReadbackIdentifier,
        on_screenshot: impl FnOnce(&[u8], glam::UVec2, T),
    ) -> Option<()> {
        let mut screenshot_was_available = None;
        ctx.gpu_readback_belt
            .lock()
            .readback_data::<ReadbackBeltMetadata<T>>(identifier, |data: &[u8], metadata| {
                screenshot_was_available = Some(());
                let buffer_info =
                    Texture2DBufferInfo::new(Self::SCREENSHOT_COLOR_FORMAT, metadata.extent);
                let texture_data = buffer_info.remove_padding(data);
                on_screenshot(
                    &texture_data,
                    glam::uvec2(metadata.extent.width, metadata.extent.height),
                    metadata.user_data,
                );
            });
        screenshot_was_available
    }
}