Repository: sotetsuk/pgx

Issue:
[Shogi] Fix king move bug
```
    actions = [304, 14, 205, 59, 358, 79, 32, 192, 367, 196, 214, 530, 14, 214, 25, 141, 177, 182, 296, 1346, 789, 196,
               187, 305, 431, 202, 7, 61, 440, 7, 31, 137, 79, 149, 41, 5, 340, 15, 305, 41, 104, 376, 124, 208, 5, 542,
               22, 304, 277, 145, 715, 58, 611, 1884, 278, 157, 42, 123, 133, 228, 141, 983, 1532, 601, 194, 163, 389,
               68, 124, 1756, 2063, 562, 444, 407, 377, 23, 1689, 467, 33, 216, 59, 1658, 211, 67, 348, 50, 38, 187, 23,
               228, 217, 123, 2112, 61, 40, 65, 159, 22, 1935, 306, 461, 2079, 449, 1858, 184, 66, 275, 1179, 55, 561,
               1697, 1682, 220, 294, 1682, 77, 39, 875, 302, 315, 207, 634, 609, 583, 1441, 313, 1636, 1645, 329, 1660,
               1958, 13, 965, 315, 389, 387, 48, 315, 76, 159, 316, 387, 544, 317, 337, 2148, 4, 378, 1793, 789, 600,
               230, 165, 220, 1772, 1685, 1628, 210, 75, 480, 1873, 360, 530, 1694, 255, 203, 554, 537, 256, 39, 2039,
               2140, 1659, 883, 627, 1624, 569, 61, 420, 1707, 149, 286, 1643, 72, 14, 203, 205, 38, 539, 12, 32, 1939,
               451, 864, 2165, 49, 196, 3, 884, 11, 1759, 820, 274, 227, 1660, 356]
```

```
before   	1ns6/+S1p+Ng1p1l/+P2pg1nNS/4k2G1/2L1PR2s/p1G2+BPR1/3Pp2+p1/1+p3B1P1/1LPK2+l1+p w 4p 206
dlaction	356
action  	 Action(is_drop=Array(False, dtype=bool), piece=Array(7, dtype=int8), to=Array(32, dtype=int8), from_=Array(41, dtype=int32), is_promotion=Array(False, dtype=bool))
after   	1ns6/+S1p+Ng1p1l/+P2pg1nNS/3k3G1/2L1PR2s/p1G2+BPR1/3Pp2+p1/1+p3B1P1/1LPK2+l1+p b 4p 207
pgx     	1ns6/+S1p+Ng1p1l/+P2pg1nNS/4k2G1/2L1PR2s/p1G2+BPR1/3Pp2+p1/1+p3B1P1/1LPK2+l1+p w 4p 1
```

<img width="206" alt="image" src="https://user-images.githubusercontent.com/5868442/226862067-c0ebdbb6-0388-489a-a902-a4567173aac7.png">



Discussion:

Resolving patch:
diff --git a/pgx/shogi.py b/pgx/shogi.py
--- a/pgx/shogi.py
+++ b/pgx/shogi.py
@@ -683,12 +683,21 @@ def _filter_pawn_drop_mate(
     )  # NOTE: 王が一番下の段にいるとき間違っているが、その場合は使われないので問題ない
     can_check_by_pawn_drop = opp_king_pos % 9 != 8
 
+    @jax.vmap
+    def filter_effect_by_pawn(p, f):
+        # 歩打によってフィルタされる利き
+        mask = IS_ON_THE_WAY[p, f, :, opp_king_head_pos]
+        return jnp.where(p >= 0, mask, FALSE)
+
     # 王が利きも味方の駒もないところへ逃げられるか
     king_escape_mask = RAW_EFFECT_BOARDS[KING, opp_king_pos, :]  # (81,)
     king_escape_mask &= ~(
         (OPP_PAWN <= pb) & (pb <= OPP_DRAGON)
     )  # 味方駒があり、逃げられない
