From d2d4e8882d10be085ab92f20c640a69c1e6edf13 Mon Sep 17 00:00:00 2001
From: Connor De Meyer <connordemeyer@gmail.com>
Date: Sun, 24 Aug 2025 19:05:16 +0900
Subject: [PATCH] bug fix

---
 demos/sudoku/sudoku_wfc.cpp | 10 +++++-----
 include/nd-wfc/wfc.hpp      | 31 +++++++++++++++++++------------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/demos/sudoku/sudoku_wfc.cpp b/demos/sudoku/sudoku_wfc.cpp
index 81c85dd..9b70bb3 100644
--- a/demos/sudoku/sudoku_wfc.cpp
+++ b/demos/sudoku/sudoku_wfc.cpp
@@ -23,12 +23,12 @@ int main()
             }
 
             // Add box constraints (3x3 box)
-            int box_x = (x / 3) * 3;
-            int box_y = (y / 3) * 3;
+            size_t box_x = (x / 3) * 3;
+            size_t box_y = (y / 3) * 3;
             for (size_t j = 0; j < 3; ++j) {
                 for (size_t k = 0; k < 3; ++k) {
-                    int cell_x = box_x + j;
-                    int cell_y = box_y + k;
+                    size_t cell_x = box_x + j;
+                    size_t cell_y = box_y + k;
                     size_t cell_index = cell_x + cell_y * 9;
                     if (cell_index != index) {
                         constrainer.Exclude(val, cell_index);
@@ -46,7 +46,7 @@ int main()
         // Print the solved sudoku
         for (size_t y = 0; y < 9; ++y) {
             for (size_t x = 0; x < 9; ++x) {
-                std::cout << static_cast<int>(sudokuWorld.at(x, y)) << " ";
+                std::cout << static_cast<int>(sudokuWorld.at(static_cast<int>(x), static_cast<int>(y))) << " ";
                 if (x == 2 || x == 5) std::cout << "| ";
             }
             std::cout << std::endl;
diff --git a/include/nd-wfc/wfc.hpp b/include/nd-wfc/wfc.hpp
index b5a88bd..bb34d3b 100644
--- a/include/nd-wfc/wfc.hpp
+++ b/include/nd-wfc/wfc.hpp
@@ -12,20 +12,23 @@
 #include <cassert>
 #include <algorithm>
 #include <concepts>
+#include <bit>
 
 namespace WFC {
 
 inline int FindNthSetBit(size_t num, int n) {
-    assert(!(n <= 0 || static_cast<size_t>(n) > std::popcount(num)));
+    auto popCount = std::popcount(num);
+    assert(n < std::popcount(num) && "index is out of range");
     int bitCount = 0;
     while (num) {
-        bitCount++;
         if (bitCount == n) {
             return std::countr_zero(num); // Index of the current set bit
         }
-        num &= (num - 1);
+        bitCount++;
+        num &= (num - 1); // turn of lowest set bit
     }
-    return -1;
+    assert(bitCount < popCount && "out of bounds");
+    return bitCount;
 }
 
 template<typename T>
@@ -103,6 +106,7 @@ public:
     }
 
     static VarT GetValue(size_t index) {
+        assert(index < sizeof...(Values));
         constexpr VarT arr[] = {Values...};
         return arr[index];
     }
@@ -182,7 +186,7 @@ public:
     }
 
     void Exclude(WorldValue<typename VariableIDMapT::Type> value, size_t cellId) {
-        ApplyMask(cellId, ~value.InternalIndex);
+        ApplyMask(cellId, ~(1 << value.InternalIndex));
     }
 
     /**
@@ -197,17 +201,17 @@ public:
     }
 
     void Only(WorldValue<typename VariableIDMapT::Type> value, size_t cellId) {
-        ApplyMask(cellId, value.InternalIndex);
+        ApplyMask(cellId, 1 << value.InternalIndex);
     }
 
 private:
     void ApplyMask(size_t cellId, MaskType mask) {
-        bool collapsedBefore = m_wave.IsCollapsed(cellId);
+        if (m_wave.IsCollapsed(cellId)) return;
 
         m_wave.Collapse(cellId, mask);
-        
-        bool collapsedAfter = m_wave.IsCollapsed(cellId);
-        if (!collapsedBefore && collapsedAfter) {
+        assert(!m_wave.HasContradiction() && "Contradiction found");
+
+        if (m_wave.IsCollapsed(cellId)) {
             m_propagationQueue.push(cellId);
         }
     }
@@ -288,7 +292,7 @@ public:
                 GetMinEntropyCell(worldSolver);
             }
         }
-        return true;
+        return false;
     }
 
     /**
@@ -335,14 +339,17 @@ private:
                 minEntropyCell = i;
             }
         }
+        assert(!worldSolver.wave.IsCollapsed(minEntropyCell));
 
         // Randomly select a value from possible values
         size_t availableValues = worldSolver.wave.Entropy(minEntropyCell);
         std::uniform_int_distribution<size_t> dist(0, availableValues - 1);
         size_t selectedValue = FindNthSetBit(worldSolver.wave.GetMask(minEntropyCell), dist(worldSolver.rng));
+        assert(selectedValue < VariableIDMapT::ValuesRegisteredAmount && "Selected Value went outside bounds");
 
         // Collapse the cell to the selected value
         worldSolver.wave.Collapse(minEntropyCell, 1 << selectedValue);
+        assert(worldSolver.wave.IsCollapsed(minEntropyCell) && "Cell was not collapsed correctly");
 
         worldSolver.propagationQueue.push(minEntropyCell);
 
@@ -356,7 +363,7 @@ private:
             size_t cellId = worldSolver.propagationQueue.front();
             worldSolver.propagationQueue.pop();
 
-            assert(worldSolver.wave.IsCollapsed(cellId));
+            assert(worldSolver.wave.IsCollapsed(cellId) && "Cell was not collapsed");
 
             uint16_t variableID = worldSolver.wave.GetVariableID(cellId);
             Constrainer<VariableIDMapT> constrainer(worldSolver.wave, worldSolver.propagationQueue);