Wave function collapse

5 files changed, 824 insertions, 59 deletions

diff --git a/dungeon/src/bsp.rs b/dungeon/src/bsp.rs
new file mode 100644
index 0000000..be0f282
--- /dev/null
+++ b/dungeon/src/bsp.rs
@@ -0,0 +1,475 @@
+//! The `bsp` module implements a Binary-Space Partitioning based dungeon generator.
+//! Produces a tile array and a player start position for a given seed.
+
+use rand::{Rng, SeedableRng};
+use std::cmp; // for min/max
+
+use crate::map::{MAP_SIZE_USIZE, TILE_COUNT, Tile};
+use crate::{const_pos, pos::Pos};
+
+/// `MIN_LEAF_SIZE` is the minimum width/height for a leaf to be considered "splittable".
+const MIN_LEAF_SIZE: usize = 8;
+/// `MIN_ROOM_SIZE` is the minimum room size (w/h).
+const MIN_ROOM_SIZE: usize = 3;
+/// `MAX_ROOM_SIZE` is the maximum room size (w/h).
+const MAX_ROOM_SIZE: usize = 12;
+
+/// The `Rect` type represents a rectangle inside the map (x,y,width,height).
+/// (x,y) is the top-left corner.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct Rect {
+	x: usize,
+	y: usize,
+	w: usize,
+	h: usize,
+}
+// since this is all "internal", you likely have to use unit tests and not doctests
+impl Rect {
+	fn new(x: usize, y: usize, w: usize, h: usize) -> Self {
+		Self { x, y, w, h }
+	}
+
+	/// Returns the center point (cx, cy) of the rectangle.
+	fn center(&self) -> (usize, usize) {
+		let cx = self.x + self.w / 2;
+		let cy = self.y + self.h / 2;
+		(cx, cy)
+	}
+}
+
+/// The `Node` type represents a node in the BSP tree.
+/// Has optional left and right children, and an optional room rectangle.
+#[derive(Debug)]
+pub(crate) struct Node {
+	rect: Rect,               // Area this node covers
+	left: Option<Box<Node>>,  // Left child (if not leaf)
+	right: Option<Box<Node>>, // Right child (if not leaf)
+	room: Option<Rect>,       // Room created in this node (if leaf)
+}
+
+impl Node {
+	fn new(rect: Rect) -> Self {
+		Self {
+			rect,
+			left: None,
+			right: None,
+			room: None,
+		}
+	}
+
+	/// Try to split the node. Returns true if split happened.
+	/// Splitting is done either horizontally or vertically,
+	/// depending on the dimensions of the node.
+	fn split<R: Rng>(&mut self, rng: &mut R) -> bool {
+		// Already split
+		if self.left.is_some() || self.right.is_some() {
+			return false;
+		}
+
+		// Determine if we can split
+		let can_split_h = self.rect.h >= (MIN_LEAF_SIZE * 2);
+		let can_split_v = self.rect.w >= (MIN_LEAF_SIZE * 2);
+
+		if !can_split_h && !can_split_v {
+			return false;
+		}
+
+		// Choose orientation: prefer the longer side
+		let split_h = if self.rect.w > self.rect.h {
+			false
+		} else if self.rect.h > self.rect.w {
+			true
+		} else {
+			rng.random_bool(0.5)
+		};
+
+		// Choose split coordinate with margin so each side can contain a room
+		if split_h {
+			// horizontal split => choose y
+			// avoid too close to borders of rect
+			let min = cmp::max(
+				self.rect.y + MIN_ROOM_SIZE + 1,
+				self.rect.y + MIN_LEAF_SIZE + 1,
+			);
+			let max = cmp::min(
+				self.rect.y + self.rect.h - MIN_ROOM_SIZE - 1,
+				self.rect.y + self.rect.h - MIN_LEAF_SIZE - 1,
+			);
+			if max <= min {
+				return false;
+			}
+			// Add some heterogeneity: choose random position between min..max
+			let split_y = rng.random_range(min..=max);
+			let top_h = split_y - self.rect.y;
+			let bottom_h = self.rect.h - top_h;
+			if top_h < MIN_LEAF_SIZE || bottom_h < MIN_LEAF_SIZE {
+				return false;
+			}
+
+			let top = Rect::new(self.rect.x, self.rect.y, self.rect.w, top_h);
+			let bottom = Rect::new(self.rect.x, split_y, self.rect.w, bottom_h);
+			self.left = Some(Box::new(Self::new(top)));
+			self.right = Some(Box::new(Self::new(bottom)));
+		} else {
+			// vertical split => choose x
+			let min = cmp::max(
+				self.rect.x + MIN_ROOM_SIZE + 1,
+				self.rect.x + MIN_LEAF_SIZE + 1,
+			);
+			let max = cmp::min(
+				self.rect.x + self.rect.w - MIN_ROOM_SIZE - 1,
+				self.rect.x + self.rect.w - MIN_LEAF_SIZE - 1,
+			);
+			if max <= min {
+				return false;
+			}
+			let split_x = rng.random_range(min..=max);
+			let left_w = split_x - self.rect.x;
+			let right_w = self.rect.w - left_w;
+			if left_w < MIN_LEAF_SIZE || right_w < MIN_LEAF_SIZE {
+				return false;
+			}
+
+			let left_rect = Rect::new(self.rect.x, self.rect.y, left_w, self.rect.h);
+			let right_rect = Rect::new(split_x, self.rect.y, right_w, self.rect.h);
+			self.left = Some(Box::new(Self::new(left_rect)));
+			self.right = Some(Box::new(Self::new(right_rect)));
+		}
+		true
+	}
+
+	/// Create a room inside this node (called for leaves).
+	/// Room size and position chosen randomly.
+	fn create_room<R: Rng>(&mut self, rng: &mut R) {
+		if self.left.is_some() || self.right.is_some() {
+			// This is not a leaf
+			if let Some(left) = &mut self.left {
+				left.create_room(rng);
+			}
+			if let Some(right) = &mut self.right {
+				right.create_room(rng);
+			}
+			return;
+		}
+
+		// Leaf: choose room size
+		// Ensure room fits inside rect with at least 1 tile border to walls
+		let max_w = (self.rect.w - 2).max(MIN_ROOM_SIZE);
+		let max_h = (self.rect.h - 2).max(MIN_ROOM_SIZE);
+		let room_w = rng.random_range(MIN_ROOM_SIZE..=max_w.min(MAX_ROOM_SIZE));
+		let room_h = rng.random_range(MIN_ROOM_SIZE..=max_h.min(MAX_ROOM_SIZE));
+
+		// Choose position inside rect, leaving at least 1 tile to the border
+		let x_range = (self.rect.x + 1)..=(self.rect.x + self.rect.w - 1 - room_w);
+		let y_range = (self.rect.y + 1)..=(self.rect.y + self.rect.h - 1 - room_h);
+
+		// The ranges should be valid; if not, fallback to center
+		let rx = if x_range.is_empty() {
+			self.rect.x + (self.rect.w.saturating_sub(room_w)) / 2
+		} else {
+			rng.random_range(x_range)
+		};
+		let ry = if y_range.is_empty() {
+			self.rect.y + (self.rect.h.saturating_sub(room_h)) / 2
+		} else {
+			rng.random_range(y_range)
+		};
+
+		self.room = Some(Rect::new(rx, ry, room_w, room_h));
+	}
+
+	/// Return a vector of references to all leaf nodes (in-order).
+	fn collect_leaves<'a>(&'a self, out: &mut Vec<&'a Self>) {
+		if self.left.is_none() && self.right.is_none() {
+			out.push(self);
+			return;
+		}
+		if let Some(left) = &self.left {
+			left.collect_leaves(out);
+		}
+		if let Some(right) = &self.right {
+			right.collect_leaves(out);
+		}
+	}
+
+	/// Return the "room center" for this subtree: if node has a room, return its center,
+	/// else try left then right.
+	fn room_center(&self) -> Option<(usize, usize)> {
+		if let Some(room) = &self.room {
+			return Some(room.center());
+		}
+		if let Some(left) = &self.left
+			&& let Some(center) = left.room_center()
+		{
+			return Some(center);
+		}
+		if let Some(right) = &self.right
+			&& let Some(center) = right.room_center()
+		{
+			return Some(center);
+		}
+		None
+	}
+
+	/// Return a point (x,y) on the border of the room in this node.
+	fn random_point_in_room<R: Rng>(&self, rng: &mut R) -> (usize, usize) {
+		if let Some(room) = &self.room {
+			let rx = rng.random_range(room.x..(room.x + room.w));
+			let ry = rng.random_range(room.y..(room.y + room.h));
+			(rx, ry)
+		} else {
+			// No room: return center of rect
+			self.rect.center()
+		}
+	}
+
+	/// Connect rooms of child nodes recursively and collect corridors to carve.
+	///
+	/// `corridors` is appended with pairs of points (x1,y1,x2,y2) describing corridors to carve.
+	fn connect_children(
+		&self,
+		corridors: &mut Vec<(usize, usize, usize, usize)>,
+		rng: &mut rand::rngs::StdRng,
+	) {
+		if let (Some(left), Some(right)) = (&self.left, &self.right) {
+			// Connect a room (or corridor endpoint) from left subtree to right subtree
+			// Select a random point from each side
+			let (x1, y1) = left.random_point_in_room(rng);
+			let (x2, y2) = right.random_point_in_room(rng);
+			corridors.push((x1, y1, x2, y2));
+
+			left.connect_children(corridors, rng);
+			right.connect_children(corridors, rng);
+		}
+		// leaf: nothing to connect!
+	}
+}
+
+/// Carve a room rectangle into the map (tile array) by setting tiles inside to Air.
+fn carve_room(tiles: &mut [Tile; TILE_COUNT], room: &Rect) {
+	for yy in room.y..(room.y + room.h) {
+		for xx in room.x..(room.x + room.w) {
+			let idx = xx + yy * MAP_SIZE_USIZE;
+			tiles[idx] = Tile::Air;
+		}
+	}
+}
+
+/// Carve a horizontal corridor from x1..=x2 at y (inclusive)
+fn carve_h_corridor(tiles: &mut [Tile; TILE_COUNT], x1: usize, x2: usize, y: usize) {
+	let (sx, ex) = if x1 <= x2 { (x1, x2) } else { (x2, x1) };
+	for x in sx..=ex {
+		let idx = x + y * MAP_SIZE_USIZE;
+		tiles[idx] = Tile::Air;
+	}
+}
+
+/// Carve a vertical corridor from y1..=y2 at x
+fn carve_v_corridor(tiles: &mut [Tile; TILE_COUNT], y1: usize, y2: usize, x: usize) {
+	let (sy, ey) = if y1 <= y2 { (y1, y2) } else { (y2, y1) };
+	for y in sy..=ey {
+		let idx = x + y * MAP_SIZE_USIZE;
+		tiles[idx] = Tile::Air;
+	}
+}
+
+/// Top-level generator function for the dungeon using BSP.
+/// Returns (tiles, player_start).
+pub fn generate(seed: u64) -> (Box<[Tile; TILE_COUNT]>, Pos) {
+	let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
+
+	// Initialize all tiles to walls
+	let mut tiles_box: Box<[Tile; TILE_COUNT]> = Box::new([Tile::Wall; TILE_COUNT]);
+
+	// Rooms and coorridors will be carved as `Air` into the tile array.
+	// Build root rect (whole map)
+	let root_rect = Rect::new(0, 0, MAP_SIZE_USIZE, MAP_SIZE_USIZE);
+	let mut root = Node::new(root_rect);
+
+	// Repeatedly try splitting nodes to produce the tree.
+	// Collect leaves and try to split them until no splits happen or we reach a cap.
+	// Number of splits proportional to map size to avoid infinite loops.
+	let max_attempts = 2000usize;
+	let mut attempts = 0usize;
+	while attempts < max_attempts {
+		attempts += 1;
+
+		// Collect mutable references to current leaves
+		let mut leaves = vec![];
+		{
+			// Collect leaves with a stack
+			let mut stack = vec![&mut root];
+			while let Some(node) = stack.pop() {
+				if node.left.is_none() && node.right.is_none() {
+					leaves.push(node as *mut Node);
+				} else {
+					if let Some(left) = &mut node.left {
+						stack.push(left);
+					}
+					if let Some(right) = &mut node.right {
+						stack.push(right);
+					}
+				}
+			}
+		}
+
+		let mut splitted_any = false;
+
+		// Try splitting each leaf in order
+		for &node_ptr in leaves.iter() {
+			// Pointers obtained from boxed root remain valid while root lives
+			// TODO: Store nodes in smart pointers to avoid unsafe?
+			let node = unsafe { &mut *node_ptr };
+			// Attempt to split if possible
+			if node.split(&mut rng) {
+				splitted_any = true;
+			}
+		}
+
+		// If no splits happened, we're done
+		if !splitted_any {
+			break;
+		}
+	}
+
+	// Create rooms in all leaves
+	root.create_room(&mut rng);
+
+	// Carve all rooms into the tile array
+	let mut leaves = vec![];
+	root.collect_leaves(&mut leaves);
+	for leaf in leaves.iter() {
+		if let Some(room) = &leaf.room {
+			carve_room(&mut tiles_box, room);
+		}
+	}
+
+	// Collect corridors (pairs of centers) by connecting children bottom-up
+	let mut corridors = Vec::new();
+	root.connect_children(&mut corridors, &mut rng);
+
+	// Carve corridors. For each corridor (x1,y1,x2,y2), carve straight or L-shape.
+	for (x1, y1, x2, y2) in corridors {
+		// If aligned horizontally or vertically, carve straight
+		if x1 == x2 {
+			carve_v_corridor(&mut tiles_box, y1, y2, x1);
+		} else if y1 == y2 {
+			carve_h_corridor(&mut tiles_box, x1, x2, y1);
+		} else {
+			// Choose whether to go horizontal then vertical or vertical then horizontal
+			if rng.random_bool(0.5) {
+				carve_h_corridor(&mut tiles_box, x1, x2, y1);
+				carve_v_corridor(&mut tiles_box, y1, y2, x2);
+			} else {
+				carve_v_corridor(&mut tiles_box, y1, y2, x1);
+				carve_h_corridor(&mut tiles_box, x1, x2, y2);
+			}
+		}
+	}
+
+	// Choose player start randomly in the center of one of the rooms
+	// TODO: Better Pos creation
+	#[expect(clippy::cast_possible_truncation)]
+	let player_start = root
+		.room_center()
+		.map(|(cx, cy)| Pos::new(cx as u16, cy as u16).unwrap_or(const_pos!(1, 1)))
+		.unwrap_or(const_pos!(1, 1));
+
+	// Choose random location among air tiles
+	/*
+	let mut player_start = Pos::new_unchecked(1, 1);
+	let mut found = false;
+	for pos in Pos::values() {
+		if tiles_box[pos.idx()] == Tile::Air {
+			player_start = pos;
+			found = true;
+			break;
+		}
+	}
+	// If we didn't find any air (extremely unlikely), put start at (1,1) and carve it
+	if !found {
+		player_start = Pos::new_unchecked(1, 1);
+		let idx = player_start.idx();
+		tiles_box[idx] = Tile::Air;
+	}*/
+
+	// Return tiles and player_start
+	(tiles_box, player_start)
+}
+
+/// Tests
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn test_rect_center() {
+		let rect = Rect::new(2, 2, 4, 4);
+		let (cx, cy) = rect.center();
+		assert_eq!(cx, 4);
+		assert_eq!(cy, 4);
+	}
+
+	#[test]
+	fn test_node_split() {
+		let rect = Rect::new(0, 0, 20, 20);
+		let mut node = Node::new(rect);
+		let mut rng = rand::rngs::StdRng::seed_from_u64(12345);
+		let splitted = node.split(&mut rng);
+		assert!(splitted);
+		assert!(node.left.is_some());
+		assert!(node.right.is_some());
+	}
+
+	#[test]
+	fn test_node_create_room() {
+		let rect = Rect::new(0, 0, 20, 20);
+		let mut node = Node::new(rect);
+		let mut rng = rand::rngs::StdRng::seed_from_u64(12345);
+		node.create_room(&mut rng);
+		assert!(node.room.is_some());
+		match &node.room {
+			Some(room) => {
+				assert!(room.w >= MIN_ROOM_SIZE && room.h >= MIN_ROOM_SIZE);
+				assert!(room.x >= 1 && room.y >= 1);
+				assert!(room.x + room.w < MAP_SIZE_USIZE);
+				assert!(room.y + room.h < MAP_SIZE_USIZE);
+			}
+			None => panic!("Room should be created"),
+		}
+	}
+
+	#[test]
+	fn test_node_collect_leaves() {
+		let rect = Rect::new(0, 0, 20, 20);
+		let mut node = Node::new(rect);
+		let mut rng = rand::rngs::StdRng::seed_from_u64(12345);
+		node.split(&mut rng);
+		let mut leaves = Vec::new();
+		node.collect_leaves(&mut leaves);
+		assert!(!leaves.is_empty());
+	}
+
+	#[test]
+	fn test_node_connect_children() {
+		let rect = Rect::new(0, 0, 20, 20);
+		let mut node = Node::new(rect);
+		let mut rng = rand::rngs::StdRng::seed_from_u64(12345);
+		node.split(&mut rng);
+		let mut corridors = Vec::new();
+		node.connect_children(&mut corridors, &mut rng);
+		assert!(!corridors.is_empty());
+	}
+
+	#[test]
+	fn test_generate() {
+		let seed = 12345u64;
+		let (tiles, player_start) = generate(seed);
+		// Check that tiles contain some Air tiles
+		let air_count = tiles.iter().filter(|&&t| t == Tile::Air).count();
+		assert!(air_count > 0);
+		// Check that player_start is within bounds
+		assert!(player_start.xy().0 < MAP_SIZE_USIZE as u16);
+		assert!(player_start.xy().1 < MAP_SIZE_USIZE as u16);
+	}
+}
diff --git a/dungeon/src/lib.rs b/dungeon/src/lib.rs
index 0e1a501..39d7c60 100644
--- a/dungeon/src/lib.rs
+++ b/dungeon/src/lib.rs
@@ -1,11 +1,12 @@
 //! The `dungon` crate contains the core functionality for
 //! interacting with a `Dungeon` and its components.
 
