1use crate::{Column, GameResult, Move, MoveList, Position, Side};
2use std::fmt;
3use std::str::FromStr;
4
5const INFINITY: i16 = 31_000;
6const FORCED_SCORE: i16 = 30_000;
7const MAX_GAME_PLIES: i16 = 42;
8const HEURISTIC_LIMIT: i16 = 4_000;
9const DEFAULT_TT_ENTRIES: usize = 1 << 16;
10
11#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
13pub enum Algorithm {
14 PlainNegamax,
15 AlphaBeta,
16 OrderedAlphaBeta,
17 #[default]
18 TranspositionTable,
19}
20
21impl Algorithm {
22 pub const ALL: [Self; 4] = [
23 Self::PlainNegamax,
24 Self::AlphaBeta,
25 Self::OrderedAlphaBeta,
26 Self::TranspositionTable,
27 ];
28
29 pub const fn name(self) -> &'static str {
30 match self {
31 Self::PlainNegamax => "plain",
32 Self::AlphaBeta => "alpha-beta",
33 Self::OrderedAlphaBeta => "ordered",
34 Self::TranspositionTable => "tt",
35 }
36 }
37
38 const fn uses_alpha_beta(self) -> bool {
39 !matches!(self, Self::PlainNegamax)
40 }
41
42 const fn uses_center_order(self) -> bool {
43 matches!(self, Self::OrderedAlphaBeta | Self::TranspositionTable)
44 }
45
46 const fn uses_table(self) -> bool {
47 matches!(self, Self::TranspositionTable)
48 }
49}
50
51impl fmt::Display for Algorithm {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 f.write_str(self.name())
54 }
55}
56
57impl FromStr for Algorithm {
58 type Err = &'static str;
59
60 fn from_str(value: &str) -> Result<Self, Self::Err> {
61 match value.to_ascii_lowercase().as_str() {
62 "plain" | "negamax" => Ok(Self::PlainNegamax),
63 "alpha-beta" | "alphabeta" | "alpha_beta" => Ok(Self::AlphaBeta),
64 "ordered" | "ordered-alpha-beta" | "ordered_alphabeta" => Ok(Self::OrderedAlphaBeta),
65 "tt" | "transposition" | "transposition-table" => Ok(Self::TranspositionTable),
66 _ => Err("algorithm must be plain, alpha-beta, ordered, or tt"),
67 }
68 }
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
75pub struct Score(i16);
76
77impl Score {
78 pub const fn raw(self) -> i16 {
79 self.0
80 }
81
82 pub const fn kind(self) -> ScoreKind {
83 if self.0 >= FORCED_SCORE - MAX_GAME_PLIES {
84 ScoreKind::ForcedWin {
85 plies: (FORCED_SCORE - self.0) as u8,
86 }
87 } else if self.0 <= -FORCED_SCORE + MAX_GAME_PLIES {
88 ScoreKind::ForcedLoss {
89 plies: (FORCED_SCORE + self.0) as u8,
90 }
91 } else {
92 ScoreKind::Estimate(self.0)
93 }
94 }
95
96 const fn negated(self) -> Self {
97 Self(-self.0)
98 }
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub enum ScoreKind {
103 ForcedWin { plies: u8 },
104 ForcedLoss { plies: u8 },
105 Estimate(i16),
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub struct SearchLimits {
110 pub depth: u8,
111 pub nodes: Option<u64>,
112}
113
114impl SearchLimits {
115 pub const fn depth(depth: u8) -> Self {
116 Self { depth, nodes: None }
117 }
118
119 pub const fn with_nodes(depth: u8, nodes: u64) -> Self {
120 Self {
121 depth,
122 nodes: Some(nodes),
123 }
124 }
125}
126
127#[derive(Debug, Clone, Copy, PartialEq, Eq)]
129pub struct SearchPreset {
130 pub name: &'static str,
131 pub limits: SearchLimits,
132}
133
134pub const SEARCH_PRESETS: [SearchPreset; 6] = [
136 SearchPreset {
137 name: "beginner",
138 limits: SearchLimits::with_nodes(3, 250),
139 },
140 SearchPreset {
141 name: "easy",
142 limits: SearchLimits::with_nodes(5, 2_000),
143 },
144 SearchPreset {
145 name: "medium",
146 limits: SearchLimits::with_nodes(7, 10_000),
147 },
148 SearchPreset {
149 name: "hard",
150 limits: SearchLimits::with_nodes(9, 50_000),
151 },
152 SearchPreset {
153 name: "expert",
154 limits: SearchLimits::with_nodes(12, 200_000),
155 },
156 SearchPreset {
157 name: "maximum",
158 limits: SearchLimits::with_nodes(42, 750_000),
159 },
160];
161
162pub fn search_preset(name: &str) -> Option<SearchPreset> {
164 SEARCH_PRESETS
165 .iter()
166 .copied()
167 .find(|preset| preset.name == name)
168}
169
170#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
171pub struct SearchStats {
172 pub nodes: u64,
173 pub leaves: u64,
174 pub cutoffs: u64,
175 pub tt_probes: u64,
176 pub tt_hits: u64,
177 pub tt_stores: u64,
178 pub max_ply: u8,
179}
180
181#[derive(Debug, Clone, PartialEq, Eq)]
182pub struct RootBranch {
183 pub mv: Move,
184 pub score: Score,
185 pub nodes: u64,
186 pub cutoffs: u64,
187 pub completed: bool,
188}
189
190#[derive(Debug, Clone, PartialEq, Eq)]
192pub struct SearchReport {
193 pub algorithm: Algorithm,
194 pub requested_depth: u8,
195 pub completed: bool,
196 pub best_move: Option<Move>,
197 pub score: Score,
198 pub principal_variation: Vec<Move>,
199 pub root_branches: Vec<RootBranch>,
200 pub stats: SearchStats,
201}
202
203#[derive(Debug, Clone, Copy, PartialEq, Eq)]
204pub struct IterationSummary {
205 pub depth: u8,
206 pub completed: bool,
207 pub best_move: Option<Move>,
208 pub score: Score,
209 pub nodes: u64,
210 pub cutoffs: u64,
211 pub tt_hits: u64,
212}
213
214#[derive(Debug, Clone, PartialEq, Eq)]
216pub struct IterativeSearchReport {
217 pub result: SearchReport,
218 pub requested_depth: u8,
219 pub completed_depth: u8,
220 pub completed: bool,
221 pub total_nodes: u64,
222 pub iterations: Vec<IterationSummary>,
223}
224
225#[derive(Debug, Clone, Copy)]
226struct Line {
227 moves: [Move; 42],
228 len: u8,
229}
230
231impl Default for Line {
232 fn default() -> Self {
233 Self {
234 moves: [Move::new(Column::new(0)); 42],
235 len: 0,
236 }
237 }
238}
239
240impl Line {
241 fn prepend(mv: Move, child: Self) -> Self {
242 let mut line = Self::default();
243 line.moves[0] = mv;
244 let child_len = usize::from(child.len).min(line.moves.len() - 1);
245 if child_len > 0 {
246 line.moves[1..=child_len].copy_from_slice(&child.moves[..child_len]);
247 }
248 line.len = child_len as u8 + 1;
249 line
250 }
251
252 fn to_vec(self) -> Vec<Move> {
253 self.moves[..usize::from(self.len)].to_vec()
254 }
255
256 fn push(&mut self, mv: Move) -> bool {
257 let index = usize::from(self.len);
258 if index == self.moves.len() {
259 return false;
260 }
261 self.moves[index] = mv;
262 self.len += 1;
263 true
264 }
265}
266
267#[derive(Debug, Clone, Copy)]
268struct NodeResult {
269 score: Score,
270 line: Line,
271 stopped: bool,
272}
273
274impl NodeResult {
275 const fn leaf(score: Score) -> Self {
276 Self {
277 score,
278 line: Line {
279 moves: [Move::new(Column::new(0)); 42],
280 len: 0,
281 },
282 stopped: false,
283 }
284 }
285
286 const fn stopped(score: Score) -> Self {
287 Self {
288 score,
289 line: Line {
290 moves: [Move::new(Column::new(0)); 42],
291 len: 0,
292 },
293 stopped: true,
294 }
295 }
296}
297
298#[derive(Debug, Clone, Copy, PartialEq, Eq)]
299enum Bound {
300 Exact,
301 Lower,
302 Upper,
303}
304
305#[derive(Debug, Clone, Copy)]
306struct TableEntry {
307 key: u64,
308 depth: u8,
309 score: Score,
310 bound: Bound,
311 best_move: Option<Move>,
312}
313
314struct TranspositionTable {
315 entries: Vec<Option<TableEntry>>,
316}
317
318impl TranspositionTable {
319 fn new(entries: usize) -> Self {
320 debug_assert!(entries.is_power_of_two());
321 Self {
322 entries: vec![None; entries],
323 }
324 }
325
326 fn probe(&self, key: u64) -> Option<TableEntry> {
327 self.entries[self.index(key)].filter(|entry| entry.key == key)
328 }
329
330 fn store(&mut self, entry: TableEntry) -> bool {
331 let index = self.index(entry.key);
332 let replace =
333 self.entries[index].is_none_or(|old| old.key != entry.key || entry.depth >= old.depth);
334 if replace {
335 self.entries[index] = Some(entry);
336 }
337 replace
338 }
339
340 fn index(&self, key: u64) -> usize {
341 let mixed = key ^ (key >> 23) ^ (key >> 41);
342 mixed as usize & (self.entries.len() - 1)
343 }
344}
345
346struct Searcher {
347 algorithm: Algorithm,
348 limits: SearchLimits,
349 stats: SearchStats,
350 table: Option<TranspositionTable>,
351}
352
353impl Searcher {
354 fn new(algorithm: Algorithm, limits: SearchLimits) -> Self {
355 Self {
356 algorithm,
357 limits,
358 stats: SearchStats::default(),
359 table: algorithm
360 .uses_table()
361 .then(|| TranspositionTable::new(DEFAULT_TT_ENTRIES)),
362 }
363 }
364
365 fn enter_node(&mut self, ply: u8) -> bool {
366 if self
367 .limits
368 .nodes
369 .is_some_and(|limit| self.stats.nodes >= limit)
370 {
371 return false;
372 }
373 self.stats.nodes += 1;
374 self.stats.max_ply = self.stats.max_ply.max(ply);
375 true
376 }
377
378 fn begin_iteration(&mut self, limits: SearchLimits) {
379 self.limits = limits;
380 self.stats = SearchStats::default();
381 }
382
383 fn table_line(&self, mut position: Position, depth: u8) -> Line {
384 let Some(table) = &self.table else {
385 return Line::default();
386 };
387 let mut line = Line::default();
388 for _ in 0..depth {
389 if position.result() != GameResult::Ongoing {
390 break;
391 }
392 let Some(mv) = table
393 .probe(position.key())
394 .and_then(|entry| entry.best_move)
395 else {
396 break;
397 };
398 if !position.can_play(mv.column()) || !line.push(mv) {
399 break;
400 }
401 position
402 .make_move(mv)
403 .expect("a stored table move must remain legal");
404 }
405 line
406 }
407
408 fn node(
409 &mut self,
410 position: &mut Position,
411 depth: u8,
412 ply: u8,
413 mut alpha: Score,
414 mut beta: Score,
415 ) -> NodeResult {
416 if !self.enter_node(ply) {
417 return NodeResult::stopped(evaluate(*position));
418 }
419 if let Some(score) = terminal_score(*position, ply) {
420 self.stats.leaves += 1;
421 return NodeResult::leaf(score);
422 }
423 if depth == 0 {
424 self.stats.leaves += 1;
425 return NodeResult::leaf(evaluate(*position));
426 }
427
428 if !self.algorithm.uses_alpha_beta() {
429 return self.plain_node(position, depth, ply);
430 }
431
432 let original_alpha = alpha;
433 let original_beta = beta;
434 let mut table_move = None;
435 if self.algorithm.uses_table() {
436 self.stats.tt_probes += 1;
437 if let Some(entry) = self
438 .table
439 .as_ref()
440 .and_then(|table| table.probe(position.key()))
441 {
442 self.stats.tt_hits += 1;
443 table_move = entry.best_move;
444 if entry.depth >= depth {
445 match entry.bound {
446 Bound::Exact => {
447 return NodeResult {
448 score: entry.score,
449 line: self.table_line(*position, depth),
450 stopped: false,
451 };
452 }
453 Bound::Lower => alpha = alpha.max(entry.score),
454 Bound::Upper => beta = beta.min(entry.score),
455 }
456 if alpha >= beta {
457 self.stats.cutoffs += 1;
458 return NodeResult {
459 score: entry.score,
460 line: self.table_line(*position, depth),
461 stopped: false,
462 };
463 }
464 }
465 }
466 }
467
468 let mut best = Score(-INFINITY);
469 let mut best_move = None;
470 let mut best_line = Line::default();
471 for mv in ordered_moves(*position, self.algorithm, table_move) {
472 let undo = position
473 .make_move(mv)
474 .expect("generated moves must be legal");
475 let child = self.node(
476 position,
477 depth - 1,
478 ply + 1,
479 beta.negated(),
480 alpha.negated(),
481 );
482 position.unmake_move(undo);
483 let score = child.score.negated();
484 if score > best || (score == best && prefer(mv, best_move)) {
485 best = score;
486 best_move = Some(mv);
487 best_line = Line::prepend(mv, child.line);
488 }
489 if child.stopped {
490 return NodeResult {
491 score: best,
492 line: best_line,
493 stopped: true,
494 };
495 }
496 alpha = alpha.max(score);
497 if alpha >= beta {
498 self.stats.cutoffs += 1;
499 break;
500 }
501 }
502
503 if self.algorithm.uses_table() {
504 let bound = if best <= original_alpha {
505 Bound::Upper
506 } else if best >= original_beta {
507 Bound::Lower
508 } else {
509 Bound::Exact
510 };
511 let entry = TableEntry {
512 key: position.key(),
513 depth,
514 score: best,
515 bound,
516 best_move,
517 };
518 if self
519 .table
520 .as_mut()
521 .expect("table algorithm owns a table")
522 .store(entry)
523 {
524 self.stats.tt_stores += 1;
525 }
526 }
527
528 NodeResult {
529 score: best,
530 line: best_line,
531 stopped: false,
532 }
533 }
534
535 fn plain_node(&mut self, position: &mut Position, depth: u8, ply: u8) -> NodeResult {
536 let mut best = Score(-INFINITY);
537 let mut best_move = None;
538 let mut best_line = Line::default();
539 for mv in ordered_moves(*position, self.algorithm, None) {
540 let undo = position
541 .make_move(mv)
542 .expect("generated moves must be legal");
543 let child = self.node(
544 position,
545 depth - 1,
546 ply + 1,
547 Score(-INFINITY),
548 Score(INFINITY),
549 );
550 position.unmake_move(undo);
551 let score = child.score.negated();
552 if score > best || (score == best && prefer(mv, best_move)) {
553 best = score;
554 best_move = Some(mv);
555 best_line = Line::prepend(mv, child.line);
556 }
557 if child.stopped {
558 return NodeResult {
559 score: best,
560 line: best_line,
561 stopped: true,
562 };
563 }
564 }
565 NodeResult {
566 score: best,
567 line: best_line,
568 stopped: false,
569 }
570 }
571}
572
573pub fn search(position: Position, algorithm: Algorithm, limits: SearchLimits) -> SearchReport {
576 let mut searcher = Searcher::new(algorithm, limits);
577 search_with(&mut searcher, position, limits, None)
578}
579
580pub fn iterative_search(
584 position: Position,
585 algorithm: Algorithm,
586 limits: SearchLimits,
587) -> IterativeSearchReport {
588 if limits.depth == 0 {
589 let result = search(position, algorithm, limits);
590 return IterativeSearchReport {
591 requested_depth: 0,
592 completed_depth: 0,
593 completed: result.completed,
594 total_nodes: result.stats.nodes,
595 iterations: Vec::new(),
596 result,
597 };
598 }
599
600 let mut searcher = Searcher::new(algorithm, SearchLimits::depth(1));
601 let mut iterations = Vec::with_capacity(usize::from(limits.depth));
602 let mut total_nodes = 0u64;
603 let mut completed_depth = 0u8;
604 let mut root_hint = None;
605 let mut best_complete = None;
606 let mut last_attempt = None;
607
608 for depth in 1..=limits.depth {
609 let remaining = limits
610 .nodes
611 .map(|budget| budget.saturating_sub(total_nodes));
612 if remaining == Some(0) {
613 break;
614 }
615 let iteration_limits = SearchLimits {
616 depth,
617 nodes: remaining,
618 };
619 let report = search_with(&mut searcher, position, iteration_limits, root_hint);
620 total_nodes += report.stats.nodes;
621 iterations.push(IterationSummary {
622 depth,
623 completed: report.completed,
624 best_move: report.best_move,
625 score: report.score,
626 nodes: report.stats.nodes,
627 cutoffs: report.stats.cutoffs,
628 tt_hits: report.stats.tt_hits,
629 });
630 if report.completed {
631 completed_depth = depth;
632 root_hint = report.best_move;
633 best_complete = Some(report);
634 } else {
635 last_attempt = Some(report);
636 break;
637 }
638 }
639
640 let result = best_complete.or(last_attempt).unwrap_or_else(|| {
641 search_with(
642 &mut searcher,
643 position,
644 SearchLimits {
645 depth: 1,
646 nodes: Some(0),
647 },
648 None,
649 )
650 });
651 IterativeSearchReport {
652 requested_depth: limits.depth,
653 completed_depth,
654 completed: completed_depth == limits.depth,
655 total_nodes,
656 iterations,
657 result,
658 }
659}
660
661fn search_with(
662 searcher: &mut Searcher,
663 position: Position,
664 limits: SearchLimits,
665 root_hint: Option<Move>,
666) -> SearchReport {
667 searcher.begin_iteration(limits);
668 let algorithm = searcher.algorithm;
669 let mut position = position;
670 if !searcher.enter_node(0) {
671 return empty_report(position, algorithm, limits, false, searcher.stats);
672 }
673 if let Some(score) = terminal_score(position, 0) {
674 searcher.stats.leaves = 1;
675 return SearchReport {
676 algorithm,
677 requested_depth: limits.depth,
678 completed: true,
679 best_move: None,
680 score,
681 principal_variation: Vec::new(),
682 root_branches: Vec::new(),
683 stats: searcher.stats,
684 };
685 }
686 if limits.depth == 0 {
687 searcher.stats.leaves = 1;
688 return empty_report(position, algorithm, limits, true, searcher.stats);
689 }
690
691 let mut best = Score(-INFINITY);
692 let mut best_move = None;
693 let mut best_line = Line::default();
694 let mut root_branches = Vec::with_capacity(Column::COUNT);
695 let mut completed = true;
696
697 for mv in ordered_moves(position, algorithm, root_hint) {
698 let nodes_before = searcher.stats.nodes;
699 let cutoffs_before = searcher.stats.cutoffs;
700 let undo = position
701 .make_move(mv)
702 .expect("generated moves must be legal");
703 let child = searcher.node(
707 &mut position,
708 limits.depth - 1,
709 1,
710 Score(-INFINITY),
711 Score(INFINITY),
712 );
713 position.unmake_move(undo);
714 let score = child.score.negated();
715 root_branches.push(RootBranch {
716 mv,
717 score,
718 nodes: searcher.stats.nodes - nodes_before,
719 cutoffs: searcher.stats.cutoffs - cutoffs_before,
720 completed: !child.stopped,
721 });
722 if score > best || (score == best && prefer(mv, best_move)) {
723 best = score;
724 best_move = Some(mv);
725 best_line = Line::prepend(mv, child.line);
726 }
727 if child.stopped {
728 completed = false;
729 break;
730 }
731 }
732
733 SearchReport {
734 algorithm,
735 requested_depth: limits.depth,
736 completed,
737 best_move,
738 score: if best_move.is_some() {
739 best
740 } else {
741 evaluate(position)
742 },
743 principal_variation: best_line.to_vec(),
744 root_branches,
745 stats: searcher.stats,
746 }
747}
748
749fn empty_report(
750 position: Position,
751 algorithm: Algorithm,
752 limits: SearchLimits,
753 completed: bool,
754 stats: SearchStats,
755) -> SearchReport {
756 SearchReport {
757 algorithm,
758 requested_depth: limits.depth,
759 completed,
760 best_move: None,
761 score: evaluate(position),
762 principal_variation: Vec::new(),
763 root_branches: Vec::new(),
764 stats,
765 }
766}
767
768fn ordered_moves(position: Position, algorithm: Algorithm, first: Option<Move>) -> MoveList {
769 let mut moves = MoveList::default();
770 if position.result() != GameResult::Ongoing {
771 return moves;
772 }
773 if first.is_some_and(|mv| position.can_play(mv.column())) {
774 moves.push(first.expect("checked above"));
775 }
776 let columns = if algorithm.uses_center_order() {
777 [3, 2, 4, 1, 5, 0, 6]
778 } else {
779 [0, 1, 2, 3, 4, 5, 6]
780 };
781 for index in columns {
782 let mv = Move::new(Column::new(index));
783 if Some(mv) != first && position.can_play(mv.column()) {
784 moves.push(mv);
785 }
786 }
787 moves
788}
789
790fn prefer(candidate: Move, current: Option<Move>) -> bool {
791 const PRIORITY: [u8; Column::COUNT] = [5, 3, 1, 0, 2, 4, 6];
792 current.is_none_or(|mv| PRIORITY[candidate.column().index()] < PRIORITY[mv.column().index()])
793}
794
795fn terminal_score(position: Position, ply: u8) -> Option<Score> {
796 match position.result() {
797 GameResult::Ongoing => None,
798 GameResult::Draw => Some(Score(0)),
799 GameResult::Win(winner) => {
800 let distance = i16::from(ply);
801 if winner == position.side_to_move() {
802 Some(Score(FORCED_SCORE - distance))
803 } else {
804 Some(Score(-FORCED_SCORE + distance))
805 }
806 }
807 }
808}
809
810fn evaluate(position: Position) -> Score {
814 let us = position.side_to_move();
815 let them = us.other();
816 let mut value = center_count(position, us) * 3 - center_count(position, them) * 3;
817 for column in 0..Column::COUNT as i8 {
818 for row in 0..crate::position::HEIGHT as i8 {
819 for (dc, dr) in [(1i8, 0i8), (0, 1), (1, 1), (1, -1)] {
820 let end_column = column + dc * 3;
821 let end_row = row + dr * 3;
822 if !(0..Column::COUNT as i8).contains(&end_column)
823 || !(0..crate::position::HEIGHT as i8).contains(&end_row)
824 {
825 continue;
826 }
827 let mut ours = 0;
828 let mut theirs = 0;
829 for step in 0..4 {
830 let cell = crate::Cell::new(
831 Column::new((column + dc * step) as u8),
832 (row + dr * step) as u8,
833 );
834 match position.side_at(cell) {
835 Some(side) if side == us => ours += 1,
836 Some(_) => theirs += 1,
837 None => {}
838 }
839 }
840 if theirs == 0 {
841 value += window_weight(ours);
842 } else if ours == 0 {
843 value -= window_weight(theirs);
844 }
845 }
846 }
847 }
848 Score(value.clamp(-HEURISTIC_LIMIT, HEURISTIC_LIMIT))
849}
850
851fn center_count(position: Position, side: Side) -> i16 {
852 let center = Column::new(3);
853 (0..crate::position::HEIGHT as u8)
854 .filter(|row| position.side_at(crate::Cell::new(center, *row)) == Some(side))
855 .count() as i16
856}
857
858const fn window_weight(discs: i16) -> i16 {
859 match discs {
860 1 => 1,
861 2 => 8,
862 3 => 40,
863 _ => 0,
864 }
865}