diff --git a/README.rst b/README.rst index d4d51a128..3a52817cb 100644 --- a/README.rst +++ b/README.rst @@ -279,6 +279,31 @@ Features >>> tablebase.close() +* Probe chesstb endgame tablebases (WDL, DTZ, DTC, DTM, DTM50). + `Docs `__. + + .. code:: python + + >>> import chess.chesstb + + >>> tablebase = chess.chesstb.open_tablebase("data/chesstb") + + >>> # White mates in 19 half moves in this KQvK endgame, and the 50-move + >>> # rule does not get in the way of it. + >>> board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1") + >>> tablebase.probe_dtm(board) + 19 + >>> tablebase.probe_dtm50(board) + (2, 19) + + >>> # DTC prices the pawn pushes a win still owes before it converts: this + >>> # KPvK win spends four of them, the first 23 half moves away. + >>> board = chess.Board("8/8/8/k7/8/8/K4P2/8 w - - 0 1") + >>> tablebase.probe_dtc(board) + (2, 4, 23) + + >>> tablebase.close() + * Communicate with UCI/XBoard engines. Based on ``asyncio``. `Docs `__. diff --git a/chess/chesstb.py b/chess/chesstb.py new file mode 100644 index 000000000..75cd4bbae --- /dev/null +++ b/chess/chesstb.py @@ -0,0 +1,3925 @@ +"""Pure-Python prober for the *chesstb* endgame tablebase format +(WDL ``.lzw`` / DTZ ``.lzdtz`` / DTC ``.lzdtc`` / DTM ``.lzdtm`` +/ DTM50 ``.lzdtm50``). + +Upstream: https://github.com/noobpwnftw/chesstb + +This is a faithful re-implementation of the C++ probe library in +``src/probe`` and is validated bit-for-bit against ``tools/probe_fen``. + +Square numbering matches python-chess exactly (a1=0 .. h8=63, rank-major), +so positions are taken as :class:`chess.Board` instances directly. + +Design mirrors :mod:`chess.syzygy`: a :class:`Tablebase` opens a directory of +table files and answers WDL / DTZ / DTC / DTM / DTM50 queries. + +Thread safety +------------- +A :class:`Tablebase` is safe to share between threads and probe concurrently, +following :mod:`chess.syzygy`. Everything a probe touches is either immutable +after construction (parsed headers, index configs, the mapping itself) or +guarded below, and locks are held only across work that is *not* the probe +itself: + +* **Lazily opened tables.** Unlike :mod:`chess.syzygy`, which registers every + ``Table`` up front in :meth:`~chess.syzygy.Tablebase.add_directory` and so + never mutates its dicts while probing, this module resolves and opens a + table the first time some material is probed. Each of the open caches + therefore gets a lock, held across the resolve-and-open only + (double-checked, so a warm probe never acquires it). Because tables are + memory-mapped rather than read (see :meth:`_TableFile._open_source`), that + open is a syscall plus a small header parse, not a multi-GB read. The lock + is held across it, so another source has to keep that property. +* **Decoded blocks.** Decoding is guarded per ``(color, block)``, not per + table: two threads that need the same cold block decode it once, while + threads needing *different* cold blocks never queue behind each other. For + DTZ and the packs that buys real parallelism, since ``lzma`` releases the GIL + (measured at ~2x on four threads); for WDL, whose LZ4 decoder here is pure + Python and so holds the GIL, it only avoids the duplicate work. +* **Closing.** :meth:`Tablebase.close` unmaps files, which must not happen + under a running probe. A read-count/condition pair (again mirroring + :mod:`chess.syzygy`) makes it wait for in-flight probes first. + +Index computation, reads of already-decoded blocks, and the derivation walk +hold no lock at all. +""" +from __future__ import annotations + +import array +import collections +import lzma +import mmap +import os +import struct +import threading +from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Type, TypeVar + +import chess + +__all__ = ["Tablebase", "ProbeResult", "MissingTableError", "open_tablebase"] + +# --------------------------------------------------------------------------- +# Chess primitive constants, mirroring src/chess/chess.h. +# python-chess: WHITE=True, BLACK=False; piece types KING..PAWN = 6..1? No: +# chess.PAWN=1, KNIGHT=2, BISHOP=3, ROOK=4, QUEEN=5, KING=6. +# The C++ enum differs (KING=1..PAWN=6) but we never serialize C++ piece ints; +# we only need: square transforms, the legal-square set per piece, the material +# key, and class ordering. Those we encode against the C++ semantics below. +# --------------------------------------------------------------------------- + +WHITE = chess.WHITE +BLACK = chess.BLACK + +# Piece "type" codes as used by the C++ side (KING=1,QUEEN=2,ROOK=3,BISHOP=4, +# KNIGHT=5,PAWN=6) and Piece = (color<<3)+type, color WHITE=0 BLACK=1. +KING, QUEEN, ROOK, BISHOP, KNIGHT, PAWN = 1, 2, 3, 4, 5, 6 + +# map C++ piece-type code -> python-chess piece type +_CPP_TO_PC = {KING: chess.KING, QUEEN: chess.QUEEN, ROOK: chess.ROOK, + BISHOP: chess.BISHOP, KNIGHT: chess.KNIGHT, PAWN: chess.PAWN} +_PC_TO_CPP = {v: k for k, v in _CPP_TO_PC.items()} + +# C++ Color: WHITE=0, BLACK=1. python-chess: WHITE=True(1), BLACK=False(0). +# We use a dedicated 0/1 color int matching C++ where indexing matters. +CPP_WHITE, CPP_BLACK = 0, 1 + + +def cpp_color(piece_color: bool) -> int: + return CPP_WHITE if piece_color == WHITE else CPP_BLACK + + +# --- square transforms (src/chess/chess.h tables) --- + +def sq_file(sq: int) -> int: + return sq & 7 + + +def sq_rank(sq: int) -> int: + return sq >> 3 + + +def sq_make(rank: int, file: int) -> int: + return (rank << 3) + file + + +def sq_file_mirror(sq: int) -> int: + return sq_make(sq_rank(sq), 7 - sq_file(sq)) + + +def sq_rank_mirror(sq: int) -> int: + return sq_make(7 - sq_rank(sq), sq_file(sq)) + + +def sq_diag_mirror(sq: int) -> int: + # transpose along a1-h8: (file f, rank r) -> (file r, rank f) + return sq_make(sq_file(sq), sq_rank(sq)) + + +def apply_transform(sq: int, t: int) -> int: + """Symmetry_Transform: bit0=file flip, bit1=rank flip, bit2=diag swap.""" + f = sq_file(sq) + r = sq_rank(sq) + if t & 1: + f = 7 - f + if t & 2: + r = 7 - r + if t & 4: + f, r = r, f + return sq_make(r, f) + + +T_IDENTITY, T_FILE, T_RANK, T_FILE_RANK = 0, 1, 2, 3 +T_DIAG, T_FILE_DIAG, T_RANK_DIAG, T_FILE_RANK_DIAG = 4, 5, 6, 7 + +SYM_FILE_MIRROR = 0 +SYM_DIHEDRAL_8 = 1 + +# Anchor square sets for king canonicalization. +_ANCHOR_FILE_MIRROR = [sq_make(r, f) for r in range(8) for f in range(4)] # files a-d +_ANCHOR_TRIANGLE = [sq_make(r, f) for r in range(4) for f in range(r, 4)] # a1,b1..d1,b2.. + +# --------------------------------------------------------------------------- +# Binomial table C(n, k) for n<=64, k<=7. C++ BINOMIAL[k][n] indexing. +# --------------------------------------------------------------------------- +_BINOM = [[0] * 8 for _ in range(65)] +for _k in range(65): + _BINOM[_k][0] = 1 + for _n in range(1, 8): + _BINOM[_k][_n] = 0 if _n > _k else _BINOM[_k - 1][_n - 1] + _BINOM[_k - 1][_n] + + +def binom(n: int, k: int) -> int: + if k < 0 or k > 7 or n < 0 or n > 64: + return 0 + return _BINOM[n][k] + + +# --------------------------------------------------------------------------- +# Material_Key (src/chess/chess.h): base-9 mixed radix, indexed by C++ Piece. +# WHITE_QUEEN..WHITE_PAWN -> 9^4..9^0 ; BLACK_QUEEN..BLACK_PAWN -> 9^9..9^5. +# --------------------------------------------------------------------------- +_MAT_WEIGHT = { # (cpp_color, type) -> weight + (CPP_WHITE, QUEEN): 9 ** 4, (CPP_WHITE, ROOK): 9 ** 3, (CPP_WHITE, BISHOP): 9 ** 2, + (CPP_WHITE, KNIGHT): 9 ** 1, (CPP_WHITE, PAWN): 9 ** 0, (CPP_WHITE, KING): 0, + (CPP_BLACK, QUEEN): 9 ** 9, (CPP_BLACK, ROOK): 9 ** 8, (CPP_BLACK, BISHOP): 9 ** 7, + (CPP_BLACK, KNIGHT): 9 ** 6, (CPP_BLACK, PAWN): 9 ** 5, (CPP_BLACK, KING): 0, +} + + +def material_key_of(pieces: List[Tuple[int, int]]) -> int: + """pieces: list of (cpp_color, type).""" + return sum(_MAT_WEIGHT[(c, t)] for c, t in pieces) + + +# --------------------------------------------------------------------------- +# Piece_Config: canonical (strength-ordered) piece list, white = stronger side. +# --------------------------------------------------------------------------- +_STRENGTH = {QUEEN: 900, ROOK: 500, BISHOP: 330, KNIGHT: 320, PAWN: 100, KING: 0} +# within-side sort order: K,Q,R,B,N,P (descending strength, kings first) +_TYPE_ORDER = {KING: 0, QUEEN: 1, ROOK: 2, BISHOP: 3, KNIGHT: 4, PAWN: 5} + + +def _composition_key(pieces: List[Tuple[int, int]], color: int) -> List[int]: + """Per-type piece counts for `color`, indexed by C++ type code so that + lexicographic comparison orders by KING, QUEEN, ROOK, BISHOP, KNIGHT, PAWN. + Mirrors the ``std::array`` tiebreak key in + ``Piece_Config::sort_pieces`` (src/chess/piece_config.cpp).""" + counts = [0] * (PAWN + 1) + for c, t in pieces: + if c == color: + counts[t] += 1 + return counts + + +class PieceConfig: + """Canonical material config. `pieces` is a list of (cpp_color, type) of the + FREE pieces. `has_pair` marks an opposing pawn pair (lowercase 'p'): + one white + one black pawn locked on a file, indexed jointly and excluded + from `pieces` (mirrors src/chess/piece_config.h).""" + + def __init__(self, pieces: List[Tuple[int, int]], has_pair: bool = False): + # Determine side ordering by total strength; stronger side -> WHITE. + # Input may be in any orientation; we canonicalize. The pair is + # strength-neutral (one pawn each side), so it never affects the swap. + ws = sum(_STRENGTH[t] for c, t in pieces if c == CPP_WHITE) + bs = sum(_STRENGTH[t] for c, t in pieces if c == CPP_BLACK) + swap = bs > ws + if bs == ws: + swap = _composition_key(pieces, CPP_BLACK) > _composition_key(pieces, CPP_WHITE) + if swap: + pieces = [(CPP_BLACK if c == CPP_WHITE else CPP_WHITE, t) for c, t in pieces] + # sort: white side first then black; within side by type order. + pieces = sorted(pieces, key=lambda ct: (ct[0], _TYPE_ORDER[ct[1]])) + self.pieces = pieces + self.has_pair = has_pair + self.base_key = material_key_of(pieces) + self.mirr_key = material_key_of( + [(CPP_BLACK if c == CPP_WHITE else CPP_WHITE, t) for c, t in pieces]) + + @property + def min_key(self) -> int: + # Pair-stripped 30-bit value, as serialized in the on-disk header. + return min(self.base_key, self.mirr_key) + + @property + def cache_key(self) -> Tuple[int, bool]: + # In-memory key distinguishing a 'p'-material from its free-pieces twin, + # which share min_key. + return (self.min_key, self.has_pair) + + def name(self) -> str: + # The pair contributes one pawn to each side; emit a 'p' at the end of + # the white side (before the 2nd king) and another at the very end, e.g. + # "KQpKp" -- so a board-derived pair config resolves the right filename. + letters = {KING: "K", QUEEN: "Q", ROOK: "R", BISHOP: "B", KNIGHT: "N", PAWN: "P"} + s = [] + seen_white_king = False + for c, t in self.pieces: + if t == KING: + if seen_white_king and self.has_pair: + s.append("p") + seen_white_king = True + s.append(letters[t]) + if self.has_pair: + s.append("p") + return "".join(s) + + @property + def num_pieces(self) -> int: + return len(self.pieces) + + @property + def is_bare_kings(self) -> bool: + """Nothing on the board but the two kings: a trivial draw, so no table + exists and no probe consults one. `num_pieces` counts only FREE pieces, + so an opposing pair sits at 2 while physically holding four (KpKp) -- a + real material with a real table. Every "is it just kings?" test goes + through here rather than comparing `num_pieces` directly + (Piece_Config::is_bare_kings in src/chess/piece_config.h).""" + return self.num_pieces <= 2 and not self.has_pair + + +def piece_config_from_board(board: chess.Board) -> Tuple["PieceConfig", bool]: + """Return (canonical PieceConfig, mirrored?) for the board's material. + + mirrored is True when the literal material had to swap colors to match the + canonical base orientation (white = stronger side). + """ + pieces = [] + for sq in range(64): + p = board.piece_at(sq) + if p: + pieces.append((cpp_color(p.color), _PC_TO_CPP[p.piece_type])) + cfg = PieceConfig(pieces) + literal = material_key_of(pieces) + mirrored = literal != cfg.base_key + return cfg, mirrored + + +def pair_config_from_board(board: chess.Board) -> Optional[Tuple["PieceConfig", bool]]: + """If `board` has an opposing pawn pair, return (its 'p' PieceConfig, + mirrored?); else None. The pair's two pawns are excluded from the config and + flagged via has_pair, mirroring src/probe/probe.cpp pair_config_from_position. + Used to prefer a 'p' table over the full material when one is on disk.""" + wp = list(board.pieces(chess.PAWN, WHITE)) + bp = list(board.pieces(chess.PAWN, BLACK)) + found = PairGroup.find_canonical(wp, bp) + if found is None: + return None + pw, pb = found + pieces = [] + for sq in range(64): + if sq == pw or sq == pb: + continue + p = board.piece_at(sq) + if p: + pieces.append((cpp_color(p.color), _PC_TO_CPP[p.piece_type])) + cfg = PieceConfig(pieces, has_pair=True) + mirrored = material_key_of(pieces) != cfg.base_key + return cfg, mirrored + + +# --------------------------------------------------------------------------- +# Piece_Class enum (src/chess/piece_config.h) +# --------------------------------------------------------------------------- +(BLACK_KINGS, BLACK_KNIGHTS, BLACK_BISHOPS, BLACK_ROOKS, BLACK_QUEENS, BLACK_PAWNS, + WHITE_KINGS, WHITE_KNIGHTS, WHITE_BISHOPS, WHITE_ROOKS, WHITE_QUEENS, WHITE_PAWNS) = range(12) +PIECE_CLASS_NB = 12 + +# Piece_Type_Class: KINGS=0,KNIGHTS=1,BISHOPS=2,ROOKS=3,QUEENS=4,PAWNS=5 +_PTCLASS = {KING: 0, KNIGHT: 1, BISHOP: 2, ROOK: 3, QUEEN: 4, PAWN: 5} + + +def make_piece_class(cpp_col: int, ptype: int) -> int: + base = WHITE_KINGS if cpp_col == CPP_WHITE else BLACK_KINGS + return base + _PTCLASS[ptype] + + +def class_to_piece(pcl: int) -> Tuple[int, int]: + """piece class -> (cpp_color, ptype).""" + cpp_col = CPP_WHITE if pcl >= WHITE_KINGS else CPP_BLACK + off = pcl - (WHITE_KINGS if cpp_col == CPP_WHITE else BLACK_KINGS) + inv = {0: KING, 1: KNIGHT, 2: BISHOP, 3: ROOK, 4: QUEEN, 5: PAWN} + return cpp_col, inv[off] + + +# --------------------------------------------------------------------------- +# Piece_Group: combinatorial ranking of `count` identical pieces over the +# legal squares of their type. +# --------------------------------------------------------------------------- +class PieceGroup: + def __init__(self, ptype: int, count: int): + self.count = count + if ptype == PAWN: + legal = list(range(chess.A2, chess.H7 + 1)) # 8..55 + else: + legal = list(range(64)) + legal.sort() + self.pos_to_sq = legal + self.sq_to_pos = {sq: i for i, sq in enumerate(legal)} + self.num_legal = len(legal) + self.table_size = binom(self.num_legal, count) + + def compound_index(self, squares: List[int]) -> int: + # squares: the placement (any order); rank a sorted combination. + sqs = sorted(squares) + rank = 0 + for i, sq in enumerate(sqs): + p = self.sq_to_pos[sq] + rank += binom(p, i + 1) + return rank + + def squares(self, idx: int) -> List[int]: + pos = [0] * self.count + rank = idx + hi = self.num_legal + for k in range(self.count, 0, -1): + p = hi - 1 + while binom(p, k) > rank: + p -= 1 + pos[k - 1] = p + rank -= binom(p, k) + hi = p + return [self.pos_to_sq[p] for p in pos] + + +# --------------------------------------------------------------------------- +# King_Slice_Manager: built once per symmetry group. +# --------------------------------------------------------------------------- +def _king_attacks(sq: int) -> List[int]: + f, r = sq_file(sq), sq_rank(sq) + out = [] + for df in (-1, 0, 1): + for dr in (-1, 0, 1): + if df == 0 and dr == 0: + continue + nf, nr = f + df, r + dr + if 0 <= nf < 8 and 0 <= nr < 8: + out.append(sq_make(nr, nf)) + return out + + +def _kings_adjacent(a: int, b: int) -> bool: + return b in _king_attacks(a) + + +def _sq_on_main_diag(sq: int) -> bool: + return sq_file(sq) == sq_rank(sq) + + +class KingSliceManager: + def __init__(self, sym: int): + self.sym = sym + n_trans = 8 if sym == SYM_DIHEDRAL_8 else 2 + anchors = _ANCHOR_TRIANGLE if sym == SYM_DIHEDRAL_8 else _ANCHOR_FILE_MIRROR + anchor_set = set(anchors) + SLICE_NONE = -1 + # pair_lookup[wk*64+bk] = [slice_id, transform, has_diag_stabilizer] + self.pair: List[List[int]] = [[SLICE_NONE, T_IDENTITY, 0] for _ in range(64 * 64)] + self.kings_of_slice: List[Tuple[int, int]] = [] + + for wk in range(64): + if wk not in anchor_set: + continue + for bk in range(64): + if bk == wk or _kings_adjacent(wk, bk): + continue + if sym == SYM_DIHEDRAL_8 and _sq_on_main_diag(wk): + bk_d = sq_diag_mirror(bk) + if bk_d != bk and bk > bk_d: + continue + sid = len(self.kings_of_slice) + self.kings_of_slice.append((wk, bk)) + stab = 1 if (sym == SYM_DIHEDRAL_8 and _sq_on_main_diag(wk) + and _sq_on_main_diag(bk)) else 0 + self.pair[wk * 64 + bk] = [sid, T_IDENTITY, stab] + self.num_slices = len(self.kings_of_slice) + + for wk in range(64): + for bk in range(64): + e = self.pair[wk * 64 + bk] + if e[0] != SLICE_NONE: + continue + if wk == bk or _kings_adjacent(wk, bk): + continue + for t in range(n_trans): + wk_t = apply_transform(wk, t) + bk_t = apply_transform(bk, t) + look = self.pair[wk_t * 64 + bk_t] + if look[0] != -1 and look[1] == T_IDENTITY: + e[0] = look[0] + e[1] = t + e[2] = look[2] + break + + def lookup(self, wk: int, bk: int) -> List[int]: + return self.pair[wk * 64 + bk] + + +_KSM_CACHE: Dict[int, KingSliceManager] = {} +# Building a manager walks 64x64 king placements over up to 8 transforms, so +# two threads first probing different materials of the same symmetry must not +# both do it. Double-checked: the cached case never acquires. Publishing the +# entry only after __init__ returns is what keeps a half-built manager (its +# `pair` table is filled in place) from ever being visible to a reader. +_KSM_LOCK = threading.Lock() + + +def king_slice_mgr(sym: int) -> KingSliceManager: + ksm = _KSM_CACHE.get(sym) + if ksm is not None: + return ksm + with _KSM_LOCK: + ksm = _KSM_CACHE.get(sym) + if ksm is None: + ksm = KingSliceManager(sym) + _KSM_CACHE[sym] = ksm + return ksm + + +# --------------------------------------------------------------------------- +# Pawn_Slice_Manager +# --------------------------------------------------------------------------- +class PairGroup: + """Opposing pawn pair (lowercase 'p'): white pawn on rank r, black on rank s, + r < s, same file. Enumeration / index_of / find_canonical must match + src/egtb/pair_group.h exactly -- the on-disk pawn-slice ids depend on them. + White ranks 2..6, black 3..7: C(6,2)=15 rank pairs x 8 files = 120.""" + + def __init__(self) -> None: + self.white: List[int] = [] + self.black: List[int] = [] + self._inv: Dict[Tuple[int, int], int] = {} + for f in range(8): + for wr in range(1, 6): # ranks 2..6 + for br in range(wr + 1, 7): # ranks 3..7 + w = sq_make(wr, f) + b = sq_make(br, f) + self._inv[(w, b)] = len(self.white) + self.white.append(w) + self.black.append(b) + + @property + def table_size(self) -> int: + return len(self.white) + + def white_square(self, i: int) -> int: + return self.white[i] + + def black_square(self, i: int) -> int: + return self.black[i] + + def index_of(self, w: int, b: int) -> int: + return self._inv[(w, b)] + + @staticmethod + def is_opposing_pair(w: int, b: int) -> bool: + return (sq_file(w) == sq_file(b) + and sq_rank(w) >= 1 and sq_rank(b) <= 6 + and sq_rank(w) < sq_rank(b)) + + @staticmethod + def find_canonical(white_sqs: List[int], black_sqs: List[int] + ) -> Optional[Tuple[int, int]]: + """The opposing pair minimal by (file, white_rank, black_rank), or None. + Both the generator's prune and this lookup use this one rule.""" + best: Optional[Tuple[Tuple[int, int, int], int, int]] = None + for w in white_sqs: + for b in black_sqs: + if not PairGroup.is_opposing_pair(w, b): + continue + key = (sq_file(w), sq_rank(w), sq_rank(b)) + if best is None or key < best[0]: + best = (key, w, b) + return None if best is None else (best[1], best[2]) + + @staticmethod + def canonical_pair(white_sqs: List[int], black_sqs: List[int]) -> Tuple[int, int]: + """For callers that know an opposing pair is present (indexing a stored + pair-table position): the canonical pair, asserting one exists.""" + found = PairGroup.find_canonical(white_sqs, black_sqs) + assert found is not None + return found + + +class PawnSliceManager: + def __init__(self, pair_group: Optional[PairGroup], + white_group: Optional[PieceGroup], black_group: Optional[PieceGroup]): + self.pair_group = pair_group + self.white_group = white_group + self.black_group = black_group + # All three None is the pawnless configuration, which falls through the + # same enumeration: the cartesian space is the single empty cell, so it + # yields exactly one slice with id 0. + self.has_pawns = (pair_group is not None + or white_group is not None or black_group is not None) + self.pair_table_size = pair_group.table_size if pair_group else 1 + self.white_table_size = white_group.table_size if white_group else 1 + self.black_table_size = black_group.table_size if black_group else 1 + # Mixed radix: cart = pair_idx + w_idx*pair_size + b_idx*pair_size*white_size. + n_cart = self.pair_table_size * self.white_table_size * self.black_table_size + + def occupancies(g: Optional[PieceGroup]) -> Tuple[List[List[int]], List[int]]: + if g is None: + return [[]], [0] + sqs = [g.squares(i) for i in range(g.table_size)] + masks = [] + for pl in sqs: + m = 0 + for s in pl: + m |= 1 << s + masks.append(m) + return sqs, masks + + white_sqs, white_masks = occupancies(white_group) + black_sqs, black_masks = occupancies(black_group) + if pair_group: + pair_masks = [(1 << pair_group.white_square(p)) | (1 << pair_group.black_square(p)) + for p in range(self.pair_table_size)] + else: + pair_masks = [0] + + self._survivor_bits = bytearray((n_cart + 7) // 8) + rank_before = array.array("q", bytes(8 * ((n_cart + 63) // 64))) + + # Storage ids are ranks in ascending cartesian order, so the walk has to + # follow that order: the pair digit is least significant, then white, + # then black. + num_slices = 0 + cart = 0 + for b_idx in range(self.black_table_size): + b_mask = black_masks[b_idx] + for w_idx in range(self.white_table_size): + # Free pawns of the two colors collide: no pair placement rescues + # the cell, so skip its whole run of pair cells at once. + if white_masks[w_idx] & b_mask: + cart += self.pair_table_size + continue + free_mask = white_masks[w_idx] | b_mask + for pair_idx in range(self.pair_table_size): + if pair_masks[pair_idx] & free_mask: + cart += 1 + continue + if pair_group: + pair_w = pair_group.white_square(pair_idx) + pair_b = pair_group.black_square(pair_idx) + cw, cb = PairGroup.canonical_pair( + [pair_w] + white_sqs[w_idx], [pair_b] + black_sqs[b_idx]) + if cw != pair_w or cb != pair_b: + cart += 1 + continue + self._survivor_bits[cart >> 3] |= 1 << (cart & 7) + num_slices += 1 + cart += 1 + assert cart == n_cart + self.num_slices = num_slices + + # Prefix the per-block survivor counts, so a storage id needs only the + # one block holding its cell. + running = 0 + for blk in range(len(rank_before)): + rank_before[blk] = running + running += chess.popcount(int.from_bytes( + self._survivor_bits[blk * 8:blk * 8 + 8], "little")) + assert running == num_slices + self._rank_before_block = rank_before + + def compose(self, pair_idx: int, w_idx: int, b_idx: int) -> int: + if not self.has_pawns: + return 0 + cart = (pair_idx + + w_idx * self.pair_table_size + + b_idx * self.pair_table_size * self.white_table_size) + blk, off = divmod(cart, 64) + word = int.from_bytes(self._survivor_bits[blk * 8:blk * 8 + 8], "little") + assert word & (1 << off) + return self._rank_before_block[blk] + chess.popcount(word & ((1 << off) - 1)) + + def lookup_from_squares(self, pair_w: int, pair_b: int, + white_pawn_sqs: List[int], black_pawn_sqs: List[int]) -> int: + if not self.has_pawns: + return 0 + pair_idx = self.pair_group.index_of(pair_w, pair_b) if self.pair_group else 0 + w_idx = self.white_group.compound_index(white_pawn_sqs) if self.white_group else 0 + b_idx = self.black_group.compound_index(black_pawn_sqs) if self.black_group else 0 + return self.compose(pair_idx, w_idx, b_idx) + + +# --------------------------------------------------------------------------- +# Index permutation (src/chess/index_permutation.h) +# --------------------------------------------------------------------------- +_FACT = [1, 1, 2, 6, 24, 120, 720, 5040, 40320] + + +def index_permutation_valid(n_classes: int, perm: int) -> bool: + return n_classes <= 8 and perm < _FACT[n_classes] + + +def storage_within_class_order(populated: List[int], perm: int) -> List[int]: + n = len(populated) + available = list(populated) + order = [] + idx = perm + for i in range(n): + f = _FACT[n - 1 - i] + pick = idx // f + idx %= f + order.append(available[pick]) + del available[pick] + return order + + +# --------------------------------------------------------------------------- +# Position_Index_Config +# --------------------------------------------------------------------------- +class PositionIndexConfig: + def __init__(self, cfg: PieceConfig): + self.cfg = cfg + counts: Dict[Tuple[int, int], int] = {} + for c, t in cfg.pieces: + counts[(c, t)] = counts.get((c, t), 0) + 1 + # An opposing pair is two pawns on the board, so it breaks symmetry to + # file-mirror like any pawn even with no free pawns. + has_pawns = (counts.get((CPP_WHITE, PAWN), 0) > 0 + or counts.get((CPP_BLACK, PAWN), 0) > 0 or cfg.has_pair) + self.sym = SYM_FILE_MIRROR if has_pawns else SYM_DIHEDRAL_8 + self.ksm = king_slice_mgr(self.sym) + + self.groups: Dict[int, PieceGroup] = {} + # make groups for Q,R,B,N,P per color (NOT kings) + for c in (CPP_WHITE, CPP_BLACK): + for t in (QUEEN, ROOK, BISHOP, KNIGHT, PAWN): + n = counts.get((c, t), 0) + if n == 0: + continue + pcl = make_piece_class(c, t) + self.groups[pcl] = PieceGroup(t, n) + + self.pair_group = PairGroup() if cfg.has_pair else None + self.psm = PawnSliceManager(self.pair_group, + self.groups.get(WHITE_PAWNS), self.groups.get(BLACK_PAWNS)) + self.num_pawn_slices = self.psm.num_slices + + # populated non-pawn classes in ascending class-id order; native weights. + self.populated: List[int] = [] + self.weights: Dict[int, int] = {} + w = 1 + for i in range(PIECE_CLASS_NB): + if i not in self.groups: + continue + if i == WHITE_PAWNS or i == BLACK_PAWNS: + continue + self.populated.append(i) + self.weights[i] = w + w *= self.groups[i].table_size + self.within_slice_size = w + self.num_king_slices = self.ksm.num_slices + self.pawn_slice_stride = self.num_king_slices * self.within_slice_size + self.num_positions = self.num_pawn_slices * self.pawn_slice_stride + + def num_populated_classes(self) -> int: + return len(self.populated) + + def make_layout(self, perm: int) -> Tuple[List[int], List[int]]: + """Return (order, radix) lists per index permutation.""" + order = storage_within_class_order(self.populated, perm) + radix = [self.groups[c].table_size for c in order] + return order, radix + + # --- canonicalization + indexing --- + def _placements_from_board(self, board: chess.Board) -> Dict[int, List[int]]: + pl: Dict[int, List[int]] = {c: [] for c in range(PIECE_CLASS_NB)} + wk, bk = board.king(WHITE), board.king(BLACK) + assert wk is not None and bk is not None # tablebase positions have both kings + pl[WHITE_KINGS] = [wk] + pl[BLACK_KINGS] = [bk] + for c in self.populated: + cc, tt = class_to_piece(c) + color = WHITE if cc == CPP_WHITE else BLACK + pl[c] = list(board.pieces(_CPP_TO_PC[tt], color)) + # Collect pawns when a free-pawn class is populated OR an opposing pair adds + # pawns with no free-pawn class (e.g. KpKp). The pair members are folded + # in here and split out by find_canonical at index time. + for c in (WHITE_PAWNS, BLACK_PAWNS): + if c in self.groups or self.pair_group is not None: + color = WHITE if c == WHITE_PAWNS else BLACK + pl[c] = list(board.pieces(chess.PAWN, color)) + return pl + + def _canonicalize(self, pl: Dict[int, List[int]]) -> bool: + wk = pl[WHITE_KINGS][0] + bk = pl[BLACK_KINGS][0] + look = self.ksm.lookup(wk, bk) + if look[0] == -1: + return False + t = look[1] + if t != T_IDENTITY: + for c in range(PIECE_CLASS_NB): + if pl[c]: + pl[c] = [apply_transform(s, t) for s in pl[c]] + if look[2]: # diagonal stabilizer tie-break (non-pawn populated only) + cur = alt = 0 + for c in self.populated: + g = self.groups[c] + cur += self.weights[c] * g.compound_index(pl[c]) + alt += self.weights[c] * g.compound_index([sq_diag_mirror(s) for s in pl[c]]) + if alt < cur: + for c in self.populated: + pl[c] = [sq_diag_mirror(s) for s in pl[c]] + return True + + def board_index(self, board: chess.Board, order: List[int], radix: List[int]) -> Optional[int]: + pl = self._placements_from_board(board) + if not self._canonicalize(pl): + return None + wk = pl[WHITE_KINGS][0] + bk = pl[BLACK_KINGS][0] + ksid = self.ksm.lookup(wk, bk)[0] + if ksid == -1: + return None + pawn_slice = 0 + if self.psm.has_pawns: + w_pl, b_pl = pl[WHITE_PAWNS], pl[BLACK_PAWNS] + if self.pair_group is not None: + # Identify the pair (canonical opposing pair); the rest are free. + pw, pb = PairGroup.canonical_pair(w_pl, b_pl) + free_w = [s for s in w_pl if s != pw] + free_b = [s for s in b_pl if s != pb] + else: + pw = pb = -1 + free_w, free_b = w_pl, b_pl + pawn_slice = self.psm.lookup_from_squares(pw, pb, free_w, free_b) + within_idx = {c: self.groups[c].compound_index(pl[c]) for c in self.populated} + within = 0 + w = 1 + for i in range(len(order)): + within += w * within_idx[order[i]] + w *= radix[i] + outer = pawn_slice * self.pawn_slice_stride + ksid * self.within_slice_size + return outer + within + + +_INDEX_CFG_CACHE: Dict[Tuple[int, bool], PositionIndexConfig] = {} +# As for _KSM_LOCK: enumerating pawn slices is expensive enough to be worth +# doing once. A config is immutable once built, so probes read it unlocked. +# Lock order is _INDEX_CFG_LOCK -> _KSM_LOCK (via PositionIndexConfig, which +# calls king_slice_mgr); nothing acquires them the other way round. +_INDEX_CFG_LOCK = threading.Lock() + + +def position_index_config(cfg: PieceConfig) -> PositionIndexConfig: + k = cfg.cache_key + icfg = _INDEX_CFG_CACHE.get(k) + if icfg is not None: + return icfg + with _INDEX_CFG_LOCK: + icfg = _INDEX_CFG_CACHE.get(k) + if icfg is None: + icfg = PositionIndexConfig(cfg) + _INDEX_CFG_CACHE[k] = icfg + return icfg + + +# =========================================================================== +# On-disk file framing (src/util/memory.h, mono_uint_vec.h, egtb_format.h) +# =========================================================================== + +WDL_MAGIC, DTZ_MAGIC, DTC_MAGIC, DTM_MAGIC, DTM50_MAGIC = ( + 0x9bd1e3a6, 0x2ec8b161, 0x7c1de4a3, 0xab57c134, 0xab57c151) +SINGULAR_FLAG = 0x80 +DROPPED_FLAG = 0x40 +LOSS_ONLY_FLAG = 0x20 +RELAXED_FLAG = 0x10 +MAX_NON_CURSED_DTZ = 100 +DTM50_HMC_COUNT = 100 +DTM50_PACK_LAYERS = 101 +# An 8-man winner has at most six pawns, each with at most five non-converting +# pushes, so its budget curve has at most 30 relevant changepoints. For clean +# W/L the terminal one is exactly DTZ: solve the preceding 29 and embed that +# endpoint at row 0. Cursed/blessed cells use row 0 only as DTZ. +DTC_BUDGET_LAYERS = 29 +DTC_PACK_LAYERS = 30 +IGNORE_50MR = -1 # sentinel (C++ uses ~0u) + +# WDL_Entry +LOSE, BLESSED_LOSS, DRAW, CURSED_WIN, WIN, ILLEGAL = 0, 1, 2, 3, 4, 7 +# WDL_Stored adds BOUNDARY_LOSS=5, BOUNDARY_WIN=6. + + +def wdl_from_storage(s: int) -> int: + if s == 6: # BOUNDARY_WIN + return WIN + if s == 5: # BOUNDARY_LOSS + return LOSE + return s + + +class _Serial: + """Sequential little-endian reader over a table's buffer, mirroring + Serial_Memory_Reader (offset tracking + align relative to begin). + + The buffer is whatever :meth:`_TableFile._open_source` returned. Multi-byte + reads materialize the few bytes they need rather than ``unpack_from`` on + the buffer, which would demand the buffer protocol; opening is cold. + + Nothing derived from the buffer outlives parsing -- a table keeps the + buffer plus offsets into it -- so :meth:`_TableFile.close` has one export + to release, a live memoryview making ``mmap.close()`` raise.""" + + def __init__(self, data: Any): + self.d = data + self.pos = 0 + + def u8(self) -> int: + v = int(self.d[self.pos]) + self.pos += 1 + return v + + def _le(self, size: int) -> int: + v = int.from_bytes(bytes(self.d[self.pos:self.pos + size]), "little") + self.pos += size + return v + + def u16(self) -> int: + return self._le(2) + + def u32(self) -> int: + return self._le(4) + + def u64(self) -> int: + return self._le(8) + + def advance(self, n: int) -> None: + self.pos += n + + def caret(self) -> int: + return self.pos + + def align(self, alignment: int) -> None: + mis = self.pos % alignment + if mis: + self.pos += alignment - mis + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +#: Precompiled struct for the 64-bit bit-window reads below, mirroring the +#: module-level UINT64_BE / UINT32 / UINT16 Structs in :mod:`chess.syzygy`. +_U64LE = struct.Struct(" int: + """The 8 bytes at `offset`, little-endian, read straight out of the + mapping. ``int.from_bytes(buf[offset:offset + 8], "little")`` would build + a throwaway memoryview slice per call, and these are the hottest reads in + a probe. Fewer than 8 bytes remain near the end of a mapping, which + ``unpack_from`` rejects and the slicing form tolerates. + + A buffer that is not a ``memoryview`` (:meth:`_TableFile._open_source`) + cannot go to ``unpack_from`` at all, so it takes the slicing form.""" + if isinstance(buf, memoryview): + try: + return _U64LE.unpack_from(buf, offset)[0] # type: ignore[no-any-return] + except struct.error: + return int.from_bytes(buf[offset:offset + 8], "little") + return int.from_bytes(bytes(buf[offset:offset + 8]), "little") + + +def _compressed_block(chunk: Any) -> Any: + """One compressed block as the decoders want it: slicing a memoryview + already gives that, any other buffer (:meth:`_TableFile._open_source`) + materializes here -- a block being the widest span a probe asks for.""" + if isinstance(chunk, (bytes, bytearray, memoryview)): + return chunk + return bytes(chunk) + + +class MonoUintVec: + """Block-sampled delta coder for a monotone uint64 sequence. + + Addressed as `base` plus a bit offset into the table's buffer, not as a + slice of it: no span wider than one read is taken + (:meth:`_TableFile._open_source`).""" + + def __init__(self, blob: Any, base: int, num_values: int, log2_bu: int, + sample_width: int, offset_width: int): + self.blob = blob + self.base = base + self.num_values = num_values + self.log2_bu = log2_bu + self.sample_width = sample_width + self.offset_width = offset_width + num_samples = _ceil_div(num_values, 1 << log2_bu) + self.delta_off = _ceil_div(num_samples * sample_width, 8) + + @staticmethod + def on_disk_bytes(num_values: int, log2_bu: int, sample_width: int, offset_width: int) -> int: + num_samples = _ceil_div(num_values, 1 << log2_bu) + return (_ceil_div(num_samples * sample_width, 8) + + _ceil_div(num_values * offset_width, 8)) + + def _read_bits(self, base_off: int, bitpos: int, width: int) -> int: + if width == 0: + return 0 + byte = self.base + base_off + (bitpos >> 3) + bit = bitpos & 7 + # read up to 16 bytes to cover width<=64 with bit offset + lo = _read_u64le(self.blob, byte) + v = lo >> bit + if bit + width > 64: + hi = _read_u64le(self.blob, byte + 8) + v |= hi << (64 - bit) + mask = (1 << width) - 1 if width < 64 else (1 << 64) - 1 + return v & mask + + def get(self, i: int) -> int: + sb = i >> self.log2_bu + base = self._read_bits(0, sb * self.sample_width, self.sample_width) + delta = self._read_bits(self.delta_off, i * self.offset_width, self.offset_width) + return base + delta + + def get2(self, i: int) -> Tuple[int, int]: + return (self.get(i), self.get(i + 1)) + + +class Min0UintVec: + """As :class:`MonoUintVec`, `base` plus a bit offset rather than a slice.""" + + def __init__(self, data: Any, base: int, size: int, width: int): + self.data = data + self.base = base + self.size = size + self.width = width + + @staticmethod + def on_disk_bytes(size: int, width: int) -> int: + return _ceil_div(size * width, 8) + + def get(self, i: int) -> int: + if self.width == 0: + return 0 + bitpos = i * self.width + byte = self.base + (bitpos >> 3) + bit = bitpos & 7 + lo = _read_u64le(self.data, byte) + v = lo >> bit + if bit + self.width > 64: + hi = _read_u64le(self.data, byte + 8) + v |= hi << (64 - bit) + mask = (1 << self.width) - 1 + return v & mask + + +# --- pure-Python LZ4 block decompression with optional dictionary prefix --- + +def lz4_decompress_block(src: memoryview, expected_size: int, dict_bytes: bytes = b"") -> bytes: + """Decompress an LZ4 *block* (not frame). A dictionary, if given, logically + precedes the output; match offsets may reference into it. Mirrors + LZ4_decompress_safe[_usingDict].""" + out = bytearray(dict_bytes) + base = len(dict_bytes) + si = 0 + n = len(src) + while si < n: + token = src[si] + si += 1 + lit_len = token >> 4 + if lit_len == 15: + while True: + b = src[si] + si += 1 + lit_len += b + if b != 255: + break + out += src[si:si + lit_len] + si += lit_len + if si >= n: + break + offset = src[si] | (src[si + 1] << 8) + si += 2 + match_len = (token & 0xF) + 4 + if (token & 0xF) == 15: + while True: + b = src[si] + si += 1 + match_len += b + if b != 255: + break + start = len(out) - offset + # overlapping copy, byte by byte + for j in range(match_len): + out.append(out[start + j]) + result = bytes(out[base:]) + if len(result) != expected_size: + raise ValueError(f"LZ4 size mismatch: got {len(result)} expected {expected_size}") + return result + + +# =========================================================================== +# Decoded-block cache +# =========================================================================== + +#: Default soft budget (bytes) for decoded blocks held resident across all +#: tables of a :class:`Tablebase`. The cache evicts least-recently-used blocks +#: once the budget is exceeded, so memory is reclaimed automatically without an +#: explicit :meth:`Tablebase.close`. +DEFAULT_BLOCK_CACHE_BYTES = 64 * 1024 * 1024 + + +class _PerColor: + """State for one color's frame of one table: the decoded-block dict plus + the locks guarding decodes into it. + + Everything else a per-color object holds (the parsed header, the buffer + and the offsets addressing it) is written once while the table is + being constructed -- before any other thread can reach the table, since + :class:`Tablebase` only publishes it into an open cache once fully built -- + and read-only afterwards. ``_blocks`` is the exception: probes fill it as + they go, so it is the only part that needs guarding. + """ + + __slots__ = ("_blocks", "_block_locks", "_meta_lock") + + _blocks: Dict[int, Any] + + def __init__(self) -> None: + self._blocks = {} + # One lock per block id, created on demand, rather than one lock for + # the whole color: the slow path guards the decode-and-insert of a + # *single* block, so two threads wanting different cold blocks have no + # reason to queue behind each other. lzma releases the GIL, so for + # DTZ/DTM50 those two decodes really do proceed in parallel. + self._block_locks: Dict[int, threading.Lock] = {} + # Guards creation of entries in `_block_locks` -- a dict lookup and + # maybe an insert. Never held across a decode. + self._meta_lock = threading.Lock() + + def lock_for(self, block_id: int) -> threading.Lock: + """The lock guarding `block_id`'s decode, created on first use. + Double-checked, so the usual case never touches ``_meta_lock``.""" + lk = self._block_locks.get(block_id) + if lk is not None: + return lk + with self._meta_lock: + lk = self._block_locks.get(block_id) + if lk is None: + lk = threading.Lock() + self._block_locks[block_id] = lk + return lk + + +class _BlockCache: + """LRU reclaimer shared by every table of a :class:`Tablebase`. + + Decoding a block is expensive, so each per-color object keeps its own + ``_blocks`` dict of decoded blocks. This cache tracks those entries in a + global least-recently-used order keyed by ``(per_color, block_id)`` and, + once the resident byte estimate exceeds ``max_bytes``, evicts the oldest + blocks by dropping them from their owning ``_blocks`` dict. Sizes are + approximate; the budget is a soft target. + """ + + def __init__(self, max_bytes: int) -> None: + self.max_bytes = max_bytes + self.cur_bytes = 0 + # (per_color, block_id) -> approximate size in bytes, ordered LRU-first. + self._lru: "collections.OrderedDict[Tuple[Any, int], int]" = collections.OrderedDict() + self._lock = threading.Lock() + + def touch(self, pc: Any, block_id: int) -> None: + """Mark an already-cached block as most-recently-used (a cache hit).""" + key = (pc, block_id) + with self._lock: + if key in self._lru: + self._lru.move_to_end(key) + + def record(self, pc: _PerColor, block_id: int, size: int) -> None: + """Register a freshly decoded block and evict until within budget. + + Re-recording a block already tracked (which two threads that decoded + it concurrently will do) subtracts the old size first, so the running + total stays right however the two interleave. + """ + key = (pc, block_id) + with self._lock: + old = self._lru.pop(key, None) + if old is not None: + self.cur_bytes -= old + self._lru[key] = size + self.cur_bytes += size + # Keep the just-added block (it is at the end); never empty fully. + while self.cur_bytes > self.max_bytes and len(self._lru) > 1: + (ev_pc, ev_id), ev_size = self._lru.popitem(last=False) + self.cur_bytes -= ev_size + self._drop(ev_pc, ev_id) + + def forget(self, pc: _PerColor) -> None: + """Stop tracking `pc`'s blocks, for a table that is being closed.""" + with self._lock: + for key in [k for k in self._lru if k[0] is pc]: + self.cur_bytes -= self._lru.pop(key) + + def clear(self) -> None: + """Drop every tracked block and reset the budget.""" + with self._lock: + for pc, block_id in self._lru: + self._drop(pc, block_id) + self._lru.clear() + self.cur_bytes = 0 + + @staticmethod + def _drop(pc: _PerColor, block_id: int) -> None: + """Evict one block from its owning per-color object. + + Also drops the block's decode lock, so `_block_locks` stays bounded by + the cache budget instead of growing once per block ever decoded. A + thread may be holding that lock right now, mid-decode: dropping the + entry only means a later thread creates a fresh lock for the same + block, so the two could decode it concurrently. Both then write the + same bytes and re-`record`, which costs duplicated work but cannot + produce a wrong answer -- whereas an unbounded lock dict is a leak in + every long-running probe. + """ + pc._blocks.pop(block_id, None) + pc._block_locks.pop(block_id, None) + + +#: What :meth:`Tablebase._find` resolves a table to and :meth:`_TableFile._open` +#: opens: a filesystem path for the default source, otherwise whatever handle +#: that transport's :meth:`_TableFile._open_source` understands. Only ever +#: interpolated into a message, so a handle wants a ``__str__`` naming the table. +TableSource = Any + + +class _TableFile: + """Shared source lifecycle for the four table kinds. + + Table files are mapped read-only rather than read into memory: a probe + touches a handful of blocks, so the OS page cache serves the file and the + resident pages stay reclaimable. Mirrors + :meth:`chess.syzygy.Table.init_mmap`. :meth:`_open_source` is where that + choice is made, and the only place a different one has to be made. + """ + + EXT: str + MAGIC: int + KIND: str + cache: _BlockCache + per_color: List[Any] + path: TableSource + + _data: Optional[Any] = None + _reader: Optional[_Serial] = None + + def _open(self, path: TableSource) -> None: + """Open `path` and parse its header, releasing the source again if + parsing fails.""" + self.path = path + try: + # Held by the reader before anything can fail, so close() releases + # the buffer however far the open got. + self._reader = _Serial(self._open_source(path)) + if (len(self._reader.d) & 63) != 8: + raise ValueError(f"Invalid {self.KIND} file size {path}") + self._parse(self._reader) + except BaseException: + self.close() + raise + + def _parse(self, r: _Serial) -> None: + raise NotImplementedError + + def _open_source(self, path: TableSource) -> Any: + """Open `path`, set ``self._data`` to what :meth:`close` releases, and + return the buffer to read the table through. Maps the file by default. + + The one seam another transport replaces; :attr:`Tablebase.WDL_FILE` and + its siblings are how the subclass carrying the override gets used. An + override may return any object with ``len()``, indexing and slicing (a + slice being ``bytes`` or another such object); nothing above needs the + buffer protocol (see :class:`_Serial`, :func:`_read_u64le`, + :func:`_compressed_block`) and no span wider than one block is ever + taken, so a source that fetches on slice costs what a probe touches. + Materializing the table to satisfy the interface gives that up, and does + it under the open lock. + + Spans are one compressed block, and far more often the 8 bytes of a + bit-window read, so a lazy source wants a page cache under it; against + a mapping the OS is one. + """ + fd = os.open(path, os.O_RDONLY | getattr(os, "O_BINARY", 0)) + try: + data = mmap.mmap(fd, 0, access=mmap.ACCESS_READ) + finally: + os.close(fd) + + try: + # Unix: probing jumps between blocks, it does not stream. + data.madvise(mmap.MADV_RANDOM) + except AttributeError: + pass + + self._data = data + return memoryview(data) + + def close(self) -> None: + """Drop decoded blocks and release the source. Idempotent. + + Requires that no thread is inside a probe of this table: releasing the + buffer invalidates it for a concurrent reader, which is why + :meth:`Tablebase.close` drains its readers before calling this. + """ + for pc in self.per_color: + if pc is not None: + pc._blocks.clear() + self.cache.forget(pc) + # The view goes before munmap or mmap.close() raises BufferError; a + # non-memoryview buffer has nothing to release. + reader, self._reader = self._reader, None + if reader is not None and isinstance(reader.d, memoryview): + reader.d.release() + # Dropping `_data` before closing is what makes this idempotent. + data, self._data = self._data, None + if data is not None: + try: + data.close() + except BufferError: + # Something still holds a slice of the mapping; + # leave the unmap to refcounting rather than raising out of a + # close(). + pass + + +_TableFileT = TypeVar("_TableFileT", bound=_TableFile) + + +# =========================================================================== +# WDL table file +# =========================================================================== + +def egtb_table_colors(table_num: int) -> List[int]: + # WHITE always; BLACK only when table_num == 2. C++ Color WHITE=0,BLACK=1. + return [CPP_WHITE] + ([CPP_BLACK] if table_num == 2 else []) + + +class _WDLPerColor(_PerColor): + __slots__ = ("order", "radix", "block_size", "tail_size", "block_cnt", + "data_size", "offsets", "buf", "data_off", "dict", "single_val", + "dict_size") + order: List[int] + radix: List[int] + block_size: int + tail_size: int + block_cnt: int + data_size: int + offsets: MonoUintVec + buf: Any + data_off: int + dict: bytes + single_val: int + dict_size: int + _blocks: Dict[int, bytes] + + def __init__(self) -> None: + super().__init__() + self.single_val = DRAW + self.dict = b"" + self.dict_size = 0 + + +class WDLFile(_TableFile): + EXT = ".lzw" + MAGIC = WDL_MAGIC + KIND = "WDL" + + def __init__(self, cfg: PieceConfig, path: TableSource, + cache: Optional[_BlockCache] = None): + self.cfg = cfg + self.index_cfg = position_index_config(cfg) + self.cache = cache if cache is not None else _BlockCache(DEFAULT_BLOCK_CACHE_BYTES) + self.is_singular = [False, False] + self.is_dropped = [False, False] + self.is_loss_only = [False, False] + self.is_relaxed = [False, False] + self.per_color: List[Optional[_WDLPerColor]] = [None, None] + self._open(path) + + def _parse(self, r: _Serial) -> None: + magic = r.u32() + if magic != self.MAGIC: + raise ValueError(f"Invalid WDL magic {self.path}") + key_and_table = r.u32() + key = key_and_table >> 2 + if key != self.cfg.min_key: + raise ValueError(f"Wrong material key in WDL {self.path}: {key} != {self.cfg.min_key}") + table_num = key_and_table & 3 + colors = egtb_table_colors(table_num) + for c in colors: + flag = r.u8() + pc = _WDLPerColor() + self.per_color[c] = pc + self.is_loss_only[c] = bool(flag & LOSS_ONLY_FLAG) + self.is_relaxed[c] = bool(flag & RELAXED_FLAG) + if flag & SINGULAR_FLAG: + self.is_singular[c] = True + pc.single_val = r.u8() + elif flag & DROPPED_FLAG: + self.is_dropped[c] = True + else: + self._parse_header(r, pc) + if table_num == 1: + # Symmetric material: BLACK is WHITE mirrored, flag byte and all. + self.is_dropped[CPP_BLACK] = True + self.is_loss_only[CPP_BLACK] = self.is_loss_only[CPP_WHITE] + self.is_relaxed[CPP_BLACK] = self.is_relaxed[CPP_WHITE] + self._finalize(r, colors) + + def _parse_header(self, r: _Serial, pc: _WDLPerColor) -> None: + perm = r.u32() + n = self.index_cfg.num_populated_classes() + if not index_permutation_valid(n, perm): + raise ValueError("Invalid WDL index permutation") + pc.order, pc.radix = self.index_cfg.make_layout(perm) + pc.tail_size = r.u16() + pc.block_size = r.u32() + pc.block_cnt = r.u64() + pc.data_size = r.u64() + + def _finalize(self, r: _Serial, colors: List[int]) -> None: + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + pc.dict_size = r.u16() + if pc.dict_size != 0: + start = r.caret() + # Copied out: it prefixes every decode of this color's blocks. + pc.dict = bytes(r.d[start:start + pc.dict_size]) + r.advance(pc.dict_size) + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + log2_bu = r.u8() + sample_width = r.u8() + offset_width = r.u8() + r.advance(1) # usz_width + mono_off = r.caret() + mono_bytes = MonoUintVec.on_disk_bytes(pc.block_cnt + 1, log2_bu, + sample_width, offset_width) + r.advance(mono_bytes) + pc.offsets = MonoUintVec(r.d, mono_off, pc.block_cnt + 1, log2_bu, + sample_width, offset_width) + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + r.align(64) + start = r.caret() + # Buffer plus offset, never a slice spanning the section. + pc.buf, pc.data_off = r.d, start + r.advance(pc.data_size) + + def _get_block(self, pc: _WDLPerColor, block_id: int) -> bytes: + blk = pc._blocks.get(block_id) + if blk is not None: + self.cache.touch(pc, block_id) + return blk + # Cold block: one thread decodes it, others wait for that result + # instead of duplicating the work. Re-checked inside the lock because + # the winner will have filled `_blocks` by the time we get in. Nothing + # else in this method touches shared state, so threads after + # *different* cold blocks proceed independently. + with pc.lock_for(block_id): + blk = pc._blocks.get(block_id) + if blk is not None: + self.cache.touch(pc, block_id) + return blk + doff, dnext = pc.offsets.get2(block_id) + dsz = dnext - doff + out_sz = (pc.tail_size if (block_id == pc.block_cnt - 1 and pc.tail_size != 0) + else pc.block_size) + blk = lz4_decompress_block( + _compressed_block(pc.buf[pc.data_off + doff:pc.data_off + doff + dsz]), + out_sz, pc.dict) + pc._blocks[block_id] = blk + self.cache.record(pc, block_id, len(blk)) + return blk + + def read(self, color: int, board: chess.Board) -> int: + """Return WDL_Stored at the board's index in `color`'s frame.""" + pc = self.per_color[color] + assert pc is not None + if self.is_singular[color]: + return pc.single_val + pos = self.index_cfg.board_index(board, pc.order, pc.radix) + assert pos is not None + packed_byte = pos // 2 + block_id = packed_byte // pc.block_size + in_block = packed_byte % pc.block_size + lo, hi = pc.offsets.get2(block_id) + if lo == hi: + return 7 # ILLEGAL + data = self._get_block(pc, block_id) + entry = data[in_block] + return (entry >> ((pos % 2) * 4)) & 0xF + + +# =========================================================================== +# Probe orchestration (subset: WDL). Mirrors src/probe/probe.cpp. +# =========================================================================== + +def mirror_for_canonical(board: chess.Board) -> chess.Board: + """Swap colors and rank-mirror every piece; flip side to move. The ep square + is rank-mirrored with everything else. Mirrors Position::mirror in + src/chess/position.h.""" + out = board.copy(stack=False) + out.apply_mirror() + return out + + +_WDL_NAME = {LOSE: "LOSE", BLESSED_LOSS: "BLESSED_LOSS", DRAW: "DRAW", + CURSED_WIN: "CURSED_WIN", WIN: "WIN", ILLEGAL: "ILLEGAL"} + + +# --- WDL semantic helpers (egtb_entry.h / probe.cpp) --- + +def invert_wdl(w: int) -> int: + return {WIN: LOSE, CURSED_WIN: BLESSED_LOSS, DRAW: DRAW, + BLESSED_LOSS: CURSED_WIN, LOSE: WIN, ILLEGAL: ILLEGAL}[w] + + +def invert_stored(s: int) -> int: + # WDL_Stored -> WDL_Entry, inverted across one quiet ply (markers tip a ply). + return {0: WIN, 1: CURSED_WIN, 2: DRAW, 3: BLESSED_LOSS, 4: LOSE, + 5: CURSED_WIN, # BOUNDARY_LOSS -> we win but only cursed + 6: BLESSED_LOSS, # BOUNDARY_WIN -> we lose but only blessed + 7: ILLEGAL}[s] + + +def wdl_rank(w: int) -> int: + return {WIN: 4, CURSED_WIN: 3, DRAW: 2, BLESSED_LOSS: 1, LOSE: 0, ILLEGAL: -1}[w] + + +def is_symmetric_material(cfg: PieceConfig) -> bool: + return cfg.base_key == cfg.mirr_key + + +def material_has_pawns(cfg: PieceConfig) -> bool: + """Whether a DTC pack exists for this material at all: with no push to + budget the stack is one layer, which is the DTZ table.""" + return cfg.has_pair or any(t == PAWN for _c, t in cfg.pieces) + + +def is_win_class(w: int) -> bool: + """Win classes are the ones a loss-only frame leaves out; CURSED_WIN counts, + since a distance cell is stored under the win flag either way.""" + return w == WIN or w == CURSED_WIN + + +def locate_frame(f: Any, cfg: PieceConfig, board: chess.Board, + wdl: int) -> Tuple[int, Optional[chess.Board], bool]: + """Which frame holds a cell, whether it needs mirroring, and whether it is + there to be read. A dropped frame is reachable through the mirror when the + material is its own mirror; a loss-only frame holds no win. Both gaps are + filled by the same one-ply derive, run on the unmirrored board.""" + color = CPP_WHITE if board.turn == WHITE else CPP_BLACK + if not f.is_dropped[color]: + return color, None, not (f.is_loss_only[color] and is_win_class(wdl)) + if not is_symmetric_material(cfg): + return color, None, False + kept = CPP_BLACK if color == CPP_WHITE else CPP_WHITE + if f.is_loss_only[kept] and is_win_class(wdl): + return kept, None, False + return kept, mirror_for_canonical(board), True + + +MAX_DERIVE_DEPTH = 16 + + +def _internal_board(board: chess.Board) -> chess.Board: + """A copy with no en-passant square (derive move-gen excludes ep; the + overlay handles ep separately).""" + if board.ep_square is not None: + board = board.copy(stack=False) + board.ep_square = None + return board + + +# =========================================================================== +# LZMA block decode + value-from-storage helpers (egtb_entry.h) +# =========================================================================== + +def lzma_raw_decompress(block: memoryview, expected_size: int) -> bytes: + """Decode a raw LZMA1 stream with 5 props bytes appended at the tail + (the LZMA SDK ``LzmaUncompress`` framing used by the C++ side).""" + if len(block) < 5: + raise ValueError("LZMA block too small") + props = bytes(block[-5:]) + raw = bytes(block[:-5]) + d0 = props[0] + lc = d0 % 9 + rem = d0 // 9 + lp = rem % 5 + pb = rem // 5 + dict_size = int.from_bytes(props[1:5], "little") + dec = lzma.LZMADecompressor(format=lzma.FORMAT_RAW, filters=[{ + "id": lzma.FILTER_LZMA1, "dict_size": dict_size, + "lc": lc, "lp": lp, "pb": pb}]) + out = dec.decompress(raw, expected_size) + if len(out) != expected_size: + out += dec.decompress(b"", expected_size - len(out)) + if len(out) != expected_size: + raise ValueError(f"LZMA size mismatch {len(out)} != {expected_size}") + return out + + +def dtz_value_from_storage(stored: int, w: int, entry_bytes: int) -> int: + if w == DRAW: + return 0 + if entry_bytes == 1 and (w == CURSED_WIN or w == BLESSED_LOSS): + return (stored << 1) - 1 + return stored + + +def dtm_value_from_storage(stored: int, w: int) -> int: + if w in (WIN, CURSED_WIN): + return (stored << 1) | 1 + if w in (LOSE, BLESSED_LOSS): + return stored << 1 + return 0 + + +def dtm50_value_from_storage(stored: int, w: int) -> int: + if w == WIN: + return (stored << 1) | 1 + if w == LOSE: + return stored << 1 + return 0 + + +# =========================================================================== +# DTZ / DTM table files — src/probe/dtz_file.cpp, src/probe/dtm_file.cpp +# =========================================================================== + +class _RankPerColor(_PerColor): + __slots__ = ("order", "radix", "entry_bytes", "block_size", "tail_size", + "block_cnt", "data_size", "offsets", "buf", "data_off", + "rank_to_value", "single_val") + order: List[int] + radix: List[int] + entry_bytes: int + block_size: int + tail_size: int + block_cnt: int + data_size: int + offsets: MonoUintVec + buf: Any + data_off: int + rank_to_value: List[int] + single_val: int + _blocks: Dict[int, bytes] + + def __init__(self) -> None: + super().__init__() + self.single_val = 0 + + +class _RankFile(_TableFile): + """The rank-coded, LZMA-blocked table body. DTM is the byte-for-byte twin + of DTZ (as ``src/probe/dtm_file.cpp`` says of its own traits); only the + magic and :meth:`_value_from_storage` differ, so both files parse and read + through this one.""" + + @staticmethod + def _value_from_storage(stored: int, wdl: int, entry_bytes: int) -> int: + raise NotImplementedError + + def __init__(self, cfg: PieceConfig, path: TableSource, + cache: Optional[_BlockCache] = None): + self.cfg = cfg + self.index_cfg = position_index_config(cfg) + self.cache = cache if cache is not None else _BlockCache(DEFAULT_BLOCK_CACHE_BYTES) + self.is_singular = [False, False] + self.is_dropped = [False, False] + self.is_loss_only = [False, False] + self.per_color: List[Optional[_RankPerColor]] = [None, None] + self._open(path) + + def _parse(self, r: _Serial) -> None: + if r.u32() != self.MAGIC: + raise ValueError(f"Invalid {self.KIND} magic {self.path}") + kat = r.u32() + if (kat >> 2) != self.cfg.min_key: + raise ValueError(f"Wrong material key in {self.KIND}") + table_num = kat & 3 + colors = egtb_table_colors(table_num) + for c in colors: + flag = r.u8() + pc = _RankPerColor() + self.per_color[c] = pc + self.is_loss_only[c] = bool(flag & LOSS_ONLY_FLAG) + if flag & SINGULAR_FLAG: + self.is_singular[c] = True + pc.single_val = r.u8() + elif flag & DROPPED_FLAG: + self.is_dropped[c] = True + else: + self._parse_header(r, pc) + if table_num == 1: + # Symmetric material: BLACK is WHITE mirrored, flag byte and all. + self.is_dropped[CPP_BLACK] = True + self.is_loss_only[CPP_BLACK] = self.is_loss_only[CPP_WHITE] + self._finalize(r, colors) + + def _parse_header(self, r: _Serial, pc: _RankPerColor) -> None: + perm = r.u32() + pc.order, pc.radix = self.index_cfg.make_layout(perm) + pc.entry_bytes = r.u8() + pc.tail_size = r.u32() + pc.block_size = r.u32() + pc.block_cnt = r.u64() + pc.data_size = r.u64() + num_ranks = r.u16() + pc.rank_to_value = [r.u16() for _ in range(num_ranks)] + + def _finalize(self, r: _Serial, colors: List[int]) -> None: + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + log2_bu = r.u8() + sample_width = r.u8() + offset_width = r.u8() + r.advance(1) # usz_width (unused here) + mono_off = r.caret() + mb = MonoUintVec.on_disk_bytes(pc.block_cnt + 1, log2_bu, sample_width, offset_width) + r.advance(mb) + pc.offsets = MonoUintVec(r.d, mono_off, pc.block_cnt + 1, log2_bu, + sample_width, offset_width) + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + r.align(64) + start = r.caret() + # Buffer plus offset, never a slice spanning the section. + pc.buf, pc.data_off = r.d, start + r.advance(pc.data_size) + + def _get_block_raw(self, pc: _RankPerColor, block_id: int) -> bytes: + blk = pc._blocks.get(block_id) + if blk is not None: + self.cache.touch(pc, block_id) + return blk + with pc.lock_for(block_id): # see WDLFile._get_block + blk = pc._blocks.get(block_id) + if blk is not None: + self.cache.touch(pc, block_id) + return blk + decode_sz = (pc.tail_size if (block_id == pc.block_cnt - 1 and pc.tail_size != 0) + else pc.block_size) + doff, dnext = pc.offsets.get2(block_id) + dsz = dnext - doff + blk = b"" if dsz == 0 else lzma_raw_decompress( + _compressed_block(pc.buf[pc.data_off + doff:pc.data_off + doff + dsz]), + decode_sz) + pc._blocks[block_id] = blk + self.cache.record(pc, block_id, len(blk)) + return blk + + def read(self, color: int, board: chess.Board, wdl: int) -> int: + assert wdl != DRAW and wdl != ILLEGAL + pc = self.per_color[color] + assert pc is not None + if self.is_singular[color]: + return self._value_from_storage(pc.single_val, wdl, 1) + pos = self.index_cfg.board_index(board, pc.order, pc.radix) + assert pos is not None + ppb = pc.block_size // pc.entry_bytes + block_id = pos // ppb + in_block = pos % ppb + lo, hi = pc.offsets.get2(block_id) + if lo == hi: + return 0 # skip block: uniform DRAW/ILLEGAL + raw = self._get_block_raw(pc, block_id) + if pc.entry_bytes == 1: + stored = raw[in_block] + else: + stored = struct.unpack_from(" int: + return dtz_value_from_storage(stored, wdl, entry_bytes) + + +class DTMFile(_RankFile): + EXT = ".lzdtm" + MAGIC = DTM_MAGIC + KIND = "DTM" + + @staticmethod + def _value_from_storage(stored: int, wdl: int, entry_bytes: int) -> int: + # Unlike DTZ, the mate distance is a plain parity encoding: the class + # supplies the low bit, so the cell width never enters the decode. + return dtm_value_from_storage(stored, wdl) + + +# =========================================================================== +# Changepoint packs — src/probe/layered_file.h, dtm50_file.cpp, dtc_file.cpp +# =========================================================================== +# +# One container, two metrics: DTM50 stacks 101 layers by halfmove clock; DTC has +# 29 separately solved budget points plus the terminal clean changepoint supplied +# by DTZ. Both open every record with that embedded endpoint -- DTM50's the DTM, +# DTC's the DTZ. Everything down to the block decode is shared +# (:class:`_LayeredFile`); each metric's ``read`` is its own, as in the C++. + +def _bit(buf: bytes, i: int) -> int: + return (buf[i >> 3] >> (i & 7)) & 1 + + +# --- block-local prefix indexes (src/probe/layered_file.h) --- +# +# A read needs a position's 2-bit state and its index among positions sharing +# it: prefix counts over a bitmap `block_positions` long, 1,048,576 in every +# table the generator writes. Answering per position costs a 1M-iteration loop +# and a ~40 MiB list per block, which dwarfs the payload and makes the decoded +# -block budget meaningless. So index every STRIDE-th position and walk at +# most one stride at read time, as the C++ does. +_LAYERED_STRIDE = 256 +_STATE_STRIDE_BYTES = _LAYERED_STRIDE // 4 # 2-bit states: 4 per byte +_HINT_STRIDE_BYTES = _LAYERED_STRIDE // 8 # hint bitmaps: 8 per byte + +#: A drawn layer, which no mate distance reaches. +DTM50_DRAWN = 0xFFFF + +#: State-bitmap byte -> how many of its four fields are state `s`, so that +#: ``translate`` + ``sum`` counts a byte range in C. +_STATE_COUNT_XLAT = tuple( + bytes(sum(1 for k in range(4) if (b >> (2 * k)) & 3 == s) for b in range(256)) + for s in range(4)) + +#: ``[b][q][s]``: fields before `q` of byte `b` in state `s` -- the partial +#: byte a walk ends on. +_STATE_FIELD_PREFIX = tuple( + tuple(tuple(sum(1 for k in range(q) if (b >> (2 * k)) & 3 == s) for s in range(4)) + for q in range(4)) + for b in range(256)) + + +def _state_stride_index(state_bits: bytes, num_positions: int) -> Tuple[array.array[int], ...]: + """Per-state cumulative counts at every stride boundary: ``cum[s][k]`` is + how many of the first ``k * _LAYERED_STRIDE`` positions are in state `s`. + + The last entry may include padding fields, which no read reaches: stride + `k` only ever consults ``cum[s][k]``.""" + n_strides = _ceil_div(num_positions, _LAYERED_STRIDE) + out = [] + for s in range(4): + per_byte = state_bits.translate(_STATE_COUNT_XLAT[s]) + cum = array.array("q", bytes(8 * (n_strides + 1))) + run = 0 + at = 0 + for k in range(n_strides): + cum[k] = run + run += sum(per_byte[at:at + _STATE_STRIDE_BYTES]) + at += _STATE_STRIDE_BYTES + cum[n_strides] = run + out.append(cum) + return tuple(out) + + +def _state_and_index(state_bits: bytes, cum: Tuple[array.array[int], ...], + pos: int) -> Tuple[int, int]: + """`pos`'s state, and its index among positions of that state before it.""" + sid, within = divmod(pos, _LAYERED_STRIDE) + byte_in_stride, field = divmod(within, 4) + base = sid * _STATE_STRIDE_BYTES + b = state_bits[base + byte_in_stride] + st = (b >> (2 * field)) & 3 + idx = cum[st][sid] + if byte_in_stride: + idx += sum(state_bits[base:base + byte_in_stride].translate(_STATE_COUNT_XLAT[st])) + return st, idx + _STATE_FIELD_PREFIX[b][field][st] + + +def _build_hints(payload: bytes, off: int, hint_byte: int, n: int, + short: int, long: int) -> bytes: + """The draw-end bits of `n` variable-width records, gathered into a bitmap. + The bit is the MSB of the record's last h byte -- its first for SINGLE, its + second for DOUBLE -- and only a walk finds the next record.""" + hints = bytearray((n + 7) // 8) + for j in range(n): + if payload[off + hint_byte] & 0x80: + hints[j >> 3] |= 1 << (j & 7) + off += short + else: + off += long + return bytes(hints) + + +def _hint_stride_index(hints: bytes, n: int) -> array.array[int]: + """Cumulative popcount of a hint bitmap at every stride boundary: entry + `k` counts the set bits below bit ``k * _LAYERED_STRIDE``.""" + n_strides = _ceil_div(n, _LAYERED_STRIDE) + cum = array.array("q", bytes(8 * (n_strides + 1))) + run = 0 + for k in range(n_strides): + cum[k] = run + chunk = hints[k * _HINT_STRIDE_BYTES:(k + 1) * _HINT_STRIDE_BYTES] + run += chess.popcount(int.from_bytes(chunk, "little")) + cum[n_strides] = run + return cum + + +def _hint_prefix(hints: bytes, cum: array.array[int], i: int) -> int: + """Popcount of bits [0, `i`) of a hint bitmap: the stride total, then at + most one stride of walking. Entry `i` itself is never counted, so the + bitmap only has to cover the entries that exist.""" + sid, within = divmod(i, _LAYERED_STRIDE) + base = sid * _HINT_STRIDE_BYTES + whole, rem = divmod(within, 8) + n = cum[sid] + if whole: + n += chess.popcount(int.from_bytes(hints[base:base + whole], "little")) + if rem: + n += chess.popcount(hints[base + whole] & ((1 << rem) - 1)) + return n + + +class _LayeredPerColor(_PerColor): + __slots__ = ("order", "radix", "entry_bytes", "block_positions", + "tail_positions", "block_cnt", "data_size", "offsets", + "usizes", "buf", "data_off", "rank_to_value") + order: List[int] + radix: List[int] + entry_bytes: int + block_positions: int + tail_positions: int + block_cnt: int + data_size: int + offsets: MonoUintVec + usizes: Min0UintVec + buf: Any + data_off: int + rank_to_value: List[int] + _blocks: Dict[int, Dict[str, Any]] + + def __init__(self) -> None: + super().__init__() + + +def _multi_last_changepoint(lo64: int, hi64: int) -> int: + """Highest set changepoint bit of a MULTI record: lo bit i is layer i, hi + bit j is layer 64 + j, the way read()'s masks run.""" + if hi64: + return 64 + hi64.bit_length() - 1 + if lo64: + return lo64.bit_length() - 1 + return 0 + + +class _LayeredFile(_TableFile): + """The changepoint container both packs are written in: header, offsets + section and block decode, down to the per-position state and its record. + What a record *means* is the metric's, so ``read`` lives in the subclass.""" + + def __init__(self, cfg: PieceConfig, path: TableSource, + cache: Optional[_BlockCache] = None): + self.cfg = cfg + self.index_cfg = position_index_config(cfg) + self.cache = cache if cache is not None else _BlockCache(DEFAULT_BLOCK_CACHE_BYTES) + self.is_singular = [False, False] + self.is_dropped = [False, False] + self.is_loss_only = [False, False] + self.per_color: List[Optional[_LayeredPerColor]] = [None, None] + self._open(path) + + def _parse(self, r: _Serial) -> None: + if r.u32() != self.MAGIC: + raise ValueError(f"Invalid {self.KIND} magic") + kat = r.u32() + if (kat >> 2) != self.cfg.min_key: + raise ValueError(f"Wrong material key in {self.KIND}") + table_num = kat & 3 + colors = egtb_table_colors(table_num) + for c in colors: + flag = r.u8() + pc = _LayeredPerColor() + self.per_color[c] = pc + self.is_loss_only[c] = bool(flag & LOSS_ONLY_FLAG) + if flag & SINGULAR_FLAG: + self.is_singular[c] = True + if r.u8() != 0: + raise ValueError(f"{self.KIND} singular value must be DRAW") + elif flag & DROPPED_FLAG: + self.is_dropped[c] = True + else: + self._parse_header(r, pc) + if table_num == 1: + # Symmetric material: BLACK is WHITE mirrored, flag byte and all. + self.is_dropped[CPP_BLACK] = True + self.is_loss_only[CPP_BLACK] = self.is_loss_only[CPP_WHITE] + self._finalize(r, colors) + + def _parse_header(self, r: _Serial, pc: _LayeredPerColor) -> None: + perm = r.u32() + pc.order, pc.radix = self.index_cfg.make_layout(perm) + pc.entry_bytes = r.u8() + pc.block_positions = r.u32() + pc.block_cnt = r.u64() + pc.tail_positions = r.u32() + pc.data_size = r.u64() + num_ranks = r.u16() + pc.rank_to_value = [r.u16() for _ in range(num_ranks)] + + def _finalize(self, r: _Serial, colors: List[int]) -> None: + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + log2_bu = r.u8() + sample_width = r.u8() + offset_width = r.u8() + usz_width = r.u8() + mono_off = r.caret() + mb = MonoUintVec.on_disk_bytes(pc.block_cnt + 1, log2_bu, sample_width, offset_width) + r.advance(mb) + usz_off = r.caret() + ub = Min0UintVec.on_disk_bytes(pc.block_cnt, usz_width) + r.advance(ub) + pc.offsets = MonoUintVec(r.d, mono_off, pc.block_cnt + 1, log2_bu, + sample_width, offset_width) + pc.usizes = Min0UintVec(r.d, usz_off, pc.block_cnt, usz_width) + for c in colors: + if self.is_singular[c] or self.is_dropped[c]: + continue + pc = self.per_color[c] + assert pc is not None + r.align(64) + start = r.caret() + # Buffer plus offset, never a slice spanning the section. + pc.buf, pc.data_off = r.d, start + r.advance(pc.data_size) + + def _get_block(self, pc: _LayeredPerColor, block_id: int) -> Dict[str, Any]: + blk = pc._blocks.get(block_id) + if blk is not None: + self.cache.touch(pc, block_id) + return blk + # See WDLFile._get_block. Besides the lzma pass, a DTM50 block builds + # the stride prefix indexes below, so it stays the heaviest decode of + # the four formats and the one most worth sharing. + with pc.lock_for(block_id): + blk = pc._blocks.get(block_id) + if blk is not None: + self.cache.touch(pc, block_id) + return blk + return self._decode_block(pc, block_id) + + def _decode_block(self, pc: _LayeredPerColor, block_id: int) -> Dict[str, Any]: + """Decode, cache and return one block. Called with its lock held.""" + doff, dnext = pc.offsets.get2(block_id) + dsz = dnext - doff + usz = pc.usizes.get(block_id) + payload = lzma_raw_decompress( + _compressed_block(pc.buf[pc.data_off + doff:pc.data_off + doff + dsz]), usz) + eb = pc.entry_bytes + (num_positions, num_single, num_double, num_multi, + single_stream_bytes, double_stream_bytes) = struct.unpack_from(" int: + return payload[off] if eb == 1 else int(struct.unpack_from(" Optional[Tuple[Dict[str, Any], int, int]]: + """A position's decoded block, its record state and its index among the + records sharing that state. ``None`` where nothing is stored, which + reads as drawn at every layer.""" + pc = self.per_color[color] + assert pc is not None + if self.is_singular[color]: + return None + pos = self.index_cfg.board_index(board, pc.order, pc.radix) + assert pos is not None + block_id, pos_in_block = divmod(pos, pc.block_positions) + lo, hi = pc.offsets.get2(block_id) + if lo == hi: + return None # skip block (uniform DRAW) + blk = self._get_block(pc, block_id) + st, idx = _state_and_index(blk["state_bits"], blk["state_cum"], pos_in_block) + return blk, st, idx + + +class DTM50File(_LayeredFile): + EXT = ".lzdtm50" + MAGIC = DTM50_MAGIC + KIND = "DTM50" + + def read(self, color: int, board: chess.Board, wdl: int, + hmc: int) -> Tuple[int, int]: + """The value at `hmc`, and the layer where the cell turns DRAW; the + terminal rank decides both. Only the draw-end hint says DRAW; a stored + 0 is a mate.""" + assert wdl != DRAW and wdl != ILLEGAL + flat = (hmc == IGNORE_50MR) + if not flat and (wdl == CURSED_WIN or wdl == BLESSED_LOSS): + return (DTM50_DRAWN, 0) + found = self._locate(color, board) + if found is None: + return (DTM50_DRAWN, 0) + blk, st, idx = found + pc = self.per_color[color] + assert pc is not None + payload = blk["payload"] + eb = blk["eb"] + r2v = pc.rank_to_value + layer = 0 if flat else (hmc + 1) + + draw_flip = 0 + if st == 0: # CONST: one rank for all layers, so it never flips + stored = r2v[self._read_rank(payload, blk["const_stream_off"] + idx * eb, eb)] + elif st == 1: # SINGLE: one transition at h + short, long = 1 + eb, 1 + 2 * eb + n_short = _hint_prefix(blk["single_hints"], blk["single_pre"], idx) + off = blk["single_stream_off"] + n_short * short + (idx - n_short) * long + draw_end = _bit(blk["single_hints"], idx) + h = payload[off] & 0x7F + if draw_end: + draw_flip = h + if layer < h: + stored = r2v[self._read_rank(payload, off + 1, eb)] + elif draw_end: + return (DTM50_DRAWN, draw_flip) + else: + stored = r2v[self._read_rank(payload, off + 1 + eb, eb)] + elif st == 2: # DOUBLE: transitions at h1 < h2 + short, long = 2 + 2 * eb, 2 + 3 * eb + n_short = _hint_prefix(blk["double_hints"], blk["double_pre"], idx) + off = blk["double_stream_off"] + n_short * short + (idx - n_short) * long + draw_end = _bit(blk["double_hints"], idx) + h1 = payload[off] + h2 = payload[off + 1] & 0x7F + if draw_end: + draw_flip = h2 + if layer < h1: + stored = r2v[self._read_rank(payload, off + 2, eb)] + elif layer < h2: + stored = r2v[self._read_rank(payload, off + 2 + eb, eb)] + elif draw_end: + return (DTM50_DRAWN, draw_flip) + else: + stored = r2v[self._read_rank(payload, off + 2 + 2 * eb, eb)] + else: # MULTI: 128-bit changepoint bitmap + entry = blk["multi_data_off"] + struct.unpack_from( + " int: + """Plies a DTC answer may still spend before the 50-move claim: the whole + band where the caller ignores the rule, and none at all once the clock has + reached the claim, where only a mate already on the board, at distance 0, + outruns it.""" + if rule50 == IGNORE_50MR: + return MAX_NON_CURSED_DTZ + return 0 if rule50 >= MAX_NON_CURSED_DTZ else MAX_NON_CURSED_DTZ - rule50 + + +class _DTCCell(NamedTuple): + """One DTC cell: pushes the winning side still owes before a conversion, and + plies to the next zeroing move on the line that owes them. Both come from + the budget layer that fits the clock the caller passed, so a fresh clock + buys the fewest pushes any line manages, and a tighter one may force more of + them or leave nothing that fits at all. + + ``dtz`` is the record's unbounded row, which is the DTZ table the pack + embeds, in that table's own plies: the same decode that picks a budget + passes it, and a cursed class has only it.""" + + order: int = 0 + value: int = DTC_DRAWN + dtz: int = DTC_DRAWN + + @property + def priced(self) -> bool: + return self.value != DTC_DRAWN + + +class _BudgetSegments(NamedTuple): + """One position's value as a function of push budget, read off its record. + The budget runs down as the pack index runs up, so `segs` is ascending in + both row and value; `draw_start` is the row a trailing DRAW segment begins + at, which is what the record's hint bit says.""" + + segs: List[Tuple[int, int]] + draw_start: Optional[int] + + def end_of(self, i: int) -> int: + """First row past segment `i`.""" + if i + 1 < len(self.segs): + return self.segs[i + 1][0] + return DTC_PACK_LAYERS if self.draw_start is None else self.draw_start + + +class DTCFile(_LayeredFile): + EXT = ".lzdtc" + MAGIC = DTC_MAGIC + KIND = "DTC" + + def _segments(self, color: int, board: chess.Board) -> Optional[_BudgetSegments]: + """Decode a whole record. ``None`` where nothing is stored, which reads + as drawn at every budget.""" + found = self._locate(color, board) + if found is None: + return None + blk, st, idx = found + pc = self.per_color[color] + assert pc is not None + payload = blk["payload"] + eb = blk["eb"] + r2v = pc.rank_to_value + rank = self._read_rank + # No value transform: unlike DTM50's parity halving, a DTC rank holds + # the plies outright. + if st == 0: # CONST: one value at every budget + off = blk["const_stream_off"] + idx * eb + return _BudgetSegments([(0, r2v[rank(payload, off, eb)])], None) + if st == 1: # SINGLE: one changepoint at h + short, long = 1 + eb, 1 + 2 * eb + n_short = _hint_prefix(blk["single_hints"], blk["single_pre"], idx) + off = blk["single_stream_off"] + n_short * short + (idx - n_short) * long + draw_end = _bit(blk["single_hints"], idx) + h = payload[off] & 0x7F + segs = [(0, r2v[rank(payload, off + 1, eb)])] + if draw_end: + return _BudgetSegments(segs, h) + segs.append((h, r2v[rank(payload, off + 1 + eb, eb)])) + return _BudgetSegments(segs, None) + if st == 2: # DOUBLE: changepoints at h1 < h2 + short, long = 2 + 2 * eb, 2 + 3 * eb + n_short = _hint_prefix(blk["double_hints"], blk["double_pre"], idx) + off = blk["double_stream_off"] + n_short * short + (idx - n_short) * long + draw_end = _bit(blk["double_hints"], idx) + h1 = payload[off] + h2 = payload[off + 1] & 0x7F + segs = [(0, r2v[rank(payload, off + 2, eb)]), + (h1, r2v[rank(payload, off + 2 + eb, eb)])] + if draw_end: + return _BudgetSegments(segs, h2) + segs.append((h2, r2v[rank(payload, off + 2 + 2 * eb, eb)])) + return _BudgetSegments(segs, None) + # MULTI: changepoint bitmap, one segment per set bit + entry = blk["multi_data_off"] + struct.unpack_from( + " _DTCCell: + """The fewest pushes whose plies-to-zeroing still fit the clock the + caller holds: the deepest budget index that stays inside `rule50`. + Fewer pushes may cost a longer wait, so a fresh clock buys the fewest any + line manages -- which is whatever the position needs, not necessarily + none.""" + assert wdl != DRAW and wdl != ILLEGAL + seg = self._segments(color, board) + if seg is None: + return _DTCCell() + + # Segment 0 is the unbounded row: the DTZ table's own plies, whatever + # the budgets above it make of the cell. Every record opens with it, and + # a decisive class is priced there whether or not any budget settles it. + assert seg.segs + dtz = seg.segs[0][1] + # The budgets hold only what 50MR settles, so a cursed class gets that + # row and nothing else. + if wdl == CURSED_WIN or wdl == BLESSED_LOSS: + return _DTCCell(dtz=dtz) + + budget_plies = dtc_budget_plies(rule50) + # Values rise with the row, so the fittable segments are a prefix and + # the last of them carries the fewest pushes. + for i in range(len(seg.segs) - 1, -1, -1): + value = seg.segs[i][1] + if value > budget_plies: + continue + return _DTCCell(DTC_BUDGET_LAYERS - (seg.end_of(i) - 1), value, dtz) + return _DTCCell(dtz=dtz) + + def read_curve(self, color: int, board: chess.Board, wdl: int) -> List[int]: + """The same decode, spread over the budgets it covers instead of + resolved against a clock: ``DTC_DRAWN`` wherever a budget point settles + nothing, which for a cursed class is everywhere. The first 29 entries + are separately solved finite points; the top entry is the embedded DTZ + terminal endpoint.""" + curve = [DTC_DRAWN] * DTC_PACK_LAYERS + assert wdl != DRAW and wdl != ILLEGAL + if wdl == CURSED_WIN or wdl == BLESSED_LOSS: + return curve + seg = self._segments(color, board) + if seg is None: + return curve + for i, (start, value) in enumerate(seg.segs): + for row in range(start, seg.end_of(i)): + curve[DTC_BUDGET_LAYERS - row] = value + return curve + + +# =========================================================================== +# Probe orchestration — src/probe/probe.cpp +# =========================================================================== + +def prefer_new(new_wdl: int, new_dtz: int, old_wdl: int, old_dtz: int) -> bool: + rn, ro = wdl_rank(new_wdl), wdl_rank(old_wdl) + if rn != ro: + return rn > ro + if new_wdl in (WIN, CURSED_WIN): + return new_dtz < old_dtz + if new_wdl in (LOSE, BLESSED_LOSS): + return new_dtz > old_dtz + return False + + +def prefer_new_dtc(new_wdl: int, new_order: int, new_value: int, + old_wdl: int, old_order: int, old_value: int) -> bool: + """Same as :func:`prefer_new` on DTC's key: pushes owed lead, and the wait + only breaks a tie in them.""" + rn, ro = wdl_rank(new_wdl), wdl_rank(old_wdl) + if rn != ro: + return rn > ro + if new_wdl in (WIN, CURSED_WIN): + return (new_order < old_order if new_order != old_order + else new_value < old_value) + if new_wdl in (LOSE, BLESSED_LOSS): + return (new_order > old_order if new_order != old_order + else new_value > old_value) + return False + + +def below_pinned_class(pinned: int, offer: int) -> bool: + """Derivation stands in for `read(color, board, wdl)`. Reads are + class-driven: the WDL companion pins the class, while the file only prices + the distance. A child that cannot reach the pinned class never wins the + minimax and never needs a probe. ILLEGAL pins nothing -- the WDL table had + no answer either.""" + return pinned != ILLEGAL and wdl_rank(offer) < wdl_rank(pinned) + + +def child_class_is_forced(pinned: int) -> bool: + """A LOSE pin makes WIN a safe child-class surrogate. + + For an exact five-class pin, every child is a clean WIN for the opponent; + otherwise the parent would be a BLESSED_LOSS. A folded DTM pin may also + represent BLESSED_LOSS, but WIN and CURSED_WIN are equivalent to the DTM + minimax, so WIN remains the right surrogate there too. + """ + return pinned == LOSE + + +def dtz_lift(my_wdl: int) -> int: + """A DTZ minimax lifts a LOSE to BLESSED_LOSS past 100, the higher rank of + the two.""" + return BLESSED_LOSS if my_wdl == LOSE else my_wdl + + +def fold_dtm_wdl(w: int) -> int: + if w == CURSED_WIN: + return WIN + if w == BLESSED_LOSS: + return LOSE + return w + + +def fold_50mr_wdl(w: int) -> int: + """5-class WDL -> the three a clocked metric holds, DTM50's layers and DTC's + budgets alike: cursed and blessed are unreachable under 50MR.""" + if w == CURSED_WIN or w == BLESSED_LOSS: + return DRAW + return w + + +def _is_checkmate(board: chess.Board) -> bool: + return board.is_checkmate() + + +def dtz_from_draw_flip(h: int, wdl: int, board: chess.Board) -> Optional[int]: + """The DRAW flip pins DTZ: a W/L cell still decisive at hmc = 100 - dtz turns + one tick later, so its flip layer (hmc + 1) is h = 102 - dtz. h == 1 is a + cell already drawn at a fresh clock -- the cursed band, whose distance the + flip says nothing about; never flipping means dtz <= 1, which only a mate + splits. DRAW/ILLEGAL answer 0, the don't-care the DTZ table decodes.""" + if wdl == DRAW or wdl == ILLEGAL: + return 0 + if wdl not in (WIN, LOSE): + return None + if h == 0: + return 0 if (wdl == LOSE and _is_checkmate(board)) else 1 + if h == 1: + return None + return MAX_NON_CURSED_DTZ + 2 - h + + +# Two bounds apply to the layer at `rule50`, one for each distance. Both apply +# only to W/L positions, since cursed/blessed positions are DRAW at every layer +# and no distance is assigned to them. +# +# Pinned: rule50 + dtm <= 100 means the entire mating line fits within the +# window, so the layer uses the flat DTM. The bound is inclusive because the +# game ends at checkmate, before a draw can be claimed. +# +# Busted: rule50 + dtz > 100 means no surviving line lets the winner reset the +# count in time, so the layer is DRAW. This one-directional bound assigns no +# distance to positions that survive. It uses the same inequality encoded by a +# cell's transition to DRAW (h = 102 - dtz), based on the DTZ returned by the +# flat probe. +# +# An unpriced distance reads as 0, so the pin needs a has_dtm at the call site +# while the bust cannot fire on one. Which case left DTZ unread does not matter +# there: the cursed class a read declines to price, or a derive whose minimax +# went unpinned. +def dtm50_layer_pinned_by_dtm(wdl: int, dtm: int, rule50: int) -> bool: + if wdl not in (WIN, LOSE): + return False + return rule50 + dtm <= MAX_NON_CURSED_DTZ + + +def layer_busted_by_dtz(wdl: int, dtz: int, rule50: int) -> bool: + if wdl not in (WIN, LOSE): + return False + return rule50 + dtz > MAX_NON_CURSED_DTZ + + +class _DTM50Result(NamedTuple): + """A DTM50 answer, with DTZ riding along: off the cell's flip when read, off + its own minimax when derived.""" + + wdl: int + dtm: int + has_dtz: bool = False + dtz: int = 0 + + +class _SkippedChildren: + """Children a derive could not price. A skip only unpins the minimax if the + class it could have offered outranks the kept best; an unknown class bounds + nothing and unpins outright. Every deriver runs this -- a best over a partial + move set is a wrong answer, not a partial one.""" + + __slots__ = ("_best_rank", "_blind") + + def __init__(self) -> None: + self._best_rank = -1 + self._blind = False + + def of_class(self, my_wdl: int) -> None: + self._best_rank = max(self._best_rank, wdl_rank(my_wdl)) + + def of_dtz_class(self, my_wdl: int) -> None: + # Same, for a DTZ minimax: the child's unknown distance may push past + # 100, which lifts a LOSE to BLESSED_LOSS -- a higher rank, so bound by + # that. + self.of_class(dtz_lift(my_wdl)) + + def unknown(self) -> None: + self._blind = True + + def unpin(self, best_wdl: int) -> bool: + return self._blind or self._best_rank >= wdl_rank(best_wdl) + + +class _DTZMinimax: + """The DTZ half of a derive: zeroing distance ranks moves its own way, so + it minimaxes beside the mate distance over the same children. Fed per child; + `finish` writes the field a read cell gets from its flip.""" + + __slots__ = ("_have", "_best_wdl", "_best_dtz", "_skipped") + + def __init__(self) -> None: + self._have = False + self._best_wdl = LOSE + self._best_dtz = 0 + self._skipped = _SkippedChildren() + + def zeroing_child(self, child_wdl: int) -> None: + """A zeroing move ends the count at this ply, whatever the child holds.""" + self._offer(child_wdl, 1) + + def quiet_child(self, child_wdl: int, child: _DTM50Result) -> None: + if child.has_dtz: + self._offer(child_wdl, 1 + child.dtz) + else: + self._skipped.of_dtz_class(invert_wdl(child_wdl)) + + def unwalked(self) -> None: + """Priced against the clock instead of visited: nothing to take from it.""" + self._skipped.unknown() + + def finish(self, wdl: int, dtm: int, any_legal: bool) -> _DTM50Result: + if not any_legal: # mate or stalemate: terminal, zeroing distance 0 + return _DTM50Result(wdl, dtm, True, 0) + if not self._have or self._skipped.unpin(self._best_wdl): + return _DTM50Result(wdl, dtm) + return _DTM50Result(wdl, dtm, True, + 0 if self._best_wdl == DRAW else self._best_dtz) + + def _offer(self, child_wdl: int, dtz: int) -> None: + my_wdl = invert_wdl(child_wdl) + if dtz > MAX_NON_CURSED_DTZ: + if my_wdl == WIN: + my_wdl = CURSED_WIN + elif my_wdl == LOSE: + my_wdl = BLESSED_LOSS + if not self._have or prefer_new(my_wdl, dtz, self._best_wdl, self._best_dtz): + self._best_wdl, self._best_dtz, self._have = my_wdl, dtz, True + + +class _UnboundedFold: + """DTC's fold for the unbounded row it carries: a zeroing move ends the + count one ply out, a quiet move waits one longer, and the class says whether + the shortest or the longest of them stands. One child unable to say leaves + the row unanswered, since a best over part of the moves is a wrong answer + rather than a partial one.""" + + __slots__ = ("_winning", "_have", "_blind", "_best") + + def __init__(self, mover_class: int) -> None: + self._winning = mover_class in (WIN, CURSED_WIN) + self._have = False + self._blind = False + self._best = 0 + + def zeroing_child(self) -> None: + self._offer(1) + + def quiet_child(self, child_dtz: int) -> None: + if child_dtz == DTC_DRAWN: + self._blind = True + else: + self._offer(child_dtz + 1) + + def value(self) -> int: + return self._best if (self._have and not self._blind) else DTC_DRAWN + + def _offer(self, v: int) -> None: + if not self._have or (v < self._best if self._winning else v > self._best): + self._best = v + self._have = True + + +def _ep_capture_moves( + board: chess.Board, ep_square: int +) -> Tuple[chess.Board, List[chess.Move]]: + """Legal en-passant captures available given `ep_square`, using python-chess + to execute the capture correctly. Returns (ep_board, [moves]).""" + bcopy = board.copy(stack=False) + bcopy.ep_square = ep_square + return bcopy, [m for m in bcopy.legal_moves if bcopy.is_en_passant(m)] + + +class ProbeResult: + """Outcome of :meth:`Tablebase.probe`. + + ``status`` is ``"ok"`` or ``"tb_not_found"``. ``wdl``, ``dtc_wdl`` and + ``dtm50_wdl`` are :data:`WIN`..:data:`LOSE` codes; ``dtz``/``dtc``/``dtm``/ + ``dtm50`` are unsigned ply counts whose sign is given by the corresponding + WDL class. + + DTC answers a pair at the caller's ``rule50``: ``dtc_order`` is the pushes + the winner still owes and ``dtc`` the plies to the next zeroing move on the + line that owes them. No budget fitting the clock is a 50MR draw, which + ``dtc_wdl`` reports rather than withholding, so ``has_dtc`` says the metric + is available and ``dtc_wdl`` what it found. + """ + + __slots__ = ("status", "wdl", "has_dtz", "dtz", "has_dtc", "dtc_wdl", + "dtc_order", "dtc", "has_dtm", "dtm", + "has_dtm50", "dtm50_wdl", "dtm50") + + status: str + wdl: int + has_dtz: bool + dtz: int + has_dtc: bool + dtc_wdl: int + dtc_order: int + dtc: int + has_dtm: bool + dtm: int + has_dtm50: bool + dtm50_wdl: int + dtm50: int + + def __init__(self) -> None: + self.status = "tb_not_found" + self.wdl = ILLEGAL + self.has_dtz = False + self.dtz = 0 + self.has_dtc = False + self.dtc_wdl = ILLEGAL + self.dtc_order = 0 + self.dtc = 0 + self.has_dtm = False + self.dtm = 0 + self.has_dtm50 = False + self.dtm50_wdl = ILLEGAL + self.dtm50 = 0 + + def __repr__(self) -> str: + if self.status != "ok": + return f"" + s = f"" + + +class MissingTableError(KeyError): + """Raised when no table is available for the queried material.""" + + +# Signed WDL convention matching chess.syzygy: +2 win, +1 cursed win, +# 0 draw, -1 blessed loss, -2 loss. +_WDL_SIGNED = {WIN: 2, CURSED_WIN: 1, DRAW: 0, BLESSED_LOSS: -1, LOSE: -2} + + +def _signed(magnitude: int, wdl: int) -> int: + if wdl in (WIN, CURSED_WIN): + return magnitude + if wdl in (LOSE, BLESSED_LOSS): + return -magnitude + return 0 + + +# =========================================================================== +# Tablebase +# =========================================================================== + +class Tablebase: + """Probe a directory tree of chesstb tables. + + `directory` may contain ``wdl/``, ``dtz/``, ``dtc/``, ``dtm/`` and + ``dtm50/`` subdirectories (the generator's layout) or the table files directly. Use + :func:`open_tablebase`. Probing is read-only and the result of each query is + derived from the canonical orientation of the board's material. + + One instance is meant to be shared by all threads that probe: tables and + decoded blocks are then opened and decoded once for everyone, rather than + once per thread. See the module docstring for what is locked. + """ + + #: The classes ``_open_wdl`` and its siblings instantiate. A transport that + #: replaces :meth:`_TableFile._open_source` names its subclasses here and + #: overrides :meth:`_find`, rather than reimplementing the + #: look-once-then-cache dance they share. + WDL_FILE: Type[WDLFile] = WDLFile + DTZ_FILE: Type[DTZFile] = DTZFile + DTC_FILE: Type[DTCFile] = DTCFile + DTM_FILE: Type[DTMFile] = DTMFile + DTM50_FILE: Type[DTM50File] = DTM50File + + #: Search-directory kinds, in the order a directory tree names them. + KINDS = ("wdl", "dtz", "dtc", "dtm", "dtm50") + + def __init__(self, directory: str, *, block_cache_bytes: int = DEFAULT_BLOCK_CACHE_BYTES): + self.dirs: Dict[str, List[str]] = {kind: [] for kind in self.KINDS} + self._wdl_cache: Dict[Tuple[int, bool], Optional[WDLFile]] = {} + self._dtz_cache: Dict[Tuple[int, bool], Optional[DTZFile]] = {} + self._dtc_cache: Dict[Tuple[int, bool], Optional[DTCFile]] = {} + self._dtm_cache: Dict[Tuple[int, bool], Optional[DTMFile]] = {} + self._dtm50_cache: Dict[Tuple[int, bool], Optional[DTM50File]] = {} + # One lock per kind, guarding that kind's open cache and the `dirs` + # entry its _find walks. Held only across a first open of some + # material, never across a probe. Lock order, where more than one is + # taken: the KINDS order (add_directory and close, both rare). + self._open_locks: Dict[str, threading.Lock] = { + kind: threading.Lock() for kind in self.KINDS} + # Lets close() wait for in-flight probes before unmapping, mirroring + # chess.syzygy.Table's read_count/read_condition pair. + self._read_condition = threading.Condition() + self._read_count = 0 + # Shared LRU so decoded blocks are reclaimed automatically once the + # budget is exceeded, rather than growing for the lifetime of the probe. + self._block_cache = _BlockCache(block_cache_bytes) + self.add_directory(directory) + + def add_directory(self, directory: str) -> None: + """Add another search directory (and its kind subdirectories). + + Safe to call while other threads probe: each kind's list is extended + under the lock its :meth:`_find` holds, so no probe sees a partly + extended search path. Tables already resolved are not re-resolved, + exactly as before -- a directory added late only affects materials not + yet looked up. + """ + for kind in self.KINDS: + with self._open_locks[kind]: + self.dirs[kind].append(os.path.join(directory, kind)) + self.dirs[kind].append(directory) + + def close(self) -> None: + """Drop all cached decoded blocks and unmap all open tables. + + Waits for probes running on other threads to finish first: unmapping + under a live probe would pull the table's memory out from under it. + Probes arriving afterwards are not turned away (as in + :mod:`chess.syzygy`) -- they simply reopen what they need -- so a + continuous stream of them can keep this waiting for a while. + """ + with self._read_condition: + while self._read_count > 0: + self._read_condition.wait() + # Readers are drained and, holding the condition, none can enter, + # so the open caches are ours: a concurrent _open_* would have had + # to register as a reader first. + self._block_cache.clear() + with self._open_locks["wdl"]: + self._close_all(self._wdl_cache) + with self._open_locks["dtz"]: + self._close_all(self._dtz_cache) + with self._open_locks["dtc"]: + self._close_all(self._dtc_cache) + with self._open_locks["dtm"]: + self._close_all(self._dtm_cache) + with self._open_locks["dtm50"]: + self._close_all(self._dtm50_cache) + + @staticmethod + def _close_all(cache: Dict[Tuple[int, bool], Optional[_TableFileT]]) -> None: + """Unmap every open table in one kind's cache and forget them all.""" + while cache: + _, table = cache.popitem() + if table is not None: + table.close() + + def __enter__(self) -> "Tablebase": + return self + + def __exit__(self, *exc: Any) -> None: + self.close() + + # --- table file resolution / caching --- + def _find(self, kind: str, name: str, ext: str) -> Optional[TableSource]: + """Resolve one table to something :meth:`_TableFile._open_source` opens, + or ``None`` if this tablebase has no such table.""" + for d in self.dirs[kind]: + p = os.path.join(d, name + ext) + if os.path.exists(p): + return p + return None + + # A cached `None` means "looked, no such table", so these subscript rather + # than `.get()`: absence of the key is the only thing that means "not looked + # up yet". The hit path runs outside the lock; only a first open takes it, + # and re-checks once inside in case another thread just did the work. + def _open_wdl(self, cfg: PieceConfig) -> Optional[WDLFile]: + k = cfg.cache_key + try: + return self._wdl_cache[k] + except KeyError: + pass + with self._open_locks["wdl"]: + try: + return self._wdl_cache[k] + except KeyError: + pass + p = self._find("wdl", cfg.name(), self.WDL_FILE.EXT) + # Published only once fully constructed, so no other thread can + # reach a table whose header is still being parsed. + table = self.WDL_FILE(cfg, p, self._block_cache) if p is not None else None + self._wdl_cache[k] = table + return table + + def _open_dtz(self, cfg: PieceConfig) -> Optional[DTZFile]: + k = cfg.cache_key + try: + return self._dtz_cache[k] + except KeyError: + pass + with self._open_locks["dtz"]: + try: + return self._dtz_cache[k] + except KeyError: + pass + p = self._find("dtz", cfg.name(), self.DTZ_FILE.EXT) + table = self.DTZ_FILE(cfg, p, self._block_cache) if p is not None else None + self._dtz_cache[k] = table + return table + + def _open_dtm(self, cfg: PieceConfig) -> Optional[DTMFile]: + k = cfg.cache_key + try: + return self._dtm_cache[k] + except KeyError: + pass + with self._open_locks["dtm"]: + try: + return self._dtm_cache[k] + except KeyError: + pass + p = self._find("dtm", cfg.name(), self.DTM_FILE.EXT) + table = self.DTM_FILE(cfg, p, self._block_cache) if p is not None else None + self._dtm_cache[k] = table + return table + + def _open_dtc(self, cfg: PieceConfig) -> Optional[DTCFile]: + k = cfg.cache_key + try: + return self._dtc_cache[k] + except KeyError: + pass + with self._open_locks["dtc"]: + try: + return self._dtc_cache[k] + except KeyError: + pass + p = self._find("dtc", cfg.name(), self.DTC_FILE.EXT) + table = self.DTC_FILE(cfg, p, self._block_cache) if p is not None else None + self._dtc_cache[k] = table + return table + + def _open_dtm50(self, cfg: PieceConfig) -> Optional[DTM50File]: + k = cfg.cache_key + try: + return self._dtm50_cache[k] + except KeyError: + pass + with self._open_locks["dtm50"]: + try: + return self._dtm50_cache[k] + except KeyError: + pass + p = self._find("dtm50", cfg.name(), self.DTM50_FILE.EXT) + table = self.DTM50_FILE(cfg, p, self._block_cache) if p is not None else None + self._dtm50_cache[k] = table + return table + + def _has_any_table(self, cfg: PieceConfig) -> bool: + return self._open_wdl(cfg) is not None + + # --- child construction for the derive / overlay paths --- + def _make_child(self, parent: chess.Board, move: chess.Move + ) -> Tuple[PieceConfig, chess.Board, Optional[int], bool]: + zeroing = parent.is_zeroing(move) + child = parent.copy(stack=False) + child.push(move) # also sets child.ep_square on a double push + + # Prefer the child's opposing-pair table when one is on disk: a move that + # keeps the pair (any non-capture, including a free-pawn push) stays in a + # 'p' material that the board-derived config below would miss -- it sees + # the pair pawns as ordinary free pawns. Falls back to the full physical + # material for captures/promotions and when no 'p' table is present. + paired = pair_config_from_board(child) + if paired is not None and self._has_any_table(paired[0]): + cfg, mirrored = paired + else: + cfg, mirrored = piece_config_from_board(child) + + # Mirror before lifting the ep off, so it can't desync from the position. + if mirrored: + child = mirror_for_canonical(child) + ep = child.ep_square + child.ep_square = None + return cfg, child, ep, zeroing + + # --- WDL --- + def _relax_bound_wdl(self, board: chess.Board, depth: int) -> int: + """Best class a move out of this material reaches. Captures and + promotions only: their children live in a sub-table, so this never + re-enters the frame it is resolving.""" + if depth >= MAX_DERIVE_DEPTH: + return ILLEGAL + b = _internal_board(board) + best = ILLEGAL + for m in b.legal_moves: + # A stored cell carries no ep rights; _probe_wdl_impl folds an ep + # capture in on top of whatever this returns. + if m.promotion is None and b.piece_at(m.to_square) is None: + continue + cfg_c, cboard, cep, _zeroing = self._make_child(b, m) + cw = DRAW if cfg_c.is_bare_kings else self._probe_wdl_impl( + cfg_c, cboard, cep, depth + 1) + if cw == ILLEGAL: + continue + mine = invert_wdl(cw) + if wdl_rank(mine) > wdl_rank(best): + best = mine + return best + + def _raise_by_bound(self, w: WDLFile, frame: int, stored: int, + board: chess.Board, depth: int) -> int: + """A relaxed frame stores no better than the truth, and stores strictly + worse only where the bound above reaches it. ILLEGAL stays ILLEGAL.""" + if not w.is_relaxed[frame] or stored == ILLEGAL: + return stored + bound = self._relax_bound_wdl(board, depth) + return bound if wdl_rank(bound) > wdl_rank(stored) else stored + + def _read_wdl_stored(self, w: Optional[WDLFile], board: chess.Board, + depth: int) -> int: + if w is None: + return 7 # WDL_Stored::ILLEGAL + color = CPP_WHITE if board.turn == WHITE else CPP_BLACK + s = w.read(color, board) + if not w.is_relaxed[color] or s == 7: + return s + # Relaxation never touches a marker cell. + if s == 6 or s == 5: # BOUNDARY_WIN / BOUNDARY_LOSS + return s + bound = self._relax_bound_wdl(board, depth) + return bound if wdl_rank(bound) > wdl_rank(s) else s + + def _probe_wdl_internal(self, w: Optional[WDLFile], cfg: PieceConfig, + board: chess.Board, depth: int) -> int: + if w is None: + return ILLEGAL + color = CPP_WHITE if board.turn == WHITE else CPP_BLACK + if w.is_dropped[color]: + if not is_symmetric_material(cfg): + return self._derive_wdl(board, depth) + mp = mirror_for_canonical(board) + mc = CPP_WHITE if mp.turn == WHITE else CPP_BLACK + return self._raise_by_bound(w, mc, wdl_from_storage(w.read(mc, mp)), + board, depth) + return self._raise_by_bound(w, color, wdl_from_storage(w.read(color, board)), + board, depth) + + def _derive_wdl(self, board: chess.Board, depth: int) -> int: + if depth >= MAX_DERIVE_DEPTH: + return ILLEGAL + b = _internal_board(board) + any_legal = have = False + best = LOSE + skipped = _SkippedChildren() + for m in b.legal_moves: + any_legal = True + cfg_c, cboard, cep, zeroing = self._make_child(b, m) + if cfg_c.is_bare_kings: + mw = DRAW + elif zeroing: + cw = self._probe_wdl_impl(cfg_c, cboard, cep, depth + 1) + # No entry, so no class either: nothing bounds this skip. + if cw == ILLEGAL: + skipped.unknown() + continue + mw = invert_wdl(cw) + else: + cs = self._read_wdl_stored(self._open_wdl(cfg_c), cboard, depth + 1) + if cs == 7: + skipped.unknown() + continue + mw = invert_stored(cs) + if wdl_rank(mw) > wdl_rank(best): + best = mw + have = True + if not any_legal: + return LOSE if b.is_check() else DRAW + if not have or skipped.unpin(best): + return ILLEGAL + return best + + def _probe_wdl_impl(self, cfg: PieceConfig, board: chess.Board, + ep_square: Optional[int], depth: int) -> int: + best = self._probe_wdl_internal(self._open_wdl(cfg), cfg, board, depth) + if best == ILLEGAL or ep_square is None: + return best + bcopy, eps = _ep_capture_moves(board, ep_square) + for m in eps: + cfg_c, cboard, _cep, _ = self._make_child(bcopy, m) + cw = DRAW if cfg_c.is_bare_kings else self._probe_wdl_internal( + self._open_wdl(cfg_c), cfg_c, cboard, depth + 1) + if cw == ILLEGAL: + continue + mine = invert_wdl(cw) + if wdl_rank(mine) > wdl_rank(best): + best = mine + return best + + # --- DTZ --- + def _probe_dtz_internal(self, d: Optional[DTZFile], cfg: PieceConfig, + board: chess.Board, wdl: int, depth: int) -> Optional[int]: + if d is None: + return None + color, mp, readable = locate_frame(d, cfg, board, wdl) + if not readable: + return self._derive_dtz(board, wdl, depth) + return d.read(color, board if mp is None else mp, wdl) + + def _derive_dtz(self, board: chess.Board, wdl: int, depth: int) -> Optional[int]: + if depth >= MAX_DERIVE_DEPTH: + return None + # A DRAW zeroes nowhere, and every caller prices it without a derive. + assert wdl != DRAW + b = _internal_board(board) + any_legal = have = False + best_wdl, best_dtz = LOSE, 0 + skipped = _SkippedChildren() + for m in b.legal_moves: + any_legal = True + cfg_c, cboard, cep, zeroing = self._make_child(b, m) + if cfg_c.is_bare_kings: + cw, my_dtz = DRAW, 1 + elif zeroing: + if child_class_is_forced(wdl): + cw = WIN + else: + cw = self._probe_wdl_impl(cfg_c, cboard, cep, depth + 1) + if cw == ILLEGAL: + skipped.unknown() + continue + my_dtz = 1 + else: + if child_class_is_forced(wdl): + cw = WIN + else: + cw = self._probe_wdl_internal(self._open_wdl(cfg_c), cfg_c, cboard, depth + 1) + if cw == ILLEGAL: + skipped.unknown() + continue + # Outranked: it loses the minimax however near it zeroes, + # so skipping its distance cannot affect the result. + if below_pinned_class(wdl, dtz_lift(invert_wdl(cw))): + continue + child_dtz = self._probe_dtz_internal(self._open_dtz(cfg_c), cfg_c, cboard, cw, depth + 1) + if child_dtz is None: + skipped.of_dtz_class(invert_wdl(cw)) + continue + my_dtz = 1 + child_dtz + my_wdl = invert_wdl(cw) + if my_dtz > MAX_NON_CURSED_DTZ: + if my_wdl == WIN: + my_wdl = CURSED_WIN + elif my_wdl == LOSE: + my_wdl = BLESSED_LOSS + if not have or prefer_new(my_wdl, my_dtz, best_wdl, best_dtz): + best_wdl, best_dtz, have = my_wdl, my_dtz, True + if not any_legal: + return 0 + if not have or skipped.unpin(best_wdl): + return None + if best_wdl == DRAW: + return 0 + return best_dtz + + # --- DTC (whose unbounded row answers DTZ too) --- + def _probe_dtc_internal(self, c: Optional[DTCFile], cfg: PieceConfig, + board: chess.Board, wdl: int, rule50: int, + depth: int) -> Optional[_DTCCell]: + if c is None: + return None + color, mp, readable = locate_frame(c, cfg, board, wdl) + if not readable: + return self._derive_dtc(board, wdl, rule50, depth) + return c.read(color, board if mp is None else mp, wdl, rule50) + + def _read_dtc_curve(self, c: Optional[DTCFile], cfg: PieceConfig, + board: chess.Board, wdl: int) -> Optional[List[int]]: + """The whole record of a position this table does hold, for a derive to + minimax over. The child keeps the physical material, so the pack its own + config names answers -- this file, or the 'p' table that re-indexes an + opposing pair.""" + if c is None: + return None + color, mp, readable = locate_frame(c, cfg, board, wdl) + if not readable: + return None + return c.read_curve(color, board if mp is None else mp, wdl) + + @staticmethod + def _dtc_move_kind(board: chess.Board, move: chess.Move) -> Tuple[bool, bool]: + """(conversion, push). A capture or promotion converts; a pawn move that + does neither spends one of the budget instead.""" + conversion = (move.promotion is not None + or board.piece_at(move.to_square) is not None) + push = not conversion and board.piece_type_at(move.from_square) == chess.PAWN + return conversion, push + + def _child_dtc_cell(self, cfg_c: PieceConfig, cboard: chess.Board, + cep: Optional[int], cw: int, child_rule50: int, + depth: int) -> Optional[_DTCCell]: + """The child's DTC answer with its ep rights folded in. A double push + leaves the opponent a capture the child's own record cannot express, and + the overlay is where that is priced, so a child carrying ep rights + answers through it.""" + if cep is None: + return self._probe_dtc_internal(self._open_dtc(cfg_c), cfg_c, cboard, + cw, child_rule50, depth) + cr = self._probe_impl(cfg_c, cboard, child_rule50, cep, depth) + if not cr.has_dtc or not cr.has_dtz: + return None + if cr.dtc_wdl == DRAW: + return _DTCCell(dtz=cr.dtz) + return _DTCCell(cr.dtc_order, cr.dtc, cr.dtz) + + def _ep_conversion_wins(self, cboard: chess.Board, cep: int, + depth: int) -> Optional[bool]: + """Whether an ep capture out of the child wins for the side holding it. + A capture is a conversion, so it owes no push and lands one ply out: the + cheapest answer DTC has, which every budget affords. Nothing answers for + a capture whose sub-table is absent, and the child's whole value turns on + it, so that is a third answer rather than a no -- though a capture that + does win settles it whatever the other one says. `depth` is the child's, + so a grandchild is one deeper, as under the overlay.""" + bcopy, eps = _ep_capture_moves(cboard, cep) + unknown = False + for mv in eps: + cfg_g, gboard, gep, _zeroing = self._make_child(bcopy, mv) + gw = DRAW if cfg_g.is_bare_kings else self._probe_wdl_impl( + cfg_g, gboard, gep, depth + 1) + if gw == ILLEGAL: + unknown = True + continue + if fold_50mr_wdl(invert_wdl(gw)) == WIN: + return True + return None if unknown else False + + def _derive_dtc(self, board: chess.Board, wdl: int, rule50: int, + depth: int) -> Optional[_DTCCell]: + """DTC by one-ply minimax, for a frame the file does not hold. It never + leaves this table: a push and a quiet move both keep the material, so + they land in the frame that was kept, and a conversion is terminal at + value 1 under its own WDL, as inside a layer's retro. A winner's push + spends one of the budget and reads one budget lower; nothing else moves + it. The same walk folds the unbounded row the pack carries, so this + derives everything the table serves. + + The class, known before any child is read, says how much of one to read: + a win takes the cheapest budget any move offers, which is the pair a + read returns at the clock behind that move, while a loss settles on the + budget the most stubborn defence forces, which no single pair names.""" + if depth >= MAX_DERIVE_DEPTH: + return None + if wdl == WIN: + return self._derive_dtc_win(board, rule50, depth) + if wdl == LOSE: + return self._derive_dtc_loss(board, rule50, depth) + return self._derive_dtc_cursed(board, wdl, depth) + + def _derive_dtc_win(self, board: chess.Board, rule50: int, + depth: int) -> Optional[_DTCCell]: + b = _internal_board(board) + have = False + best_order = best_value = 0 + budget_plies = dtc_budget_plies(rule50) + unbounded = _UnboundedFold(WIN) + for m in b.legal_moves: + conversion, push = self._dtc_move_kind(b, m) + cfg_c, cboard, cep, _zeroing = self._make_child(b, m) + if cfg_c.is_bare_kings: + continue + # A double push hands the opponent an ep capture, which the class has + # to carry: the generator prices the same reply into the push it + # evaluates. + cw = self._probe_wdl_impl(cfg_c, cboard, cep, depth + 1) + if cw == ILLEGAL: + return None + # Only a clean loss for the other side carries a clean win, and + # anything else leaves a budget drawn however the rest are priced. + if cw != LOSE: + continue + order, value = 0, 1 + if conversion: + unbounded.zeroing_child() + else: + # A push zeroes the clock and spends one of the budget; a quiet + # move spends a ply of the clock and none of the budget. + child_rule50 = 0 if push else ( + IGNORE_50MR if rule50 == IGNORE_50MR else rule50 + 1) + cell = self._child_dtc_cell(cfg_c, cboard, cep, cw, + child_rule50, depth + 1) + if cell is None: + return None + # The row is clock-free, so it takes this child whatever the + # clock made of its budgets. + if push: + unbounded.zeroing_child() + else: + unbounded.quiet_child(cell.dtz) + if not cell.priced: + continue # that clock has taken this line + order = cell.order + (1 if push else 0) + value = 1 if push else cell.value + 1 + # The clock this position holds: a child behind a zeroing move + # answered against a fresh one. + if value > budget_plies: + continue + if not have or order < best_order or (order == best_order + and value < best_value): + best_order, best_value, have = order, value, True + dtz = unbounded.value() + return _DTCCell(best_order, best_value, dtz) if have else _DTCCell(dtz=dtz) + + def _derive_dtc_loss(self, board: chess.Board, rule50: int, + depth: int) -> Optional[_DTCCell]: + """Every move prices a loss, so this one walks the children's whole + records: the budget it settles on is the highest any defence needs, and + the value there is the longest wait among them, which a child's own + cheapest budget does not report. Every child is a win for the other + side, so none needs its class read.""" + b = _internal_board(board) + worst = [0] * DTC_PACK_LAYERS + drawn = [False] * DTC_PACK_LAYERS + any_legal = False + unbounded = _UnboundedFold(LOSE) + + # Past the band the budget answers nothing, the ply ceiling being the + # layer's own draw. + def raise_at(k: int, val: int) -> None: + if val > MAX_NON_CURSED_DTZ: + drawn[k] = True + elif val > worst[k]: + worst[k] = val + + for m in b.legal_moves: + any_legal = True + conversion, push = self._dtc_move_kind(b, m) + cfg_c, cboard, cep, _zeroing = self._make_child(b, m) + if conversion or cfg_c.is_bare_kings: + # A conversion ends the count here, whatever budget the rest + # settle on. + for k in range(DTC_PACK_LAYERS): + raise_at(k, 1) + unbounded.zeroing_child() + continue + # A double push leaves the winner an ep capture the child's own + # record cannot express. It converts, so where it wins it settles + # every budget one ply out and this defence gains nothing by the + # push. + if cep is not None: + ep_wins = self._ep_conversion_wins(cboard, cep, depth + 1) + if ep_wins is None: + return None + if ep_wins: + assert push # only a double push leaves ep rights behind + for k in range(DTC_PACK_LAYERS): + raise_at(k, 1) + unbounded.zeroing_child() + continue + + curve = self._read_dtc_curve(self._open_dtc(cfg_c), cfg_c, cboard, WIN) + if curve is None: + return None + # A push is the zeroing move itself and spends none of the winner's + # budget; a quiet move waits one ply more at the same budget. + if push: + unbounded.zeroing_child() + else: + unbounded.quiet_child(curve[DTC_BUDGET_LAYERS]) + for k in range(DTC_PACK_LAYERS): + v = curve[k] + if v == DTC_DRAWN: + drawn[k] = True + else: + raise_at(k, 1 if push else v + 1) + + if not any_legal: + return _DTCCell(0, 0, 0) # mate: converted, nothing owed + dtz = unbounded.value() + budget_plies = dtc_budget_plies(rule50) + for k in range(DTC_PACK_LAYERS): + if not drawn[k] and worst[k] <= budget_plies: + return _DTCCell(k, worst[k], dtz) + return _DTCCell(dtz=dtz) + + def _derive_dtc_cursed(self, board: chess.Board, wdl: int, + depth: int) -> Optional[_DTCCell]: + """A cursed class has no budget to look for, none of them settling it, + so the unbounded row is the whole of what this derives.""" + b = _internal_board(board) + winning = (wdl == CURSED_WIN) + unbounded = _UnboundedFold(wdl) + any_legal = False + for m in b.legal_moves: + any_legal = True + conversion, push = self._dtc_move_kind(b, m) + cfg_c, cboard, cep, _zeroing = self._make_child(b, m) + if cfg_c.is_bare_kings: + if winning: + continue # a draw no cursed win would take + return None + cw = self._probe_wdl_impl(cfg_c, cboard, cep, depth + 1) + if cw == ILLEGAL: + return None + # A win runs over the moves that win and nothing else can beat them; + # a loss is priced by every move it has. + if winning and cw not in (LOSE, BLESSED_LOSS): + continue + if conversion or push: + unbounded.zeroing_child() + continue + cell = self._probe_dtc_internal(self._open_dtc(cfg_c), cfg_c, cboard, + cw, IGNORE_50MR, depth + 1) + if cell is None: + return None + unbounded.quiet_child(cell.dtz) + return _DTCCell(dtz=unbounded.value()) if any_legal else _DTCCell() + + # --- DTM (the standalone `dtm/` table) --- + def _probe_dtm_internal(self, d: Optional[DTMFile], cfg: PieceConfig, + board: chess.Board, wdl: int, depth: int) -> Optional[int]: + if d is None: + return None + color, mp, readable = locate_frame(d, cfg, board, wdl) + if not readable: + return self._derive_dtm(board, wdl, depth) + return d.read(color, board if mp is None else mp, wdl) + + def _derive_dtm(self, board: chess.Board, wdl: int, depth: int) -> Optional[int]: + """DTM by one-ply minimax for a dropped frame of the DTM table. 50MR-free, + so a cursed win mates like any other -- which is what folding the pin + expresses, and why no child needs its clock tracked.""" + if depth >= MAX_DERIVE_DEPTH: + return None + pinned = fold_dtm_wdl(wdl) + # A DRAW mates nowhere, and every caller prices it without a derive. + assert pinned != DRAW + b = _internal_board(board) + any_legal = have = False + best_wdl, best_dtm = LOSE, 0 + skipped = _SkippedChildren() + for mv in b.legal_moves: + any_legal = True + cfg_c, cboard, cep, _zeroing = self._make_child(b, mv) + if cfg_c.is_bare_kings: + cw, cd = DRAW, 0 + elif cep is not None: + cr = self._probe_impl(cfg_c, cboard, IGNORE_50MR, cep, depth + 1) + if cr.status != "ok" or cr.wdl == ILLEGAL: + skipped.unknown() + continue + if not cr.has_dtm: + skipped.of_class(invert_wdl(fold_dtm_wdl(cr.wdl))) + continue + cw, cd = cr.wdl, cr.dtm + else: + if child_class_is_forced(pinned): + cw = WIN + else: + cw = self._probe_wdl_internal(self._open_wdl(cfg_c), cfg_c, + cboard, depth + 1) + if cw == ILLEGAL: + skipped.unknown() + continue + # Outranked: no mate distance it holds can win the minimax, + # and a skip of it would bound nothing -- of_class ranks it + # this far. + if below_pinned_class(pinned, invert_wdl(fold_dtm_wdl(cw))): + continue + # The `dtm/` table alone: each derive stays on the file it + # rebuilds. + child_dtm = self._probe_dtm_internal(self._open_dtm(cfg_c), cfg_c, + cboard, cw, depth + 1) + if child_dtm is None: + skipped.of_class(invert_wdl(fold_dtm_wdl(cw))) + continue + cd = child_dtm + my_wdl = invert_wdl(fold_dtm_wdl(cw)) + my_dtm = 1 + cd + if not have or prefer_new(my_wdl, my_dtm, best_wdl, best_dtm): + best_wdl, best_dtm, have = my_wdl, my_dtm, True + if not any_legal: + return 0 + if not have or skipped.unpin(best_wdl): + return None + if best_wdl in (WIN, LOSE): + return best_dtm + return 0 + + # --- DTM50 (whose flat layer answers DTM too) --- + def _probe_dtm50_internal(self, m: Optional[DTM50File], cfg: PieceConfig, + board: chess.Board, wdl: int, rule50: int, + depth: int) -> _DTM50Result: + flat = (rule50 == IGNORE_50MR) + if not flat and rule50 >= DTM50_HMC_COUNT: + return _DTM50Result(DRAW, 0) + if m is None: + return _DTM50Result(ILLEGAL, 0) + + # A cell carries the DRAW flip that pins DTZ; a derived one prices it + # itself. Only layer 0 bothers: zeroing is clock-free, so _probe_impl + # reads DTZ off the flat probe and a layered one would only recompute + # what it drops. + def from_cell(cell: Tuple[int, int]) -> _DTM50Result: + value, flip = cell + if value == DTM50_DRAWN: + cls, value = DRAW, 0 + else: + cls = wdl if flat else fold_50mr_wdl(wdl) + if not flat: + return _DTM50Result(cls, value) + dtz = dtz_from_draw_flip(flip, wdl, board) + if dtz is None: + return _DTM50Result(cls, value) + return _DTM50Result(cls, value, True, dtz) + + color, mp, readable = locate_frame(m, cfg, board, wdl) + if not readable: + return (self._derive_dtm50_flat(board, wdl, depth) if flat + else self._derive_dtm50(board, wdl, rule50, depth)) + return from_cell(m.read(color, board if mp is None else mp, wdl, rule50)) + + def _derive_dtm50_flat(self, board: chess.Board, wdl: int, + depth: int) -> _DTM50Result: + """Reconstruct the pack's dropped layer-0 frame with an unbounded, + one-ply DTM minimax. The result carries DTZ, just as a cell read derives + it from the flip.""" + if depth >= MAX_DERIVE_DEPTH: + return _DTM50Result(ILLEGAL, 0) + # Layer 0 is 50MR-free, so a cursed win still mates; only a true DRAW + # stays one, and _probe_impl prices that without reaching a derive. + pinned = fold_dtm_wdl(wdl) + assert pinned != DRAW + b = _internal_board(board) + any_legal = have = False + best_wdl, best_dtm = LOSE, 0 + dtz = _DTZMinimax() + skipped = _SkippedChildren() + for mv in b.legal_moves: + any_legal = True + cfg_c, cboard, cep, zeroing = self._make_child(b, mv) + if cfg_c.is_bare_kings: + cw, cd = DRAW, 0 + dtz.zeroing_child(DRAW) + elif cep is not None: + cr = self._probe_impl(cfg_c, cboard, IGNORE_50MR, cep, depth + 1) + if cr.status != "ok" or cr.wdl == ILLEGAL: + skipped.unknown() + dtz.unwalked() + continue + dtz.zeroing_child(cr.wdl) # a double push zeroes the clock + if not cr.has_dtm: + skipped.of_class(invert_wdl(fold_dtm_wdl(cr.wdl))) + continue + cw, cd = cr.wdl, cr.dtm + else: + # The pack alone: a child shipping only `dtm/` goes unpriced here. + if child_class_is_forced(wdl): + # Keep this exact rather than folded: the DTZ minimax running + # beside the mate one still needs all five WDL classes. + raw_wdl = WIN + else: + raw_wdl = self._probe_wdl_internal(self._open_wdl(cfg_c), cfg_c, + cboard, depth + 1) + if raw_wdl == ILLEGAL: + skipped.unknown() + dtz.unwalked() + continue + # Outranked for the mate distance, and so for zeroing too: + # neither accumulator can take it, so the cell goes unread. + if below_pinned_class(pinned, invert_wdl(fold_dtm_wdl(raw_wdl))): + continue + child = self._probe_dtm50_internal(self._open_dtm50(cfg_c), cfg_c, + cboard, raw_wdl, IGNORE_50MR, depth + 1) + if zeroing: + dtz.zeroing_child(raw_wdl) + else: + dtz.quiet_child(raw_wdl, child) + if child.wdl == ILLEGAL: + skipped.of_class(invert_wdl(fold_dtm_wdl(raw_wdl))) + continue + cw, cd = child.wdl, child.dtm + my_wdl = invert_wdl(fold_dtm_wdl(cw)) + my_dtm = 1 + cd + if not have or prefer_new(my_wdl, my_dtm, best_wdl, best_dtm): + best_wdl, best_dtm, have = my_wdl, my_dtm, True + if not any_legal: + return dtz.finish(LOSE if b.is_check() else DRAW, 0, False) + if not have or skipped.unpin(best_wdl): + return dtz.finish(ILLEGAL, 0, True) + if best_wdl in (WIN, LOSE): + return dtz.finish(best_wdl, best_dtm, True) + return dtz.finish(DRAW, 0, True) + + def _derive_dtm50(self, board: chess.Board, wdl: int, rule50: int, + depth: int) -> _DTM50Result: + """rule50-aware derive: per-child hmc (zeroing resets, quiet increments); + once >=100, the move is DRAW unless it mates. No DTZ rides along: + zeroing is clock-free, so _probe_impl prices it once off the layer-0 + probe and drops whatever a layered one finds.""" + if depth >= MAX_DERIVE_DEPTH: + return _DTM50Result(ILLEGAL, 0) + # Under the clock, cursed wins and blessed losses collapse to DRAW along + # with a DRAW itself; _probe_impl settles all three before the layer. + assert fold_50mr_wdl(wdl) != DRAW + b = _internal_board(board) + any_legal = have = False + best_wdl, best_dtm = LOSE, 0 + skipped = _SkippedChildren() + for mv in b.legal_moves: + any_legal = True + cfg_c, cboard, cep, zeroing = self._make_child(b, mv) + child_rule50 = 0 if zeroing else rule50 + 1 + if cfg_c.is_bare_kings: + cd_wdl, cd_dtm = DRAW, 0 + elif child_rule50 >= DTM50_HMC_COUNT: + # The mate test is the only thing that prices it, so it runs + # unscreened. + cd_wdl, cd_dtm = (LOSE, 0) if _is_checkmate(cboard) else (DRAW, 0) + elif cep is not None: + cr = self._probe_impl(cfg_c, cboard, child_rule50, cep, depth + 1) + if cr.status != "ok" or cr.wdl == ILLEGAL: + skipped.unknown() + continue + if not cr.has_dtm50: + skipped.of_class(invert_wdl(fold_50mr_wdl(cr.wdl))) + continue + cd_wdl, cd_dtm = cr.dtm50_wdl, cr.dtm50 + else: + if child_class_is_forced(wdl): + cw = WIN + else: + cw = self._probe_wdl_internal(self._open_wdl(cfg_c), cfg_c, cboard, depth + 1) + if cw == ILLEGAL: + skipped.unknown() + continue + # Outranked: the clock only walks an offer further toward + # DRAW, never up to the pinned class, so the cell goes unread. + if below_pinned_class(wdl, invert_wdl(cw)): + continue + child = self._probe_dtm50_internal(self._open_dtm50(cfg_c), cfg_c, cboard, + cw, child_rule50, depth + 1) + if child.wdl == ILLEGAL: + skipped.of_class(invert_wdl(fold_50mr_wdl(cw))) + continue + cd_wdl, cd_dtm = child.wdl, child.dtm + my_wdl = invert_wdl(fold_50mr_wdl(cd_wdl)) + my_dtm = 1 + cd_dtm + if not have or prefer_new(my_wdl, my_dtm, best_wdl, best_dtm): + best_wdl, best_dtm, have = my_wdl, my_dtm, True + if not any_legal: + return _DTM50Result(LOSE if b.is_check() else DRAW, 0) + if not have or skipped.unpin(best_wdl): + return _DTM50Result(ILLEGAL, 0) + if best_wdl in (WIN, LOSE): + return _DTM50Result(best_wdl, best_dtm) + return _DTM50Result(DRAW, 0) + + # --- combined probe (mirrors probe.cpp's probe_impl, with ep overlay) --- + def _probe_impl(self, cfg: PieceConfig, board: chess.Board, rule50: int, + ep_square: Optional[int], depth: int) -> ProbeResult: + r = ProbeResult() + w = self._open_wdl(cfg) + rule50_drawn = (rule50 != IGNORE_50MR and rule50 >= DTM50_HMC_COUNT) + # DTZ/DTM/DTM50 reads are all gated on WDL, so WDL absent (and not a + # rule50 auto-draw) means there is nothing to return. + if w is None and not rule50_drawn: + return r # tb_not_found + r.status = "ok" + if w is not None: + r.wdl = self._probe_wdl_internal(w, cfg, board, depth) + if r.wdl == ILLEGAL: + return r + # A DRAW is priced without opening anything: no mate to count, + # nowhere to zero, and no clock that can turn it decisive. Every + # reader answers 0 for it anyway. + if r.wdl == DRAW: + r.has_dtm = True + r.has_dtz = True + r.has_dtc = True + r.dtc_wdl = DRAW + if rule50 != IGNORE_50MR: + r.dtm50_wdl = DRAW + r.has_dtm50 = True + m50 = None + else: + m50 = self._open_dtm50(cfg) + if m50 is None: + # No pack, so the mate distance comes from the `dtm/` table + # the pack otherwise makes redundant. DTZ still follows + # below, and nothing here can answer the rule-true layer. + m = self._open_dtm(cfg) + if m is not None: + dtm = self._probe_dtm_internal(m, cfg, board, r.wdl, depth) + r.has_dtm = dtm is not None + if dtm is not None: + r.dtm = dtm + if m50 is not None: # never for a DRAW, which is already answered + d50 = self._probe_dtm50_internal(m50, cfg, board, r.wdl, IGNORE_50MR, depth) + r.dtm = d50.dtm + r.has_dtm = (d50.wdl != ILLEGAL) + r.has_dtz = d50.has_dtz + r.dtz = d50.dtz + if rule50_drawn: + # Mate outruns the claim, and a LOSE at flat distance 0 is + # one. An unpriced distance tells them apart from neither, + # so the layer goes unanswered there. + if r.has_dtm or r.wdl != LOSE: + r.dtm50_wdl = LOSE if (r.wdl == LOSE and r.dtm == 0) else DRAW + r.dtm50 = 0 + r.has_dtm50 = True + elif rule50 != IGNORE_50MR: + # Cursed/blessed are DRAW at every layer; the clock is never read. + if fold_50mr_wdl(r.wdl) == DRAW: + r.dtm50_wdl = DRAW + r.dtm50 = 0 + r.has_dtm50 = True + elif r.has_dtm and dtm50_layer_pinned_by_dtm(r.wdl, r.dtm, rule50): + r.dtm50_wdl = r.wdl # plain W/L, so the DTM50 fold is the identity + r.dtm50 = r.dtm + r.has_dtm50 = True + elif layer_busted_by_dtz(r.wdl, r.dtz, rule50): + r.dtm50_wdl = DRAW + r.dtm50 = 0 + r.has_dtm50 = True + else: + rr = self._probe_dtm50_internal(m50, cfg, board, r.wdl, rule50, depth) + r.dtm50_wdl = rr.wdl + r.dtm50 = rr.dtm + r.has_dtm50 = (rr.wdl != ILLEGAL) + # DTC next: its pack answers both metrics, the budget the caller's + # clock picks and the unbounded row that is the DTZ table it embeds. + # A cursed class or an outrun clock leaves no budget to pick, which + # the read reports by pricing nothing, and a DRAW is priced above + # without opening anything. + if not r.has_dtc and material_has_pawns(cfg): + c = self._open_dtc(cfg) + if c is not None: + cell = self._probe_dtc_internal(c, cfg, board, r.wdl, rule50, depth) + if cell is not None: + r.has_dtc = True + r.dtc_wdl = r.wdl if cell.priced else DRAW + r.dtc_order = cell.order if cell.priced else 0 + r.dtc = cell.value if cell.priced else 0 + # A record carries that row and so does the derive, so a + # decisive class always has one. + assert cell.dtz != DTC_DRAWN + r.has_dtz = True + r.dtz = cell.dtz + # What the packs above left: the DTM50 one stops at the cursed band, + # and a DTC one answers only the materials it is built for. + if not r.has_dtz: + d = self._open_dtz(cfg) + if d is not None: + dtz = self._probe_dtz_internal(d, cfg, board, r.wdl, depth) + if dtz is not None: + r.has_dtz = True + r.dtz = dtz + # Pawnless materials carry no pack: with no push to budget the stack + # is one layer, every zeroing move is a capture and so a conversion, + # which leaves DTZ's own number at order 0. A cursed class is drawn + # without reading it. + if not r.has_dtc and not material_has_pawns(cfg): + if r.wdl not in (WIN, LOSE): + r.has_dtc = True + r.dtc_wdl = DRAW + elif r.has_dtz: + busted = (rule50 != IGNORE_50MR + and layer_busted_by_dtz(r.wdl, r.dtz, rule50)) + r.has_dtc = True + r.dtc_wdl = DRAW if busted else r.wdl + r.dtc = 0 if busted else r.dtz + + if ep_square is None: + return r + bcopy, eps = _ep_capture_moves(board, ep_square) + if not eps: + return r + + best = r + best_dtz_wdl = r.wdl + best_dtz = r.dtz if r.has_dtz else 0 + best_dtm_wdl = fold_dtm_wdl(r.wdl) + best_dtm = r.dtm if r.has_dtm else 0 + best_dtm50_wdl = r.dtm50_wdl if r.has_dtm50 else fold_50mr_wdl(r.wdl) + best_dtm50 = r.dtm50 if r.has_dtm50 else 0 + # DTC compares on its own class, not the clock-independent one: a base + # this clock has drawn must lose to an ep conversion that still wins. + best_dtc_wdl = r.dtc_wdl if r.has_dtc else r.wdl + best_dtc_order = r.dtc_order if r.has_dtc else 0 + best_dtc = r.dtc if r.has_dtc else 0 + for mv in eps: + cfg_c, cboard, _cep, _ = self._make_child(bcopy, mv) + if cfg_c.is_bare_kings: + cr = ProbeResult() + cr.status = "ok" + cr.wdl = DRAW + cr.has_dtz = best.has_dtz; cr.dtz = 0 + cr.has_dtm = best.has_dtm; cr.dtm = 0 + cr.has_dtm50 = best.has_dtm50; cr.dtm50_wdl = DRAW; cr.dtm50 = 0 + else: + cr = self._probe_impl(cfg_c, cboard, 0, None, depth + 1) # ep is zeroing + if cr.status != "ok" or cr.wdl == ILLEGAL: + return ProbeResult() + my_wdl = invert_wdl(cr.wdl) + if wdl_rank(my_wdl) > wdl_rank(best.wdl): + best.wdl = my_wdl + # An ep capture zeroes, so its dtz is 1 whatever the child + # holds. The base value has to be known only to break a tie in class + # -- outranked, it does not enter, which is how a cursed base the + # pack cannot pin still reports a dtz once ep lifts it. + ep_outranks_dtz = wdl_rank(my_wdl) > wdl_rank(best_dtz_wdl) + if ep_outranks_dtz or (best.has_dtz + and prefer_new(my_wdl, 1, best_dtz_wdl, best_dtz)): + best_dtz_wdl, best_dtz = my_wdl, 1 + best.dtz = 0 if my_wdl == DRAW else 1 + best.has_dtz = True + if best.has_dtm and cr.has_dtm: + my_dtm_wdl = fold_dtm_wdl(my_wdl) + my_dtm = 1 + cr.dtm + if prefer_new(my_dtm_wdl, my_dtm, best_dtm_wdl, best_dtm): + best_dtm_wdl, best_dtm = my_dtm_wdl, my_dtm + best.dtm = my_dtm if my_dtm_wdl in (WIN, LOSE) else 0 + # An ep capture is a conversion: it owes no push and lands one ply + # out, so (0, 1) needs no table to know, and the base value enters + # only to break a tie in class. A conversion is terminal to DTC, so + # its class is the child's WDL folded to what a layer holds, as the + # generator classifies one. + my_dtc_wdl = fold_50mr_wdl(invert_wdl(cr.wdl)) + if (wdl_rank(my_dtc_wdl) > wdl_rank(best_dtc_wdl) + or (best.has_dtc and prefer_new_dtc(my_dtc_wdl, 0, 1, + best_dtc_wdl, best_dtc_order, + best_dtc))): + best_dtc_wdl, best_dtc_order, best_dtc = my_dtc_wdl, 0, 1 + best.has_dtc = True + best.dtc_wdl = my_dtc_wdl + best.dtc_order = 0 + best.dtc = 1 if my_dtc_wdl != DRAW else 0 + if best.has_dtm50 and cr.has_dtm50: + my_dtm50_wdl = invert_wdl(cr.dtm50_wdl) + my_dtm50 = 1 + cr.dtm50 + if prefer_new(my_dtm50_wdl, my_dtm50, best_dtm50_wdl, best_dtm50): + best_dtm50_wdl, best_dtm50 = my_dtm50_wdl, my_dtm50 + best.dtm50_wdl = my_dtm50_wdl + best.dtm50 = my_dtm50 if my_dtm50_wdl in (WIN, LOSE) else 0 + return best + + # --- public API --- + def probe(self, board: chess.Board, rule50: int = 0) -> ProbeResult: + """Full probe of `board`. `rule50` (the halfmove clock) selects the DTM50 + layer. Returns a :class:`ProbeResult`; its ``status`` is ``"ok"`` or + ``"tb_not_found"`` (no table for the material). + + Raises :class:`MissingTableError` if `board` has castling rights: the + tables are built without them, so no cell describes such a position. + + Safe to call concurrently from several threads. This is the one entry + point every other probe method funnels through, so registering as a + reader here covers the whole walk, including the tables it opens on the + way; the registration itself only ever blocks a concurrent + :meth:`close`, never another probe.""" + if board.castling_rights: + raise MissingTableError( + f"chesstb tables do not contain positions with castling rights: {board.fen()}") + with self._read_condition: + self._read_count += 1 + try: + # Prefer the pair table: if the board has an opposing pawn pair and + # that 'p' table is on disk, route there; else fall back to the + # full material. The pair table is partial and a position has the + # same value in either, so preferring the pair table when present is + # safe. + paired = pair_config_from_board(board) + if paired is not None and self._has_any_table(paired[0]): + cfg, mirrored = paired + else: + cfg, mirrored = piece_config_from_board(board) + cboard = mirror_for_canonical(board) if mirrored else board.copy(stack=False) + # ep travels as an overlay, not on the board. + ep = cboard.ep_square + cboard.ep_square = None + return self._probe_impl(cfg, cboard, rule50, ep, 0) + finally: + with self._read_condition: + self._read_count -= 1 + self._read_condition.notify_all() + + def _require(self, board: chess.Board) -> ProbeResult: + r = self.probe(board) + if r.status == "tb_not_found": + cfg, _ = piece_config_from_board(board) + raise MissingTableError(f"no chesstb table for {cfg.name()}") + if r.wdl == ILLEGAL: + # Tables exist for the material but cannot resolve this cell: a + # dropped color is rebuilt by one-ply minimax, and that needs every + # capture/promotion sub-table the walk reaches. + cfg, _ = piece_config_from_board(board) + raise MissingTableError( + f"chesstb tables for {cfg.name()} cannot resolve this position") + return r + + def probe_wdl(self, board: chess.Board) -> int: + """5-class WDL as a signed int: +2 win, +1 cursed win, 0 draw, + -1 blessed loss, -2 loss. Raises :class:`MissingTableError` if no table + for the position.""" + return _WDL_SIGNED[self._require(board).wdl] + + def get_wdl(self, board: chess.Board, default: Any = None) -> Any: + try: + return self.probe_wdl(board) + except MissingTableError: + return default + + def probe_dtz(self, board: chess.Board) -> int: + """Signed distance to zeroing: +N = side to move reaches a capture, + promotion or pawn move toward a win in N plies, -N toward a loss, + 0 = draw. + + Measures what :meth:`chess.syzygy.Tablebase.probe_dtz` measures, but the + two are not interchangeable. Syzygy bases the cursed band off 100 -- + ``n > 100`` is a cursed win whose zeroing move is n or n - 100 plies + away -- so a DTZ magnitude carries the class along with the count. Here + the class stays in WDL and the count is plain at every value. + """ + r = self._require(board) + if not r.has_dtz: + raise MissingTableError("DTZ table unavailable") + return _signed(r.dtz, r.wdl) + + def get_dtz(self, board: chess.Board, default: Any = None) -> Any: + try: + return self.probe_dtz(board) + except MissingTableError: + return default + + def probe_dtc(self, board: chess.Board, + rule50: Optional[int] = None) -> Tuple[int, int, int]: + """Distance to conversion at the board's halfmove clock (or `rule50` if + given), with pawn pushes priced separately from waiting. Returns + ``(signed_wdl, pushes_owed, plies)``: the pushes the winning side still + owes before a capture or promotion, and the plies to the next zeroing + move on the line that owes them. + + The key is pushes first, then plies: a line that converts without + touching a pawn beats one that spends a push, even when the push is + quicker. Fewer pushes cost a longer wait, so a tight clock forces the + trade the other way, and where no budget fits the clock the answer is a + 50MR draw -- ``(0, 0, 0)`` -- since that clock has taken the win. + + Pawnful materials carry a ``.lzdtc`` pack; for a pawnless one every + zeroing move is already a conversion, so the answer is DTZ's own number + at nothing owed.""" + hmc = board.halfmove_clock if rule50 is None else rule50 + r = self.probe(board, hmc) + if r.status == "tb_not_found": + cfg, _ = piece_config_from_board(board) + raise MissingTableError(f"no chesstb table for {cfg.name()}") + if not r.has_dtc: + raise MissingTableError("DTC unavailable") + return (_WDL_SIGNED[r.dtc_wdl], r.dtc_order, r.dtc) + + def probe_dtm(self, board: chess.Board) -> int: + """Signed distance-to-mate, ignoring the 50-move rule.""" + r = self._require(board) + if not r.has_dtm: + raise MissingTableError("DTM unavailable") + return _signed(r.dtm, r.wdl) + + def get_dtm(self, board: chess.Board, default: Any = None) -> Any: + try: + return self.probe_dtm(board) + except MissingTableError: + return default + + def probe_dtm50(self, board: chess.Board, rule50: Optional[int] = None) -> Tuple[int, int]: + """50-move-rule-aware distance to mate at the board's halfmove clock (or + `rule50` if given). Returns ``(signed_wdl, plies)``; cursed/blessed both + collapse to draw under the 50-move rule.""" + hmc = board.halfmove_clock if rule50 is None else rule50 + cfg, _ = piece_config_from_board(board) + r = self.probe(board, hmc) + if r.status == "tb_not_found": + raise MissingTableError(f"no chesstb table for {cfg.name()}") + if not r.has_dtm50: + raise MissingTableError("DTM50 unavailable") + return (_WDL_SIGNED[r.dtm50_wdl], r.dtm50) + + +def open_tablebase(directory: str, *, + block_cache_bytes: int = DEFAULT_BLOCK_CACHE_BYTES) -> Tablebase: + """Open a directory tree of chesstb tables (``wdl/``, ``dtz/``, ``dtc/``, + ``dtm/``, ``dtm50/`` subdirectories, or table files directly under + `directory`). + + Decoded blocks are kept in a least-recently-used cache bounded by + `block_cache_bytes`; older blocks are reclaimed automatically once the + budget is exceeded.""" + return Tablebase(directory, block_cache_bytes=block_cache_bytes) diff --git a/data/chesstb/dtc/KPK.lzdtc b/data/chesstb/dtc/KPK.lzdtc new file mode 100644 index 000000000..6082b74c1 Binary files /dev/null and b/data/chesstb/dtc/KPK.lzdtc differ diff --git a/data/chesstb/dtm/KBK.lzdtm b/data/chesstb/dtm/KBK.lzdtm new file mode 100644 index 000000000..3d591942c Binary files /dev/null and b/data/chesstb/dtm/KBK.lzdtm differ diff --git a/data/chesstb/dtm/KBNK.lzdtm b/data/chesstb/dtm/KBNK.lzdtm new file mode 100644 index 000000000..ab0ee11e5 Binary files /dev/null and b/data/chesstb/dtm/KBNK.lzdtm differ diff --git a/data/chesstb/dtm/KNK.lzdtm b/data/chesstb/dtm/KNK.lzdtm new file mode 100644 index 000000000..d938a6ef6 Binary files /dev/null and b/data/chesstb/dtm/KNK.lzdtm differ diff --git a/data/chesstb/dtm/KPK.lzdtm b/data/chesstb/dtm/KPK.lzdtm new file mode 100644 index 000000000..a2fe75dec Binary files /dev/null and b/data/chesstb/dtm/KPK.lzdtm differ diff --git a/data/chesstb/dtm/KQK.lzdtm b/data/chesstb/dtm/KQK.lzdtm new file mode 100644 index 000000000..6921af47d Binary files /dev/null and b/data/chesstb/dtm/KQK.lzdtm differ diff --git a/data/chesstb/dtm/KRK.lzdtm b/data/chesstb/dtm/KRK.lzdtm new file mode 100644 index 000000000..2a97654fe Binary files /dev/null and b/data/chesstb/dtm/KRK.lzdtm differ diff --git a/data/chesstb/dtm/KRKR.lzdtm b/data/chesstb/dtm/KRKR.lzdtm new file mode 100644 index 000000000..de537422d Binary files /dev/null and b/data/chesstb/dtm/KRKR.lzdtm differ diff --git a/data/chesstb/dtm50/KBK.lzdtm50 b/data/chesstb/dtm50/KBK.lzdtm50 new file mode 100644 index 000000000..c5506547a Binary files /dev/null and b/data/chesstb/dtm50/KBK.lzdtm50 differ diff --git a/data/chesstb/dtm50/KBNK.lzdtm50 b/data/chesstb/dtm50/KBNK.lzdtm50 new file mode 100644 index 000000000..70a0acc07 Binary files /dev/null and b/data/chesstb/dtm50/KBNK.lzdtm50 differ diff --git a/data/chesstb/dtm50/KNK.lzdtm50 b/data/chesstb/dtm50/KNK.lzdtm50 new file mode 100644 index 000000000..54e3e398b Binary files /dev/null and b/data/chesstb/dtm50/KNK.lzdtm50 differ diff --git a/data/chesstb/dtm50/KPK.lzdtm50 b/data/chesstb/dtm50/KPK.lzdtm50 new file mode 100644 index 000000000..fb4a58759 Binary files /dev/null and b/data/chesstb/dtm50/KPK.lzdtm50 differ diff --git a/data/chesstb/dtm50/KQK.lzdtm50 b/data/chesstb/dtm50/KQK.lzdtm50 new file mode 100644 index 000000000..92b7ddfd1 Binary files /dev/null and b/data/chesstb/dtm50/KQK.lzdtm50 differ diff --git a/data/chesstb/dtm50/KRK.lzdtm50 b/data/chesstb/dtm50/KRK.lzdtm50 new file mode 100644 index 000000000..280b4bcfb Binary files /dev/null and b/data/chesstb/dtm50/KRK.lzdtm50 differ diff --git a/data/chesstb/dtm50/KRKR.lzdtm50 b/data/chesstb/dtm50/KRKR.lzdtm50 new file mode 100644 index 000000000..59df7a692 Binary files /dev/null and b/data/chesstb/dtm50/KRKR.lzdtm50 differ diff --git a/data/chesstb/dtz/KBK.lzdtz b/data/chesstb/dtz/KBK.lzdtz new file mode 100644 index 000000000..5618b1ced Binary files /dev/null and b/data/chesstb/dtz/KBK.lzdtz differ diff --git a/data/chesstb/dtz/KBNK.lzdtz b/data/chesstb/dtz/KBNK.lzdtz new file mode 100644 index 000000000..a3c0c7b72 Binary files /dev/null and b/data/chesstb/dtz/KBNK.lzdtz differ diff --git a/data/chesstb/dtz/KNK.lzdtz b/data/chesstb/dtz/KNK.lzdtz new file mode 100644 index 000000000..00ccec828 Binary files /dev/null and b/data/chesstb/dtz/KNK.lzdtz differ diff --git a/data/chesstb/dtz/KPK.lzdtz b/data/chesstb/dtz/KPK.lzdtz new file mode 100644 index 000000000..caf575bba Binary files /dev/null and b/data/chesstb/dtz/KPK.lzdtz differ diff --git a/data/chesstb/dtz/KQK.lzdtz b/data/chesstb/dtz/KQK.lzdtz new file mode 100644 index 000000000..566c1019d Binary files /dev/null and b/data/chesstb/dtz/KQK.lzdtz differ diff --git a/data/chesstb/dtz/KRK.lzdtz b/data/chesstb/dtz/KRK.lzdtz new file mode 100644 index 000000000..479ebb001 Binary files /dev/null and b/data/chesstb/dtz/KRK.lzdtz differ diff --git a/data/chesstb/dtz/KRKR.lzdtz b/data/chesstb/dtz/KRKR.lzdtz new file mode 100644 index 000000000..1ec41ae35 Binary files /dev/null and b/data/chesstb/dtz/KRKR.lzdtz differ diff --git a/data/chesstb/wdl/KBK.lzw b/data/chesstb/wdl/KBK.lzw new file mode 100644 index 000000000..be565c5b3 Binary files /dev/null and b/data/chesstb/wdl/KBK.lzw differ diff --git a/data/chesstb/wdl/KBNK.lzw b/data/chesstb/wdl/KBNK.lzw new file mode 100644 index 000000000..7e28a0427 Binary files /dev/null and b/data/chesstb/wdl/KBNK.lzw differ diff --git a/data/chesstb/wdl/KNK.lzw b/data/chesstb/wdl/KNK.lzw new file mode 100644 index 000000000..69ba166b4 Binary files /dev/null and b/data/chesstb/wdl/KNK.lzw differ diff --git a/data/chesstb/wdl/KPK.lzw b/data/chesstb/wdl/KPK.lzw new file mode 100644 index 000000000..ccd46f2b6 Binary files /dev/null and b/data/chesstb/wdl/KPK.lzw differ diff --git a/data/chesstb/wdl/KQK.lzw b/data/chesstb/wdl/KQK.lzw new file mode 100644 index 000000000..8b12a27b8 Binary files /dev/null and b/data/chesstb/wdl/KQK.lzw differ diff --git a/data/chesstb/wdl/KRK.lzw b/data/chesstb/wdl/KRK.lzw new file mode 100644 index 000000000..0caa26a31 Binary files /dev/null and b/data/chesstb/wdl/KRK.lzw differ diff --git a/data/chesstb/wdl/KRKR.lzw b/data/chesstb/wdl/KRKR.lzw new file mode 100644 index 000000000..3a00e8c13 Binary files /dev/null and b/data/chesstb/wdl/KRKR.lzw differ diff --git a/docs/chesstb.rst b/docs/chesstb.rst new file mode 100644 index 000000000..17dbcf523 --- /dev/null +++ b/docs/chesstb.rst @@ -0,0 +1,73 @@ +chesstb endgame tablebase probing +================================== + +`chesstb `_ tablebases provide +50-move-rule-aware **WDL** (win/draw/loss, with cursed/blessed classes), +**DTZ** (distance to zeroing -- plies to the next capture, promotion or pawn +move), a **DTC** pack pricing pawn pushes separately from waiting, and a +**DTM50** pack giving both the unbounded **DTM** (depth to mate) and the exact +50-move-rule DTM at any halfmove clock. Where no pack is present, a standalone +**DTM** table answers the unbounded mate distance on its own. Positions with +castling rights are not included. + +Table files are looked up in the ``wdl/``, ``dtz/``, ``dtc/``, ``dtm/`` and +``dtm50/`` subdirectories of each search directory, and in the directory +itself. Each pack embeds an unbounded table and so makes it redundant -- the +DTM50 pack the standalone DTM, the DTC pack the DTZ -- and a material shipping +both is read from the pack alone. + +DTZ measures the same thing as syzygy's, but encodes it differently and the two +are not interchangeable: syzygy bases its cursed band off 100, so a DTZ magnitude +carries the WDL class along with the count, whereas a chesstb DTZ is a plain +distance in every class, the class being WDL's alone. + +**DTC** answers a pair at a given halfmove clock: how many pawn pushes the +winning side still owes before a conversion (a capture or a promotion), and the +plies to the next zeroing move on the line that owes them. The key is pushes +first, then plies, which is what DTZ cannot express -- it prices every push at 1 +and so reads 1 almost everywhere a pawn can move. Fewer pushes cost a longer +wait, so a tighter clock forces the trade the other way, and where no budget +fits the clock the answer is a 50-move draw. Packs exist for pawnful materials; +for a pawnless one every zeroing move is already a conversion, so the answer is +DTZ's own number at nothing owed. + +Probes assume a legal position -- both kings present, the side not to move not +in check, no pawn outside ranks 2-7 -- and do not validate it. Probing +anything else is undefined: it reads whatever cell the placement maps to, or +raises out of the indexer. Screen untrusted input with +:func:`chess.Board.is_valid` first. Castling rights are the one excluded input +that is reported, as :class:`~chess.chesstb.MissingTableError`. + +This is a pure-Python prober (it depends only on the standard library); no +native extension is required. + +.. code-block:: python + + import chess + import chess.chesstb + + with chess.chesstb.open_tablebase("data/chesstb") as tablebase: + board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1") + print(tablebase.probe_wdl(board)) # 2 (+2 win .. -2 loss) + print(tablebase.probe_dtz(board)) # 19 (signed distance to zeroing) + print(tablebase.probe_dtm(board)) # 19 (signed distance to mate) + print(tablebase.probe_dtm50(board)) # (2, 19): rule-true (wdl, plies) + print(tablebase.probe_dtc(board)) # (2, 0, 19): (wdl, pushes, plies) + + board = chess.Board("8/8/8/k7/8/8/K4P2/8 w - - 0 1") + print(tablebase.probe_dtc(board)) # (2, 4, 23): four pushes, 23 plies + print(tablebase.probe_dtc(board, 80)) # (2, 5, 19): tighter clock, one push more + print(tablebase.probe_dtc(board, 95)) # (0, 0, 0): that clock has taken the win + +.. warning:: + Maliciously crafted tablebase files may cause denial of service. + +.. autofunction:: chess.chesstb.open_tablebase + +.. autoclass:: chess.chesstb.Tablebase + :members: probe_wdl, get_wdl, probe_dtz, get_dtz, probe_dtc, probe_dtm, get_dtm, probe_dtm50, probe, add_directory, close + +.. autoclass:: chess.chesstb.ProbeResult + :members: + +.. autoexception:: chess.chesstb.MissingTableError diff --git a/docs/index.rst b/docs/index.rst index da8558da6..efee8528f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -11,6 +11,7 @@ Contents polyglot gaviota syzygy + chesstb engine svg variant diff --git a/test.py b/test.py index cc8a95b7c..699860f2b 100755 --- a/test.py +++ b/test.py @@ -6,13 +6,16 @@ import os import os.path import platform +import struct import sys import tempfile import textwrap +import threading import unittest import io import chess +import chess.chesstb import chess.gaviota import chess.engine import chess.pgn @@ -4955,6 +4958,603 @@ def test_antichess_pgn(self): self.assertEqual(game.end().board().fen(), "8/6k1/3K4/8/8/3k4/8/8 w - - 4 33") +class ChesstbTestCase(unittest.TestCase): + + def open_without_pack(self): + # A tablebase with the DTM50 pack out of the search path, which is how + # material shipping no pack is probed: the mate distance has to come + # from the standalone dtm/ table instead of the pack's flat layer. + tables = chess.chesstb.open_tablebase("data/chesstb") + tables.dirs["dtm50"].clear() + return tables + + # --- The four table kinds, and how a directory is searched for them. --- + + def test_table_kinds(self): + # One file class and one search subdirectory per kind. DTZ and DTM share + # one on-disk layout and so one reader (_RankFile); the value decode is + # the whole of what separates them, so it is asserted directly rather + # than through a probe that could pass for other reasons. DTC and DTM50 + # likewise share the changepoint container (_LayeredFile) and differ in + # what a record means, which the DTC tests below cover. + classes = [chess.chesstb.WDLFile, chess.chesstb.DTZFile, + chess.chesstb.DTCFile, chess.chesstb.DTMFile, + chess.chesstb.DTM50File] + self.assertEqual([c.KIND.lower() for c in classes], + list(chess.chesstb.Tablebase.KINDS)) + self.assertEqual(len({c.MAGIC for c in classes}), len(classes)) + self.assertEqual([c.EXT for c in classes], + [".lzw", ".lzdtz", ".lzdtc", ".lzdtm", ".lzdtm50"]) + # A stored 9 under a WIN: DTZ keeps the count, DTM doubles it and adds + # the ply the class implies. + self.assertEqual(chess.chesstb.DTZFile._value_from_storage(9, chess.chesstb.WIN, 1), 9) + self.assertEqual(chess.chesstb.DTMFile._value_from_storage(9, chess.chesstb.WIN, 1), 19) + self.assertEqual(chess.chesstb.DTMFile._value_from_storage(9, chess.chesstb.LOSE, 1), 18) + # Cell width changes the DTZ decode and never the DTM one, which is why + # only one of the two reads it. + self.assertEqual(chess.chesstb.DTZFile._value_from_storage(9, chess.chesstb.CURSED_WIN, 1), 17) + self.assertEqual(chess.chesstb.DTZFile._value_from_storage(9, chess.chesstb.CURSED_WIN, 2), 9) + + def test_search_directories(self): + with chess.chesstb.open_tablebase("data/chesstb") as tables: + # Every kind looks in its own subdirectory first, then in the + # directory itself, so a flat dump of table files is also probeable. + for kind in chess.chesstb.Tablebase.KINDS: + self.assertEqual(tables.dirs[kind], + [os.path.join("data/chesstb", kind), "data/chesstb"]) + tables.add_directory("data") + for kind in chess.chesstb.Tablebase.KINDS: + self.assertEqual(tables.dirs[kind][-2:], [os.path.join("data", kind), "data"]) + + # A directory holding no tables at all is not an error, just an absence. + with tempfile.TemporaryDirectory() as empty: + with chess.chesstb.open_tablebase(empty) as tables: + board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1") + self.assertIsNone(tables.get_wdl(board)) + self.assertEqual(tables.probe(board).status, "tb_not_found") + + # --- What a probe answers, and in what convention. --- + + def test_probe(self): + with chess.chesstb.open_tablebase("data/chesstb") as tables: + # KQK: mate distances and the signed WDL convention. + board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1") + self.assertEqual(tables.probe_wdl(board), 2) + self.assertEqual(tables.probe_dtz(board), 19) + self.assertEqual(tables.probe_dtm(board), 19) + self.assertEqual(tables.probe_dtm50(board), (2, 19)) + + board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 b - - 0 1") + self.assertEqual(tables.probe_wdl(board), -2) + self.assertEqual(tables.probe_dtz(board), -20) + self.assertEqual(tables.probe_dtm(board), -20) # signed: losing side + self.assertEqual(tables.probe_dtm50(board), (-2, 20)) + + def test_mirrored_material(self): + # Stronger side is Black: internally mirrored to the canonical KQK table. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + self.assertEqual(tables.probe_wdl(chess.Board("k7/1q6/8/5K2/8/8/8/8 b - - 0 1")), 2) + + def test_opposing_pair_child_is_not_bare_kings(self): + # An opposing-pair config counts only its free pieces, so KpKp sits at + # num_pieces == 2 while holding four: it must not be read as KK (an + # unconditional draw) when the walk routes a child into a 'p' table. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + tables._has_any_table = lambda cfg: cfg.has_pair # pretend KpKp is on disk + board = chess.Board("8/8/8/4p3/4P3/8/8/K6k w - - 0 1") + cfg, _, _, _ = tables._make_child(board, chess.Move.from_uci("a1a2")) + self.assertEqual(cfg.name(), "KpKp") + self.assertEqual(cfg.num_pieces, 2) + self.assertFalse(cfg.is_bare_kings) + # Genuine bare kings still short-circuit. + tables._has_any_table = lambda cfg: False + board = chess.Board("8/8/8/8/8/8/1r6/K6k w - - 0 1") + cfg, _, _, _ = tables._make_child(board, chess.Move.from_uci("a1b2")) + self.assertTrue(cfg.is_bare_kings) + + # --- Frames a table does not hold, recovered two ways. --- + + def test_dropped_frame_symmetric(self): + # KRKR ships one frame dropped; the missing side is reconstructed by the + # symmetric color mirror. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + board = chess.Board("8/2r5/8/8/8/1k6/8/K1R5 b - - 0 1") + self.assertEqual(tables.probe_wdl(board), 2) + self.assertEqual(tables.probe_dtz(board), 1) + self.assertEqual(tables.probe_dtm(board), 1) + + def test_dropped_frame_minimax(self): + # KBNK ships an asymmetric dropped frame, reconstructed by one-ply minimax. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + board = chess.Board("8/8/8/8/8/2k5/2N5/KB6 w - - 0 1") + self.assertEqual(tables.probe_wdl(board), 2) + self.assertEqual(tables.probe_dtm(board), 61) + board = chess.Board("8/8/8/8/8/2k5/2N5/KB6 b - - 0 1") + self.assertEqual(tables.probe_wdl(board), -2) + self.assertEqual(tables.probe_dtm(board), -60) # signed: losing side + + # --- The DTM50 pack, and the standalone DTM table it stands in for. --- + + def test_dtm_table_without_pack(self): + # Without the DTM50 pack the mate distance comes from the standalone + # dtm/ table, which must answer exactly what the pack does. + with self.open_without_pack() as tables: + board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1") + self.assertEqual(tables.probe_dtm(board), 19) + # The pack alone answers the rule-true layer. + with self.assertRaises(chess.chesstb.MissingTableError): + tables.probe_dtm50(board) + # DTZ still comes from its own table. + self.assertEqual(tables.probe_dtz(board), 19) + # A DRAW is priced without opening the table at all. + self.assertEqual(tables.probe_dtm(chess.Board("8/8/8/8/4k3/8/4P3/4K3 w - - 0 1")), 0) + # Each shipped dtm/ table drops one frame (KQK black, KBNK white), + # so the other side comes back the same way the pack's frames do. + self.assertEqual(tables.probe_dtm(chess.Board("8/8/8/5k2/8/8/1Q6/K7 b - - 0 1")), -20) + self.assertEqual(tables.probe_dtm(chess.Board("8/8/8/8/8/2k5/2N5/KB6 w - - 0 1")), 61) + + def test_dtm50_pack_preferred_over_dtm(self): + # The pack carries the flat DTM, which is what makes the dtm/ table + # redundant: with a pack on disk it must never be opened. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + def fail(cfg): + self.fail("dtm/ opened while the DTM50 pack is available") + tables._open_dtm = fail + board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1") + self.assertEqual(tables.probe_dtm(board), 19) + self.assertEqual(tables.probe_dtm50(board), (2, 19)) + + def test_dtm50_layer_decoded_at_the_clock(self): + # The pack stores a layer per halfmove clock, so a rule-true probe + # decodes a different cell than the flat one. Reaching that decode at + # all needs material that converts sooner than it mates: where dtz and + # dtm agree, a bound settles the layer before any cell is read. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + board = chess.Board("k7/8/8/1K6/8/8/6P1/8 b - - 0 1") # dtz 2, dtm 48 + cfg, _ = chess.chesstb.piece_config_from_board(board) + pack = tables._open_dtm50(cfg) + layers = [] + original = pack.read + pack.read = lambda color, b, wdl, hmc: (layers.append(hmc), + original(color, b, wdl, hmc))[1] + + self.assertEqual(tables.probe_dtm(board), -48) + self.assertEqual(tables.probe_dtz(board), -2) + # rule50 + dtm <= 100: the flat layer already answers, nothing else + # is decoded. + layers.clear() + self.assertEqual(tables.probe_dtm50(board, 0), (-2, 48)) + self.assertEqual(layers, [chess.chesstb.IGNORE_50MR]) + # Past that bound the layer at the clock is decoded on its own. + layers.clear() + self.assertEqual(tables.probe_dtm50(board, 60), (-2, 48)) + self.assertIn(60, layers) + # rule50 + dtz > 100: no line resets the clock in time, so the layer + # is a draw and again no cell is read for it. Past the window itself + # the flat mate distance still stands while 50MR calls it a draw. + layers.clear() + self.assertEqual(tables.probe_dtm50(board, 99), (0, 0)) + self.assertEqual(layers, [chess.chesstb.IGNORE_50MR]) + self.assertEqual(tables.probe_dtm50(board, 100), (0, 0)) + self.assertEqual(tables.probe_dtm(board), -48) + + def test_dtm50_stride_index(self): + # The stride prefix index has to answer exactly what walking every + # position from the start of the block would. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + board = chess.Board("8/8/8/8/8/2k5/8/K1N1B3 w - - 0 1") + cfg, _ = chess.chesstb.piece_config_from_board(board) + dtm50 = tables._open_dtm50(cfg) + color = next(c for c in range(2) + if dtm50.per_color[c] is not None + and not dtm50.is_singular[c] and not dtm50.is_dropped[c]) + pc = dtm50.per_color[color] + blk = dtm50._get_block(pc, 0) + state_bits, cum = blk["state_bits"], blk["state_cum"] + + counts = [0, 0, 0, 0] + for pos in range(4096): + state = (state_bits[pos // 4] >> (2 * (pos % 4))) & 3 + self.assertEqual(chess.chesstb._state_and_index(state_bits, cum, pos), + (state, counts[state])) + counts[state] += 1 + + hints, pre = blk["single_hints"], blk["single_pre"] + popcount = 0 + for i in range(min(4096, (len(hints) - 1) * 8)): + self.assertEqual(chess.chesstb._hint_prefix(hints, pre, i), popcount) + popcount += (hints[i >> 3] >> (i & 7)) & 1 + + def test_dtm50_hint_bitmaps(self): + # SINGLE and DOUBLE records vary in width and the file stores no bitmap + # saying which, so the prober rebuilds one by reading each record's + # draw-end bit -- the MSB of byte 0 for SINGLE, of byte 1 for DOUBLE. + # Read the wrong byte and every record after the first short one is + # located at the wrong offset, which no single probe reveals. Both + # walks have to land exactly on the stream sizes the header declares. + with chess.chesstb.open_tablebase("data/chesstb") as tables: + board = chess.Board("8/8/8/8/4k3/8/4P3/4K3 w - - 0 1") + cfg, _ = chess.chesstb.piece_config_from_board(board) + pack = tables._open_dtm50(cfg) + color = next(c for c in range(2) + if pack.per_color[c] is not None + and not pack.is_singular[c] and not pack.is_dropped[c]) + blk = pack._get_block(pack.per_color[color], 0) + payload, eb = blk["payload"], blk["eb"] + num_single, num_double = struct.unpack_from("