Wave function collapse

1 files changed, 275 insertions, 36 deletions

diff --git a/dungeon/src/wfc.rs b/dungeon/src/wfc.rs
index 0309546..e0d820b 100644
--- a/dungeon/src/wfc.rs
+++ b/dungeon/src/wfc.rs
@@ -1,23 +1,32 @@
 //! The `wfc` module contains the implementation of the
 //! Wave Function Collapse algorithm for procedural generation of the `Floor`.
 
+// TODO: Add pattern size input
+
+// Reference map's map size and tile count constants
+use crate::map::TILE_COUNT;
+
 use crate::map::Tile;
+use crate::pos::Direction;
 use crate::pos::Pos;
 
+use rand::prelude::IndexedRandom;
+
 /// The `State` struct represents each possible state of a tile in the WFC algorithm.
-struct State {
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub(crate) struct State {
 	/// Position of the tile
 	#[expect(dead_code)]
 	pos: Pos,
 	/// Possible tiles for this state
 	possible_tiles: Vec<Tile>,
 	/// Entropy of the state
-	entropy: f32,
+	entropy: usize,
 }
 impl State {
 	/// Creates a new State instance
-	fn new(pos: Pos, possible_tiles: Vec<Tile>) -> Self {
-		let entropy = 0.0;
+	pub fn new(pos: Pos, possible_tiles: Vec<Tile>) -> Self {
+		let entropy = possible_tiles.len();
 		Self {
 			pos,
 			possible_tiles,
@@ -26,47 +35,129 @@ impl State {
 	}
 
 	/// Updates the possible tiles and recalculates entropy
-	#[expect(dead_code)]
-	fn update_possible_tiles(&mut self, possible_tiles: Vec<Tile>) {
+	pub fn set_possible_tiles(&mut self, possible_tiles: Vec<Tile>) {
 		self.possible_tiles = possible_tiles;
 		self.entropy = self.entropy();
 	}
 
 	/// Calculates the entropy of the state
 	/// TODO: Implement shannon entropy calculation
-	fn entropy(&self) -> f32 {
-		self.possible_tiles.len() as f32
+	pub fn entropy(&self) -> usize {
+		self.possible_tiles.len()
+	}
+
+	/// Checks if the state has been collapsed to a single tile
+	pub fn is_collapsed(&self) -> bool {
+		self.possible_tiles.len() == 1
+	}
+
+	/// Returns a reference to the possible tiles
+	pub fn possible_tiles(&self) -> &Vec<Tile> {
+		&self.possible_tiles
+	}
+
+	/// Collapses the state to a single tile
+	pub fn collapse(&mut self, tile: Tile) {
+		self.possible_tiles = vec![tile];
+		self.entropy = 1;
+	}
+}
+
+/// The `Compatability` struct represents the compatibility rules between tiles.
+pub(crate) struct Compatability {
+	/// The first tile in the compatibility pair
+	tile1: Tile,
+	/// The second tile in the compatibility pair
+	tile2: Tile,
+	/// The direction from tile1 -> tile2
+	direction: Direction,
+}
+
+impl Compatability {
+	/// Creates a new Compatability instance
+	pub fn new(tile1: Tile, tile2: Tile, direction: Direction) -> Self {
+		Self {
+			tile1,
+			tile2,
+			direction,
+		}
+	}
+}
+
+/// The `CompatabilityChecker` struct manages compatibility rules between tiles.
+pub(crate) struct CompatabilityChecker {
+	compatabilities: Vec<Compatability>,
+}
+
+impl CompatabilityChecker {
+	/// Creates a new CompatabilityChecker instance
+	pub fn new(input: &[Vec<Tile>]) -> Self {
+		let mut compatabilities = Vec::new();
+		// Which pairs of tiles can be placed next to each other and in which directions
+		for row in input {
+			for window in row.windows(2) {
+				if let [tile1, tile2] = window {
+					compatabilities.push(Compatability::new(
+						*tile1,
+						*tile2,
+						Direction::East,
+					));
+					compatabilities.push(Compatability::new(
+						*tile2,
+						*tile1,
+						Direction::West,
+					));
+				}
+			}
+		}
+
+		for col in 0..input[0].len() {
+			for row in 0..input.len() - 1 {
+				let tile1 = input[row][col];
+				let tile2 = input[row + 1][col];
+				compatabilities.push(Compatability::new(tile1, tile2, Direction::North));
+				compatabilities.push(Compatability::new(tile2, tile1, Direction::South));
+			}
+		}
+
+		Self { compatabilities }
+	}
+
+	/// Checks if two tiles are compatible in the given direction
+	pub fn is_valid(&self, tile1: Tile, tile2: Tile, direction: Direction) -> bool {
+		self.compatabilities
+			.iter()
+			.any(|c| c.tile1 == tile1 && c.tile2 == tile2 && c.direction == direction)
 	}
 }
 
 /// The `Wfc` struct encapsulates the Wave Function Collapse algorithm.
-pub(crate) struct Wfc {
-	/// The seed used for randomness in the algorithm
-	#[expect(dead_code)]
-	seed: u64,
-	/// The size of the map to generate
-	#[expect(dead_code)]
-	mapsize: u16,
-	/// The current state of the WFC algorithm (smart pointer for interior mutability)
-	states: Vec<std::cell::RefCell<State>>,
+pub struct Wfc {
+	/// The compatibility checker for tile constraints
+	compatabilities: CompatabilityChecker,
 	/// The random number generator
-	#[expect(dead_code)]
 	rng: rand::rngs::StdRng,
+	/// The current state of the WFC algorithm (smart pointer for interior mutability)
+	states: Vec<std::cell::RefCell<State>>,
+	/// The States that have been collapsed (as positions)
+	collapsed_states: Vec<Pos>,
+	// Note: uses map's map size and tile count constants
 }
 impl Wfc {
 	/// Creates a new Wfc instance with the given seed and map size
-	pub(crate) fn new(seed: u64, mapsize: u16) -> Self {
+	pub fn new(input: &[Vec<Tile>], seed: u64) -> Self {
 		let rng = rand::SeedableRng::seed_from_u64(seed);
 		Self {
-			seed,
-			mapsize,
-			states: Vec::with_capacity(mapsize as usize),
+			compatabilities: CompatabilityChecker::new(input),
 			rng,
+			states: Vec::with_capacity(TILE_COUNT),
+			collapsed_states: Vec::new(),
 		}
 	}
 
 	/// Initializes the states vector with default states
-	pub(crate) fn initialize_states(&mut self) {
+	fn initialize_states(&mut self) {
+		// Store a superposition of all possible tiles for each position
 		for pos in Pos::values() {
 			let possible_tiles = Tile::values().collect();
 			let state = State::new(pos, possible_tiles);
@@ -74,29 +165,177 @@ impl Wfc {
 		}
 	}
 
+	/// Find the state with the lowest entropy that has not been collapsed yet
+	fn find_lowest_entropy_state(&self) -> Option<State> {
+		self.states
+			.iter()
+			.filter_map(|state| {
+				let borrowed_state = state.borrow();
+				if !borrowed_state.is_collapsed() {
+					Some(borrowed_state.clone())
+				} else {
+					None
+				}
+			})
+			.min_by_key(State::entropy)
+	}
+
+	/// Collapses the given state by selecting one of its possible tiles (using the rng)
+	fn collapse_state(&mut self, state: &State) {
+		// Select a random tile from the possible tiles
+		if let Some(&selected_tile) = state.possible_tiles().choose(&mut self.rng) {
+			// Update the state to be collapsed with the selected tile
+			if let Some(state_ref) = self.states.get(state.pos.idx()) {
+				let mut state_mut = state_ref.borrow_mut();
+				state_mut.collapse(selected_tile);
+				self.collapsed_states.push(state.pos);
+			}
+		}
+	}
+
+	/// Propagates constraints after a state has been collapsed, starting from the given state
+	fn propagate(&mut self, state: &State) {
+		// Keep a queue of states to propagate constraints from
+		let mut queue = vec![state.pos];
+
+		while let Some(current_pos) = queue.pop() {
+			// For each direction, check neighboring states
+			for direction in Direction::values() {
+				if let Some(neighbor_pos) = current_pos.step(direction)
+					&& let Some(neighbor_state_ref) = self.states.get(neighbor_pos.idx())
+				{
+					let mut neighbor_state = neighbor_state_ref.borrow_mut();
+					let valid_tiles: Vec<Tile> = neighbor_state
+						.possible_tiles()
+						.iter()
+						.cloned()
+						.filter(|&tile| {
+							self.compatabilities.is_valid(
+								state.possible_tiles()[0],
+								tile,
+								direction,
+							)
+						})
+						.collect();
+
+					if valid_tiles.len() < neighbor_state.possible_tiles().len() {
+						neighbor_state.set_possible_tiles(valid_tiles);
+						queue.push(neighbor_pos);
+					}
+				}
+			}
+		}
+	}
+
 	/// Runs the Wave Function Collapse algorithm to generate the map
 	#[expect(clippy::unused_self)]
 	pub(crate) fn run(&mut self) {
 		// ----- Step 1: Read in input sample and define constraints -----
-		// TODO: add support for sample image input
+		// (This step is currently handled in the CompatabilityChecker::new method)
 
 		// ----- Step 2: Initialize the grid of possible states -----
+		self.initialize_states();
 
-		// Start with an empty box of tiles
+		// Keep running until all states are collapsed
+		while self.collapsed_states.len() != TILE_COUNT {
+			// ----- Step 3: Collapse the superpositions -----
+			// Find the shannon entropy of all superpositions
+			if let Some(state) = self.find_lowest_entropy_state() {
+				// Collapse the tile with the lowest entropy
+				// (Assign a seeded random superposition)
+				self.collapse_state(&state);
 
-		// Store a superposition of all possible tiles for each position
+				// ----- Step 4: Propagate -----
+				// Update neighboring superpositions based on the collapsed tile
+				// (Remove any that violate constraints)
+				self.propagate(&state);
+			}
 
-		// ----- Step 3: Collapse the superpositions -----
-		// Calculate the shannon entropy of all superpositions
-		// Collapse the tile with the lowest entropy
-		// (Assign a seeded random superposition)
+			// ----- Step 5: Repeat until all superpositions are collapsed -----
+			// If any tiles are left with no possible states, restart
+			// TODO: backtrack
+		}
+	}
+
+	/// Returns the generated tiles as a boxed array
+	pub fn to_tiles(&self) -> Box<[Tile; TILE_COUNT]> {
+		let mut arr = [Tile::Air; TILE_COUNT];
+
+		for (i, state_ref) in self.states.iter().enumerate() {
+			let state = state_ref.borrow();
+			if state.is_collapsed() {
+				arr[i] = state.possible_tiles()[0];
+			}
+		}
+
+		Box::new(arr)
+	}
+}
+
+/// Tests
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	use crate::map::MAP_SIZE_USIZE;
+
+	#[test]
+	fn test_wfc_generation() {
+		// 16x16 input
+		// 0 is wall, 1 is air
+		let input = vec![
+			vec![0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+			vec![0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
+			vec![0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
+			vec![0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
+			vec![0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
+			vec![1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
+			vec![0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
+			vec![0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0],
+			vec![0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
+			vec![0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
+			vec![0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0],
+			vec![1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1],
+			vec![0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
+			vec![0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0],
+			vec![0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 0],
+			vec![0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+		];
+
+		// Convert input to Tiles
+		let input_tiles: Vec<Vec<Tile>> = input
+			.iter()
+			.map(|row| {
+				row.iter()
+					.map(|&cell| if cell == 0 { Tile::Wall } else { Tile::Air })
+					.collect()
+			})
+			.collect();
+
+		let seed = 12345;
+		let mut wfc = Wfc::new(&input_tiles, seed);
+		wfc.run();
+		let tiles = wfc.to_tiles();
+
+		// Print the generated tiles
+		for y in 0..MAP_SIZE_USIZE {
+			for x in 0..MAP_SIZE_USIZE {
+				// Print air as "██", walls as "  " and stairs as ">"
+				let tile = tiles[y * MAP_SIZE_USIZE + x];
+				let symbol = match tile {
+					Tile::Air => "██",
+					Tile::Wall => "  ",
+					Tile::Stairs => ">",
+				};
+				print!("{symbol}");
+			}
+			println!();
+		}
 
-		// ----- Step 4: Propogate -----
-		// Update neighboring superpositions based on the collapsed tile
-		// (Remove any that violate constraints)
+		// Check that the tiles array has the correct length
+		assert_eq!(tiles.len(), TILE_COUNT);
 
-		// ----- Step 5: Repeat until all superpositions are collapsed -----
-		// If any tiles are left with no possible states, restart
-		// TODO: backtrack
+		// Check that there is at least one Air tile
+		assert!(tiles.contains(&Tile::Air));
 	}
 }