1 files changed, 27 insertions, 1 deletions
diff --git a/include/negamax.c b/include/negamax.c
index c86eeab..f9a52e6 100644
--- a/include/negamax.c
+++ b/include/negamax.c
@@ -52,6 +52,30 @@ static inline enum WIN_TYPE check_win_cached(tak_state_p state,
 }
 
 // ===================================================================
+// leaf evaluation memoisation
+// ===================================================================
+
+typedef struct eval_memo_s {
+  uint64_t key;
+  nn_function eval;
+  float value;
+} eval_memo_t;
+
+#define EVAL_MEMO_SIZE 262144
+static eval_memo_t eval_memo[EVAL_MEMO_SIZE];
+
+static inline float evaluate_memo(tak_state_p state, const uint64_t key) {
+  const uint32_t midx = key & (EVAL_MEMO_SIZE - 1);
+  if (eval_memo[midx].key == key && eval_memo[midx].eval == evaluate)
+    return eval_memo[midx].value;
+  const float value = evaluate(state);
+  eval_memo[midx].key = key;
+  eval_memo[midx].eval = evaluate;
+  eval_memo[midx].value = value;
+  return value;
+}
+
+// ===================================================================
 // Helper declarations
 // ===================================================================
 
@@ -204,6 +228,8 @@ static float negamax(tak_state_p state, const uint8_t cur_depth,
     float node_value;
     if (state->ply >= 2 * state->board_size - 2 &&
         (w = check_win_cached(state, child_hash)) < 0xFF) {
+      if (w == WIN_DRAGON)
+        w = (colour > 0) ? WIN_ROAD_BLACK : WIN_ROAD_WHITE;
       node_value = -colour * infty;
       // Check win if far enough into the game
       if (w == WIN_ROAD_BLACK || w == WIN_FLAT_BLACK) {
@@ -223,7 +249,7 @@ static float negamax(tak_state_p state, const uint8_t cur_depth,
 #endif
       );
     } else {
-      node_value = colour * evaluate(state);
+      node_value = colour * evaluate_memo(state, child_hash);
 #ifdef RECORD_PV
       pv_capture_leaf(&candidate_leaf, state, (enum WIN_TYPE)0xFF);
 #endif