+mod bsp;
 mod entity;
 mod map;
 mod pos;
-mod wfc;
 
+pub use bsp::*;
 pub use entity::*;
 pub use map::*;
 pub use pos::*;
diff --git a/dungeon/src/map.rs b/dungeon/src/map.rs
index d185547..f4a8729 100644
--- a/dungeon/src/map.rs
+++ b/dungeon/src/map.rs
@@ -4,16 +4,16 @@
 use strum::IntoEnumIterator;
 use strum_macros::EnumIter;
 
-use crate::wfc::Wfc;
 use std::{
 	cell::RefCell,
 	hash::{DefaultHasher, Hash, Hasher},
 };
 
+use crate::bsp;
 use crate::{const_pos, pos::Pos};
 
 /// `MAP_SIZE` is the size of the size of the dungeon grid.
-pub const MAP_SIZE: u16 = 100;
+pub const MAP_SIZE: u16 = 48;
 
 /// `MAP_SIZE` as a usize
 pub const MAP_SIZE_USIZE: usize = MAP_SIZE as usize;
@@ -77,7 +77,7 @@ impl Floor {
 		}
 	}
 
-	/// Generates a dungeon `Floor` using wave function collapse.
+	/// Generates a dungeon `Floor` using binary space partitioning.
 	///
 	/// # Examples
 	///
