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("tests/terrain"),
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 when a heightmap arrives.
58/// - When in save mode, also write the tile to disk.
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        // Always feed heightmaps back into the physics system.
78        if *map_type == 0 {
79            terrain.dirty_physics.push((*pos, 0, 0, data.clone()));
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            match map_type {
87                0 => {
88                    let path = format!("{}/heightmap_{}_{}.png", full_path, pos.x, pos.y);
89                    if let Err(e) = save_png_from_channels(
90                        &path,
91                        data,
92                        1,
93                        (CHUNK_RENDER_SUBDIVISIONS, CHUNK_RENDER_SUBDIVISIONS)
94                    ) {
95                        error!("Failed to save heightmap ({}, {}): {}", pos.x, pos.y, e);
96                    }
97                }
98                1 => {
99                    let path = format!(
100                        "{}/splatmap_{}_{}-{}.png",
101                        full_path, pos.x, pos.y, splat_map_index
102                    );
103                    if let Err(e) = save_png_from_channels(
104                        &path,
105                        data,
106                        4,
107                        (CHUNK_RENDER_SUBDIVISIONS, CHUNK_RENDER_SUBDIVISIONS)
108                    ) {
109                        error!("Failed to save splatmap ({}, {}): {}", pos.x, pos.y, e);
110                    }
111                }
112                _ => continue
113            }
114            save_manager
115                .tiles_to_save
116                .retain(|(p, t, s)| !(p == pos && t == map_type && s == splat_map_index));
117        }
118    }
119
120    if save_manager.saving && save_manager.tiles_to_save.is_empty() {
121        info!("Finished saving terrain.");
122        save_manager.saving = false;
123    }
124}
125
126fn save_png_from_channels(
127    path: &str,
128    data: &[u8],
129    channels: usize,
130    dimensions: (u32, u32)
131) -> Result<(), String> {
132    let rgba_data: Vec<u8> = match channels {
133        1 => data.iter().flat_map(|&v| [v, v, v, 255]).collect(),
134        4 => data.to_vec(),
135        _ => return Err("Unsupported number of channels".to_string())
136    };
137
138    let file =
139        std::fs::File::create(path).map_err(|e| format!("Failed to create file {path}: {e}"))?;
140    let w = std::io::BufWriter::new(file);
141    let mut encoder = png::Encoder::new(w, dimensions.0, dimensions.1);
142    encoder.set_color(png::ColorType::Rgba);
143    encoder.set_depth(png::BitDepth::Eight);
144    let mut writer = encoder
145        .write_header()
146        .map_err(|e| format!("Failed to write PNG header: {e}"))?;
147    writer
148        .write_image_data(&rgba_data)
149        .map_err(|e| format!("Failed to write PNG data: {e}"))
150}
wde_terrain_editor/lib.rs

wde_terrain_editor/
lib.rs
