Skip to main content

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/// Message sent from the render world to the main world when a tile is read back from GPU.
10#[derive(Message)]
11pub struct ExtractedTileMessage {
12    pub pos: ChunkPos,
13    pub map_type: u32,        // 0 = heightmap, 1 = splatmap
14    pub splat_map_index: u32, // only relevant when map_type == 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::<StagingBuffers>()
22            .init_resource::<TerrainExtractor>()
23            .add_systems(Startup, init_staging_buffers);
24
25        let render_app = app.get_sub_app_mut(RenderApp).unwrap();
26        render_app
27            .init_resource::<TerrainExtractor>()
28            .init_resource::<StagingBuffers>()
29            .add_systems(Extract, extract_staging_buffers)
30            .add_systems(Extract, extract_messages)
31            .add_systems(Render, extract_tiles_from_gpu.in_set(RenderSet::Render));
32    }
33}
34
35/// Staging buffers for GPU → CPU readback.
36#[derive(Resource, Default)]
37pub struct StagingBuffers {
38    pub heightmap: Option<Handle<Buffer>>,
39    pub splatmap: Option<Handle<Buffer>>
40}
41
42/// Queue of tiles to read back from the GPU, shared between main and render worlds.
43#[derive(Resource, Default)]
44pub struct TerrainExtractor {
45    /// Tiles queued for GPU extraction.
46    pub tiles_to_extract: Vec<(ChunkPos, u32, u32)>,
47    /// Tiles already extracted and ready to send to the main world.
48    extracted_tiles: Vec<TerrainDirtyTile>
49}
50
51impl TerrainExtractor {
52    /// Queue a tile for GPU readback.
53    pub fn queue_tile_extraction(&mut self, pos: ChunkPos, map_type: u32, splat_index: u32) {
54        self.tiles_to_extract.push((pos, map_type, splat_index));
55    }
56}
57
58fn init_staging_buffers(asset_server: Res<AssetServer>, mut staging: ResMut<StagingBuffers>) {
59    let usage = BufferUsage::COPY_DST | BufferUsage::MAP_READ;
60    let px = CHUNK_RENDER_SUBDIVISIONS as usize;
61
62    staging.heightmap = Some(asset_server.add(Buffer {
63        label: "staging_heightmap".to_string(),
64        size: px * px * std::mem::size_of::<u8>(),
65        usage,
66        content: None
67    }));
68    staging.splatmap = Some(asset_server.add(Buffer {
69        label: "staging_splatmap".to_string(),
70        size: px * px * std::mem::size_of::<[u8; 4]>(),
71        usage,
72        content: None
73    }));
74}
75
76/// Drain the main world's extraction queue into the render world's queue.
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
81fn extract_staging_buffers(
82    main: ExtractWorld<Res<StagingBuffers>>,
83    mut render: ResMut<StagingBuffers>
84) {
85    if render.heightmap.is_some() {
86        return;
87    }
88    if let (Some(h), Some(s)) = (main.heightmap.clone(), main.splatmap.clone()) {
89        render.heightmap = Some(h);
90        render.splatmap = Some(s);
91    }
92}
93
94fn extract_tiles_from_gpu(
95    mut extractor: ResMut<TerrainExtractor>,
96    gpu: Res<TerrainRendererGPU>,
97    textures: Res<RenderAssets<GpuTexture>>,
98    render_instance: Res<RenderInstance>,
99    mut buffers: ResMut<RenderAssets<GpuBuffer>>,
100    staging: Res<StagingBuffers>
101) {
102    if extractor.tiles_to_extract.is_empty() {
103        return;
104    }
105
106    let render_instance = render_instance.0.read().unwrap();
107    let tiles = std::mem::take(&mut extractor.tiles_to_extract);
108
109    for (pos, map_type, splat_index) in tiles {
110        let layer = match gpu.pos_to_layer.get(&pos) {
111            Some(&l) => l,
112            None => continue
113        };
114
115        let (texture_handle, staging_handle) = match map_type {
116            0 => (gpu.heightmap_array.as_ref(), staging.heightmap.as_ref()),
117            1 => (gpu.splatmap_array.as_ref(), staging.splatmap.as_ref()),
118            _ => continue
119        };
120        let (Some(tex_h), Some(stag_h)) = (texture_handle, staging_handle) else {
121            continue;
122        };
123        let tex = match textures.get(tex_h) {
124            Some(t) => t,
125            None => continue
126        };
127        let staging_buf = match buffers.get_mut(stag_h) {
128            Some(b) => b,
129            None => continue
130        };
131
132        staging_buf
133            .buffer
134            .copy_from_texture_layered(&render_instance, &tex.texture.texture, layer);
135
136        let mut data = Vec::new();
137        staging_buf.buffer.map_read(&render_instance, |view| {
138            data = view.to_vec();
139        });
140
141        extractor
142            .extracted_tiles
143            .push((pos, map_type, splat_index, data));
144    }
145}
146
147fn extract_messages(
148    mut render_extractor: ResMut<TerrainExtractor>,
149    mut main_world: ResMut<MainWorld>
150) {
151    for (pos, map_type, splat_map_index, data) in
152        std::mem::take(&mut render_extractor.extracted_tiles)
153    {
154        main_world.write_message(ExtractedTileMessage {
155            pos,
156            map_type,
157            splat_map_index,
158            data
159        });
160    }
161}