1 files changed, 97 insertions, 105 deletions

diff --git a/include/tak.c b/include/tak.c
index 78d3854..6505bb1 100644
--- a/include/tak.c
+++ b/include/tak.c
@@ -32,9 +32,6 @@ uint8_t white_count, black_count, ply;
 // Helpers
 // ===================================================================
 
-#define DFS_MASK       0x0C // 0b00001100
-#define NOT_DFS_MASK   0xF3 // 0b11110011
-
 #define NUM_SQUARES   (board_size * board_size)
 
 // ===================================================================
@@ -233,110 +230,108 @@ try_move(const int8_t location, const enum MOVE_DIRECTION direction,
 // Checking win
 // ===================================================================
 
-// Depth-first search of the board for a road with a given
-// directionality:
-// direction = 0 --> left-to-right, direction = 1 --> bottom-to-top
-static uint8_t
-dfs_road(uint8_t dfs_stack[NUM_SQUARES], uint8_t dfs_pntr,
-				 const enum COLOUR colour, const uint8_t direction) {
-	while (dfs_pntr > 0) {
-		const uint8_t cur = dfs_stack[--dfs_pntr];
-		// Made it to the other side?
-		if ( (direction == 1 && cur >= board_size * (board_size - 1))
-				 || (direction == 0 && cur % board_size + 1 == board_size) )
-			return 1;
-
-		// Check the four neighbours of this cell, provided they exist,
-		// are inhabited, and are of the appropriate colour
-
-		// Direction: > (same row)
-		if ( ((cur % board_size) + 1 < board_size)
-				 && (COUNT_AT(cur + 1)) // wont ever be out of bounds
-				 && ((colours[cur+1] & 1) == colour)
-				 && ((celldat[cur+1] & DFS_MASK) == 0)
-				 && (STONE_AT(cur+1) != STONE_STANDING)) {
-			dfs_stack[dfs_pntr++] = cur + 1;
-			celldat[cur+1] |= DFS_MASK;
-		}
-
-		// Direction: <  (same row)
-		if ( (cur % board_size > 0)
-				 && (COUNT_AT(cur - 1)) // wont ever be out of bounds
-				 && ((colours[cur-1] & 1) == colour)
-				 && ((celldat[cur-1] & DFS_MASK) == 0)
-				 && (STONE_AT(cur-1) != STONE_STANDING)) {
-			dfs_stack[dfs_pntr++] = cur - 1;
-			celldat[cur-1] |= DFS_MASK;
-		}
-
-		// Direction: +
-		if ( (cur + board_size < NUM_SQUARES)
-				 && (COUNT_AT(cur + board_size))
-				 && ((colours[cur+board_size] & 1) == colour)
-				 && ((celldat[cur+board_size] & DFS_MASK) == 0)
-				 && (STONE_AT(cur+board_size) != STONE_STANDING)) {
-			dfs_stack[dfs_pntr++] = cur + board_size;
-			celldat[cur+board_size] |= DFS_MASK;
-		}
-
-		// Direction: -
-		if ( (cur >= board_size)
-				 && (COUNT_AT(cur - board_size))
-				 && ((colours[cur-board_size] & 1) == colour)
-				 && ((celldat[cur-board_size] & DFS_MASK) == 0)
-				 && (STONE_AT(cur-board_size) != STONE_STANDING)) {
-			dfs_stack[dfs_pntr++] = cur - board_size;
-			celldat[cur-board_size] |= DFS_MASK;
-		}
-	}
-	return 0;
-}
-
+// Check for the presence of a road connecting opposite sides
 static enum WIN_TYPE
 check_road_colour(const enum COLOUR colour) {
-	uint8_t dfs_stack[NUM_SQUARES];
-	uint8_t dfs_pntr;
-
-	// Prime the depth-first-search stack with all boundary cells of
-	// colour COLOUR.
-
-	const enum WIN_TYPE winner =
-		(colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
-
-	// We're using the two left-over bits in data_t to track whether
-	// we've seen it. Reset those before anything.
-
-	for (uint8_t k=0; k<NUM_SQUARES; k++)
-		celldat[k] &= NOT_DFS_MASK;
-
-	// First left-to-right
-	dfs_pntr = 0;
-	for (uint8_t y=0; y<board_size; y++) {
-		if ( COUNT_AT(THE_COORDS(0, y))
-				 && (colours[THE_COORDS(0, y)] & 1) == colour
-				 && STONE_AT(THE_COORDS(0, y)) != STONE_STANDING ) {
-			dfs_stack[dfs_pntr++] = THE_COORDS(0, y);
-			celldat[THE_COORDS(0, y)] |= DFS_MASK;
-		}
+	int component[NUM_SQUARES], touching[NUM_SQUARES];
+
+	/*
+	  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.
+	*/
+	for (int k=0; k<NUM_SQUARES; k++) {
+		component[k] = k; // every square is in its own connected
+											// component initially
+		touching[k] = 0;  // and not connected to any sides
 	}
-	if (dfs_road(dfs_stack, dfs_pntr, colour, 0)) return winner;
-
-	// Clear visited squares
-	for (uint8_t k=0; k<NUM_SQUARES; k++)
-		celldat[k] &= NOT_DFS_MASK;
-
-	// Then bottom-to-top
-	dfs_pntr = 0;
-	for (uint8_t x=0; x<board_size; x++) {
-		if ( COUNT_AT(THE_COORDS(x, 0))
-				 && (colours[THE_COORDS(x, 0)] & 1) == colour
-				 && STONE_AT(THE_COORDS(x, 0)) != STONE_STANDING ) {
-			dfs_stack[dfs_pntr++] = THE_COORDS(x, 0);
-			celldat[THE_COORDS(x, 0)] |= DFS_MASK;
+
+	// touching is the bit mask for connectivity,
+	// bottom | top | left | right
+	//    1      2     4       8
+
+	int touch = 5;
+	for (int row = 0; row < board_size; row++) {
+		for (int col = 0; col < board_size; col++) {
+			const int cur = THE_COORDS(col, row);
+			if (COUNT_AT(cur)
+			    && (colours[cur] & 1) == colour
+			    && STONE_AT(cur) != STONE_STANDING) {
+
+				// do we have any neighbours to the left and below?
+				const int left_neighbour = ((cur % board_size > 0)
+				                            && (COUNT_AT(cur - 1)) // wont ever be out of bounds
+				                            && ((colours[cur - 1] & 1) == colour)
+				                            && (STONE_AT(cur - 1) != STONE_STANDING));
+
+				const int lowr_neighbour = ((cur >= board_size)
+				                           && (COUNT_AT(cur - board_size))
+				                           && ((colours[cur - board_size] & 1) == colour)
+				                           && (STONE_AT(cur - 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 - 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] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC)
+							return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
+					}
+					// 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] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC)
+								return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
+						}
+						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] & 0x3) == 0x3 || (touching[root] & 0xC) == 0xC)
+							return (colour == C_BLACK) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
+					}
+				} else if (touch) {
+					// we had no left or lower neighbour, so we're on our own
+					touching[cur] = touch;
+				}
+			}
+			if (col + 2 == board_size) touch |= 8;
+			else touch &= 0x3;
 		}
+		if (row + 2 == board_size) touch = 6;
+		else touch = 4;
 	}
-	if (dfs_road(dfs_stack, dfs_pntr, colour, 1)) return winner;
-
 	return 0xFF;
 }
 
@@ -346,9 +341,6 @@ check_win(void) {
 	enum WIN_TYPE rb, rw;
 	rb = check_road_colour(C_BLACK);
 	rw = check_road_colour(C_WHITE);
-	for (uint8_t k=0; k<NUM_SQUARES; k++)
-		celldat[k] &= NOT_DFS_MASK;
-
 
 	if (rb == WIN_ROAD_BLACK && rw == WIN_ROAD_WHITE) {
 		return (ply & 1) ? rb : rw; // Dragons