Add chess.chesstb: pure-Python prober for chesstb endgame tablebases

docs/chesstb.rst

@@ -0,0 +1,44 @@
+chesstb endgame tablebase probing
+==================================
+
+`chesstb <https://github.com/noobpwnftw/chesstb>`_ tablebases provide
+50-move-rule-aware **WDL** (win/draw/loss, with cursed/blessed classes),
+**DTC** (distance to conversion -- plies to the next zeroing move), and a
+**DTM50** pack giving both the unbounded **DTM** (depth to mate) and the exact
+50-move-rule DTM at any halfmove clock. Positions with castling rights are not
+included.
+
+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 conversion)
+        print(tablebase.probe_dtm(board))    # 19 (signed distance to mate)
+        print(tablebase.probe_dtm50(board))  # (2, 19): rule-true (wdl, plies)
+
+.. 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_dtm, get_dtm, probe_dtm50, probe, add_directory, close
+
+.. autoclass:: chess.chesstb.ProbeResult
+    :members:
+
+.. autoexception:: chess.chesstb.MissingTableError

docs/index.rst

@@ -11,6 +11,7 @@ Contents
     polyglot
     gaviota
     syzygy
+    chesstb
     engine
     svg
     variant

test.py

@@ -9,10 +9,12 @@
 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 +4957,295 @@ 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 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_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
