Files
nd-wfc/include/nd-wfc/wave.hpp
2025-08-21 13:00:00 +09:00

277 lines
8.6 KiB
C++

#pragma once
#include "types.hpp"
#include "grid.hpp"
#include <vector>
#include <random>
namespace nd_wfc {
// Wave function for tracking possible states at each grid cell
template<std::size_t N>
class Wave {
private:
Size<N> size_;
Index total_size_;
std::vector<TileMask> possibilities_;
std::vector<std::uint8_t> entropy_cache_;
std::vector<bool> is_collapsed_;
std::vector<TileId> tiles_;
Index num_tiles_;
bool entropy_dirty_;
public:
Wave(const Size<N>& size, const std::vector<TileId>& tiles)
: size_(size), tiles_(tiles), num_tiles_(tiles.size()), entropy_dirty_(true) {
assert(num_tiles_ <= MAX_TILES);
// Calculate total size
total_size_ = 1;
for (std::size_t i = 0; i < N; ++i) {
total_size_ *= size_[i];
}
// Initialize all cells with all possibilities
TileMask all_possible;
for (Index i = 0; i < num_tiles_; ++i) {
all_possible.set(i);
}
possibilities_.resize(total_size_, all_possible);
entropy_cache_.resize(total_size_, static_cast<std::uint8_t>(num_tiles_));
is_collapsed_.resize(total_size_, false);
}
// Convert coordinate to linear index
Index coordToIndex(const Coord<N>& coord) const {
Index index = 0;
Index stride = 1;
for (std::size_t i = N; i > 0; --i) {
assert(coord[i-1] < size_[i-1]);
index += coord[i-1] * stride;
stride *= size_[i-1];
}
return index;
}
// Convert linear index to coordinate
Coord<N> indexToCoord(Index index) const {
assert(index < total_size_);
Coord<N> coord;
for (std::size_t i = 0; i < N; ++i) {
coord[i] = index % size_[i];
index /= size_[i];
}
return coord;
}
// Get possible tiles at a coordinate
const TileMask& getPossibilities(const Coord<N>& coord) const {
return possibilities_[coordToIndex(coord)];
}
// Get possible tiles at an index
const TileMask& getPossibilities(Index index) const {
assert(index < total_size_);
return possibilities_[index];
}
// Check if a tile is possible at a coordinate
bool isPossible(const Coord<N>& coord, TileId tile_id) const {
Index tile_index = getTileIndex(tile_id);
return tile_index < num_tiles_ && possibilities_[coordToIndex(coord)][tile_index];
}
// Remove a tile possibility from a coordinate
bool removePossibility(const Coord<N>& coord, TileId tile_id) {
Index tile_index = getTileIndex(tile_id);
if (tile_index >= num_tiles_) return false;
Index cell_index = coordToIndex(coord);
if (!possibilities_[cell_index][tile_index]) {
return false; // Already removed
}
possibilities_[cell_index][tile_index] = false;
entropy_dirty_ = true;
// Check for contradiction
return possibilities_[cell_index].any();
}
// Collapse a cell to a specific tile
bool collapse(const Coord<N>& coord, TileId tile_id) {
Index tile_index = getTileIndex(tile_id);
if (tile_index >= num_tiles_) return false;
Index cell_index = coordToIndex(coord);
if (!possibilities_[cell_index][tile_index]) {
return false; // Tile not possible
}
// Clear all possibilities except the chosen one
possibilities_[cell_index].reset();
possibilities_[cell_index][tile_index] = true;
is_collapsed_[cell_index] = true;
entropy_dirty_ = true;
return true;
}
// Get entropy (number of possible tiles) at a coordinate
std::uint8_t getEntropy(const Coord<N>& coord) const {
Index cell_index = coordToIndex(coord);
if (entropy_dirty_) {
// Recalculate entropy cache
const_cast<Wave*>(this)->updateEntropyCache();
}
return entropy_cache_[cell_index];
}
// Check if a cell is collapsed
bool isCollapsed(const Coord<N>& coord) const {
return is_collapsed_[coordToIndex(coord)];
}
// Get the collapsed tile at a coordinate (only valid if collapsed)
TileId getCollapsedTile(const Coord<N>& coord) const {
Index cell_index = coordToIndex(coord);
assert(is_collapsed_[cell_index]);
const auto& mask = possibilities_[cell_index];
for (Index i = 0; i < num_tiles_; ++i) {
if (mask[i]) {
return tiles_[i];
}
}
assert(false); // Should never reach here if properly collapsed
return static_cast<TileId>(-1);
}
// Find the cell with minimum entropy (excluding collapsed cells)
bool findMinEntropyCell(Coord<N>& coord, std::mt19937& rng) const {
if (entropy_dirty_) {
const_cast<Wave*>(this)->updateEntropyCache();
}
std::uint8_t min_entropy = 255;
std::vector<Index> candidates;
// Find minimum entropy among uncollapsed cells
for (Index i = 0; i < total_size_; ++i) {
if (!is_collapsed_[i]) {
std::uint8_t entropy = entropy_cache_[i];
if (entropy == 0) {
return false; // Contradiction found
}
if (entropy < min_entropy) {
min_entropy = entropy;
candidates.clear();
candidates.push_back(i);
} else if (entropy == min_entropy) {
candidates.push_back(i);
}
}
}
if (candidates.empty()) {
return false; // All cells collapsed or contradiction
}
// Randomly select among candidates with same entropy
std::uniform_int_distribution<std::size_t> dist(0, candidates.size() - 1);
Index chosen_index = candidates[dist(rng)];
coord = indexToCoord(chosen_index);
return true;
}
// Select a random tile from possibilities at a coordinate
TileId selectRandomTile(const Coord<N>& coord, std::mt19937& rng) const {
Index cell_index = coordToIndex(coord);
const auto& mask = possibilities_[cell_index];
std::vector<TileId> available_tiles;
for (Index i = 0; i < num_tiles_; ++i) {
if (mask[i]) {
available_tiles.push_back(tiles_[i]);
}
}
if (available_tiles.empty()) {
return static_cast<TileId>(-1); // No possibilities
}
std::uniform_int_distribution<std::size_t> dist(0, available_tiles.size() - 1);
return available_tiles[dist(rng)];
}
// Check if the wave function is completely collapsed
bool isComplete() const {
for (Index i = 0; i < total_size_; ++i) {
if (!is_collapsed_[i]) {
return false;
}
}
return true;
}
// Check for contradictions (cells with no possibilities)
bool hasContradiction() const {
for (Index i = 0; i < total_size_; ++i) {
if (!is_collapsed_[i] && !possibilities_[i].any()) {
return true;
}
}
return false;
}
// Get the size of the grid
const Size<N>& getSize() const { return size_; }
Index getTotalSize() const { return total_size_; }
// Reset the wave function to initial state
void reset() {
TileMask all_possible;
for (Index i = 0; i < num_tiles_; ++i) {
all_possible.set(i);
}
std::fill(possibilities_.begin(), possibilities_.end(), all_possible);
std::fill(entropy_cache_.begin(), entropy_cache_.end(), static_cast<std::uint8_t>(num_tiles_));
std::fill(is_collapsed_.begin(), is_collapsed_.end(), false);
entropy_dirty_ = true;
}
private:
// Get the index of a tile ID in the tiles vector
Index getTileIndex(TileId tile_id) const {
for (Index i = 0; i < num_tiles_; ++i) {
if (tiles_[i] == tile_id) {
return i;
}
}
return static_cast<Index>(-1);
}
// Update the entropy cache
void updateEntropyCache() {
for (Index i = 0; i < total_size_; ++i) {
if (is_collapsed_[i]) {
entropy_cache_[i] = 1;
} else {
entropy_cache_[i] = static_cast<std::uint8_t>(possibilities_[i].count());
}
}
entropy_dirty_ = false;
}
};
// Type aliases for common dimensions
using Wave2D = Wave<2>;
using Wave3D = Wave<3>;
using Wave4D = Wave<4>;
} // namespace nd_wfc