1 files changed, 11 insertions, 3 deletions
diff --git a/include/nn.c b/include/nn.c
index 705f258..c1a806d 100644
--- a/include/nn.c
+++ b/include/nn.c
@@ -45,9 +45,17 @@ static float fast_tanh(const float x) {
 
 inline uint8_t cell_idx_at(const tak_state_p state, const uint8_t loc) {
   const uint8_t d = state->celldat[loc];
-  if ((d & NUM_MASK) == 0)
-    return 0; // empty
-  return (uint8_t)(1u + 3u * (state->colours[loc] & 1u) + (d & STONE_MASK));
+  switch (d & NUM_MASK) {
+    case 0: {
+      // index 0
+      return 0;
+    }
+    default: {
+      // |{flat, standing, cap} * {black, white}| = 6
+      // index 1 .. 6
+      return (uint8_t)(1u + 3u * (state->colours[loc] & 1u) + (d & STONE_MASK));
+    }
+  }
 }
 
 inline void nn_cell_idxs(uint8_t *out, tak_state_p state) {