Implement saving map as image file (Astrabit-ST#131)

* fix: fix 1 frame delay when grid shader viewport size changes

This fixes a bug where when the size of the map editor changes, the grid
shader's internally stored value for the size of the map editor's
viewport doesn't get updated until one frame later if the shader is
using push constants. This caused the thicknesses of the grid lines to
be inaccurate for a frame after resizing.

* feat: add "Save map preview" button for saving map as image

* style: replace `ScreenshotState` with an async block

* perf: don't use padding in the texture when screenshotting

* feat: handle screenshotting for maps larger than max texture size

* fix: fix `width_padded` alignment value

* style: remove unnecessary rounding from `max_texture_width`

crates/components/Cargo.toml

@@ -46,6 +46,9 @@ lexical-sort = "0.3.1"
 
 fragile.workspace = true
 parking_lot.workspace = true
+oneshot.workspace = true
 
 fuzzy-matcher = "0.3.7"
 murmur3.workspace = true
+
+image.workspace = true

crates/components/src/map_view.rs

@@ -15,7 +15,9 @@
 // You should have received a copy of the GNU General Public License
 // along with Luminol.  If not, see <http://www.gnu.org/licenses/>.
 
+use color_eyre::eyre::{ContextCompat, WrapErr};
 use itertools::Itertools;
+use std::io::Write;
 
 pub struct MapView {
     /// Toggle to display the visible region in-game.
@@ -697,4 +699,284 @@ impl MapView {
 
         response
     }
+
+    /// Saves the current state of the map to an image file of the user's choice (will prompt the
+    /// user with a file picker).
+    /// This function returns a future that you need to `.await` to finish saving the image, but
+    /// the future doesn't borrow anything so you don't need to worry about lifetime-related issues.
+    pub fn save_as_image(
+        &self,
+        graphics_state: &std::sync::Arc<luminol_graphics::GraphicsState>,
+        map: &luminol_data::rpg::Map,
+    ) -> impl std::future::Future<Output = color_eyre::Result<()>> {
+        let c = "While screenshotting the map";
+
+        let max_texture_dimension_2d = graphics_state
+            .render_state
+            .device
+            .limits()
+            .max_texture_dimension_2d
+            / wgpu::COPY_BYTES_PER_ROW_ALIGNMENT
+            * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
+        let max_buffer_size = graphics_state.render_state.device.limits().max_buffer_size as u32
+            / wgpu::COPY_BYTES_PER_ROW_ALIGNMENT
+            * wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
+
+        let screenshot_width = map.width as u32 * 32;
+        let screenshot_height = map.height as u32 * 32;
+
+        let max_texture_width = screenshot_width
+            .min(max_texture_dimension_2d)
+            .min(max_buffer_size);
+        let max_texture_height = screenshot_height
+            .min(max_texture_dimension_2d)
+            .min(max_buffer_size / (max_texture_width * 4));
+
+        let buffers = (0..screenshot_height)
+            .step_by(max_texture_height as usize)
+            .cartesian_product((0..screenshot_width).step_by(max_texture_width as usize))
+            .map(|(y_offset, x_offset)| {
+                let width = max_texture_width.min(screenshot_width - x_offset);
+                let height = max_texture_height.min(screenshot_height - y_offset);
+                let width_padded = width.next_multiple_of(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT / 4);
+                let viewport_rect = egui::Rect::from_min_size(
+                    egui::pos2(x_offset as f32, y_offset as f32),
+                    egui::vec2(width as f32, height as f32),
+                );
+
+                let texture =
+                    graphics_state
+                        .render_state
+                        .device
+                        .create_texture(&wgpu::TextureDescriptor {
+                            label: Some("map editor screenshot texture"),
+                            size: wgpu::Extent3d {
+                                width,
+                                height,
+                                depth_or_array_layers: 1,
+                            },
+                            mip_level_count: 1,
+                            sample_count: 1,
+                            dimension: wgpu::TextureDimension::D2,
+                            format: graphics_state.render_state.target_format,
+                            usage: wgpu::TextureUsages::COPY_SRC
+                                | wgpu::TextureUsages::RENDER_ATTACHMENT,
+                            view_formats: &[],
+                        });
+                let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
+                let buffer =
+                    graphics_state
+                        .render_state
+                        .device
+                        .create_buffer(&wgpu::BufferDescriptor {
+                            label: Some("map editor screenshot buffer"),
+                            size: width_padded as u64 * height as u64 * 4,
+                            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+                            mapped_at_creation: false,
+                        });
+
+                self.map.set_proj(
+                    &graphics_state.render_state,
+                    glam::Mat4::orthographic_rh(
+                        x_offset as f32,
+                        (x_offset + width) as f32,
+                        (y_offset + height) as f32,
+                        y_offset as f32,
+                        -1.,
+                        1.,
+                    ),
+                );
+
+                let mut command_encoder = graphics_state
+                    .render_state
+                    .device
+                    .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
+                {
+                    let map_callback = self.map.callback(
+                        graphics_state.clone(),
+                        match self.selected_layer {
+                            SelectedLayer::Events => None,
+                            SelectedLayer::Tiles(selected_layer)
+                                if self.darken_unselected_layers =>
+                            {
+                                Some(selected_layer)
+                            }
+                            SelectedLayer::Tiles(_) => None,
+                        },
+                    );
+                    let map_overlay_callback =
+                        self.map.overlay_callback(graphics_state.clone(), 1.);
+
+                    let event_callbacks = map
+                        .events
+                        .iter()
+                        .filter_map(|(_, event)| {
+                            let sprites = self.events.get(event.id);
+                            let tile_size = 32.;
+                            let event_size = sprites
+                                .map(|e| e.0.sprite_size)
+                                .unwrap_or(egui::vec2(32., 32.));
+
+                            if let Some((sprite, _)) = sprites {
+                                sprite.sprite().graphic.set_opacity_multiplier(
+                                    &graphics_state.render_state,
+                                    if self.darken_unselected_layers
+                                        && !matches!(self.selected_layer, SelectedLayer::Events)
+                                    {
+                                        0.5
+                                    } else {
+                                        1.
+                                    },
+                                );
+                            }
+
+                            let rect = egui::Rect::from_min_size(
+                                egui::pos2(
+                                    (event.x as f32 * tile_size) + (tile_size - event_size.x) / 2.,
+                                    (event.y as f32 * tile_size) + (tile_size - event_size.y),
+                                ),
+                                event_size,
+                            );
+
+                            sprites.and_then(|(sprite, _)| {
+                                viewport_rect.intersects(rect).then(|| {
+                                    let x = event.x as f32 * 32. + (32. - event_size.x) / 2.;
+                                    let y = event.y as f32 * 32. + (32. - event_size.y);
+                                    sprite.set_proj(
+                                        &graphics_state.render_state,
+                                        glam::Mat4::orthographic_rh(
+                                            x_offset as f32 - x,
+                                            (x_offset + width) as f32 - x,
+                                            (y_offset + height) as f32 - y,
+                                            y_offset as f32 - y,
+                                            -1.,
+                                            1.,
+                                        ),
+                                    );
+                                    sprite.callback(graphics_state.clone())
+                                })
+                            })
+                        })
+                        .collect_vec();
+
+                    let mut render_pass =
+                        command_encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                            label: Some("map editor screenshot render pass"),
+                            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                                view: &view,
+                                resolve_target: None,
+                                ops: wgpu::Operations::default(),
+                            })],
+                            depth_stencil_attachment: None,
+                            timestamp_writes: None,
+                            occlusion_query_set: None,
+                        });
+
+                    map_callback.paint(&mut render_pass);
+                    for event_callback in event_callbacks.iter() {
+                        event_callback.paint(&mut render_pass);
+                    }
+                    map_overlay_callback.paint(
+                        egui::PaintCallbackInfo {
+                            viewport: viewport_rect,
+                            clip_rect: viewport_rect,
+                            pixels_per_point: 1.,
+                            screen_size_px: [screenshot_width, screenshot_height],
+                        },
+                        &mut render_pass,
+                    );
+                }
+
+                command_encoder.copy_texture_to_buffer(
+                    wgpu::ImageCopyTexture {
+                        texture: &texture,
+                        mip_level: 0,
+                        origin: wgpu::Origin3d::ZERO,
+                        aspect: wgpu::TextureAspect::All,
+                    },
+                    wgpu::ImageCopyBuffer {
+                        buffer: &buffer,
+                        layout: wgpu::ImageDataLayout {
+                            offset: 0,
+                            bytes_per_row: Some(width_padded * 4),
+                            rows_per_image: Some(height),
+                        },
+                    },
+                    wgpu::Extent3d {
+                        width,
+                        height,
+                        depth_or_array_layers: 1,
+                    },
+                );
+                graphics_state
+                    .render_state
+                    .queue
+                    .submit(Some(command_encoder.finish()));
+
+                buffer
+            })
+            .collect_vec();
+
+        let graphics_state = graphics_state.clone();
+        let mut vec = vec![0; screenshot_width as usize * screenshot_height as usize * 4];
+        async move {
+            for ((y_offset, x_offset), buffer) in (0..screenshot_height)
+                .step_by(max_texture_height as usize)
+                .cartesian_product((0..screenshot_width).step_by(max_texture_width as usize))
+                .zip(buffers)
+            {
+                let width = max_texture_width.min(screenshot_width - x_offset);
+                let width_padded = width.next_multiple_of(wgpu::COPY_BYTES_PER_ROW_ALIGNMENT / 4);
+
+                let (tx, rx) = oneshot::channel();
+                buffer
+                    .slice(..)
+                    .map_async(wgpu::MapMode::Read, move |result| {
+                        let _ = tx.send(result);
+                    });
+                if !graphics_state
+                    .render_state
+                    .device
+                    .poll(wgpu::Maintain::Wait)
+                    .is_queue_empty()
+                {
+                    return Err(color_eyre::eyre::eyre!("wgpu::Device::poll timed out").wrap_err(c));
+                }
+                rx.await.unwrap().wrap_err(c)?;
+
+                for (i, row) in buffer
+                    .slice(..)
+                    .get_mapped_range()
+                    .chunks_exact(width_padded as usize * 4)
+                    .enumerate()
+                {
+                    let offset = ((y_offset as usize + i) * screenshot_width as usize
+                        + x_offset as usize)
+                        * 4;
+                    vec[offset..offset + width as usize * 4]
+                        .copy_from_slice(&row[..width as usize * 4]);
+                }
+            }
+
+            if graphics_state.render_state.target_format == wgpu::TextureFormat::Bgra8Unorm {
+                for (b, _g, r, _a) in vec.iter_mut().tuples() {
+                    std::mem::swap(b, r);
+                }
+            }
+
+            let screenshot =
+                image::RgbaImage::from_raw(screenshot_width, screenshot_height, vec).wrap_err(c)?;
+            let mut file = luminol_filesystem::host::File::new().wrap_err(c)?;
+            screenshot
+                .write_to(
+                    &mut std::io::BufWriter::new(&mut file),
+                    image::ImageOutputFormat::Png,
+                )
+                .wrap_err(c)?;
+            file.flush().wrap_err(c)?;
+            file.save("map.png", "Portable Network Graphics")
+                .await
+                .wrap_err(c)
+        }
+    }
 }