+
+    def test_layered_dtm50(self):
+        with chess.chesstb.open_tablebase("data/chesstb") as tables:
+            board = chess.Board("8/8/8/4k3/8/8/Q7/K7 w - - 10 30")
+            # Halfmove clock selects the DTM50 layer.
+            self.assertEqual(tables.probe_dtm50(board, 10), (2, 17))
+            # Past the 50-move window the win still has a flat DTM but draws under 50MR.
+            board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 b - - 0 1")
+            self.assertEqual(tables.probe_dtm(board), -20)
+            self.assertEqual(tables.probe_dtm50(board, 100), (0, 0))
+
+    def test_missing_table(self):
+        with chess.chesstb.open_tablebase("data/chesstb") as tables:
+            # Legal position whose material (KQQQQK) has no table on disk.
+            self.assertIsNone(tables.get_wdl(chess.Board("7k/8/8/8/8/8/3QQQQ1/K7 w - - 0 1")))
+            with self.assertRaises(chess.chesstb.MissingTableError):
+                tables.probe_wdl(chess.Board("7k/8/8/8/8/8/3QQQQ1/K7 w - - 0 1"))
+
+    def test_castling_rights_rejected(self):
+        # The tables are built without castling rights, so such a position has
+        # no cell: reject it rather than answer as if the rights were absent.
+        with chess.chesstb.open_tablebase("data/chesstb") as tables:
+            board = chess.Board("4k3/8/8/8/8/8/8/R3K3 w Q - 0 1")
+            with self.assertRaises(chess.chesstb.MissingTableError):
+                tables.probe_wdl(board)
+            self.assertIsNone(tables.get_wdl(board))
+            self.assertIsNone(tables.get_dtz(board))
+            with self.assertRaises(chess.chesstb.MissingTableError):
+                tables.probe_dtm50(board)
+            # The same material without the rights probes normally.
+            self.assertEqual(tables.probe_wdl(chess.Board("4k3/8/8/8/8/8/8/R3K3 w - - 0 1")), 2)
+
+    def test_frozen_pair_child_is_not_bare_kings(self):
+        # A frozen-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)
+
+    def test_block_cache_reclaim(self):
+        # A tiny budget forces decoded blocks to be evicted from their owning
+        # per-color dicts, while answers stay correct.
+        with chess.chesstb.open_tablebase("data/chesstb", block_cache_bytes=1) as tables:
+            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_dtm(board), 19)
+            cache = tables._block_cache
+            cfg, _ = chess.chesstb.piece_config_from_board(board)
+            wdl = tables._open_wdl(cfg)
+            # At most a single decoded WDL block stays resident under the budget.
+            self.assertLessEqual(sum(len(pc._blocks) for pc in wdl.per_color if pc), 1)
+            # close() drops everything tracked.
+            tables.close()
+            self.assertEqual(cache.cur_bytes, 0)
+
+    def test_concurrent_probe(self):
+        # Threads sharing one Tablebase must agree with a single-threaded
+        # reference. The barrier puts them all on the *first* open of these
+        # five materials at once (the lazy-open race), and the 1-byte block
+        # budget makes eviction race decoding on essentially every probe.
+        # KRKR is left out on purpose: under that budget its blocks re-decode
+        # per probe at ~100 ms a go, and test_dropped_frame_symmetric covers it.
+        queries = [
+            ("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1", 0),
+            ("8/8/8/5k2/8/8/1Q6/K7 b - - 0 1", 0),
+            ("8/8/8/3k4/8/8/4P3/K7 w - - 0 1", 0),
+            ("8/8/8/4k3/8/8/Q7/K7 w - - 10 30", 10),
+            ("8/8/8/4k3/8/8/2R5/K7 b - - 0 1", 0),
+            ("8/8/8/4k3/8/8/2B5/K7 w - - 0 1", 0),
+            ("8/8/8/4k3/8/8/2N5/K7 w - - 0 1", 0),
+        ]
+
+        def answers(tables):
+            return [tables.probe(chess.Board(fen), rule50).__repr__()
+                    for fen, rule50 in queries]
+
+        with chess.chesstb.open_tablebase("data/chesstb") as tables:
+            expected = answers(tables)
+
+        num_threads = 4
+        barrier = threading.Barrier(num_threads)
+        results: list = [None] * num_threads
+        errors: list = []
+
+        with chess.chesstb.open_tablebase("data/chesstb", block_cache_bytes=1) as tables:
+            def worker(i):
+                try:
+                    barrier.wait(timeout=30)
+                    results[i] = answers(tables)
+                except BaseException as e:
+                    errors.append(e)
+
+            threads = [threading.Thread(target=worker, args=(i,)) for i in range(num_threads)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join(timeout=60)
+            self.assertEqual([t.name for t in threads if t.is_alive()], [], "deadlocked")
+            self.assertEqual(errors, [])
+            for got in results:
+                self.assertEqual(got, expected)
+
+    def test_close_waits_for_probes(self):
+        # close() unmaps files and releases the views into them, so it must
+        # block until in-flight probes are done. Asserted on the read-count
+        # registration directly -- a thread standing in for a probe that is
+        # mid-read -- so the ordering is checked rather than raced for.
+        tables = chess.chesstb.open_tablebase("data/chesstb")
+        tables.probe(chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1"))  # warm and mapped
+
+        registered = threading.Event()
+        may_finish = threading.Event()
+        closed = threading.Event()
+
+        def in_flight_probe():
+            with tables._read_condition:
+                tables._read_count += 1
+            registered.set()
+            may_finish.wait(timeout=30)
+            with tables._read_condition:
+                tables._read_count -= 1
+                tables._read_condition.notify_all()
+
+        reader = threading.Thread(target=in_flight_probe)
+        reader.start()
+        self.assertTrue(registered.wait(timeout=30))
+
+        closer = threading.Thread(target=lambda: (tables.close(), closed.set()))
+        closer.start()
+        self.assertFalse(closed.wait(timeout=0.05), "close() did not wait for the probe")
+        may_finish.set()
+        self.assertTrue(closed.wait(timeout=30), "close() did not resume")
+
+        reader.join(timeout=30)
+        closer.join(timeout=30)
+        self.assertEqual([t.name for t in (reader, closer) if t.is_alive()], [])
+
+    def test_mmap_lifecycle(self):
+        with chess.chesstb.open_tablebase("data/chesstb") as tables:
+            board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1")
+            self.assertEqual(tables.probe_dtm(board), 19)
+            cfg, _ = chess.chesstb.piece_config_from_board(board)
+            dtm50 = tables._open_dtm50(cfg)
+            # Table data is read through the mapping, not copied into memory.
+            pc = next(pc for pc in dtm50.per_color if pc)
+            self.assertIs(pc.buf.obj, dtm50._data)
+            self.assertIs(pc.offsets.blob.obj, dtm50._data)
+
+            # close() has to release every exported view before unmapping,
+            # otherwise mmap.close() raises BufferError.
+            tables.close()
+            self.assertIsNone(dtm50._data)
+            # A closed tablebase maps the files again on the next probe.
+            self.assertEqual(tables.probe_dtm(board), 19)
+
+    def test_open_source_seam(self):
+        # A source that is not a buffer at all: len(), indexing and slicing to
+        # bytes, nothing more. Everything above _open_source has to work
+        # through just that, and must never ask for more than one block at a
+        # time -- a source that fetches on slice (a remote table, say) would
+        # otherwise pull the whole file to answer one probe.
+        opened = []
+
+        class LazySource:
+            def __init__(self, path):
+                with open(path, "rb") as f:
+                    self.data = f.read()
+                self.widest = 0
+                self.closed = False
+
+            def __len__(self):
+                return len(self.data)
+
+            def __getitem__(self, key):
+                if isinstance(key, slice):
+                    start, stop, _ = key.indices(len(self.data))
+                    self.widest = max(self.widest, max(0, stop - start))
+                    return self.data[key]  # bytes, not a view
+                return self.data[key]
+
+            def close(self):
+                self.closed = True
+
+        def open_source(table, path):
+            src = LazySource(path)
+            opened.append(src)
+            table._data = src
+            return src
+
+        board = chess.Board("8/8/8/5k2/8/8/1Q6/K7 w - - 0 1")
+        with chess.chesstb.open_tablebase("data/chesstb") as tables:
+            expected = tables.probe(board, 0)
+
+        original = chess.chesstb._TableFile._open_source
+        chess.chesstb._TableFile._open_source = open_source
+        try:
+            with chess.chesstb.open_tablebase("data/chesstb") as tables:
+                got = tables.probe(board, 0)
+                dtm50 = tables._open_dtm50(chess.chesstb.piece_config_from_board(board)[0])
+                block_positions = next(pc for pc in dtm50.per_color if pc).block_positions
+        finally:
+            chess.chesstb._TableFile._open_source = original
+
+        self.assertTrue(opened)
+        self.assertEqual(got.wdl, expected.wdl)
+        self.assertEqual(got.dtm, expected.dtm)
+        self.assertEqual(got.dtc, expected.dtc)
+        self.assertTrue(all(src.closed for src in opened))
+        # Nothing asked for a span wider than one block.
+        self.assertLessEqual(max(src.widest for src in opened), block_positions)
+
+    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
+
+
 if __name__ == "__main__":
     verbosity = sum(arg.count("v") for arg in sys.argv if all(c == "v" for c in arg.lstrip("-")))
     verbosity += sys.argv.count("--verbose")