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#[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 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 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 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}