@@ -92,10 +92,10 @@ impl Floor {
 		Self::generate_seeded(seed)
 	}
 
-	/// Genreates a dungeon `Floor` using wave function collapse provided with a seed.
+	/// Generates a dungeon `Floor` using binary space partitioning provided with a seed.
 	///
-	/// The provided seed is used for randomness in the wave function
-	/// collapse algorithm.
+	/// The provided seed is used for randomness in the binary space partitioning
+	/// algorithm.
 	///
 	/// # Examples
 	///
@@ -104,16 +104,15 @@ impl Floor {
 	///
 	/// /// here is our very seedy seed
 	/// let seed = 2893249402u64;
-	/// let floor = Floor::generate_seeded(seed);
+	/// let floor_1 = Floor::generate_seeded(seed);
+	/// let floor_2 = Floor::generate_seeded(seed);
+	/// assert_eq!(floor_1, floor_2); // both floors will be identical
 	/// ```
 	#[must_use]
 	pub fn generate_seeded(seed: u64) -> Self {
-		let mut wfc = Wfc::new(seed, MAP_SIZE);
-		wfc.initialize_states();
-		wfc.run();
+		let (tiles, player_start) = bsp::generate(seed);
 
-		// TODO
-		unimplemented!()
+		Self::new(tiles, player_start, seed)
 	}
 
 	/// Returns the start position of the player
@@ -178,7 +177,41 @@ impl Floor {
 		*dirty = false;
 		*hash
 	}
+
+	/// Display the floor as a string for debugging
+	///
+	/// # Examples
+	/// ```no_run
+	/// use dungeon::Floor;
+	/// let floor = Floor::generate();
+	/// println!("{}", floor.display());
+	/// ```
+	#[must_use]
+	pub fn display(&self) -> String {
+		let mut output = String::new();
+		for pos in Pos::values() {
+			// If it's the player start, show 'P'
+			if self.player_start == pos {
+				output.push('P');
+				continue;
+			}
+			// Otherwise, show the tile character
+			let tile = self.get(pos);
+			let char = match tile {
+				Tile::Wall => '#',
+				Tile::Air => '.',
+				Tile::Stairs => '>',
+			};
+			output.push(char);
+			// Newline at the end of each row
+			if pos.xy().0 == MAP_SIZE - 1 {
+				output.push('\n');
+			}
+		}
+		output
+	}
 }
