wde_terrain_editor/
lib.rs

1//! Terrain editor plugin for WaterDropEngine.
2use wde_logger::prelude::*;
3
4use bevy::prelude::*;
5use wde_terrain::prelude::*;
6
7use crate::{
8    paint::{brush::PaintBrush, paint_manager::PaintManagerPlugin},
9    processor::PaintProcessorPlugin,
10    ui_painter::TerrainEditorUIPlugin
11};
12
13mod paint;
14mod processor;
15mod ui_painter;
16
17pub struct TerrainEditorPlugin;
18impl Plugin for TerrainEditorPlugin {
19    fn build(&self, app: &mut App) {
20        app.add_plugins((
21            PaintManagerPlugin,
22            PaintProcessorPlugin,
23            TerrainEditorUIPlugin
24        ))
25        .init_resource::<SaveManager>()
26        .add_systems(PreStartup, init)
27        .add_systems(Update, handle_extracted_tiles)
28        .add_message::<ExtractedTileMessage>();
29    }
30}
31
32fn init(mut commands: Commands, asset_server: Res<AssetServer>) {
33    let entity = commands
34        .spawn((
35            Name::new("Terrain"),
36            Transform::default(),
37            Terrain::load(DEFAULT_TERRAIN_PATH),
38            TerrainRenderer::new(&asset_server),
39            TerrainPhysics::default()
40        ))
41        .id();
42
43    commands.spawn((
44        Name::new("Terrain Default Brush"),
45        PaintBrush::default(),
46        ChildOf(entity)
47    ));
48}
49
50#[derive(Resource, Default)]
51struct SaveManager {
52    saving: bool,
53    tiles_to_save: Vec<(ChunkPos, u32, u32)>
54}
55
56/// Handle all tile readbacks from the GPU:
57/// - Always update physics and the CPU heights when a heightmap arrives.
58/// - When in save mode, also write the tile to disk (without its border, as loaded).
59fn handle_extracted_tiles(
60    mut terrain: Query<&mut Terrain>,
61    mut save_manager: ResMut<SaveManager>,
62    mut extracted_tiles: MessageReader<ExtractedTileMessage>
63) {
64    let mut terrain = match terrain.single_mut() {
65        Ok(t) => t,
66        Err(_) => return
67    };
68
69    for message in extracted_tiles.read() {
70        let ExtractedTileMessage {
71            pos,
72            map_type,
73            splat_map_index,
74            data
75        } = message;
76
77        if *map_type == 0 {
78            terrain.dirty_physics.push((*pos, 0, 0, data.clone()));
79            terrain.set_chunk_heights(*pos, data);
80        }
81
82        // Optionally save to disk.
83        if save_manager.saving {
84            let cur_dir = std::env::current_dir().unwrap();
85            let full_path = format!("{}/res/{}", cur_dir.display(), terrain.path);
86            let result = match map_type {
87                0 => save_heightmap(
88                    &format!("{}/heightmap_{}_{}.png", full_path, pos.x, pos.y),
89                    data
90                ),
91                1 => save_splatmap(
92                    &format!(
93                        "{}/splatmap_{}_{}-{}.png",
94                        full_path, pos.x, pos.y, splat_map_index
95                    ),
96                    data
97                ),
98                _ => continue
99            };
100            if let Err(e) = result {
101                error!("Failed to save terrain tile ({}, {}): {}", pos.x, pos.y, e);
102            }
103            save_manager
104                .tiles_to_save
105                .retain(|(p, t, s)| !(p == pos && t == map_type && s == splat_map_index));
106        }
107    }
108
109    if save_manager.saving && save_manager.tiles_to_save.is_empty() {
110        info!("Finished saving terrain.");
111        save_manager.saving = false;
112    }
113}
114
115/// Saves the chunk samples of a GPU heightmap as a 16-bit grayscale PNG.
116fn save_heightmap(path: &str, map: &[u8]) -> Result<(), String> {
117    let data: Vec<u8> = chunk_samples(map)
118        .iter()
119        .flat_map(|h| ((h.clamp(0.0, 1.0) * u16::MAX as f32).round() as u16).to_be_bytes())
120        .collect();
121    write_png(
122        path,
123        &data,
124        png::ColorType::Grayscale,
125        png::BitDepth::Sixteen
126    )
127}
128
129/// Saves the chunk samples of a GPU splat map as an RGBA PNG.
130fn save_splatmap(path: &str, map: &[u8]) -> Result<(), String> {
131    let size = CHUNK_MAP_SIZE as usize;
132    let data: Vec<u8> = (1..size - 1)
133        .flat_map(|z| (1..size - 1).map(move |x| (z * size + x) * 4))
134        .flat_map(|i| map[i..i + 4].to_vec())
135        .collect();
136    write_png(path, &data, png::ColorType::Rgba, png::BitDepth::Eight)
137}
138
139fn write_png(
140    path: &str,
141    data: &[u8],
142    color: png::ColorType,
143    depth: png::BitDepth
144) -> Result<(), String> {
145    let file =
146        std::fs::File::create(path).map_err(|e| format!("Failed to create file {path}: {e}"))?;
147    let mut encoder =
148        png::Encoder::new(std::io::BufWriter::new(file), CHUNK_SAMPLES, CHUNK_SAMPLES);
149    encoder.set_color(color);
150    encoder.set_depth(depth);
151    let mut writer = encoder
152        .write_header()
153        .map_err(|e| format!("Failed to write PNG header: {e}"))?;
154    writer
155        .write_image_data(data)
156        .map_err(|e| format!("Failed to write PNG data: {e}"))
157}
wde_terrain_editor/lib.rs

wde_terrain_editor/
lib.rs
