1 files changed, 78 insertions, 127 deletions

diff --git a/include/tak.c b/include/tak.c
index e9f27ff..f31d4a9 100644
--- a/include/tak.c
+++ b/include/tak.c
@@ -16,6 +16,7 @@
 */
 
 #include "tak.h"
+#include <stdio.h>
 #include <stdlib.h>
 
 // ===================================================================
@@ -240,154 +241,100 @@ enum ACT_RESULT try_move(tak_state_p state, const int8_t location,
 // Checking win
 // ===================================================================
 
-// Sides of the board a connected component can touch
-#define SIDE_BOTTOM 0x1 // row 0
-#define SIDE_TOP 0x2    // last row
-#define SIDE_LEFT 0x4   // column 0
-#define SIDE_RIGHT 0x8  // last column
-#define SIDES_BOTTOM_TOP (SIDE_BOTTOM | SIDE_TOP)
-#define SIDES_LEFT_RIGHT (SIDE_LEFT | SIDE_RIGHT)
-
 // One bit per colour, set once a road for that colour has been found
 #define ROAD_WHITE (1u << C_WHITE)
 #define ROAD_BLACK (1u << C_BLACK)
 #define ROAD_BOTH (ROAD_WHITE | ROAD_BLACK)
 
+// One bit per square, marking the cells on a given side of the board.
+// Indexed by board_size - 5 (0 for 5x5, 1 for 6x6):
+//   5x5: bottom bits 0-4, top bits 20-24,
+//        left bits 0,5,10,15,20, right bits 4,9,14,19,24
+//   6x6: bottom bits 0-5, top bits 30-35,
+//        left bits 0,6,12,18,24,30, right bits 5,11,17,23,29,35
+static const uint64_t side_bottom_row[2] = {0x000000000000001FULL,
+                                            0x000000000000003FULL};
+static const uint64_t side_top_row[2] = {0x0000000001F00000ULL,
+                                         0x00000FC0000000ULL};
+static const uint64_t side_left_col[2] = {0x0000000000108421ULL,
+                                          0x000000041041041ULL};
+static const uint64_t side_right_col[2] = {0x000000001084210ULL,
+                                           0x000820820820ULL};
+/*
+  Flood fill, restricted to the squares set in `mask', starting from the
+  squares set in `start'. Each iteration grows the fill by one square in all
+  four directions (the left/right edge masks stop the horizontal shifts from
+  wrapping to the next row), so the fill reaches a fixed point within the
+  board's diameter, 2*(board_size - 1), iterations.
+*/
+static inline uint64_t side_flood(const uint64_t mask, const uint64_t start,
+                                  const uint8_t board_size,
+                                  const uint64_t left,
+                                  const uint64_t right) {
+  uint64_t fill = start & mask;
+  for (uint8_t i = 0; i < 2 * (board_size - 1); i++) {
+    // << 1 moves toward the right column (and would wrap to the next row's
+    // left edge, hence the ~right mask); >> 1 is the mirror image
+    const uint64_t grow =
+        ((fill & ~right) << 1) | ((fill & ~left) >> 1) |
+        (fill << board_size) | (fill >> board_size);
+    const uint64_t nf = fill | (grow & mask);
+    if (nf == fill)
+      break;
+    fill = nf;
+  }
+  return fill;
+}
+
 enum WIN_TYPE check_win(tak_state_p state) {
-  uint8_t component[NUM_SQUARES(state->board_size)],
-      touching[NUM_SQUARES(state->board_size)];
+  const uint8_t board_size = state->board_size;
+  const uint8_t si = board_size - 5;
+  const uint64_t bottom = side_bottom_row[si], top = side_top_row[si],
+      left = side_left_col[si], right = side_right_col[si];
 
   /*
-    We're doing a poor version of a disjoint set data structure to track and
-    merge connected components. The array `component' stores indices to the
-    representative cells of each connected component. A cell is representative
-    if component[cell] = cell. We track only whether representatives are
-    touching sides, and do a 2-dimensional DP approach to forming these from the
-    board.
-
-    Note: we don't track the rank/size of each tree, because we're not
-    interested in good asymptotic complexity in the size of the board, merely
-    good performance for a single board size in practice. For the same reason we
-    also don't do path flattening/halving or anything.
-
-    A cell is road-eligible only for the colour of its top stone, so black and
-    white components are disjoint and can share one component/touching pair,
-    letting us check both colours (and count top-flat stones) in a single pass.
-    The dragon rule needs the road status of both colours.
+    A cell is road-eligible only for the colour of its top stone, and only
+    if that stone is not a standing stone. We build one bitset of
+    road-eligible squares per colour, then flood fill from each pair of
+    opposite sides: a road exists for a colour if the fill from the left
+    column (or the bottom row) reaches the opposite side. Both colours are
+    checked in a single pass over the board, which also counts top-flat
+    stones. The dragon rule needs the road status of both colours.
   */
-  for (int k = 0; k < NUM_SQUARES(state->board_size); k++) {
-    component[k] = k; // every square is in its own connected
-    // component initially
-    touching[k] = 0; // and not connected to any sides
-  }
-
-  uint8_t road = 0;
+  uint64_t m[2] = {0, 0};
 
   // flat count: +1 per top-flat black stone, -1 per top-flat white
   int8_t total = 0;
   uint8_t board_full = 1;
 
-  int touch = SIDE_BOTTOM | SIDE_LEFT;
-  for (int row = 0; row < state->board_size; row++) {
-    for (int col = 0; col < state->board_size; col++) {
-      const int cur = THE_COORDS(state->board_size, col, row);
-      const uint8_t count = COUNT_AT(state, cur);
-      if (count == 0) {
-        board_full = 0;
-      } else if (STONE_AT(state, cur) == STONE_FLAT) {
-        total += (TOP_COLOUR_AT(state, cur) == C_BLACK) ? +1 : -1;
-      }
-
-      if (count && STONE_AT(state, cur) != STONE_STANDING) {
-        const uint8_t colour = TOP_COLOUR_AT(state, cur);
-        const uint8_t done = 1u << colour;
-        if ((road & done) == 0) {
-
-          // do we have any neighbours to the left and below?
-          const int left_neighbour =
-              ((cur % state->board_size > 0) &&
-               (COUNT_AT(state, cur - 1)) // wont ever be out of bounds
-               && (TOP_COLOUR_AT(state, cur - 1) == colour) &&
-               (STONE_AT(state, cur - 1) != STONE_STANDING));
-
-          const int lowr_neighbour =
-              ((cur >= state->board_size) &&
-               (COUNT_AT(state, cur - state->board_size)) &&
-               (TOP_COLOUR_AT(state, cur - state->board_size) == colour) &&
-               (STONE_AT(state, cur - state->board_size) != STONE_STANDING));
-
-          // always take the component of the lower neighbour if
-          // possible, failing that take the left neighbour, otherwise
-          // we're not yet connected, so update our own component.
-          if (lowr_neighbour) {
-            // look up the root of the lower neighbour
-            int root = cur - state->board_size;
-            while (root != component[root])
-              root = component[root];
-            // join the set
-            component[cur] = root;
-            if (touch) {
-              // something new
-              touching[root] |= touch;
-              // are we done?
-              if ((touching[root] & SIDES_BOTTOM_TOP) == SIDES_BOTTOM_TOP ||
-                  (touching[root] & SIDES_LEFT_RIGHT) == SIDES_LEFT_RIGHT)
-                road |= done;
-            }
-            // if we also have a left neighbour then we should `merge'
-            // sets, and here we assume that the left neighbour set is
-            // always smaller (may not be) for the direction of merge
-            if (left_neighbour) {
-              int left_root = cur - 1;
-              while (left_root != component[left_root])
-                left_root = component[left_root];
-              // merge
-              const int left_touch = touching[left_root];
-              if (left_touch) {
-                touching[root] |= left_touch;
-                if ((touching[root] & SIDES_BOTTOM_TOP) == SIDES_BOTTOM_TOP ||
-                    (touching[root] & SIDES_LEFT_RIGHT) == SIDES_LEFT_RIGHT)
-                  road |= done;
-              }
-              component[left_root] = root;
-            }
-          } else if (left_neighbour) {
-            int root = cur - 1;
-            while (root != component[root])
-              root = component[root];
-            component[cur] = root;
-            if (touch) {
-              touching[root] |= touch;
-              if ((touching[root] & SIDES_BOTTOM_TOP) == SIDES_BOTTOM_TOP ||
-                  (touching[root] & SIDES_LEFT_RIGHT) == SIDES_LEFT_RIGHT)
-                road |= done;
-            }
-          } else if (touch) {
-            // we had no left or lower neighbour, so we're on our own
-            touching[cur] = touch;
-          }
-          if (road == ROAD_BOTH)
-            break; // both colours decided, the scan is done
-        }
-      }
-      if (col + 2 == state->board_size)
-        touch |= SIDE_RIGHT;
-      else
-        touch &= (SIDE_BOTTOM | SIDE_TOP);
+  for (uint32_t k = 0; k < NUM_SQUARES(board_size); k++) {
+    const uint8_t cd = state->celldat[k];
+    const uint8_t count = cd >> NUM_SHIFT;
+    if (count == 0) {
+      board_full = 0;
+    } else {
+      const uint8_t stone = cd & STONE_MASK;
+      if (stone == STONE_FLAT)
+        total += (state->colours[k] & 1) ? +1 : -1;
+      if (stone != STONE_STANDING)
+        m[state->colours[k] & 1] |= 1ULL << k;
     }
-    if (row + 2 == state->board_size)
-      touch = SIDE_TOP | SIDE_LEFT;
-    else
-      touch = SIDE_LEFT;
-    if (road == ROAD_BOTH)
-      break;
+  }
+
+  uint8_t road = 0;
+  for (uint8_t c = 0; c < 2; c++) {
+    if (m[c] == 0)
+      continue;
+    if ((side_flood(m[c], left, board_size, left, right) & right) != 0 ||
+        (side_flood(m[c], bottom, board_size, left, right) & top) != 0)
+      road |= 1u << c;
   }
 
   const enum WIN_TYPE rw = (road & ROAD_WHITE) ? WIN_ROAD_WHITE : 0xFF,
                       rb = (road & ROAD_BLACK) ? WIN_ROAD_BLACK : 0xFF;
 
   if (rb == WIN_ROAD_BLACK && rw == WIN_ROAD_WHITE) {
-    return (state->ply & 1) ? rb : rw; // Dragons
+    return WIN_DRAGON;
   } else if (rw == WIN_ROAD_WHITE) {
     return rw;
   } else if (rb == WIN_ROAD_BLACK) {
@@ -696,6 +643,10 @@ enum ACT_RESULT do_ptn(tak_state_p state, char *ptn) {
     // more conservative here :)
     if (state->ply >= state->board_size) {
       state->won = check_win(state);
+      if (state->won == WIN_DRAGON)
+        state->won = (state->current_colour == C_BLACK)
+                         ? WIN_ROAD_BLACK
+                         : WIN_ROAD_WHITE;
       if (state->won < 0xFF) {
         // Winning move, but no need to update current colour
         state->ply++;