Skip to content

Commit 7ac8170

Browse files
authored
Merge pull request #201 from smly/fix/bugfix-python-select-api
fix: disambiguate select_action_from_mjai with tsumogiri/consumed (#195)
2 parents 5829f7a + 7299d46 commit 7ac8170

10 files changed

Lines changed: 363 additions & 191 deletions

File tree

riichienv-core/src/observation/encode.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -616,6 +616,7 @@ mod tests {
616616
[None; 4], // riichi_sutehais
617617
[None; 4], // last_tedashis
618618
None, // last_discard
619+
None, // drawn_tile
619620
)
620621
}
621622

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
//! Shared helpers for mapping Mjai messages to a legal `Action`.
2+
//!
3+
//! Used by both `Observation::select_action_from_mjai` (4P) and
4+
//! `Observation3P::select_action_from_mjai` (3P).
5+
6+
use pyo3::prelude::*;
7+
use pyo3::types::{PyDict, PyDictMethods};
8+
9+
use crate::action::{Action, ActionType};
10+
use crate::parser::tid_to_mjai;
11+
12+
pub(crate) struct ParsedMjai {
13+
pub type_str: String,
14+
pub tile_str: String,
15+
pub tsumogiri: Option<bool>,
16+
pub consumed: Option<Vec<String>>,
17+
}
18+
19+
pub(crate) fn parse_mjai_message(mjai_data: &Bound<'_, PyAny>) -> Option<ParsedMjai> {
20+
if let Ok(s) = mjai_data.extract::<String>() {
21+
let v: serde_json::Value = serde_json::from_str(&s).ok()?;
22+
let type_str = v["type"].as_str()?.to_string();
23+
let tile_str = v["pai"].as_str().unwrap_or("").to_string();
24+
let tsumogiri = v.get("tsumogiri").and_then(|x| x.as_bool());
25+
let consumed = v.get("consumed").and_then(|x| x.as_array()).map(|arr| {
26+
arr.iter()
27+
.filter_map(|e| e.as_str().map(|s| s.to_string()))
28+
.collect::<Vec<_>>()
29+
});
30+
Some(ParsedMjai {
31+
type_str,
32+
tile_str,
33+
tsumogiri,
34+
consumed,
35+
})
36+
} else if let Ok(dict) = mjai_data.cast::<PyDict>() {
37+
let type_str: String = dict
38+
.get_item("type")
39+
.ok()
40+
.flatten()
41+
.and_then(|x| x.extract::<String>().ok())
42+
.unwrap_or_default();
43+
let tile_str: String = dict
44+
.get_item("pai")
45+
.ok()
46+
.flatten()
47+
.or_else(|| dict.get_item("tile").ok().flatten())
48+
.and_then(|x| x.extract::<String>().ok())
49+
.unwrap_or_default();
50+
let tsumogiri = dict
51+
.get_item("tsumogiri")
52+
.ok()
53+
.flatten()
54+
.and_then(|x| x.extract::<bool>().ok());
55+
let consumed = dict
56+
.get_item("consumed")
57+
.ok()
58+
.flatten()
59+
.and_then(|x| x.extract::<Vec<String>>().ok());
60+
Some(ParsedMjai {
61+
type_str,
62+
tile_str,
63+
tsumogiri,
64+
consumed,
65+
})
66+
} else {
67+
None
68+
}
69+
}
70+
71+
fn consumed_matches(action_consume: &[u8], expected: &[String]) -> bool {
72+
if action_consume.len() != expected.len() {
73+
return false;
74+
}
75+
let mut a: Vec<String> = action_consume.iter().map(|&t| tid_to_mjai(t)).collect();
76+
let mut b: Vec<String> = expected.to_vec();
77+
a.sort();
78+
b.sort();
79+
a == b
80+
}
81+
82+
/// Select a matching `Action` from a slice of legal actions for a parsed Mjai
83+
/// message.
84+
///
85+
/// `three_player` controls whether 3P-only types (`kita`) are recognized; chi
86+
/// is rejected when set.
87+
pub(crate) fn select_action<'a>(
88+
legal_actions: &'a [Action],
89+
parsed: &ParsedMjai,
90+
drawn_tile: Option<u8>,
91+
three_player: bool,
92+
) -> Option<&'a Action> {
93+
let atype = parsed.type_str.as_str();
94+
95+
if atype == "hora" {
96+
return legal_actions
97+
.iter()
98+
.find(|a| matches!(a.action_type, ActionType::Tsumo | ActionType::Ron));
99+
}
100+
101+
if atype == "none" {
102+
return legal_actions
103+
.iter()
104+
.find(|a| a.action_type == ActionType::Pass);
105+
}
106+
107+
let target_type = match atype {
108+
"dahai" => Some(ActionType::Discard),
109+
"chi" if !three_player => Some(ActionType::Chi),
110+
"pon" => Some(ActionType::Pon),
111+
"kakan" => Some(ActionType::Kakan),
112+
"daiminkan" => Some(ActionType::Daiminkan),
113+
"ankan" => Some(ActionType::Ankan),
114+
"kita" if three_player => Some(ActionType::Kita),
115+
"reach" => Some(ActionType::Riichi),
116+
"ryukyoku" => Some(ActionType::KyushuKyuhai),
117+
_ => None,
118+
};
119+
120+
let tt = target_type?;
121+
122+
// Special-case Discard: filter by mjai pai (or any Discard if pai is
123+
// omitted) then disambiguate via tsumogiri.
124+
//
125+
// NOTE: An mjai `dahai` message without a `pai` field is malformed per
126+
// the protocol, but we still return a non-empty Action (the first
127+
// legal Discard) instead of `None` to preserve backward compatibility
128+
// with the previous implementation; bailing out here would silently
129+
// break callers that rely on the old lenient behavior.
130+
if tt == ActionType::Discard {
131+
let candidates: Vec<&Action> = legal_actions
132+
.iter()
133+
.filter(|a| {
134+
a.action_type == ActionType::Discard
135+
&& (parsed.tile_str.is_empty()
136+
|| a.tile.is_some_and(|t| tid_to_mjai(t) == parsed.tile_str))
137+
})
138+
.collect();
139+
140+
if candidates.is_empty() {
141+
return None;
142+
}
143+
144+
if let (Some(tsumogiri), Some(drawn)) = (parsed.tsumogiri, drawn_tile) {
145+
let preferred = candidates.iter().find(|a| {
146+
let is_drawn = a.tile == Some(drawn);
147+
if tsumogiri { is_drawn } else { !is_drawn }
148+
});
149+
if let Some(a) = preferred {
150+
return Some(*a);
151+
}
152+
}
153+
154+
return Some(candidates[0]);
155+
}
156+
157+
legal_actions.iter().find(|a| {
158+
if a.action_type != tt {
159+
return false;
160+
}
161+
162+
if let Some(consumed) = parsed.consumed.as_ref() {
163+
if !consumed_matches(&a.consume_tiles, consumed) {
164+
return false;
165+
}
166+
// If pai is also given, double-check tile match for actions that
167+
// carry a meaningful tile (chi/pon/daiminkan/kakan).
168+
if !parsed.tile_str.is_empty()
169+
&& matches!(
170+
tt,
171+
ActionType::Chi | ActionType::Pon | ActionType::Daiminkan | ActionType::Kakan
172+
)
173+
{
174+
if let Some(t) = a.tile {
175+
if tid_to_mjai(t) != parsed.tile_str {
176+
return false;
177+
}
178+
} else {
179+
return false;
180+
}
181+
}
182+
return true;
183+
}
184+
185+
// No consumed field: fall back to pai-based match.
186+
if !parsed.tile_str.is_empty() {
187+
if let Some(t) = a.tile {
188+
return tid_to_mjai(t) == parsed.tile_str;
189+
}
190+
return false;
191+
}
192+
true
193+
})
194+
}

