wde_terrain/render/
extractor.rs

1use bevy::prelude::*;
2use wde_renderer::prelude::*;
3
4use crate::{
5    manager::{CHUNK_RENDER_SUBDIVISIONS, ChunkPos, TerrainDirtyTile},
6    prelude::TerrainRendererGPU
7};
8
9/// Type of the message sent from the render world to the main world when a tile is extracted from the GPU.
10#[derive(Message)]
11pub struct ExtractedTileMessage {
12    pub pos: ChunkPos,
13    pub map_type: u32,        // 0 for heightmap, 1 for splatmap
14    pub splat_map_index: u32, // only relevant if map_type is 1, should be between 0 and SPLAT_MAP_COUNT/4 - 1
15    pub data: Vec<u8>
16}
17
18pub struct TerrainExtractorPlugin;
19impl Plugin for TerrainExtractorPlugin {
20    fn build(&self, app: &mut App) {
21        app.init_resource::<StaggingBuffers>()
22            .init_resource::<TerrainExtractor>()
23            .add_systems(Startup, init_stagging_buffers);
24        app.get_sub_app_mut(RenderApp)
25            .unwrap()
26            .init_resource::<TerrainExtractor>()
27            .init_resource::<StaggingBuffers>()
28            .add_systems(Extract, extract_staging_buffers)
29            .add_systems(Render, extract_tiles_from_gpu.in_set(RenderSet::Render));
30
31        // Add the extract system
32        app.get_sub_app_mut(RenderApp)
33            .unwrap()
34            .add_systems(Extract, extract_messages);
35    }
36}
37
38/// Resource storing individual texture handles for each material
39#[derive(Resource, Default)]
40pub struct StaggingBuffers {
41    pub heightmap: Option<Handle<Buffer>>,
42    pub splatmap: Option<Handle<Buffer>>
43}
44
45#[derive(Resource, Default)]
46pub struct TerrainExtractor {
47    /// Liste of tiles that needs to be extracted from the gpu.
48    /// First the tile position, then 0 for heightmap, 1 for splatmap, and then the index of the splatmap (0 to SPLAT_MAP_COUNT/4 - 1)
49    tiles_to_extract: Vec<Option<(ChunkPos, u32, u32)>>,
50    /// List of tiles that were extracted from the gpu and are ready to be sent to the main world
51    extracted_tiles: Vec<Option<TerrainDirtyTile>>
52}
53impl TerrainExtractor {
54    /// Queues a tile for extraction from the GPU.
55    /// The tile will be processed in the next render pass and the extracted data will be available in the main world in the next frame.
56    ///
57    /// # Arguments
58    /// * `tile_pos` - The position of the tile to extract (x, z)
59    /// * `map_type` - The type of map to extract (0 for heightmap, 1 for splatmap)
60    /// * `splat_map_index` - The index of the splat map to extract (only relevant if map_type is 1, should be between 0 and SPLAT_MAP_COUNT/4 - 1)
61    pub fn queue_tile_extraction(
62        &mut self,
63        tile_pos: ChunkPos,
64        map_type: u32,
65        splat_map_index: u32
66    ) {
67        self.tiles_to_extract
68            .push(Some((tile_pos, map_type, splat_map_index)));
69    }
70
71    /// Clears the list of tiles to extract. Should be called after processing the extracted tiles to avoid extracting the same tiles multiple times.
72    pub fn clear_tiles_to_extract(&mut self) {
73        self.tiles_to_extract.clear();
74    }
75}
76
77/// Init stagging buffers
78fn init_stagging_buffers(asset_server: Res<AssetServer>, mut buffers: ResMut<StaggingBuffers>) {
79    // Build texture arrays for each type
80    let usage = BufferUsage::COPY_DST | BufferUsage::MAP_READ;
81
82    let stagging_heightmap = asset_server.add(Buffer {
83        label: "stagging_heightmap".to_string(),
84        size: (CHUNK_RENDER_SUBDIVISIONS as usize * CHUNK_RENDER_SUBDIVISIONS as usize)
85            * std::mem::size_of::<u8>(),
86        usage,
87        content: None
88    });
89    let stagging_splatmap = asset_server.add(Buffer {
90        label: "stagging_splatmap".to_string(),
91        size: (CHUNK_RENDER_SUBDIVISIONS as usize * CHUNK_RENDER_SUBDIVISIONS as usize)
92            * std::mem::size_of::<[u8; 4]>(), // RGBA8Unorm
93        usage,
94        content: None
95    });
96
97    buffers.heightmap = Some(stagging_heightmap);
98    buffers.splatmap = Some(stagging_splatmap);
99}
100
101/// Extract the tiles data from the GPU
102fn extract_tiles_from_gpu(
103    mut extractor: ResMut<TerrainExtractor>,
104    gpu_terrain_renderer: Res<TerrainRendererGPU>,
105    textures: Res<RenderAssets<GpuTexture>>,
106    render_instance: Res<RenderInstance>,
107    mut buffers: ResMut<RenderAssets<GpuBuffer>>,
108    stagging_buffers: Res<StaggingBuffers>
109) {
110    // Check if there are tiles to extract
111    if extractor.tiles_to_extract.is_empty() {
112        return;
113    }
114
115    // Process the tiles to extract
116    let render_instance = render_instance.0.read().unwrap();
117    for i in 0..extractor.tiles_to_extract.len() {
118        if let Some((pos, map_type, splat_map_index)) = extractor.tiles_to_extract[i].take() {
119            // Get the right texture handle
120            let render_tile = &gpu_terrain_renderer.tiles[gpu_terrain_renderer.pos_to_tile[&pos]];
121            let texture_handle = match map_type {
122                0 => render_tile.heightmap.clone(),
123                1 => render_tile.splatmaps[splat_map_index as usize].clone(),
124                _ => continue
125            };
126
127            // Get the GPU texture
128            let texture = match textures.get(&texture_handle) {
129                Some(texture) => texture,
130                None => continue
131            };
132
133            // Get the buffer
134            let staging_buffer_handle = match map_type {
135                0 => &stagging_buffers.heightmap,
136                1 => &stagging_buffers.splatmap,
137                _ => continue
138            };
139            let staging_buffer_handle = match staging_buffer_handle {
140                Some(handle) => handle,
141                None => continue
142            };
143            let stagging_buffer = match buffers.get_mut(staging_buffer_handle) {
144                Some(buffer) => buffer,
145                None => continue
146            };
147
148            // Copy the texture data to the staging buffer
149            stagging_buffer
150                .buffer
151                .copy_from_texture(&render_instance, &texture.texture.texture);
152
153            // Read the buffer data and create a dirty tile
154            let mut data = Vec::new();
155            stagging_buffer.buffer.map_read(&render_instance, |view| {
156                data = view.to_vec();
157            });
158
159            // Create the dirty tile and add it to the extracted tiles list
160            extractor.extracted_tiles.push(Some((
161                pos,
162                if map_type == 0 { 0 } else { 1 },
163                if map_type == 1 { splat_map_index } else { 0 },
164                data
165            )));
166            extractor.tiles_to_extract[i] = None;
167        }
168    }
169
170    // Remove None entries from the tiles to extract list
171    extractor.tiles_to_extract.retain(|entry| entry.is_some());
172}
173
174fn extract_staging_buffers(
175    main_stagging_buffers: ExtractWorld<Res<StaggingBuffers>>,
176    mut render_stagging_buffers: ResMut<StaggingBuffers>
177) {
178    if render_stagging_buffers.heightmap.is_some() {
179        return;
180    }
181
182    if let (Some(heightmap), Some(splatmap)) = (
183        main_stagging_buffers.heightmap.clone(),
184        main_stagging_buffers.splatmap.clone()
185    ) {
186        render_stagging_buffers.heightmap = Some(heightmap);
187        render_stagging_buffers.splatmap = Some(splatmap);
188    }
189}
190
191pub fn extract_dirty(
192    main_terrain_extractor: &TerrainExtractor,
193    render_terrain_extractor: &mut TerrainExtractor
194) {
195    // Extract the tiles_to_extract from the main world and move them to the render world extractor resource
196    render_terrain_extractor.tiles_to_extract = main_terrain_extractor.tiles_to_extract.clone();
197}
198
199pub fn extract_messages(
200    mut render_terrain_extractor: ResMut<TerrainExtractor>,
201    mut main_world: ResMut<MainWorld>
202) {
203    // Send events for each extracted tile to the main world
204    for (pos, map_type, splat_map_index, data) in render_terrain_extractor
205        .extracted_tiles
206        .clone()
207        .into_iter()
208        .flatten()
209    {
210        main_world.write_message(ExtractedTileMessage {
211            pos,
212            map_type,
213            splat_map_index,
214            data
215        });
216    }
217
218    // Clear the extracted tiles list after processing
219    render_terrain_extractor.extracted_tiles.clear();
220}