use std::fs;
use std::path::Path;
use crate::board::{PackedSfen, Ply, Position};
use crate::book::BookError;
use crate::types::{Move, Square};
const MAGIC: &[u8; 16] = b"YANE-BINBOOK-V1\0";
const HEADER_BYTES: usize = 32;
const INDEX_BYTES: usize = 44;
#[derive(Debug, Clone)]
pub struct YbbBook {
bytes: Vec<u8>,
rows: Vec<IndexRow>,
move_area: usize,
has_depth: bool,
options: YbbBookOpenOptions,
}
#[derive(Debug, Clone, Copy)]
struct IndexRow {
packed_sfen: PackedSfen,
moves_offset: u64,
ply: Ply,
move_count: u16,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct YbbBookOpenOptions {
ignore_ply: bool,
flipped: bool,
}
impl YbbBookOpenOptions {
#[must_use]
pub const fn new() -> Self {
Self { ignore_ply: false, flipped: false }
}
#[must_use]
pub const fn with_ignore_ply(mut self, ignore_ply: bool) -> Self {
self.ignore_ply = ignore_ply;
self
}
#[must_use]
pub const fn with_flipped(mut self, flipped: bool) -> Self {
self.flipped = flipped;
self
}
#[must_use]
pub const fn ignore_ply(self) -> bool {
self.ignore_ply
}
#[must_use]
pub const fn flipped(self) -> bool {
self.flipped
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct YbbEntry {
packed_sfen: PackedSfen,
ply: Ply,
flipped: bool,
moves: Vec<YbbMove>,
}
impl YbbEntry {
#[must_use]
pub const fn packed_sfen(&self) -> PackedSfen {
self.packed_sfen
}
#[must_use]
pub const fn ply(&self) -> Ply {
self.ply
}
#[must_use]
pub const fn flipped(&self) -> bool {
self.flipped
}
#[must_use]
pub fn moves(&self) -> &[YbbMove] {
&self.moves
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct YbbMove {
mv: Move,
raw_move: u16,
eval: i16,
depth: Option<u16>,
}
impl YbbMove {
#[must_use]
pub const fn mv(self) -> Move {
self.mv
}
#[must_use]
pub const fn raw_move(self) -> u16 {
self.raw_move
}
#[must_use]
pub const fn eval(self) -> i16 {
self.eval
}
#[must_use]
pub const fn depth(self) -> Option<u16> {
self.depth
}
}
impl YbbBook {
pub fn open(path: impl AsRef<Path>) -> Result<Self, BookError> {
Self::open_with_options(path, YbbBookOpenOptions::new())
}
pub fn open_with_options(
path: impl AsRef<Path>,
options: YbbBookOpenOptions,
) -> Result<Self, BookError> {
let bytes = fs::read(path)?;
if bytes.len() < HEADER_BYTES {
return Err(BookError::InvalidFormat("truncated YBB header"));
}
if &bytes[..16] != MAGIC {
return Err(BookError::InvalidFormat("invalid YBB magic"));
}
let count = read_u64(&bytes[16..24])?;
let flags = read_u64(&bytes[24..32])?;
if flags & !1 != 0 {
return Err(BookError::Unsupported("unknown YBB header flags"));
}
let count_usize = usize::try_from(count)
.map_err(|_| BookError::InvalidFormat("YBB record count overflows usize"))?;
let index_bytes = count_usize
.checked_mul(INDEX_BYTES)
.ok_or(BookError::InvalidFormat("YBB index size overflow"))?;
let move_area = HEADER_BYTES
.checked_add(index_bytes)
.ok_or(BookError::InvalidFormat("YBB index end overflow"))?;
if move_area > bytes.len() {
return Err(BookError::InvalidFormat("truncated YBB index"));
}
let has_depth = flags & 1 != 0;
let move_record_bytes = if has_depth { 6usize } else { 4usize };
let mut rows: Vec<IndexRow> = Vec::with_capacity(count_usize);
for index in 0..count_usize {
let start = HEADER_BYTES + index * INDEX_BYTES;
let packed_sfen = PackedSfen::from_bytes(
bytes[start..start + 32].try_into().expect("fixed YBB index slice"),
);
let moves_offset = read_u64(&bytes[start + 32..start + 40])?;
let ply = read_u16(&bytes[start + 40..start + 42])?;
let move_count = read_u16(&bytes[start + 42..start + 44])?;
let byte_count = usize::from(move_count)
.checked_mul(move_record_bytes)
.ok_or(BookError::InvalidFormat("YBB move count overflow"))?;
let offset = usize::try_from(moves_offset)
.map_err(|_| BookError::InvalidFormat("YBB move offset overflows usize"))?;
let end = offset
.checked_add(byte_count)
.ok_or(BookError::InvalidFormat("YBB move range overflow"))?;
if end > bytes.len() - move_area {
return Err(BookError::InvalidFormat("truncated YBB move data"));
}
let row = IndexRow { packed_sfen, moves_offset, ply, move_count };
if let Some(previous) = rows.last()
&& previous.packed_sfen.cmp_bytes(&row.packed_sfen).is_gt()
{
return Err(BookError::InvalidFormat("YBB index is not sorted"));
}
rows.push(row);
}
let book = Self { bytes, rows, move_area, has_depth, options };
for row in &book.rows {
let _ = book.decode_moves(*row, false)?;
}
Ok(book)
}
#[must_use]
pub const fn len(&self) -> u64 {
self.rows.len() as u64
}
#[must_use]
pub const fn is_empty(&self) -> bool {
self.rows.is_empty()
}
#[must_use]
pub const fn has_depth(&self) -> bool {
self.has_depth
}
#[must_use]
pub const fn options(&self) -> YbbBookOpenOptions {
self.options
}
pub fn lookup_position(&self, pos: &Position) -> Result<Option<YbbEntry>, BookError> {
let packed = packed_sfen_for_lookup(pos, "position")?;
if let Some(entry) = self.lookup(packed, pos.game_ply(), false)? {
return Ok(Some(entry));
}
if !self.options.flipped() {
return Ok(None);
}
let flipped_sfen = pos.to_sfen_flipped(Some(i32::from(pos.game_ply())));
let flipped = Position::from_sfen(&flipped_sfen).map_err(|error| {
BookError::InvalidData(format!("cannot flip YBB position: {error}"))
})?;
self.lookup(packed_sfen_for_lookup(&flipped, "flipped position")?, pos.game_ply(), true)
}
pub fn lookup_packed_sfen(
&self,
packed_sfen: PackedSfen,
ply: Ply,
) -> Result<Option<YbbEntry>, BookError> {
if let Some(entry) = self.lookup(packed_sfen, ply, false)? {
return Ok(Some(entry));
}
if !self.options.flipped() {
return Ok(None);
}
let mut position = Position::empty();
position.set_packed_sfen(&packed_sfen, false, ply).map_err(|error| {
BookError::InvalidData(format!("invalid YBB PackedSfen lookup key: {error}"))
})?;
let flipped_sfen = position.to_sfen_flipped(Some(i32::from(ply)));
let flipped = Position::from_sfen(&flipped_sfen).map_err(|error| {
BookError::InvalidData(format!("cannot flip YBB lookup key: {error}"))
})?;
self.lookup(flipped.to_packed_sfen(), ply, true)
}
fn lookup(
&self,
packed_sfen: PackedSfen,
ply: Ply,
flipped: bool,
) -> Result<Option<YbbEntry>, BookError> {
let Ok(mut index) =
self.rows.binary_search_by(|row| row.packed_sfen.cmp_bytes(&packed_sfen))
else {
return Ok(None);
};
while index > 0 && self.rows[index - 1].packed_sfen == packed_sfen {
index -= 1;
}
let end = self.rows[index..]
.iter()
.position(|row| row.packed_sfen != packed_sfen)
.map_or(self.rows.len(), |relative| index + relative);
let row = if self.options.ignore_ply() {
self.rows[index]
} else {
let Some(row) = self.rows[index..end].iter().find(|row| row.ply == ply) else {
return Ok(None);
};
*row
};
Ok(Some(YbbEntry {
packed_sfen: row.packed_sfen,
ply: row.ply,
flipped,
moves: self.decode_moves(row, flipped)?,
}))
}
fn decode_moves(&self, row: IndexRow, flipped: bool) -> Result<Vec<YbbMove>, BookError> {
let record_bytes = if self.has_depth { 6usize } else { 4usize };
let offset = usize::try_from(row.moves_offset)
.map_err(|_| BookError::InvalidFormat("YBB move offset overflows usize"))?;
let start = self.move_area + offset;
let mut moves = Vec::with_capacity(usize::from(row.move_count));
for index in 0..usize::from(row.move_count) {
let base = start + index * record_bytes;
let raw_move = read_u16(&self.bytes[base..base + 2])?;
if !is_valid_raw_move(raw_move) {
return Err(BookError::InvalidData(format!(
"invalid YBB raw move: {raw_move:#06x}"
)));
}
let eval = read_i16(&self.bytes[base + 2..base + 4])?;
let depth = if self.has_depth {
Some(read_u16(&self.bytes[base + 4..base + 6])?)
} else {
None
};
let mv = if flipped {
flip_move(Move::from_raw(raw_move))
} else {
Move::from_raw(raw_move)
};
moves.push(YbbMove { mv, raw_move, eval, depth });
}
Ok(moves)
}
}
fn is_valid_raw_move(raw: u16) -> bool {
let mv = Move::from_raw(raw);
if !mv.is_normal() || mv.to_sq().raw() < 0 || mv.to_sq().raw() >= Square::COUNT as i8 {
return false;
}
if mv.is_drop() {
!mv.is_promotion() && matches!((raw >> 7) & 0x7f, 1..=7)
} else {
mv.from_sq().raw() >= 0
&& mv.from_sq().raw() < Square::COUNT as i8
&& mv.from_sq() != mv.to_sq()
}
}
fn packed_sfen_for_lookup(pos: &Position, description: &str) -> Result<PackedSfen, BookError> {
pos.try_to_packed_sfen().map_err(|error| {
BookError::InvalidData(format!("invalid YBB lookup {description}: {error}"))
})
}
fn flip_move(mv: Move) -> Move {
if mv.is_drop() {
Move::drop(mv.dropped_piece().expect("validated YBB drop"), mv.to_sq().flip())
} else if mv.is_promotion() {
Move::promotion(mv.from_sq().flip(), mv.to_sq().flip())
} else {
Move::normal(mv.from_sq().flip(), mv.to_sq().flip())
}
}
fn read_u16(bytes: &[u8]) -> Result<u16, BookError> {
Ok(u16::from_le_bytes(
bytes.try_into().map_err(|_| BookError::InvalidFormat("truncated YBB u16"))?,
))
}
fn read_i16(bytes: &[u8]) -> Result<i16, BookError> {
Ok(i16::from_le_bytes(
bytes.try_into().map_err(|_| BookError::InvalidFormat("truncated YBB i16"))?,
))
}
fn read_u64(bytes: &[u8]) -> Result<u64, BookError> {
Ok(u64::from_le_bytes(
bytes.try_into().map_err(|_| BookError::InvalidFormat("truncated YBB u64"))?,
))
}
#[cfg(test)]
mod tests {
use std::fs;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::board::hirate_position;
use super::*;
fn temp_path(suffix: &str) -> PathBuf {
let nonce =
SystemTime::now().duration_since(UNIX_EPOCH).expect("clock after epoch").as_nanos();
std::env::temp_dir().join(format!("rsshogi-ybb-{}-{nonce}-{suffix}", std::process::id()))
}
fn write_single_entry(path: &Path, position: &Position) {
let packed = position.try_to_packed_sfen().expect("packed position");
let mv = Move::from_usi("7g7f").expect("move");
let mut bytes = Vec::new();
bytes.extend_from_slice(MAGIC);
bytes.extend_from_slice(&1u64.to_le_bytes());
bytes.extend_from_slice(&0u64.to_le_bytes());
bytes.extend_from_slice(packed.as_bytes());
bytes.extend_from_slice(&0u64.to_le_bytes());
bytes.extend_from_slice(&position.game_ply().to_le_bytes());
bytes.extend_from_slice(&1u16.to_le_bytes());
bytes.extend_from_slice(&mv.raw().to_le_bytes());
bytes.extend_from_slice(&(-25i16).to_le_bytes());
fs::write(path, bytes).expect("write YBB fixture");
}
#[test]
fn test_ybb_lookup_position_rejects_unencodable_positions_and_keeps_valid_lookup() {
let path = temp_path("invalid-lookup-position");
let position = hirate_position();
write_single_entry(&path, &position);
let book = YbbBook::open_with_options(&path, YbbBookOpenOptions::new().with_flipped(true))
.expect("open YBB fixture");
assert!(matches!(book.lookup_position(&Position::empty()), Err(BookError::InvalidData(_))));
let one_king = Position::from_sfen("4k4/9/9/9/9/9/9/9/9 b - 1").expect("one king");
assert!(matches!(book.lookup_position(&one_king), Err(BookError::InvalidData(_))));
let entry = book.lookup_position(&position).expect("valid lookup").expect("entry");
assert_eq!(entry.packed_sfen(), position.to_packed_sfen());
assert_eq!(entry.moves()[0].mv(), Move::from_usi("7g7f").unwrap());
let _ = fs::remove_file(path);
}
}