Skip to main content

wde_terrain/render/
extractor.rs

1use bevy::prelude::*;
2use wde_renderer::{prelude::*, wgpu_utils::AsyncReadback};
3
4use crate::{manager::ChunkPos, prelude::TerrainRendererGPU};
5
6/// Message sent from the render world to the main world when a tile is read back from GPU.
7#[derive(Message)]
8pub struct ExtractedTileMessage {
9    pub pos: ChunkPos,
10    pub map_type: u32,        // 0 = heightmap, 1 = splatmap
11    pub splat_map_index: u32, // only relevant when map_type == 1
12    pub data: Vec<u8>
13}
14
15/// System sets for ordering terrain render-world work.
16/// External plugins that write terrain textures (e.g. `wde-terrain-editor`) should place
17/// their compute dispatch systems inside [`TerrainRenderSets::TextureWrite`].
18/// [`initiate_tile_readbacks`] is configured to run **after** this set so it always
19/// captures post-compute texture data.
20#[derive(SystemSet, Debug, Clone, PartialEq, Eq, Hash)]
21pub enum TerrainRenderSets {
22    /// Any compute pass that writes to the terrain heightmap or splatmap textures.
23    TextureWrite
24}
25
26pub struct TerrainExtractorPlugin;
27impl Plugin for TerrainExtractorPlugin {
28    fn build(&self, app: &mut App) {
29        app.init_resource::<TerrainExtractor>();
30
31        let render_app = app.get_sub_app_mut(RenderApp).unwrap();
32        render_app
33            .init_resource::<TerrainExtractor>()
34            .init_resource::<PendingReadbacks>()
35            // TextureWrite set must finish before tile readbacks start.
36            .configure_sets(
37                Render,
38                TerrainRenderSets::TextureWrite.before(initiate_tile_readbacks)
39            )
40            // Poll pending readbacks and forward completions to the main world.
41            .add_systems(Extract, flush_completed_readbacks)
42            // Kick off GPU copies for the tiles queued this frame.
43            .add_systems(Render, initiate_tile_readbacks.in_set(RenderSet::Render));
44    }
45}
46
47/// Queue of tiles to read back from the GPU.
48/// Lives in both the main world (as a write target for game systems) and the render world
49/// (consumed by [`initiate_tile_readbacks`]).
50#[derive(Resource, Default)]
51pub struct TerrainExtractor {
52    pub tiles_to_extract: Vec<(ChunkPos, u32, u32)>
53}
54
55impl TerrainExtractor {
56    pub fn queue_tile_extraction(&mut self, pos: ChunkPos, map_type: u32, splat_index: u32) {
57        self.tiles_to_extract.push((pos, map_type, splat_index));
58    }
59}
60
61/// In-flight GPU→CPU readback for one tile.
62struct PendingReadback {
63    pos: ChunkPos,
64    map_type: u32,
65    splat_index: u32,
66    readback: AsyncReadback
67}
68
69/// Render-world store of all in-flight readbacks.
70#[derive(Resource, Default)]
71pub struct PendingReadbacks {
72    pending: Vec<PendingReadback>
73}
74
75/// Transfer the main world's extraction queue into the render world.
76/// Called by [`TerrainRendererGPU::extract_dirty`] during the Extract phase.
77pub fn extract_dirty(main: &mut TerrainExtractor, render: &mut TerrainExtractor) {
78    render.tiles_to_extract = std::mem::take(&mut main.tiles_to_extract);
79}
80
81/// Render phase: copy each queued tile to a staging buffer and start async mapping.
82/// Must run **after** the paint compute pass so the copy captures post-compute data.
83/// This ordering is enforced in [`PaintProcessorPlugin`].
84pub fn initiate_tile_readbacks(
85    mut extractor: ResMut<TerrainExtractor>,
86    mut pending: ResMut<PendingReadbacks>,
87    gpu: Res<TerrainRendererGPU>,
88    textures: Res<RenderAssets<GpuTexture>>,
89    render_instance: Res<RenderInstance>
90) {
91    if extractor.tiles_to_extract.is_empty() {
92        return;
93    }
94
95    let ri = render_instance.0.read().unwrap();
96    let tiles = std::mem::take(&mut extractor.tiles_to_extract);
97
98    for (pos, map_type, splat_index) in tiles {
99        let Some(&layer) = gpu.pos_to_layer.get(&pos) else {
100            continue;
101        };
102        let texture_handle = match map_type {
103            0 => gpu.heightmap_array.as_ref(),
104            1 => gpu.splatmap_array.as_ref(),
105            _ => continue
106        };
107        let Some(tex_h) = texture_handle else {
108            continue;
109        };
110        let Some(tex) = textures.get(tex_h) else {
111            continue;
112        };
113
114        pending.pending.push(PendingReadback {
115            pos,
116            map_type,
117            splat_index,
118            readback: AsyncReadback::from_texture_layer(&ri, &tex.texture.texture, layer)
119        });
120    }
121}
122
123/// Extract phase: non-blocking poll + collect any readbacks the GPU completed last frame.
124/// Forwards completed tile data to the main world as [`ExtractedTileMessage`]s.
125fn flush_completed_readbacks(
126    mut pending: ResMut<PendingReadbacks>,
127    render_instance: Res<RenderInstance>,
128    mut main_world: ResMut<MainWorld>
129) {
130    if pending.pending.is_empty() {
131        return;
132    }
133
134    // Non-blocking poll: triggers any map_async callbacks the GPU has finished.
135    render_instance.0.read().unwrap().poll_non_blocking();
136
137    let drained = std::mem::take(&mut pending.pending);
138    for PendingReadback {
139        pos,
140        map_type,
141        splat_index,
142        readback
143    } in drained
144    {
145        match readback.try_collect() {
146            Ok(data) if !data.is_empty() => {
147                main_world.write_message(ExtractedTileMessage {
148                    pos,
149                    map_type,
150                    splat_map_index: splat_index,
151                    data
152                });
153            }
154            Ok(_) => {} // empty data means a mapping error was already logged
155            Err(readback) => {
156                // Mapping not done yet — return to queue for next frame.
157                pending.pending.push(PendingReadback {
158                    pos,
159                    map_type,
160                    splat_index,
161                    readback
162                });
163            }
164        }
165    }
166}