crates/graphics/src/event.rs

@@ -33,22 +33,28 @@ pub struct Event {
 }
 
 // wgpu types are not Send + Sync on webassembly, so we use fragile to make sure we never access any wgpu resources across thread boundaries
-struct Callback {
+pub struct Callback {
     sprite: Fragile<Arc<Sprite>>,
     graphics_state: Fragile<Arc<GraphicsState>>,
 }
 
+impl Callback {
+    pub fn paint<'a>(&'a self, render_pass: &mut wgpu::RenderPass<'a>) {
+        let sprite = self.sprite.get();
+        let graphics_state = self.graphics_state.get();
+
+        sprite.draw(graphics_state, render_pass);
+    }
+}
+
 impl luminol_egui_wgpu::CallbackTrait for Callback {
     fn paint<'a>(
         &'a self,
         _info: egui::PaintCallbackInfo,
         render_pass: &mut wgpu::RenderPass<'a>,
         _callback_resources: &'a luminol_egui_wgpu::CallbackResources,
     ) {
-        let sprite = self.sprite.get();
-        let graphics_state = self.graphics_state.get();
-
-        sprite.draw(graphics_state, render_pass);
+        self.paint(render_pass)
     }
 }
 
@@ -183,6 +189,13 @@ impl Event {
         self.viewport.set_proj(render_state, proj);
     }
 
+    pub fn callback(&self, graphics_state: Arc<GraphicsState>) -> Callback {
+        Callback {
+            sprite: Fragile::new(self.sprite.clone()),
+            graphics_state: Fragile::new(graphics_state),
+        }
+    }
+
     pub fn paint(
         &self,
         graphics_state: Arc<GraphicsState>,
@@ -191,10 +204,7 @@ impl Event {
     ) {
         painter.add(luminol_egui_wgpu::Callback::new_paint_callback(
             rect,
-            Callback {
-                sprite: Fragile::new(self.sprite.clone()),
-                graphics_state: Fragile::new(graphics_state),
-            },
+            self.callback(graphics_state),
         ));
     }
 }