Skip to main content

terrain_forge/algorithms/
prefab.rs

1use crate::{Algorithm, Grid, Rng, Tile};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::path::Path;
5
6#[derive(Debug, Clone)]
7pub struct PrefabConfig {
8    pub max_prefabs: usize,
9    pub min_spacing: usize,
10    pub allow_rotation: bool,
11    pub allow_mirroring: bool,
12    pub weighted_selection: bool,
13    pub placement_mode: PrefabPlacementMode,
14    pub tags: Option<Vec<String>>,
15}
16
17impl Default for PrefabConfig {
18    fn default() -> Self {
19        Self {
20            max_prefabs: 3,
21            min_spacing: 5,
22            allow_rotation: true,
23            allow_mirroring: false,
24            weighted_selection: true,
25            placement_mode: PrefabPlacementMode::Overwrite,
26            tags: None,
27        }
28    }
29}
30
31#[derive(Debug, Clone, Copy)]
32pub enum PrefabPlacementMode {
33    Overwrite,
34    Merge,
35    PaintFloor,
36    PaintWall,
37}
38
39#[derive(Clone, Debug, Serialize, Deserialize)]
40pub struct PrefabLegendEntry {
41    pub tile: Option<String>,
42    pub marker: Option<String>,
43    pub mask: Option<String>,
44}
45
46#[derive(Clone, Debug, Serialize, Deserialize)]
47pub struct PrefabData {
48    pub name: String,
49    pub width: usize,
50    pub height: usize,
51    pub pattern: Vec<String>,
52    pub weight: f32,
53    pub tags: Vec<String>,
54    #[serde(default)]
55    pub legend: Option<HashMap<String, PrefabLegendEntry>>,
56}
57
58#[derive(Clone, Debug, Default)]
59pub struct PrefabCell {
60    pub tile: Option<Tile>,
61    pub marker: Option<String>,
62    pub mask: Option<String>,
63}
64
65#[derive(Clone)]
66pub struct Prefab {
67    pub name: String,
68    pub width: usize,
69    pub height: usize,
70    pub cells: Vec<PrefabCell>,
71    pub symbols: Vec<char>,
72    pub legend: Option<HashMap<char, PrefabLegendEntry>>,
73    pub weight: f32,
74    pub tags: Vec<String>,
75}
76
77impl Prefab {
78    pub fn new(pattern: &[&str]) -> Self {
79        let height = pattern.len();
80        let width = pattern.first().map(|s| s.len()).unwrap_or(0);
81        let (cells, symbols) = parse_pattern_with_legend(pattern, None);
82        Self {
83            name: "unnamed".to_string(),
84            width,
85            height,
86            cells,
87            symbols,
88            legend: None,
89            weight: 1.0,
90            tags: Vec::new(),
91        }
92    }
93
94    pub fn from_data(data: PrefabData) -> Self {
95        let width = data.width;
96        let height = data.height;
97        let legend = data.legend.as_ref().map(convert_legend);
98        let (cells, symbols) = parse_pattern_with_legend(&data.pattern, legend.as_ref());
99
100        Self {
101            name: data.name,
102            width,
103            height,
104            cells,
105            symbols,
106            legend,
107            weight: data.weight,
108            tags: data.tags,
109        }
110    }
111
112    pub fn rect(w: usize, h: usize) -> Self {
113        Self {
114            name: format!("rect_{}x{}", w, h),
115            width: w,
116            height: h,
117            cells: vec![
118                PrefabCell {
119                    tile: Some(Tile::Floor),
120                    marker: None,
121                    mask: None,
122                };
123                w * h
124            ],
125            symbols: vec!['.'; w * h],
126            legend: None,
127            weight: 1.0,
128            tags: vec!["rectangle".to_string()],
129        }
130    }
131
132    /// Rotate prefab 90 degrees clockwise
133    pub fn rotated(&self) -> Self {
134        let mut rotated_cells = vec![PrefabCell::default(); self.width * self.height];
135        let mut rotated_symbols = vec!['#'; self.width * self.height];
136        for y in 0..self.height {
137            for x in 0..self.width {
138                let old_idx = y * self.width + x;
139                let new_x = self.height - 1 - y;
140                let new_y = x;
141                let new_idx = new_y * self.height + new_x;
142                rotated_cells[new_idx] = self.cells[old_idx].clone();
143                rotated_symbols[new_idx] = self.symbols[old_idx];
144            }
145        }
146
147        Self {
148            name: format!("{}_rot90", self.name),
149            width: self.height,
150            height: self.width,
151            cells: rotated_cells,
152            symbols: rotated_symbols,
153            legend: self.legend.clone(),
154            weight: self.weight,
155            tags: self.tags.clone(),
156        }
157    }
158
159    /// Mirror prefab horizontally
160    pub fn mirrored_horizontal(&self) -> Self {
161        let mut mirrored_cells = vec![PrefabCell::default(); self.width * self.height];
162        let mut mirrored_symbols = vec!['#'; self.width * self.height];
163        for y in 0..self.height {
164            for x in 0..self.width {
165                let old_idx = y * self.width + x;
166                let new_x = self.width - 1 - x;
167                let new_idx = y * self.width + new_x;
168                mirrored_cells[new_idx] = self.cells[old_idx].clone();
169                mirrored_symbols[new_idx] = self.symbols[old_idx];
170            }
171        }
172
173        Self {
174            name: format!("{}_mirror_h", self.name),
175            width: self.width,
176            height: self.height,
177            cells: mirrored_cells,
178            symbols: mirrored_symbols,
179            legend: self.legend.clone(),
180            weight: self.weight,
181            tags: self.tags.clone(),
182        }
183    }
184
185    /// Mirror prefab vertically
186    pub fn mirrored_vertical(&self) -> Self {
187        let mut mirrored_cells = vec![PrefabCell::default(); self.width * self.height];
188        let mut mirrored_symbols = vec!['#'; self.width * self.height];
189        for y in 0..self.height {
190            for x in 0..self.width {
191                let old_idx = y * self.width + x;
192                let new_y = self.height - 1 - y;
193                let new_idx = new_y * self.width + x;
194                mirrored_cells[new_idx] = self.cells[old_idx].clone();
195                mirrored_symbols[new_idx] = self.symbols[old_idx];
196            }
197        }
198
199        Self {
200            name: format!("{}_mirror_v", self.name),
201            width: self.width,
202            height: self.height,
203            cells: mirrored_cells,
204            symbols: mirrored_symbols,
205            legend: self.legend.clone(),
206            weight: self.weight,
207            tags: self.tags.clone(),
208        }
209    }
210
211    pub fn get(&self, x: usize, y: usize) -> bool {
212        self.cell_tile(x, y) == Some(Tile::Floor)
213    }
214
215    pub fn cell_tile(&self, x: usize, y: usize) -> Option<Tile> {
216        if x < self.width && y < self.height {
217            self.cells[y * self.width + x].tile
218        } else {
219            None
220        }
221    }
222
223    pub fn cell_marker(&self, x: usize, y: usize) -> Option<&str> {
224        if x < self.width && y < self.height {
225            self.cells[y * self.width + x].marker.as_deref()
226        } else {
227            None
228        }
229    }
230
231    pub fn cell_mask(&self, x: usize, y: usize) -> Option<&str> {
232        if x < self.width && y < self.height {
233            self.cells[y * self.width + x].mask.as_deref()
234        } else {
235            None
236        }
237    }
238
239    pub fn has_tag(&self, tag: &str) -> bool {
240        self.tags.contains(&tag.to_string())
241    }
242}
243
244#[derive(Debug, Clone, Serialize, Deserialize)]
245pub struct PrefabLibraryData {
246    pub prefabs: Vec<PrefabData>,
247}
248
249#[derive(Clone)]
250pub struct PrefabLibrary {
251    prefabs: Vec<Prefab>,
252    by_tag: HashMap<String, Vec<usize>>,
253}
254
255impl PrefabLibrary {
256    pub fn new() -> Self {
257        Self {
258            prefabs: Vec::new(),
259            by_tag: HashMap::new(),
260        }
261    }
262
263    pub fn add_prefab(&mut self, prefab: Prefab) {
264        let index = self.prefabs.len();
265
266        for tag in &prefab.tags {
267            self.by_tag.entry(tag.clone()).or_default().push(index);
268        }
269
270        self.prefabs.push(prefab);
271    }
272
273    pub fn load_from_json<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
274        let content = std::fs::read_to_string(path)?;
275        let data: PrefabLibraryData = serde_json::from_str(&content)?;
276
277        let mut library = Self::new();
278        for prefab_data in data.prefabs {
279            library.add_prefab(Prefab::from_data(prefab_data));
280        }
281
282        Ok(library)
283    }
284
285    pub fn load_from_paths<I, P>(paths: I) -> Result<Self, Box<dyn std::error::Error>>
286    where
287        I: IntoIterator<Item = P>,
288        P: AsRef<Path>,
289    {
290        let mut library = Self::new();
291        for path in paths {
292            let loaded = Self::load_from_json(path)?;
293            library.extend_from(loaded);
294        }
295        Ok(library)
296    }
297
298    pub fn load_from_dir<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
299        let mut entries: Vec<std::path::PathBuf> = std::fs::read_dir(path)?
300            .filter_map(|entry| entry.ok().map(|e| e.path()))
301            .filter(|p| p.extension().and_then(|e| e.to_str()) == Some("json"))
302            .collect();
303        entries.sort();
304        if entries.is_empty() {
305            return Ok(Self::new());
306        }
307        Self::load_from_paths(entries)
308    }
309
310    pub fn save_to_json<P: AsRef<Path>>(&self, path: P) -> Result<(), Box<dyn std::error::Error>> {
311        let data = PrefabLibraryData {
312            prefabs: self
313                .prefabs
314                .iter()
315                .map(|p| PrefabData {
316                    name: p.name.clone(),
317                    width: p.width,
318                    height: p.height,
319                    pattern: self.prefab_to_pattern(p),
320                    weight: p.weight,
321                    tags: p.tags.clone(),
322                    legend: p.legend.as_ref().map(convert_legend_to_strings),
323                })
324                .collect(),
325        };
326
327        let content = serde_json::to_string_pretty(&data)?;
328        std::fs::write(path, content)?;
329        Ok(())
330    }
331
332    fn prefab_to_pattern(&self, prefab: &Prefab) -> Vec<String> {
333        let mut pattern = Vec::new();
334        for y in 0..prefab.height {
335            let mut row = String::new();
336            for x in 0..prefab.width {
337                let idx = y * prefab.width + x;
338                let symbol = prefab.symbols.get(idx).copied().unwrap_or_else(|| {
339                    if prefab.get(x, y) {
340                        '.'
341                    } else {
342                        '#'
343                    }
344                });
345                row.push(symbol);
346            }
347            pattern.push(row);
348        }
349        pattern
350    }
351
352    pub fn get_prefabs(&self) -> &[Prefab] {
353        &self.prefabs
354    }
355
356    pub fn extend_from(&mut self, other: PrefabLibrary) {
357        for prefab in other.prefabs {
358            self.add_prefab(prefab);
359        }
360    }
361
362    pub fn get_by_tag(&self, tag: &str) -> Vec<&Prefab> {
363        self.by_tag
364            .get(tag)
365            .map(|indices| indices.iter().map(|&i| &self.prefabs[i]).collect())
366            .unwrap_or_default()
367    }
368
369    pub fn select_weighted(&self, rng: &mut Rng, tag: Option<&str>) -> Option<&Prefab> {
370        let tags = tag.map(|t| vec![t.to_string()]);
371        self.select_with_tags(rng, tags.as_deref(), true)
372    }
373
374    pub fn select_with_tags(
375        &self,
376        rng: &mut Rng,
377        tags: Option<&[String]>,
378        weighted: bool,
379    ) -> Option<&Prefab> {
380        let candidates = if let Some(tags) = tags {
381            self.get_by_any_tag(tags)
382        } else {
383            self.prefabs.iter().collect()
384        };
385
386        if candidates.is_empty() {
387            return None;
388        }
389
390        if !weighted {
391            return rng.pick(&candidates).copied();
392        }
393
394        let total_weight: f32 = candidates.iter().map(|p| p.weight).sum();
395        if total_weight <= 0.0 {
396            return rng.pick(&candidates).copied();
397        }
398
399        let mut target = rng.random() as f32 * total_weight;
400        for prefab in &candidates {
401            target -= prefab.weight;
402            if target <= 0.0 {
403                return Some(prefab);
404            }
405        }
406
407        candidates.last().copied()
408    }
409
410    pub fn get_by_any_tag(&self, tags: &[String]) -> Vec<&Prefab> {
411        if tags.is_empty() {
412            return Vec::new();
413        }
414        let mut indices = Vec::new();
415        for tag in tags {
416            if let Some(found) = self.by_tag.get(tag) {
417                indices.extend_from_slice(found);
418            }
419        }
420        if indices.is_empty() {
421            return Vec::new();
422        }
423        indices.sort_unstable();
424        indices.dedup();
425        indices.iter().map(|&i| &self.prefabs[i]).collect()
426    }
427
428    pub fn create_default() -> Self {
429        let mut library = Self::new();
430
431        // Add some basic prefabs
432        let mut room = Prefab::rect(5, 5);
433        room.name = "small_room".to_string();
434        room.tags = vec!["room".to_string(), "small".to_string()];
435        library.add_prefab(room);
436
437        let mut corridor = Prefab::new(&[".....", "#####"]);
438        corridor.name = "corridor".to_string();
439        corridor.tags = vec!["corridor".to_string()];
440        library.add_prefab(corridor);
441
442        let mut cross = Prefab::new(&["#.#", "...", "#.#"]);
443        cross.name = "cross".to_string();
444        cross.tags = vec!["junction".to_string()];
445        library.add_prefab(cross);
446
447        library
448    }
449}
450
451impl Default for PrefabLibrary {
452    fn default() -> Self {
453        Self::create_default()
454    }
455}
456
457#[derive(Debug, Clone, Default)]
458pub struct PrefabTransform {
459    pub rotation: u8, // 0, 1, 2, 3 for 0°, 90°, 180°, 270°
460    pub mirror_h: bool,
461    pub mirror_v: bool,
462}
463
464impl PrefabTransform {
465    pub fn apply(&self, prefab: &Prefab) -> Prefab {
466        let mut result = prefab.clone();
467
468        // Apply mirroring first
469        if self.mirror_h {
470            result = result.mirrored_horizontal();
471        }
472        if self.mirror_v {
473            result = result.mirrored_vertical();
474        }
475
476        // Apply rotation
477        for _ in 0..self.rotation {
478            result = result.rotated();
479        }
480
481        result
482    }
483
484    pub fn random(rng: &mut Rng, allow_rotation: bool, allow_mirroring: bool) -> Self {
485        Self {
486            rotation: if allow_rotation {
487                rng.range(0, 4) as u8
488            } else {
489                0
490            },
491            mirror_h: allow_mirroring && rng.chance(0.5),
492            mirror_v: allow_mirroring && rng.chance(0.5),
493        }
494    }
495}
496pub struct PrefabPlacer {
497    config: PrefabConfig,
498    library: PrefabLibrary,
499}
500
501impl PrefabPlacer {
502    pub fn new(config: PrefabConfig, library: PrefabLibrary) -> Self {
503        Self { config, library }
504    }
505
506    pub fn with_library(library: PrefabLibrary) -> Self {
507        Self::new(PrefabConfig::default(), library)
508    }
509}
510
511impl Default for PrefabPlacer {
512    fn default() -> Self {
513        Self::with_library(PrefabLibrary::default())
514    }
515}
516
517impl Algorithm<Tile> for PrefabPlacer {
518    fn generate(&self, grid: &mut Grid<Tile>, seed: u64) {
519        self.generate_internal(grid, seed, None);
520    }
521
522    fn name(&self) -> &'static str {
523        "PrefabPlacer"
524    }
525}
526
527impl PrefabPlacer {
528    pub fn generate_with_semantic(
529        &self,
530        grid: &mut Grid<Tile>,
531        seed: u64,
532        semantic: &mut crate::semantic::SemanticLayers,
533    ) {
534        self.generate_internal(grid, seed, Some(semantic));
535    }
536
537    fn generate_internal(
538        &self,
539        grid: &mut Grid<Tile>,
540        seed: u64,
541        mut semantic: Option<&mut crate::semantic::SemanticLayers>,
542    ) {
543        let mut rng = Rng::new(seed);
544        let mut placed: Vec<(usize, usize, usize, usize)> = Vec::new();
545
546        for _ in 0..self.config.max_prefabs * 10 {
547            if placed.len() >= self.config.max_prefabs {
548                break;
549            }
550
551            let base_prefab = if let Some(prefab) = self.library.select_with_tags(
552                &mut rng,
553                self.config.tags.as_deref(),
554                self.config.weighted_selection,
555            ) {
556                prefab
557            } else {
558                continue;
559            };
560
561            // Apply random transformations
562            let transform = PrefabTransform::random(
563                &mut rng,
564                self.config.allow_rotation,
565                self.config.allow_mirroring,
566            );
567            let prefab = transform.apply(base_prefab);
568
569            if prefab.width + 2 >= grid.width() || prefab.height + 2 >= grid.height() {
570                continue;
571            }
572
573            let x = rng.range_usize(1, grid.width() - prefab.width - 1);
574            let y = rng.range_usize(1, grid.height() - prefab.height - 1);
575
576            let overlaps = placed.iter().any(|&(px, py, pw, ph)| {
577                let s = self.config.min_spacing;
578                !(x + prefab.width + s < px
579                    || px + pw + s < x
580                    || y + prefab.height + s < py
581                    || py + ph + s < y)
582            });
583
584            if overlaps {
585                continue;
586            }
587
588            for py in 0..prefab.height {
589                for px in 0..prefab.width {
590                    let cell_tile = prefab.cell_tile(px, py);
591                    let cell_marker = prefab.cell_marker(px, py);
592                    let cell_mask = prefab.cell_mask(px, py);
593                    let gx = (x + px) as i32;
594                    let gy = (y + py) as i32;
595
596                    let mut applied = false;
597                    if let Some(tile) = cell_tile {
598                        let current = *grid.get(gx, gy).unwrap_or(&Tile::Wall);
599                        let should_place = match self.config.placement_mode {
600                            PrefabPlacementMode::Overwrite => true,
601                            PrefabPlacementMode::Merge => matches!(current, Tile::Wall),
602                            PrefabPlacementMode::PaintFloor => matches!(current, Tile::Floor),
603                            PrefabPlacementMode::PaintWall => matches!(current, Tile::Wall),
604                        };
605                        if should_place {
606                            grid.set(gx, gy, tile);
607                            applied = true;
608                        }
609                    }
610
611                    if let Some(layers) = semantic.as_deref_mut() {
612                        let marker_allowed = cell_tile.is_none() || applied;
613                        if marker_allowed {
614                            if let Some(tag) = cell_marker {
615                                layers.markers.push(crate::semantic::Marker::with_tag(
616                                    gx as u32,
617                                    gy as u32,
618                                    tag.to_string(),
619                                ));
620                            }
621                            if let Some(mask) = cell_mask {
622                                apply_prefab_mask(&mut layers.masks, gx, gy, mask);
623                            }
624                        }
625                    }
626                }
627            }
628            placed.push((x, y, prefab.width, prefab.height));
629        }
630    }
631}
632
633fn parse_pattern_with_legend(
634    pattern: &[impl AsRef<str>],
635    legend: Option<&HashMap<char, PrefabLegendEntry>>,
636) -> (Vec<PrefabCell>, Vec<char>) {
637    let mut cells = Vec::new();
638    let mut symbols = Vec::new();
639    for row in pattern {
640        for ch in row.as_ref().chars() {
641            let cell = if let Some(legend) = legend {
642                legend
643                    .get(&ch)
644                    .map(parse_legend_entry)
645                    .unwrap_or_else(|| default_cell_from_symbol(ch))
646            } else {
647                default_cell_from_symbol(ch)
648            };
649            cells.push(cell);
650            symbols.push(ch);
651        }
652    }
653    (cells, symbols)
654}
655
656fn default_cell_from_symbol(ch: char) -> PrefabCell {
657    match ch {
658        '.' => PrefabCell {
659            tile: Some(Tile::Floor),
660            marker: None,
661            mask: None,
662        },
663        '#' => PrefabCell::default(),
664        _ => PrefabCell::default(),
665    }
666}
667
668fn parse_legend_entry(entry: &PrefabLegendEntry) -> PrefabCell {
669    PrefabCell {
670        tile: parse_tile_name(entry.tile.as_deref()),
671        marker: entry.marker.clone(),
672        mask: entry.mask.clone(),
673    }
674}
675
676fn parse_tile_name(value: Option<&str>) -> Option<Tile> {
677    let value = value?;
678    match value.trim().to_ascii_lowercase().as_str() {
679        "floor" | "f" => Some(Tile::Floor),
680        "wall" | "w" => Some(Tile::Wall),
681        "empty" | "none" | "skip" => None,
682        _ => None,
683    }
684}
685
686fn convert_legend(legend: &HashMap<String, PrefabLegendEntry>) -> HashMap<char, PrefabLegendEntry> {
687    let mut out = HashMap::new();
688    for (key, value) in legend {
689        let mut chars = key.chars();
690        if let Some(ch) = chars.next() {
691            out.insert(ch, value.clone());
692        }
693    }
694    out
695}
696
697fn convert_legend_to_strings(
698    legend: &HashMap<char, PrefabLegendEntry>,
699) -> HashMap<String, PrefabLegendEntry> {
700    legend
701        .iter()
702        .map(|(ch, entry)| (ch.to_string(), entry.clone()))
703        .collect()
704}
705
706fn apply_prefab_mask(masks: &mut crate::semantic::Masks, x: i32, y: i32, mask: &str) {
707    if x < 0 || y < 0 {
708        return;
709    }
710    let (x, y) = (x as usize, y as usize);
711    if y >= masks.height || x >= masks.width {
712        return;
713    }
714    match mask.trim().to_ascii_lowercase().as_str() {
715        "no_spawn" | "nospawn" | "reserved" => {
716            if let Some(row) = masks.no_spawn.get_mut(y) {
717                if let Some(cell) = row.get_mut(x) {
718                    *cell = true;
719                }
720            }
721        }
722        _ => {}
723    }
724}