~25% speedup in selfplay due to: incremental zobrist hashing, smarter road checking, re-using the input computed for the NN

10 files changed, 234 insertions, 118 deletions

diff --git a/include/actions.c b/include/actions.c
index 541b388..981214f 100644
--- a/include/actions.c
+++ b/include/actions.c
@@ -17,6 +17,7 @@
 
 #include "actions.h"
 #include "tak.h"
+#include "zobrist.h"
 
 #include <stdio.h>
 #include <stdlib.h>
@@ -135,7 +136,7 @@ action_t *actions_generate_at(tak_state_p state, uint32_t *num_actions,
       // Only try moves after CPS and if the colour is correct
       if (count) {
         if (state->ply >= 2 &&
-            ((state->colours[loc] & 1) == state->current_colour)) {
+            (TOP_COLOUR_AT(state, loc) == state->current_colour)) {
 
           // Pre-compute end-stops and crushes
           uint8_t end_stops[4], crushes[4] = {0, 0, 0, 0};
@@ -245,7 +246,18 @@ action_t *actions_generate_at(tak_state_p state, uint32_t *num_actions,
 
 void action_take(tak_state_p state, const action_t action) {
   const int8_t loc = A_GET_LOC(action);
+
+  // Incremental hash: record the affected cells and their current (pre-move)
+  // contributions, apply the move, then fold in the change.
+  uint8_t affected[7];
+  uint64_t old_contrib[7];
+  int n_aff;
+
   if (A_GET_TYPE(action) == A_PLACE) {
+    affected[0] = (uint8_t)loc;
+    n_aff = 1;
+    old_contrib[0] = zobrist_cell(state, affected[0]);
+
     const uint8_t black = (state->current_colour == C_BLACK);
     switch (A_GET_DATA0(action)) {
     case STONE_FLAT: {
@@ -293,6 +305,15 @@ void action_take(tak_state_p state, const action_t action) {
     for (uint8_t _gaps = gaps; _gaps; steps++)
       _gaps &= _gaps - 1;
 
+    // Record the affected cells (source + every drop square) and their
+    // current contributions
+    affected[0] = (uint8_t)loc;
+    for (int8_t k = 1; k <= steps; k++)
+      affected[k] = (uint8_t)(loc + k * delta);
+    n_aff = steps + 1;
+    for (int i = 0; i < n_aff; i++)
+      old_contrib[i] = zobrist_cell(state, affected[i]);
+
     // Move top stone type to destination
     state->celldat[loc + steps * delta] &= CLR_STONE; // necessary for crushing
     state->celldat[loc + steps * delta] |= STONE_AT(state, loc);
@@ -327,6 +348,14 @@ void action_take(tak_state_p state, const action_t action) {
   }
   // Next ply
   inline_next_ply(state);
+
+  // Fold the move into the incremental hash: remove each affected cell's old
+  // contribution, add its new one, and flip the ply-parity term.
+  for (int i = 0; i < n_aff; i++) {
+    state->hash ^= old_contrib[i];
+    state->hash ^= zobrist_cell(state, affected[i]);
+  }
+  state->hash ^= PLY_PARITY_MAGIC;
 }
 
 void action_undo(tak_state_p state, const action_t action) {
@@ -334,7 +363,18 @@ void action_undo(tak_state_p state, const action_t action) {
   inline_prev_ply(state);
 
   const int8_t loc = A_GET_LOC(action);
+
+  // Incremental hash: record the affected cells and their current
+  // (pre-undo) contributions, reverse the move, then fold in the change.
+  uint8_t affected[7];
+  uint64_t old_contrib[7];
+  int n_aff;
+
   if (A_GET_TYPE(action) == A_PLACE) {
+    affected[0] = (uint8_t)loc;
+    n_aff = 1;
+    old_contrib[0] = zobrist_cell(state, affected[0]);
+
     const uint8_t black = (state->current_colour == C_BLACK);
     state->celldat[loc] = 0;
     state->colours[loc] = 0;
@@ -358,6 +398,18 @@ void action_undo(tak_state_p state, const action_t action) {
                   dir = A_GET_DATA1(action) >> 4;
     const int8_t delta = move_deltas[dir];
 
+    // Record the affected cells (source + every drop square) and their
+    // current contributions
+    int8_t n_steps = 1;
+    for (uint8_t _gaps = gaps; _gaps; n_steps++)
+      _gaps &= _gaps - 1;
+    affected[0] = (uint8_t)loc;
+    for (int8_t k = 1; k <= n_steps; k++)
+      affected[k] = (uint8_t)(loc + k * delta);
+    n_aff = n_steps + 1;
+    for (int i = 0; i < n_aff; i++)
+      old_contrib[i] = zobrist_cell(state, affected[i]);
+
     int8_t steps = 1;
     uint8_t gap_bit = 1, total = 1;
     for (int8_t d = 1; d < num; d++, total++, gap_bit <<= 1) {
@@ -389,6 +441,14 @@ void action_undo(tak_state_p state, const action_t action) {
           STONE_FLAT; // should be optimised out
     }
   }
+
+  // Fold the undo into the incremental hash: remove each affected cell's
+  // current contribution, add its restored one, and flip the ply-parity term.
+  for (int i = 0; i < n_aff; i++) {
+    state->hash ^= old_contrib[i];
+    state->hash ^= zobrist_cell(state, affected[i]);
+  }
+  state->hash ^= PLY_PARITY_MAGIC;
 }
 
 void action_to_ptn(const action_t action, char *out_ptn) {
diff --git a/include/negamax.c b/include/negamax.c
index 6c29481..c86eeab 100644
--- a/include/negamax.c
+++ b/include/negamax.c
@@ -100,6 +100,7 @@ float negamax_generate(tak_state_p state) {
 
   // this is constant across iterative deepening
   const uint64_t root_hash = zobrist_compute(state);
+  state->hash = root_hash; // anchor the incremental hash at the root
 #ifdef RECORD_PV
   pv_leaf_t iteration_leaf;
 #endif
@@ -197,7 +198,7 @@ static float negamax(tak_state_p state, const uint8_t cur_depth,
 
     action_take(state, actions[idx]);
 
-    const uint64_t child_hash = zobrist_compute(state);
+    const uint64_t child_hash = state->hash;
 
     // Compute the value of the node and retain the leaf that produced it.
     float node_value;
diff --git a/include/nn.c b/include/nn.c
index c1a806d..f228297 100644
--- a/include/nn.c
+++ b/include/nn.c
@@ -15,6 +15,8 @@
   along with ctak. If not, see <https://www.gnu.org/licenses/>.
 */
 
+#include <string.h>
+
 #include "nn.h"
 
 #include "weights_5.h"
@@ -53,7 +55,7 @@ inline uint8_t cell_idx_at(const tak_state_p state, const uint8_t loc) {
     default: {
       // |{flat, standing, cap} * {black, white}| = 6
       // index 1 .. 6
-      return (uint8_t)(1u + 3u * (state->colours[loc] & 1u) + (d & STONE_MASK));
+      return (uint8_t)(1u + 3u * TOP_COLOUR_AT(state, loc) + (d & STONE_MASK));
     }
   }
 }
@@ -72,14 +74,30 @@ inline void nn_prepare_head(float out_head[3], tak_state_p state) {
   out_head[2] = (float)(state->black_count & 127) * inv_max_flats;
 }
 
+// The cell indices computed by the most recent nn_prepare_input, and the hash
+// of the state that produced them.
+static uint8_t prepared_idxs[36]; // max board is 6x6
+static uint64_t prepared_hash = 0;
+
 void nn_prepare_input(float *cur_board, tak_state_p state,
                       const float embedding[NN_N_EMB]) {
   const unsigned int n_cells = state->board_size * state->board_size;
   nn_prepare_head(cur_board, state);
-  uint8_t idxs[36];
-  nn_cell_idxs(idxs, state);
+  nn_cell_idxs(prepared_idxs, state);
+  prepared_hash = state->hash;
   for (unsigned int loc = 0; loc < n_cells; loc++)
-    cur_board[3 + loc] = embedding[idxs[loc]];
+    cur_board[3 + loc] = embedding[prepared_idxs[loc]];
+}
+
+// Copy the cell indices produced by the most recent nn_prepare_input into out
+// and return 1, but only if that call was made on exactly this state, otherwise
+// return 0 and leave out unfilled.
+int nn_last_cell_idxs(uint8_t *out, tak_state_p state) {
+  const unsigned int n_cells = state->board_size * state->board_size;
+  if (prepared_hash != state->hash)
+    return 0;
+  memcpy(out, prepared_idxs, n_cells);
+  return 1;
 }
 
 // ===================================================================
diff --git a/include/nn.h b/include/nn.h
index 9e576c1..871bd0f 100644
--- a/include/nn.h
+++ b/include/nn.h
@@ -35,6 +35,8 @@ void nn_prepare_head(float out_head[3], tak_state_p state);
 
 void nn_prepare_input(float *cur_board, tak_state_p state,
                       const float embedding[NN_N_EMB]);
+
+int nn_last_cell_idxs(uint8_t *out, tak_state_p state);
 // ===================================================================
 // Network
 // ===================================================================
diff --git a/include/pv.c b/include/pv.c
index 9c120f5..94c00d1 100644
--- a/include/pv.c
+++ b/include/pv.c
@@ -38,9 +38,10 @@ static float terminal_value(const enum WIN_TYPE winner) {
 
 void pv_capture_leaf(pv_leaf_t *leaf, tak_state_p state,
                      const enum WIN_TYPE winner) {
-  nn_prepare_head(leaf->head, state);
-  nn_cell_idxs(leaf->cells, state);
   const uint8_t n_cells = state->board_size * state->board_size;
+  nn_prepare_head(leaf->head, state);
+  if (!nn_last_cell_idxs(leaf->cells, state))
+    nn_cell_idxs(leaf->cells, state);
   for (uint8_t loc = n_cells; loc < PV_MAX_CELLS; loc++)
     leaf->cells[loc] = 0;
   leaf->current_colour = (uint8_t)state->current_colour;
diff --git a/include/pv.h b/include/pv.h
index 4be7b24..06455fe 100644
--- a/include/pv.h
+++ b/include/pv.h
@@ -34,9 +34,11 @@ typedef struct pv_leaf_s {
   float terminal_value;
 } pv_leaf_t;
 
-// Capture a horizon leaf (winner == 0xFF) or an exact terminal leaf
-void pv_capture_leaf(pv_leaf_t *leaf, tak_state_p state,
-                     enum WIN_TYPE winner);
+// Capture a horizon leaf (winner == 0xFF) or an exact terminal leaf.
+//
+// The cell indices are taken from the encoding of the most recent
+// nn_prepare_input.
+void pv_capture_leaf(pv_leaf_t *leaf, tak_state_p state, enum WIN_TYPE winner);
 
 // Publish and retrieve the leaf selected by the full-depth root search
 void pv_publish_leaf(const pv_leaf_t *leaf);
diff --git a/include/tak.c b/include/tak.c
index 406f8b5..e9f27ff 100644
--- a/include/tak.c
+++ b/include/tak.c
@@ -50,6 +50,7 @@ void reset_state(tak_state_p state, const uint8_t new_board_size) {
   state->ply = 0;
   state->won = 0xFF; // i may live to regret this hack
   state->current_colour = C_BLACK;
+  state->hash = 0; // empty board, ply 0
 
   for (uint8_t k = 0; k < NUM_SQUARES(new_board_size); k++) {
     state->celldat[k] = 0;
@@ -154,7 +155,7 @@ enum ACT_RESULT try_move(tak_state_p state, const int8_t location,
   if (avail == 0)
     return ACT_ILLEGAL;
   // Does the current player own the pile?
-  if ((state->colours[location] & 1) != state->current_colour)
+  if (TOP_COLOUR_AT(state, location) != state->current_colour)
     return ACT_ILLEGAL;
   // Is the desired direction and count on the board?
   int8_t delta = 0;
@@ -239,25 +240,40 @@ enum ACT_RESULT try_move(tak_state_p state, const int8_t location,
 // Checking win
 // ===================================================================
 
-// Check for the presence of a road connecting opposite sides
-static enum WIN_TYPE check_road_colour(tak_state_p state,
-                                       const enum COLOUR colour) {
-  int component[NUM_SQUARES(state->board_size)],
+// 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)
+
+enum WIN_TYPE check_win(tak_state_p state) {
+  uint8_t component[NUM_SQUARES(state->board_size)],
       touching[NUM_SQUARES(state->board_size)];
 
   /*
-    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.
+    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.
+    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.
   */
   for (int k = 0; k < NUM_SQUARES(state->board_size); k++) {
     component[k] = k; // every square is in its own connected
@@ -265,97 +281,110 @@ static enum WIN_TYPE check_road_colour(tak_state_p state,
     touching[k] = 0; // and not connected to any sides
   }
 
-  // touching is the bit mask for connectivity,
-  // bottom | top | left | right
-  //    1      2     4       8
+  uint8_t road = 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 = 5;
+  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);
-      if (COUNT_AT(state, cur) && (state->colours[cur] & 1) == colour &&
-          STONE_AT(state, cur) != STONE_STANDING) {
+      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;
+      }
 
-        // 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
-             && ((state->colours[cur - 1] & 1) == colour) &&
-             (STONE_AT(state, cur - 1) != STONE_STANDING));
+      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) {
 
-        const int lowr_neighbour =
-            ((cur >= state->board_size) &&
-             (COUNT_AT(state, cur - state->board_size)) &&
-             ((state->colours[cur - state->board_size] & 1) == colour) &&
-             (STONE_AT(state, cur - state->board_size) != STONE_STANDING));
+          // 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));
 
-        // 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] & 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;
+          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;
             }
-            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;
+            // 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;
           }
-        } 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 |= 8;
+        touch |= SIDE_RIGHT;
       else
-        touch &= 0x3;
+        touch &= (SIDE_BOTTOM | SIDE_TOP);
     }
     if (row + 2 == state->board_size)
-      touch = 6;
+      touch = SIDE_TOP | SIDE_LEFT;
     else
-      touch = 4;
+      touch = SIDE_LEFT;
+    if (road == ROAD_BOTH)
+      break;
   }
-  return 0xFF;
-}
 
-enum WIN_TYPE check_win(tak_state_p state) {
-  // Road?
-  enum WIN_TYPE rb, rw;
-  rb = check_road_colour(state, C_BLACK);
-  rw = check_road_colour(state, C_WHITE);
+  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
@@ -365,15 +394,6 @@ enum WIN_TYPE check_win(tak_state_p state) {
     return rb;
   }
 
-  // Do we do a flat count?
-  int8_t total = 0, board_full = 1;
-  for (uint8_t k = 0; k < NUM_SQUARES(state->board_size); k++) {
-    if (COUNT_AT(state, k) == 0) {
-      board_full = 0;
-    } else if (STONE_AT(state, k) == STONE_FLAT) {
-      total += ((state->colours[k] & 1) == C_BLACK) ? +1 : -1;
-    }
-  }
   if (state->black_count == 0 || state->white_count == 0 || board_full) {
     // Decide based on count
     if (total > 0)
diff --git a/include/tak.h b/include/tak.h
index d60c221..324c41f 100644
--- a/include/tak.h
+++ b/include/tak.h
@@ -55,6 +55,9 @@ typedef struct tak_state_s {
   enum WIN_TYPE won;
   enum COLOUR current_colour;
   uint8_t white_count, black_count, ply, board_size;
+  // incremental zobrist hash; maintained by action_take/action_undo and
+  // re-anchored at the search root (see zobrist.h)
+  uint64_t hash;
 } *tak_state_p;
 
 #define NUM_SHIFT 4
@@ -63,6 +66,8 @@ typedef struct tak_state_s {
 #define STONE_MASK 3                // 0b00000011
 #define STONE_AT(state, l) (state->celldat[(l)] & STONE_MASK)
 #define COUNT_AT(state, l) (state->celldat[(l)] >> NUM_SHIFT)
+// the colour stack holds one bit per stone, the top stone in the LSB
+#define TOP_COLOUR_AT(state, l) (state->colours[(l)] & 1)
 #define THE_COORDS(board_size, col, row) ((col) + (row)*board_size)
 
 // ===================================================================
diff --git a/include/zobrist.c b/include/zobrist.c
index 2922abf..5076fe1 100644
--- a/include/zobrist.c
+++ b/include/zobrist.c
@@ -50,18 +50,19 @@ void zobrist_free(void) {
 }
 
 
-// a position's value is from the perspective of the player to move, which flips
-// with parity, so we must track this
-#define PLY_PARITY_MAGIC 0x9E3779B97F4A7C15ULL
+uint64_t zobrist_cell(tak_state_p state, const uint8_t l) {
+  colour_stack_t c = state->colours[l];
+  const uint8_t count = COUNT_AT(state, l);
+  enum STONE_VARIANT s = STONE_AT(state, l);
+  uint64_t h = 0;
+  for (uint8_t hh = 0; hh < count; hh++, c >>= 1)
+    h ^= zobrist[l * (15 * 2 * 3) + hh * 2 * 3 + (c & 1) * 3 + s];
+  return h;
+}
 
 uint64_t zobrist_compute(tak_state_p state) {
   uint64_t hash = 0;
-  for (uint8_t l = 0; l < state->board_size * state->board_size; l++) {
-    colour_stack_t c = state->colours[l];
-    const uint8_t count = COUNT_AT(state, l);
-    enum STONE_VARIANT s = STONE_AT(state, l);
-    for (uint8_t h = 0; h < count; h++, c >>= 1)
-      hash ^= zobrist[l * (15 * 2 * 3) + h * 2 * 3 + (c & 1) * 3 + s];
-  }
+  for (uint8_t l = 0; l < state->board_size * state->board_size; l++)
+    hash ^= zobrist_cell(state, l);
   return hash ^ ((state->ply & 1) ? PLY_PARITY_MAGIC : 0);
 }
diff --git a/include/zobrist.h b/include/zobrist.h
index 5fd5ace..27053a6 100644
--- a/include/zobrist.h
+++ b/include/zobrist.h
@@ -19,8 +19,14 @@
 #include <tak.h>
 #include <xorshift64.h>
 
+#define PLY_PARITY_MAGIC 0x9E3779B97F4A7C15ULL
+
 void zobrist_free(void);
 
 int zobrist_init(const uint8_t board_size);
 
+// hash for whole state
 uint64_t zobrist_compute(tak_state_p state);
+
+// hash contribution from a single board cell at location l
+uint64_t zobrist_cell(tak_state_p state, const uint8_t l);