Repository: sotetsuk/pgx

Issue:
 [Mahjong] Determine the next player after Meld



Discussion:

Resolving patch:
diff --git a/pgx/_mahjong/_mahjong2.py b/pgx/_mahjong/_mahjong2.py
--- a/pgx/_mahjong/_mahjong2.py
+++ b/pgx/_mahjong/_mahjong2.py
@@ -80,9 +80,10 @@ class State(v1.State):
     # 最後のプレイヤー.  ron,pon,chiの対象
     last_player: jnp.ndarray = jnp.int8(0)
 
-    # 打牌の後に競合する副露が生じた場合、どこまで調べたかを保存
-    # 0:none 1:ron判定済み 2:pon/kan判定済み
-    furo_check_num: jnp.ndarray = jnp.int8(0)
+    # 打牌の後に競合する副露が生じた場合用
+    #     pon player, kan player, chi player
+    # b0       00         00         00
+    furo_check_num: jnp.ndarray = jnp.uint8(0)
 
     # state.current_player がリーチ宣言してから, その直後の打牌が通るまでTrue
     riichi_declared: jnp.ndarray = FALSE
@@ -154,7 +155,7 @@ def decode_state(data: dict):
                 target=jnp.int8(data["target"]),
                 last_draw=jnp.int8(data["last_draw"]),
                 last_player=jnp.int8(data["last_player"]),
-                furo_check_num=jnp.int8(data["furo_check_num"]),
+                furo_check_num=jnp.uint8(data["furo_check_num"]),
                 riichi_declared=jnp.bool_(data["riichi_declared"]),
                 riichi=jnp.array(data["riichi"], dtype=jnp.bool_),
                 n_meld=jnp.array(data["n_meld"], dtype=jnp.int8),
@@ -381,10 +382,13 @@ def _discard(state: State, tile: jnp.ndarray):
     meld_type = 0  # none < chi_L = chi_M = chi_R < pon = kan < ron の中で最大のものを探す
     pon_player = kan_player = ron_player = c_p
     chi_player = (c_p + 1) % 4
-    meld_type = jax.lax.cond(
+    can_chi = (
         Hand.can_chi(state.hand[chi_player], tile, Action.CHI_L)
         | Hand.can_chi(state.hand[chi_player], tile, Action.CHI_M)
-        | Hand.can_chi(state.hand[chi_player], tile, Action.CHI_R),
+        | Hand.can_chi(state.hand[chi_player], tile, Action.CHI_R)
+    )
+    meld_type = jax.lax.cond(
+        can_chi,
         lambda: jnp.max(jnp.array([1, meld_type])),
         lambda: meld_type,
     )
@@ -394,12 +398,12 @@ def search(i, tpl):
         player = (c_p + 1 + i) % 4
         pon_player, meld_type = jax.lax.cond(
             Hand.can_pon(state.hand[player], tile),
-            lambda: (player, jnp.max(jnp.array([2, meld_type]))),
+            lambda: (i, jnp.max(jnp.array([2, meld_type]))),
             lambda: (pon_player, meld_type),
         )
         kan_player, meld_type = jax.lax.cond(
             Hand.can_minkan(hand[player], tile),
-            lambda: (player, jnp.max(jnp.array([3, meld_type]))),
+            lambda: (i, jnp.max(jnp.array([3, meld_type]))),
             lambda: (kan_player, meld_type),
         )
         ron_player, meld_type = jax.lax.cond(
@@ -412,13 +416,24 @@ def search(i, tpl):
                 state.riichi[c_p],
                 FALSE,
             )[0].any(),
-            lambda: (player, jnp.max(jnp.array([4, meld_type]))),
+            lambda: (i, jnp.max(jnp.array([4, meld_type]))),
             lambda: (ron_player, meld_type),
         )
         return (meld_type, pon_player, kan_player, ron_player)
 
     meld_type, pon_player, kan_player, ron_player = jax.lax.fori_loop(
-        0, 3, search, (meld_type, pon_player, kan_player, ron_player)
+        jnp.int8(0),
+        jnp.int8(3),
+        search,
+        (meld_type, pon_player, kan_player, ron_player),
+    )
+    furo_num = jnp.uint8(
+        c_p << 6 | kan_player << 4 | pon_player << 2 | chi_player
+    )
+    pon_player, kan_player, ron_player = (
+        (c_p + 1 + pon_player) % 4,
+        (c_p + 1 + kan_player) % 4,
+        (c_p + 1 + ron_player) % 4,
     )
 
     rewards = jnp.float32([Hand.is_tenpai(hand) * 100 for hand in state.hand])