riichienv-core/src/observation/mod.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ mod encode;
33
#[cfg(feature = "python")]
44
pub(crate) mod helpers;
55
#[cfg(feature = "python")]
6+
pub(crate) mod mjai_select;
7+
#[cfg(feature = "python")]
68
mod python;
79
#[cfg(feature = "python")]
810
pub(crate) mod sequence_features;
@@ -49,6 +51,8 @@ pub struct Observation {
4951
pub riichi_sutehais: [Option<u8>; 4],
5052
pub last_tedashis: [Option<u8>; 4],
5153
pub last_discard: Option<u32>,
54+
#[serde(default)]
55+
pub drawn_tile: Option<u8>,
5256
}
5357

5458
/// Pure Rust methods (no PyO3 dependency).
@@ -74,6 +78,7 @@ impl Observation {
7478
riichi_sutehais: [Option<u8>; 4],
7579
last_tedashis: [Option<u8>; 4],
7680
last_discard: Option<u32>,
81+
drawn_tile: Option<u8>,
7782
) -> Self {
7883
let hands_u32 = hands.map(|h| h.into_iter().map(|x| x as u32).collect());
7984
let discards_u32 = discards.map(|d| d.into_iter().map(|x| x as u32).collect());
@@ -101,6 +106,7 @@ impl Observation {
101106
riichi_sutehais,
102107
last_tedashis,
103108
last_discard,
109+
drawn_tile,
104110
}
105111
}
106112

