Skip to main content

ai_connect4/
position.rs

1use crate::mv::{Cell, Column, Move, MoveList};
2use std::fmt;
3
4pub const WIDTH: usize = Column::COUNT;
5pub const HEIGHT: usize = Cell::ROWS;
6const STRIDE: usize = HEIGHT + 1;
7const COLUMN_BITS: u64 = (1u64 << HEIGHT) - 1;
8
9const fn board_mask() -> u64 {
10    let mut mask = 0;
11    let mut column = 0;
12    while column < WIDTH {
13        mask |= COLUMN_BITS << (column * STRIDE);
14        column += 1;
15    }
16    mask
17}
18
19const fn bottom_mask() -> u64 {
20    let mut mask = 0;
21    let mut column = 0;
22    while column < WIDTH {
23        mask |= 1u64 << (column * STRIDE);
24        column += 1;
25    }
26    mask
27}
28
29pub(crate) const BOARD_MASK: u64 = board_mask();
30pub(crate) const BOTTOM_MASK: u64 = bottom_mask();
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33pub enum Side {
34    Red,
35    Yellow,
36}
37
38impl Side {
39    pub const fn other(self) -> Self {
40        match self {
41            Self::Red => Self::Yellow,
42            Self::Yellow => Self::Red,
43        }
44    }
45
46    pub const fn as_char(self) -> char {
47        match self {
48            Self::Red => 'R',
49            Self::Yellow => 'Y',
50        }
51    }
52}
53
54impl fmt::Display for Side {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        write!(f, "{}", self.as_char())
57    }
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum GameResult {
62    Ongoing,
63    Draw,
64    Win(Side),
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub enum MoveError {
69    GameOver,
70    Full(Column),
71}
72
73impl fmt::Display for MoveError {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        match self {
76            Self::GameOver => write!(f, "the game is already over"),
77            Self::Full(column) => write!(f, "column {column} is full"),
78        }
79    }
80}
81
82impl std::error::Error for MoveError {}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub struct Undo {
86    mv: Move,
87    side: Side,
88    bit: u64,
89}
90
91impl Undo {
92    pub const fn mv(self) -> Move {
93        self.mv
94    }
95
96    pub const fn side(self) -> Side {
97        self.side
98    }
99
100    pub const fn cell(self) -> Cell {
101        let row = (self.bit.trailing_zeros() as usize % STRIDE) as u8;
102        Cell::new(self.mv.column(), row)
103    }
104}
105
106/// Two 42-bit player boards inside Connect Four's 7×7 sentinel layout.
107///
108/// Each column occupies seven adjacent bits: six playable cells and one zero
109/// sentinel. The spare row makes horizontal and diagonal shifts unable to wrap
110/// from one column into the next.
111#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
112pub struct Position {
113    red: u64,
114    yellow: u64,
115    side_to_move: Side,
116}
117
118impl Default for Position {
119    fn default() -> Self {
120        Self::start()
121    }
122}
123
124impl Position {
125    pub const fn start() -> Self {
126        Self {
127            red: 0,
128            yellow: 0,
129            side_to_move: Side::Red,
130        }
131    }
132
133    pub const fn side_to_move(self) -> Side {
134        self.side_to_move
135    }
136
137    pub const fn bits(self, side: Side) -> u64 {
138        match side {
139            Side::Red => self.red,
140            Side::Yellow => self.yellow,
141        }
142    }
143
144    pub const fn occupied(self) -> u64 {
145        self.red | self.yellow
146    }
147
148    pub const fn ply(self) -> u8 {
149        self.occupied().count_ones() as u8
150    }
151
152    /// Pascal Pons' compact, collision-free key for legal Connect Four boards.
153    ///
154    /// The side-to-move board is a subset of occupancy; adding the bottom mask
155    /// encodes each column height in its sentinel bit.
156    pub const fn key(self) -> u64 {
157        self.bits(self.side_to_move) + self.occupied() + BOTTOM_MASK
158    }
159
160    pub fn side_at(self, cell: Cell) -> Option<Side> {
161        let bit = 1u64 << cell.bit_index();
162        if self.red & bit != 0 {
163            Some(Side::Red)
164        } else if self.yellow & bit != 0 {
165            Some(Side::Yellow)
166        } else {
167            None
168        }
169    }
170
171    /// One set bit per non-full column, computed with one addition.
172    pub const fn playable_bits(self) -> u64 {
173        self.occupied().wrapping_add(BOTTOM_MASK) & BOARD_MASK
174    }
175
176    pub const fn can_play(self, column: Column) -> bool {
177        self.playable_bits() & Self::column_mask(column) != 0
178    }
179
180    /// Returns legal moves in deterministic center-first order.
181    pub fn legal_moves(self) -> MoveList {
182        const ORDER: [u8; WIDTH] = [3, 2, 4, 1, 5, 0, 6];
183        let mut moves = MoveList::default();
184        if self.result() != GameResult::Ongoing {
185            return moves;
186        }
187        for index in ORDER {
188            let column = Column::new(index);
189            if self.can_play(column) {
190                moves.push(Move::new(column));
191            }
192        }
193        moves
194    }
195
196    pub fn is_winning_move(self, side: Side, mv: Move) -> bool {
197        let bit = self.playable_bits() & Self::column_mask(mv.column());
198        bit != 0 && Self::has_four(self.bits(side) | bit)
199    }
200
201    pub fn result(self) -> GameResult {
202        if Self::has_four(self.red) {
203            GameResult::Win(Side::Red)
204        } else if Self::has_four(self.yellow) {
205            GameResult::Win(Side::Yellow)
206        } else if self.occupied() == BOARD_MASK {
207            GameResult::Draw
208        } else {
209            GameResult::Ongoing
210        }
211    }
212
213    pub fn winning_cells(self) -> Vec<Cell> {
214        let winning_bits = match self.result() {
215            GameResult::Win(side) => self.bits(side),
216            GameResult::Ongoing | GameResult::Draw => return Vec::new(),
217        };
218        let directions = [(1i8, 0i8), (0, 1), (1, 1), (1, -1)];
219        let mut cells = Vec::with_capacity(7);
220        for column in 0..WIDTH as i8 {
221            for row in 0..HEIGHT as i8 {
222                for (dc, dr) in directions {
223                    let end_column = column + dc * 3;
224                    let end_row = row + dr * 3;
225                    if !(0..WIDTH as i8).contains(&end_column)
226                        || !(0..HEIGHT as i8).contains(&end_row)
227                    {
228                        continue;
229                    }
230                    let line = (0..4).map(|step| {
231                        Cell::new(
232                            Column::new((column + dc * step) as u8),
233                            (row + dr * step) as u8,
234                        )
235                    });
236                    let line: Vec<Cell> = line.collect();
237                    if line
238                        .iter()
239                        .all(|cell| winning_bits & (1u64 << cell.bit_index()) != 0)
240                    {
241                        for cell in line {
242                            if !cells.contains(&cell) {
243                                cells.push(cell);
244                            }
245                        }
246                    }
247                }
248            }
249        }
250        cells.sort_unstable();
251        cells
252    }
253
254    pub fn make_move(&mut self, mv: Move) -> Result<Undo, MoveError> {
255        if self.result() != GameResult::Ongoing {
256            return Err(MoveError::GameOver);
257        }
258        let bit = self.playable_bits() & Self::column_mask(mv.column());
259        if bit == 0 {
260            return Err(MoveError::Full(mv.column()));
261        }
262        let side = self.side_to_move;
263        match side {
264            Side::Red => self.red |= bit,
265            Side::Yellow => self.yellow |= bit,
266        }
267        self.side_to_move = side.other();
268        debug_assert!(self.is_consistent());
269        Ok(Undo { mv, side, bit })
270    }
271
272    pub fn unmake_move(&mut self, undo: Undo) {
273        debug_assert_eq!(self.side_to_move, undo.side.other());
274        debug_assert_ne!(self.bits(undo.side) & undo.bit, 0);
275        self.side_to_move = undo.side;
276        match undo.side {
277            Side::Red => self.red &= !undo.bit,
278            Side::Yellow => self.yellow &= !undo.bit,
279        }
280        debug_assert!(self.is_consistent());
281    }
282
283    pub fn from_moves(moves: &[Move]) -> Result<Self, MoveError> {
284        let mut position = Self::start();
285        for &mv in moves {
286            position.make_move(mv)?;
287        }
288        Ok(position)
289    }
290
291    pub fn mirrored(self) -> Self {
292        let mut mirrored = Self::start();
293        mirrored.side_to_move = self.side_to_move;
294        for column in Column::all() {
295            let shift = column.index() * STRIDE;
296            let target_shift = column.mirrored().index() * STRIDE;
297            let red = (self.red >> shift) & COLUMN_BITS;
298            let yellow = (self.yellow >> shift) & COLUMN_BITS;
299            mirrored.red |= red << target_shift;
300            mirrored.yellow |= yellow << target_shift;
301        }
302        mirrored
303    }
304
305    pub const fn has_four(bits: u64) -> bool {
306        let vertical = bits & (bits >> 1);
307        if vertical & (vertical >> 2) != 0 {
308            return true;
309        }
310        let horizontal = bits & (bits >> STRIDE);
311        if horizontal & (horizontal >> (2 * STRIDE)) != 0 {
312            return true;
313        }
314        let rising = bits & (bits >> (STRIDE + 1));
315        if rising & (rising >> (2 * (STRIDE + 1))) != 0 {
316            return true;
317        }
318        let falling = bits & (bits >> (STRIDE - 1));
319        falling & (falling >> (2 * (STRIDE - 1))) != 0
320    }
321
322    const fn column_mask(column: Column) -> u64 {
323        COLUMN_BITS << (column.index() * STRIDE)
324    }
325
326    fn is_consistent(self) -> bool {
327        if self.red & self.yellow != 0 || self.occupied() & !BOARD_MASK != 0 {
328            return false;
329        }
330        let red_count = self.red.count_ones();
331        let yellow_count = self.yellow.count_ones();
332        let turn_is_consistent = match self.side_to_move {
333            Side::Red => red_count == yellow_count,
334            Side::Yellow => red_count == yellow_count + 1,
335        };
336        if !turn_is_consistent {
337            return false;
338        }
339        for column in Column::all() {
340            let occupied = (self.occupied() >> (column.index() * STRIDE)) & COLUMN_BITS;
341            if occupied != (1u64 << occupied.count_ones()) - 1 {
342                return false;
343            }
344        }
345        !(Self::has_four(self.red) && Self::has_four(self.yellow))
346    }
347}