@@ -444,6 +459,7 @@ def search(i, tpl):
                 current_player=chi_player,
                 last_player=c_p,
                 target=jnp.int8(tile),
+                furo_check_num=furo_num & 0b11111100,
                 legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
                 .at[Action.CHI_L]
                 .set(Hand.can_chi(state.hand[chi_player], tile, Action.CHI_L))
@@ -458,6 +474,7 @@ def search(i, tpl):
                 current_player=pon_player,
                 last_player=c_p,
                 target=jnp.int8(tile),
+                furo_check_num=furo_num & 0b11110011,
                 legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
                 .at[Action.PON]
                 .set(TRUE)
@@ -468,6 +485,7 @@ def search(i, tpl):
                 current_player=kan_player,
                 last_player=c_p,
                 target=jnp.int8(tile),
+                furo_check_num=furo_num & 0b11001111,
                 legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
                 .at[Action.MINKAN]
                 .set(Hand.can_minkan(hand[kan_player], tile))
@@ -478,6 +496,7 @@ def search(i, tpl):
                 current_player=ron_player,
                 last_player=c_p,
                 target=jnp.int8(tile),
+                furo_check_num=furo_num,
                 legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
                 .at[Action.RON]
                 .set(TRUE)
@@ -656,10 +675,75 @@ def _chi(state: State, action):
 
 
 def _pass(state: State):
-    # TODO
-    # ron -> pon/kan -> chiへの遷移
-
-    return _draw(state)
+    last_player = (state.furo_check_num & 0b11000000) >> 6
+    kan_player = (state.furo_check_num & 0b00110000) >> 4
+    pon_player = (state.furo_check_num & 0b00001100) >> 2
+    chi_player = state.furo_check_num & 0b00000011
+    return jax.lax.cond(
+        kan_player > 0,
+        lambda: state.replace(  # type:ignore
+            current_player=jnp.int8(last_player + 1 + kan_player) % 4,
+            furo_check_num=jnp.uint8(state.furo_check_num & 0b11001111),
+            legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
+            .at[Action.MINKAN]
+            .set(Hand.can_minkan(state.hand[kan_player], state.target))
+            .at[Action.PASS]
+            .set(TRUE),
+        ),
+        lambda: jax.lax.cond(
+            pon_player > 0,
+            lambda: state.replace(  # type:ignore
+                current_player=jnp.int8(last_player + 1 + pon_player) % 4,
+                furo_check_num=jnp.uint8(state.furo_check_num & 0b11110011),
+                legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
+                .at[Action.PON]
+                .set(Hand.can_pon(state.hand[pon_player], state.target))
+                .at[Action.PASS]
+                .set(TRUE),
+            ),
+            lambda: jax.lax.cond(
+                chi_player > 0,
+                lambda: state.replace(  # type:ignore
+                    current_player=jnp.int8(last_player + 1 + chi_player) % 4,
+                    furo_check_num=jnp.uint8(
+                        state.furo_check_num & 0b11111100
+                    ),
+                    legal_action_mask=jnp.zeros(NUM_ACTION, dtype=jnp.bool_)
+                    .at[Action.CHI_L]
+                    .set(
+                        Hand.can_chi(
+                            state.hand[(last_player + 1 + chi_player) % 4],
+                            state.target,
+                            Action.CHI_L,
+                        )
+                    )
+                    .at[Action.CHI_M]
+                    .set(
+                        Hand.can_chi(
+                            state.hand[(last_player + 1 + chi_player) % 4],
+                            state.target,
+                            Action.CHI_M,
+                        )
+                    )
+                    .at[Action.CHI_R]
+                    .set(
+                        Hand.can_chi(
+                            state.hand[(last_player + 1 + chi_player) % 4],
+                            state.target,
+                            Action.CHI_R,
+                        )
+                    )
+                    .at[Action.PASS]
+                    .set(TRUE),
+                ),
+                lambda: _draw(
+                    state.replace(  # type: ignore
+                        current_player=jnp.int8(last_player + 1) % 4,
+                    )
+                ),
+            ),
+        ),
+    )
 
 
 def _riichi(state: State):


Test patch: