Repository: sotetsuk/pgx

Issue:
[Go] Bug in Ko
<img width="103" alt="image" src="https://user-images.githubusercontent.com/5868442/224210707-5a72e43b-0d60-4634-8630-0c4a04592be0.png">
<img width="108" alt="image" src="https://user-images.githubusercontent.com/5868442/224210719-28af27a5-6428-4bd9-ba07-71fefb3f68e3.png">

このときに `action=21` が誤ってKoに判定され、置けない

```

1
16
9
11
14
6
24
4
25
17
21
18
13
23
8
0
5
25
15
19
22
25
12
10
2
25
3
20
```


Discussion:

Resolving patch:
diff --git a/pgx/go.py b/pgx/go.py
--- a/pgx/go.py
+++ b/pgx/go.py
@@ -203,7 +203,7 @@ def _not_pass_move(_state: State, _action: int, size) -> State:
     my_color_ix = _my_color_ix(state)
     agehama_before = state.agehama[my_color_ix]
 
-    kou_occurred = _kou_occurred(state, xy)
+    kou_may_occur = _kou_may_occur(state, xy)
 
     # 周囲の連から敵石を除く
     adj_xy = _neighbour(xy, size)
@@ -219,7 +219,7 @@ def _not_pass_move(_state: State, _action: int, size) -> State:
         0, 4,
         lambda i, s: jax.lax.cond(
             is_killed[i],
-            lambda: _remove_stones(s, ren_id[i], adj_xy[i]),
+            lambda: _remove_stones(s, ren_id[i], adj_xy[i], kou_may_occur),
             lambda: s,
         ),
         state,
@@ -234,7 +234,7 @@ def _not_pass_move(_state: State, _action: int, size) -> State:
 
     # コウの確認
     state = jax.lax.cond(
-        kou_occurred & state.agehama[my_color_ix] - agehama_before == 1,
+        kou_may_occur & state.agehama[my_color_ix] - agehama_before == 1,
         lambda: state,
         lambda: state.replace(kou=jnp.int32(-1)),  # type:ignore
     )
@@ -281,15 +281,21 @@ def _merge_ren(_state: State, _xy: int, _adj_xy: int):
     )
 
 
-def _remove_stones(_state: State, _rm_ren_id, _rm_stone_xy) -> State:
+def _remove_stones(
+    _state: State, _rm_ren_id, _rm_stone_xy, kou_may_occur
+) -> State:
     surrounded_stones = _state.ren_id_board == _rm_ren_id
     agehama = jnp.count_nonzero(surrounded_stones)
     ren_id_board = jnp.where(surrounded_stones, 0, _state.ren_id_board)
-
+    kou = jax.lax.cond(
+        kou_may_occur & (agehama == 1),
+        lambda: jnp.int32(_rm_stone_xy),
+        lambda: _state.kou,
+    )
     return _state.replace(  # type:ignore
         ren_id_board=ren_id_board,
         agehama=_state.agehama.at[_my_color_ix(_state)].add(agehama),
-        kou=jnp.int32(_rm_stone_xy),  # type:ignore
+        kou=kou,
     )
 
 
@@ -410,7 +416,7 @@ def _opponent_color_ix(_state: State):
     return (_state.turn + 1) % 2
 
 
-def _kou_occurred(_state: State, xy: int) -> jnp.ndarray:
+def _kou_may_occur(_state: State, xy: int) -> jnp.ndarray:
     size = _state.size
     x = xy // size
     y = xy % size


Test patch:
diff --git a/tests/test_go.py b/tests/test_go.py
--- a/tests/test_go.py
+++ b/tests/test_go.py
@@ -134,6 +134,73 @@ def test_kou():
     assert not state2.terminated
     assert state2.kou == -1
 
+    # see #468
+    state: State = env.init(key=key)
+    state = env.step(state, action=2)
+    state = env.step(state, action=9)
+    state = env.step(state, action=18)
+    state = env.step(state, action=5)
+    state = env.step(state, action=11)
+    state = env.step(state, action=22)
+    state = env.step(state, action=8)
+    state = env.step(state, action=14)
+    state = env.step(state, action=25)
+    state = env.step(state, action=1)
+    state = env.step(state, action=24)
+    state = env.step(state, action=23)
+    state = env.step(state, action=7)
+    state = env.step(state, action=4)
+    state = env.step(state, action=16)
+    state = env.step(state, action=15)
+    state = env.step(state, action=19)
+    state = env.step(state, action=6)
+    state = env.step(state, action=20)
+    state = env.step(state, action=25)
+    state = env.step(state, action=12)
+    state = env.step(state, action=3)
+    state = env.step(state, action=21)
+    state = env.step(state, action=10)
+    state = env.step(state, action=17)
+    state = env.step(state, action=25)
+    state = env.step(state, action=13)
+    state = env.step(state, action=4)
+    state = env.step(state, action=14)
+    state = env.step(state, action=23)
+    state = env.step(state, action=0)
+    assert state.kou == -1
+
+    # see #468
+    state: State = env.init(key=key)
+    state = env.step(state, action=1)
+    state = env.step(state, action=16)
+    state = env.step(state, action=9)
+    state = env.step(state, action=11)
+    state = env.step(state, action=14)
+    state = env.step(state, action=6)
+    state = env.step(state, action=24)
+    state = env.step(state, action=4)
+    state = env.step(state, action=25)
+    state = env.step(state, action=17)
+    state = env.step(state, action=21)
+    state = env.step(state, action=18)
+    state = env.step(state, action=13)
+    state = env.step(state, action=23)
+    state = env.step(state, action=8)
+    state = env.step(state, action=0)
+    state = env.step(state, action=5)
+    state = env.step(state, action=25)
+    state = env.step(state, action=15)
+    state = env.step(state, action=19)
+    state = env.step(state, action=22)
+    state = env.step(state, action=25)
+    state = env.step(state, action=12)
+    state = env.step(state, action=10)
+    state = env.step(state, action=2)
+    state = env.step(state, action=25)
+    state = env.step(state, action=3)
+    state = env.step(state, action=20)
+    assert state.kou == 15
+
 
 def test_observe():
     key = jax.random.PRNGKey(0)