riichienv-core/src/observation/python.rs

Lines changed: 6 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use super::helpers::get_next_tile;
1414
impl Observation {
1515
#[new]
1616
#[allow(clippy::too_many_arguments)]
17-
#[pyo3(signature = (player_id, hands, melds, discards, dora_indicators, scores, riichi_declared, legal_actions, events, honba, riichi_sticks, round_wind, oya, kyoku_index, waits, is_tenpai, riichi_sutehais, last_tedashis, last_discard))]
17+
#[pyo3(signature = (player_id, hands, melds, discards, dora_indicators, scores, riichi_declared, legal_actions, events, honba, riichi_sticks, round_wind, oya, kyoku_index, waits, is_tenpai, riichi_sutehais, last_tedashis, last_discard, drawn_tile=None))]
1818
pub fn py_new(
1919
player_id: u8,
2020
hands: Vec<Vec<u8>>,
@@ -35,6 +35,7 @@ impl Observation {
3535
riichi_sutehais: Vec<Option<u8>>,
3636
last_tedashis: Vec<Option<u8>>,
3737
last_discard: Option<u32>,
38+
drawn_tile: Option<u8>,
3839
) -> Self {
3940
let hands: [Vec<u8>; 4] = hands.try_into().expect("expected 4 hands");
4041
let melds: [Vec<Meld>; 4] = melds.try_into().expect("expected 4 melds");
@@ -68,6 +69,7 @@ impl Observation {
6869
riichi_sutehais,
6970
last_tedashis,
7071
last_discard,
72+
drawn_tile,
7173
)
7274
}
7375

@@ -120,100 +122,9 @@ impl Observation {
120122

121123
#[pyo3(signature = (mjai_data))]
122124
pub fn select_action_from_mjai(&self, mjai_data: &Bound<'_, PyAny>) -> Option<Action> {
123-
let (atype, tile_str) = if let Ok(s) = mjai_data.extract::<String>() {
124-
let v: serde_json::Value = serde_json::from_str(&s).ok()?;
125-
(
126-
v["type"].as_str()?.to_string(),
127-
v["pai"].as_str().unwrap_or("").to_string(),
128-
)
129-
} else if let Ok(dict) = mjai_data.cast::<PyDict>() {
130-
let type_str: String = dict
131-
.get_item("type")
132-
.ok()
133-
.flatten()
134-
.and_then(|x| x.extract::<String>().ok())
135-
.unwrap_or_default();
136-
let _args_list: Vec<String> = dict
137-
.get_item("args")
138-
.ok()
139-
.flatten()
140-
.and_then(|x| x.extract::<Vec<String>>().ok())
141-
.unwrap_or_default();
142-
let _who: i8 = dict
143-
.get_item("who")
144-
.ok()
145-
.flatten()
146-
.and_then(|x| x.extract::<i8>().ok())
147-
.unwrap_or(-1);
148-
let tile_str: String = dict
149-
.get_item("pai")
150-
.ok()
151-
.flatten()
152-
.or_else(|| dict.get_item("tile").ok().flatten())
153-
.and_then(|x| x.extract::<String>().ok())
154-
.unwrap_or_default();
155-
(type_str, tile_str)
156-
} else {
157-
return None;
158-
};
159-
160-
let target_type = match atype.as_str() {
161-
"dahai" => Some(crate::action::ActionType::Discard),
162-
"chi" => Some(crate::action::ActionType::Chi),
163-
"pon" => Some(crate::action::ActionType::Pon),
164-
"kakan" => Some(crate::action::ActionType::Kakan),
165-
"daiminkan" => Some(crate::action::ActionType::Daiminkan),
166-
"ankan" => Some(crate::action::ActionType::Ankan),
167-
"reach" => Some(crate::action::ActionType::Riichi),
168-
"hora" => None,
169-
"ryukyoku" => Some(crate::action::ActionType::KyushuKyuhai),
170-
_ => None,
171-
};
172-
173-
if atype == "hora" {
174-
return self
175-
._legal_actions
176-
.iter()
177-
.find(|a| {
178-
a.action_type == crate::action::ActionType::Tsumo
179-
|| a.action_type == crate::action::ActionType::Ron
180-
})
181-
.cloned();
182-
}
183-
184-
if let Some(tt) = target_type {
185-
return self
186-
._legal_actions
187-
.iter()
188-
.find(|a| {
189-
if a.action_type != tt {
190-
return false;
191-
}
192-
if !tile_str.is_empty() {
193-
if let Some(t) = a.tile {
194-
let t_str = crate::parser::tid_to_mjai(t);
195-
if t_str == tile_str {
196-
return true;
197-
}
198-
return false;
199-
} else {
200-
return false;
201-
}
202-
}
203-
true
204-
})
205-
.cloned();
206-
}
207-
208-
if atype == "none" {
209-
return self
210-
._legal_actions
211-
.iter()
212-
.find(|a| a.action_type == crate::action::ActionType::Pass)
213-
.cloned();
214-
}
215-
216-
None
125+
use super::mjai_select::{parse_mjai_message, select_action};
126+
let parsed = parse_mjai_message(mjai_data)?;
127+
select_action(&self._legal_actions, &parsed, self.drawn_tile, false).cloned()
217128
}
218129

219130
#[pyo3(name = "new_events")]

riichienv-core/src/observation_3p/encode.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -650,6 +650,7 @@ mod tests {
650650
[None; 3], // riichi_sutehais
651651
[None; 3], // last_tedashis
652652
None, // last_discard
653+
None, // drawn_tile
653654
)
654655
}
655656

riichienv-core/src/observation_3p/mod.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ pub struct Observation3P {
4141
pub riichi_sutehais: [Option<u8>; 3],
4242
pub last_tedashis: [Option<u8>; 3],
4343
pub last_discard: Option<u32>,
44+
#[serde(default)]
45+
pub drawn_tile: Option<u8>,
4446
}
4547

4648
/// Pure Rust methods (no PyO3 dependency).
@@ -66,6 +68,7 @@ impl Observation3P {
6668
riichi_sutehais: [Option<u8>; 3],
6769
last_tedashis: [Option<u8>; 3],
6870
last_discard: Option<u32>,
71+
drawn_tile: Option<u8>,
6972
) -> Self {
7073
let hands_u32 = hands.map(|h| h.into_iter().map(|x| x as u32).collect());
7174
let discards_u32 = discards.map(|d| d.into_iter().map(|x| x as u32).collect());
@@ -95,6 +98,7 @@ impl Observation3P {
9598
riichi_sutehais,
9699
last_tedashis,
97100
last_discard,
101+
drawn_tile,
98102
}
99103
}
100104

0 commit comments

Comments
 (0)