Skip to main content

kvas_engine/
paint.rs

1//! Texture painting: blend several surface textures with a painted mask, so dirt can be brushed
2//! along a road edge and fade into grass without a hard seam.
3//!
4//! A [`TexturePaint`] component names four layer textures and a mask. The mask is an RGBA image:
5//! red, green and blue are the weights of layers 2, 3 and 4; whatever is left is layer 1. The
6//! editor's brush edits the mask in memory ([`PaintMaps`]) and saves it to `assets/paint/<name>.png`.
7//!
8//! ```
9//! use kvas_engine::paint::{PaintMaps, BrushMode};
10//! let maps = PaintMaps::global();
11//! maps.create("doc-example", 32, 32);
12//! maps.brush("doc-example", [0.5, 0.5], 0.2, 1.0, 0.5, BrushMode::Layer(1));
13//! let (_, rgba, _) = maps.snapshot("doc-example").unwrap();
14//! assert!(rgba[(16 * 32 + 16) * 4] > 200); // the red channel (layer 2) at the centre
15//! ```
16
17use crate::assets::Assets;
18use crate::components::*;
19use crate::reflect::{AssetKind, Field};
20use serde::{Deserialize, Serialize};
21use std::collections::HashMap;
22use std::sync::{Mutex, OnceLock};
23
24/// What the brush does.
25#[derive(Clone, Copy, PartialEq, Eq, Debug)]
26pub enum BrushMode {
27    /// Paint layer 0..=3 (0 is the base layer: it erases the others).
28    Layer(usize),
29    /// Soften the mask (blur under the brush).
30    Smooth,
31}
32
33/// One painted mask kept in memory.
34#[derive(Clone)]
35pub struct PaintMap {
36    /// Width in pixels.
37    pub width: u32,
38    /// Height in pixels.
39    pub height: u32,
40    /// RGBA8 pixels, row-major.
41    pub rgba: Vec<u8>,
42    /// Counter that grows with every edit.
43    pub version: u64,
44    /// Edited since the last save.
45    pub dirty: bool,
46}
47
48/// All painted masks of the running editor / game, by name.
49pub struct PaintMaps {
50    maps: Mutex<HashMap<String, PaintMap>>,
51}
52
53impl PaintMaps {
54    /// The process-wide registry (the renderer reads masks from here).
55    pub fn global() -> &'static PaintMaps {
56        static MAPS: OnceLock<PaintMaps> = OnceLock::new();
57        MAPS.get_or_init(|| PaintMaps { maps: Mutex::new(HashMap::new()) })
58    }
59
60    /// Creates an empty mask (all base layer) unless one exists.
61    pub fn create(&self, name: &str, width: u32, height: u32) {
62        let mut maps = self.maps.lock().unwrap();
63        maps.entry(name.to_string()).or_insert_with(|| PaintMap { width, height, rgba: vec![0; (width * height * 4) as usize], version: 1, dirty: false });
64        // The alpha channel stays opaque so the image previews correctly.
65        if let Some(m) = maps.get_mut(name) {
66            for px in m.rgba.chunks_exact_mut(4) {
67                px[3] = 255;
68            }
69        }
70    }
71
72    /// Whether a mask is loaded.
73    pub fn exists(&self, name: &str) -> bool {
74        self.maps.lock().unwrap().contains_key(name)
75    }
76
77    /// Loads `paint/<name>.png` from the project, or creates a fresh mask.
78    pub fn load_or_create(&self, name: &str, assets: &Assets, size: u32) {
79        if self.exists(name) {
80            return;
81        }
82        let loaded = assets.read(&format!("paint/{name}.png")).ok().and_then(|b| image::load_from_memory(&b).ok());
83        match loaded {
84            Some(img) => {
85                let rgba = img.to_rgba8();
86                self.maps.lock().unwrap().insert(name.to_string(), PaintMap { width: rgba.width(), height: rgba.height(), rgba: rgba.into_raw(), version: 1, dirty: false });
87            }
88            None => self.create(name, size, size),
89        }
90    }
91
92    /// `(size, pixels, version)` of a mask, for uploading to the GPU.
93    pub fn snapshot(&self, name: &str) -> Option<([u32; 2], Vec<u8>, u64)> {
94        let maps = self.maps.lock().unwrap();
95        let m = maps.get(name)?;
96        Some(([m.width, m.height], m.rgba.clone(), m.version))
97    }
98
99    /// Restores an undo snapshot, incrementing the version so the renderer uploads it again.
100    /// Rejects empty/oversized images and mismatched RGBA lengths without changing the map.
101    pub fn restore(&self, name: &str, size: [u32; 2], rgba: Vec<u8>) -> bool {
102        if size.iter().any(|&v| v == 0 || v > 8192) || rgba.len() != size[0] as usize * size[1] as usize * 4 { return false; }
103        let mut maps = self.maps.lock().unwrap();
104        let version = maps.get(name).map_or(1, |m| m.version.saturating_add(1));
105        maps.insert(name.into(), PaintMap { width: size[0], height: size[1], rgba, version, dirty: true });
106        true
107    }
108
109    /// Current version of a mask (0 when missing).
110    pub fn version(&self, name: &str) -> u64 {
111        self.maps.lock().unwrap().get(name).map_or(0, |m| m.version)
112    }
113
114    /// Applies one brush dab at `uv` (0..1). `radius` is in uv units, `strength` 0..1 per dab,
115    /// `hardness` 0 (soft falloff) .. 1 (hard edge).
116    pub fn brush(&self, name: &str, uv: [f32; 2], radius: f32, strength: f32, hardness: f32, mode: BrushMode) {
117        let mut maps = self.maps.lock().unwrap();
118        let Some(m) = maps.get_mut(name) else { return };
119        let (w, h) = (m.width as i32, m.height as i32);
120        let (cx, cy) = (uv[0] * w as f32, uv[1] * h as f32);
121        let r = (radius * w as f32).max(1.0);
122        let (x0, x1) = (((cx - r).floor() as i32).max(0), ((cx + r).ceil() as i32).min(w - 1));
123        let (y0, y1) = (((cy - r).floor() as i32).max(0), ((cy + r).ceil() as i32).min(h - 1));
124        for y in y0..=y1 {
125            for x in x0..=x1 {
126                let d = (((x as f32 + 0.5 - cx).powi(2) + (y as f32 + 0.5 - cy).powi(2)).sqrt()) / r;
127                if d >= 1.0 {
128                    continue;
129                }
130                let edge = hardness.clamp(0.0, 0.98);
131                let fall = if d <= edge { 1.0 } else { let t = (d - edge) / (1.0 - edge); let s = 1.0 - t; s * s * (3.0 - 2.0 * s) };
132                let amount = fall * strength.clamp(0.0, 1.0);
133                let i = ((y * w + x) * 4) as usize;
134                let mut weights = [m.rgba[i] as f32 / 255.0, m.rgba[i + 1] as f32 / 255.0, m.rgba[i + 2] as f32 / 255.0];
135                match mode {
136                    BrushMode::Layer(0) => weights.iter_mut().for_each(|v| *v = (*v - amount).max(0.0)),
137                    BrushMode::Layer(k) => {
138                        let k = (k - 1).min(2);
139                        weights[k] = (weights[k] + amount).min(1.0);
140                        let total: f32 = weights.iter().sum();
141                        if total > 1.0 {
142                            // Take the surplus from the other layers in proportion to what they hold.
143                            let others: f32 = total - weights[k];
144                            let keep = (1.0 - weights[k]).max(0.0);
145                            for (j, v) in weights.iter_mut().enumerate() {
146                                if j != k && others > 0.0 {
147                                    *v *= keep / others;
148                                }
149                            }
150                        }
151                    }
152                    BrushMode::Smooth => {
153                        let mut avg = [0.0f32; 3];
154                        let mut n = 0.0;
155                        for (dx, dy) in [(-2, 0), (2, 0), (0, -2), (0, 2), (-1, -1), (1, 1), (-1, 1), (1, -1)] {
156                            let (nx, ny) = ((x + dx).clamp(0, w - 1), (y + dy).clamp(0, h - 1));
157                            let j = ((ny * w + nx) * 4) as usize;
158                            for c in 0..3 {
159                                avg[c] += m.rgba[j + c] as f32 / 255.0;
160                            }
161                            n += 1.0;
162                        }
163                        for c in 0..3 {
164                            weights[c] += (avg[c] / n - weights[c]) * amount;
165                        }
166                    }
167                }
168                for c in 0..3 {
169                    m.rgba[i + c] = (weights[c] * 255.0 + 0.5) as u8;
170                }
171            }
172        }
173        m.version += 1;
174        m.dirty = true;
175    }
176
177    /// Fills the whole mask with one layer (a quick base coat).
178    pub fn fill(&self, name: &str, layer: usize) {
179        let mut maps = self.maps.lock().unwrap();
180        if let Some(m) = maps.get_mut(name) {
181            for px in m.rgba.chunks_exact_mut(4) {
182                px[..3].copy_from_slice(&[0, 0, 0]);
183                if layer >= 1 {
184                    px[(layer - 1).min(2)] = 255;
185                }
186            }
187            m.version += 1;
188            m.dirty = true;
189        }
190    }
191
192    /// Writes every edited mask to `<assets>/paint/<name>.png`. Returns how many were saved.
193    pub fn save_all(&self, assets_dir: &std::path::Path) -> usize {
194        let mut maps = self.maps.lock().unwrap();
195        let mut saved = 0;
196        for (name, m) in maps.iter_mut().filter(|(_, m)| m.dirty) {
197            let dir = assets_dir.join("paint");
198            let _ = std::fs::create_dir_all(&dir);
199            if let Some(img) = image::RgbaImage::from_raw(m.width, m.height, m.rgba.clone()) {
200                if img.save(dir.join(format!("{name}.png"))).is_ok() {
201                    m.dirty = false;
202                    saved += 1;
203                }
204            }
205        }
206        saved
207    }
208
209    /// Names of masks with unsaved edits.
210    pub fn dirty_names(&self) -> Vec<String> {
211        self.maps.lock().unwrap().iter().filter(|(_, m)| m.dirty).map(|(n, _)| n.clone()).collect()
212    }
213}
214
215/// Four textures blended by a painted mask. Put it on an object with a mesh (a plane, a terrain
216/// chunk, a road strip). The component takes over the object's material.
217#[derive(Clone, Serialize, Deserialize)]
218#[serde(default)]
219pub struct TexturePaint {
220    /// Base layer (where nothing is painted), usually grass.
221    pub layer1: String,
222    /// Layer painted with the red channel (dirt, sand...).
223    pub layer2: String,
224    /// Layer painted with the green channel.
225    pub layer3: String,
226    /// Layer painted with the blue channel.
227    pub layer4: String,
228    /// Name of the painted mask (saved as `paint/<name>.png`). Empty = the object's name.
229    pub mask: String,
230    /// Texture repeats per metre of the surface.
231    pub tiling: f32,
232    /// Noise that roughens the borders between layers, 0..1.
233    pub breakup: f32,
234    /// Mask resolution when a new mask is created.
235    pub resolution: i32,
236}
237
238impl Default for TexturePaint {
239    fn default() -> Self {
240        Self { layer1: String::new(), layer2: String::new(), layer3: String::new(), layer4: String::new(), mask: String::new(), tiling: 0.25, breakup: 0.35, resolution: 512 }
241    }
242}
243
244impl Component for TexturePaint {
245    fn fields(&mut self) -> Vec<Field<'_>> {
246        vec![
247            Field::asset("layer1", &mut self.layer1, AssetKind::Texture),
248            Field::asset("layer2", &mut self.layer2, AssetKind::Texture),
249            Field::asset("layer3", &mut self.layer3, AssetKind::Texture),
250            Field::asset("layer4", &mut self.layer4, AssetKind::Texture),
251            Field::text("mask", &mut self.mask),
252            Field::f32("tiling", &mut self.tiling),
253            Field::f32("breakup", &mut self.breakup),
254            Field::i32("resolution", &mut self.resolution),
255        ]
256    }
257}
258
259impl TexturePaint {
260    /// The mask name for an entity called `entity_name`.
261    pub fn mask_name(&self, entity_name: &str) -> String {
262        let n = if self.mask.trim().is_empty() { entity_name } else { self.mask.trim() };
263        n.chars().map(|c| if c.is_ascii_alphanumeric() || c == '-' || c == '_' { c } else { '_' }).collect()
264    }
265}
266
267/// Prefix of the mask texture keys the renderer understands (`paint:<name>#<version>`).
268pub const PAINT_PREFIX: &str = "paint:";
269/// Prefix of the layer atlas keys (`splat:<a>|<b>|<c>|<d>`).
270pub const SPLAT_PREFIX: &str = "splat:";
271
272/// Texture key of a mask at its current version (changes with every edit so the GPU copy refreshes).
273pub fn paint_key(name: &str) -> String {
274    format!("{PAINT_PREFIX}{name}#{}", PaintMaps::global().version(name))
275}
276
277/// Keeps the material of every painted object in sync with its component and mask.
278pub fn update(world: &mut kvas_ecs::World, assets: Option<&Assets>) {
279    for e in world.entities_with::<TexturePaint>() {
280        let Some(tp) = world.get::<TexturePaint>(e).cloned() else { continue };
281        let name = tp.mask_name(&world.get::<Name>(e).map(|n| n.0.clone()).unwrap_or_else(|| "mask".into()));
282        let maps = PaintMaps::global();
283        if !maps.exists(&name) {
284            match assets {
285                Some(a) => maps.load_or_create(&name, a, tp.resolution.clamp(64, 2048) as u32),
286                None => maps.create(&name, tp.resolution.clamp(64, 2048) as u32, tp.resolution.clamp(64, 2048) as u32),
287            }
288        }
289        if !world.has::<MeshRenderer>(e) {
290            world.insert(e, MeshRenderer::default());
291        }
292        let key = paint_key(&name);
293        let atlas = format!("{SPLAT_PREFIX}{}|{}|{}|{}", tp.layer1, tp.layer2, tp.layer3, tp.layer4);
294        let m = world.get_mut::<MeshRenderer>(e).unwrap();
295        if m.texture != key {
296            m.texture = key;
297        }
298        if m.normal_map != atlas {
299            m.normal_map = atlas;
300        }
301        m.shader = "builtin:splat".into();
302        m.tiling = tp.tiling;
303        m.user[0] = tp.breakup;
304        m.cast_shadows = false;
305    }
306}