+
 impl Default for Floor {
 	/// Returns a floor with a single set of walls on the map border
 	fn default() -> Self {
@@ -195,3 +228,18 @@ impl Default for Floor {
 		Self::new(tiles, player_start, seed)
 	}
 }
+
+/// Tests
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	// Test floor printing
+	#[test]
+	fn test_floor_display() {
+		let floor = Floor::generate();
+		let display = floor.display();
+		// Print the display for visual inspection
+		println!("{display}");
+	}
+}
diff --git a/dungeon/src/pos.rs b/dungeon/src/pos.rs
index 8f1f8dd..2682261 100644
--- a/dungeon/src/pos.rs
+++ b/dungeon/src/pos.rs
@@ -170,18 +170,18 @@ impl Pos {
 	/// # Examples
 	///
 	/// ```
-	/// use dungeon::{Pos};
+	/// use dungeon::{Pos, MAP_SIZE_USIZE};
 	///
-	/// let idx_pos = Pos::from_idx(17);
-	/// let pos = Pos::new(17, 0);
+	/// let idx_pos = Pos::from_idx(MAP_SIZE_USIZE);
+	/// let pos = Pos::new(0, 1);
 	///
 	/// assert_eq!(idx_pos, pos);
 	/// ```
 	///
 	/// ```
-	/// use dungeon::{Pos};
+	/// use dungeon::{Pos, MAP_SIZE_USIZE};
 	///
-	/// let idx_pos = Pos::from_idx(170);
+	/// let idx_pos = Pos::from_idx(MAP_SIZE_USIZE * 70 + 1);
 	/// let pos = Pos::new(70, 1);
 	///
 	/// assert_eq!(idx_pos, pos);
@@ -285,11 +285,13 @@ impl Pos {
 	/// ```
 	/// use dungeon::{Pos, MAP_SIZE};
 	///
-	/// let pos1 = Pos::new(0, 17).unwrap();
-	/// let pos2 = Pos::new(1, 17).unwrap();
-	/// let pos3 = Pos::new(MAP_SIZE - 1, 17).unwrap();
-	/// let pos4 = Pos::new(55, MAP_SIZE - 1).unwrap();
-	/// let pos5 = Pos::new(55, 0).unwrap();
+	/// // Assuming MAP_SIZE is at least 2
+	///
+	/// let pos1 = Pos::new(0, MAP_SIZE - 1).unwrap();
+	/// let pos2 = Pos::new(1, MAP_SIZE - 2).unwrap();
+	/// let pos3 = Pos::new(MAP_SIZE - 1, MAP_SIZE - 1).unwrap();
+	/// let pos4 = Pos::new(MAP_SIZE - 1, MAP_SIZE - 1).unwrap();
+	/// let pos5 = Pos::new(MAP_SIZE - 1, 0).unwrap();
 	///
 	/// assert!(pos1.is_border());
 	/// assert!(!pos2.is_border());
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));
 	}
 }