-    king_escape_mask &= ~effect_boards.any(axis=0)  # 利きがあり逃げられない
+    effects = effect_boards & ~filter_effect_by_pawn(
+        _to_large_piece_ix(state.piece_board), jnp.arange(81)
+    )
+    king_escape_mask &= ~effects.any(axis=0)  # 利きがあり逃げられない
     can_king_escape = king_escape_mask.any()
 
     # 反転したボードで処理していることに注意


Test patch:
diff --git a/tests/assets/shogi/buggy_samples_011.svg b/tests/assets/shogi/buggy_samples_011.svg
new file mode 100644
--- /dev/null
+++ b/tests/assets/shogi/buggy_samples_011.svg
@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="500.0" version="1.1" width="550.0" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink"><defs /><rect fill="#1e1e1e" height="500" width="575.0" x="0" y="0" /><g transform="scale(1.0)"><rect fill="#1e1e1e" height="500" width="550" x="0" y="0" /><g transform="translate(15,0) translate(25.0,25.0)"><g id="hlines" stroke="gray"><line stroke-width="2px" x1="0" x2="450" y1="50" y2="50" /><line stroke-width="2px" x1="0" x2="450" y1="100" y2="100" /><line stroke-width="2px" x1="0" x2="450" y1="150" y2="150" /><line stroke-width="2px" x1="0" x2="450" y1="200" y2="200" /><line stroke-width="2px" x1="0" x2="450" y1="250" y2="250" /><line stroke-width="2px" x1="0" x2="450" y1="300" y2="300" /><line stroke-width="2px" x1="0" x2="450" y1="350" y2="350" /><line stroke-width="2px" x1="0" x2="450" y1="400" y2="400" /></g><g id="vline" stroke="gray"><line stroke-width="2px" x1="50" x2="50" y1="0" y2="450" /><line stroke-width="2px" x1="100" x2="100" y1="0" y2="450" /><line stroke-width="2px" x1="150" x2="150" y1="0" y2="450" /><line stroke-width="2px" x1="200" x2="200" y1="0" y2="450" /><line stroke-width="2px" x1="250" x2="250" y1="0" y2="450" /><line stroke-width="2px" x1="300" x2="300" y1="0" y2="450" /><line stroke-width="2px" x1="350" x2="350" y1="0" y2="450" /><line stroke-width="2px" x1="400" x2="400" y1="0" y2="450" /></g><rect fill="none" height="450" stroke="gray" stroke-width="4px" width="450" x="0" y="0" /><g fill="gray" id="cord"><text font-family="Serif" font-size="20px" x="455.0" y="30.0">一</text><text font-family="Serif" font-size="20px" x="420.0" y="-5.0">1</text><text font-family="Serif" font-size="20px" x="455.0" y="80.0">二</text><text font-family="Serif" font-size="20px" x="370.0" y="-5.0">2</text><text font-family="Serif" font-size="20px" x="455.0" y="130.0">三</text><text font-family="Serif" font-size="20px" x="320.0" y="-5.0">3</text><text font-family="Serif" font-size="20px" x="455.0" y="180.0">四</text><text font-family="Serif" font-size="20px" x="270.0" y="-5.0">4</text><text font-family="Serif" font-size="20px" x="455.0" y="229.99999999999997">五</text><text font-family="Serif" font-size="20px" x="220.00000000000003" y="-5.0">5</text><text font-family="Serif" font-size="20px" x="455.0" y="280.0">六</text><text font-family="Serif" font-size="20px" x="170.0" y="-5.0">6</text><text font-family="Serif" font-size="20px" x="455.0" y="330.0">七</text><text font-family="Serif" font-size="20px" x="120.0" y="-5.0">7</text><text font-family="Serif" font-size="20px" x="455.0" y="380.0">八</text><text font-family="Serif" font-size="20px" x="70.0" y="-5.0">8</text><text font-family="Serif" font-size="20px" x="455.0" y="430.0">九</text><text font-family="Serif" font-size="20px" x="20.0" y="-5.0">9</text></g><g><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="442.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="242.49999999999997">香</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="142.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="142.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="392.5">角</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="292.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="242.49999999999997">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="192.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="142.5">と</text><text fill="gray" font-family="Serif" font-size="26px" x="161.0" y="72.5">成</text><text fill="gray" font-family="Serif" font-size="26px" x="161.0" y="97.5">桂</text><text fill="gray" font-family="Serif" font-size="26px" x="11.0" y="72.5">成</text><text fill="gray" font-family="Serif" font-size="26px" x="11.0" y="97.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="292.5">馬</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="300.0">☗</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="335.0">先</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="370.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="405.0" /><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="440.00000000000006">歩</text></g><g transform="rotate(180,225.0,225.0)"><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="392.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="142.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="392.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="192.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="392.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="342.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="242.49999999999997">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="392.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="292.5">玉</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="42.5">と</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="142.5">と</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="92.5">と</text><text fill="gray" font-family="Serif" font-size="26px" x="111.00000000000001" y="22.5">成</text><text fill="gray" font-family="Serif" font-size="26px" x="111.00000000000001" y="47.5">香</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="265.00000000000006">☖</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="300.0">後</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="335.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="370.0" /><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="405.0">歩</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="440.00000000000006">四</text></g></g></g></svg>
\ No newline at end of file
diff --git a/tests/assets/shogi/buggy_samples_012.svg b/tests/assets/shogi/buggy_samples_012.svg
new file mode 100644
--- /dev/null
+++ b/tests/assets/shogi/buggy_samples_012.svg
@@ -0,0 +1,2 @@
+<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="500.0" version="1.1" width="550.0" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink"><defs /><rect fill="#1e1e1e" height="500" width="575.0" x="0" y="0" /><g transform="scale(1.0)"><rect fill="#1e1e1e" height="500" width="550" x="0" y="0" /><g transform="translate(15,0) translate(25.0,25.0)"><g id="hlines" stroke="gray"><line stroke-width="2px" x1="0" x2="450" y1="50" y2="50" /><line stroke-width="2px" x1="0" x2="450" y1="100" y2="100" /><line stroke-width="2px" x1="0" x2="450" y1="150" y2="150" /><line stroke-width="2px" x1="0" x2="450" y1="200" y2="200" /><line stroke-width="2px" x1="0" x2="450" y1="250" y2="250" /><line stroke-width="2px" x1="0" x2="450" y1="300" y2="300" /><line stroke-width="2px" x1="0" x2="450" y1="350" y2="350" /><line stroke-width="2px" x1="0" x2="450" y1="400" y2="400" /></g><g id="vline" stroke="gray"><line stroke-width="2px" x1="50" x2="50" y1="0" y2="450" /><line stroke-width="2px" x1="100" x2="100" y1="0" y2="450" /><line stroke-width="2px" x1="150" x2="150" y1="0" y2="450" /><line stroke-width="2px" x1="200" x2="200" y1="0" y2="450" /><line stroke-width="2px" x1="250" x2="250" y1="0" y2="450" /><line stroke-width="2px" x1="300" x2="300" y1="0" y2="450" /><line stroke-width="2px" x1="350" x2="350" y1="0" y2="450" /><line stroke-width="2px" x1="400" x2="400" y1="0" y2="450" /></g><rect fill="none" height="450" stroke="gray" stroke-width="4px" width="450" x="0" y="0" /><g fill="gray" id="cord"><text font-family="Serif" font-size="20px" x="455.0" y="30.0">一</text><text font-family="Serif" font-size="20px" x="420.0" y="-5.0">1</text><text font-family="Serif" font-size="20px" x="455.0" y="80.0">二</text><text font-family="Serif" font-size="20px" x="370.0" y="-5.0">2</text><text font-family="Serif" font-size="20px" x="455.0" y="130.0">三</text><text font-family="Serif" font-size="20px" x="320.0" y="-5.0">3</text><text font-family="Serif" font-size="20px" x="455.0" y="180.0">四</text><text font-family="Serif" font-size="20px" x="270.0" y="-5.0">4</text><text font-family="Serif" font-size="20px" x="455.0" y="229.99999999999997">五</text><text font-family="Serif" font-size="20px" x="220.00000000000003" y="-5.0">5</text><text font-family="Serif" font-size="20px" x="455.0" y="280.0">六</text><text font-family="Serif" font-size="20px" x="170.0" y="-5.0">6</text><text font-family="Serif" font-size="20px" x="455.0" y="330.0">七</text><text font-family="Serif" font-size="20px" x="120.0" y="-5.0">7</text><text font-family="Serif" font-size="20px" x="455.0" y="380.0">八</text><text font-family="Serif" font-size="20px" x="70.0" y="-5.0">8</text><text font-family="Serif" font-size="20px" x="455.0" y="430.0">九</text><text font-family="Serif" font-size="20px" x="20.0" y="-5.0">9</text></g><g><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="392.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="292.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="242.49999999999997">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="442.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="242.49999999999997">香</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="442.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="142.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="142.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="392.5">角</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="292.5">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="242.49999999999997">飛</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="192.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="292.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="152.5" y="442.5">玉</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="142.5">と</text><text fill="gray" font-family="Serif" font-size="26px" x="161.0" y="72.5">成</text><text fill="gray" font-family="Serif" font-size="26px" x="161.0" y="97.5">桂</text><text fill="gray" font-family="Serif" font-size="26px" x="11.0" y="72.5">成</text><text fill="gray" font-family="Serif" font-size="26px" x="11.0" y="97.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="292.5">馬</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="335.0">☗</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="370.0">先</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="405.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="475.0" y="440.00000000000006" /></g><g transform="rotate(180,225.0,225.0)"><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="392.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="142.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="252.5" y="342.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="392.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="402.50000000000006" y="192.5">歩</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="392.5">香</text><text fill="gray" font-family="Serif" font-size="45px" x="102.49999999999999" y="342.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="442.5">桂</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="242.49999999999997">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="302.5" y="442.5">銀</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="392.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="342.5">金</text><text fill="gray" font-family="Serif" font-size="45px" x="202.5" y="292.5">玉</text><text fill="gray" font-family="Serif" font-size="45px" x="2.5" y="42.5">と</text><text fill="gray" font-family="Serif" font-size="45px" x="52.5" y="142.5">と</text><text fill="gray" font-family="Serif" font-size="45px" x="352.5" y="92.5">と</text><text fill="gray" font-family="Serif" font-size="26px" x="111.00000000000001" y="22.5">成</text><text fill="gray" font-family="Serif" font-size="26px" x="111.00000000000001" y="47.5">香</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="265.00000000000006">☖</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="300.0">後</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="335.0">手</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="370.0" /><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="405.0">歩</text><text fill="gray" font-family="Serif" font-size="34px" x="455.0" y="440.00000000000006">四</text></g></g></g></svg>
\ No newline at end of file
diff --git a/tests/test_shogi.py b/tests/test_shogi.py
--- a/tests/test_shogi.py
+++ b/tests/test_shogi.py
@@ -603,6 +603,14 @@ def test_buggy_samples():
     expected_state = _from_sfen("2+P+P2G1+S/1P2+P+P1+Pn/+S1GK2P2/1b2PP3/1nl4PP/3k2lRL/1pg4L1/p2+s2p2/P+n+B+p+ng1+s+p b Pr 1")
     assert (state.piece_board == expected_state.piece_board).all()
 
+    # 629
+    state = _from_sfen("1ns6/+S1p+Ng1p1l/+P2pg1nNS/4k2G1/2L2R2s/p1G2+BPR1/3Pp2+p1/1+p3B1P1/1LPK2+l1+p b P4p 1")
+    visualize(state, "tests/assets/shogi/buggy_samples_011.svg")
+    dlshogi_action = 1660  # 歩打
+    state = step(state, dlshogi_action)
+    visualize(state, "tests/assets/shogi/buggy_samples_012.svg")
+    assert not state.terminated  # 打ち歩詰でない
+
 
 def test_observe():
     key = jax.random.PRNGKey(0)