1use super::models::{Color, GameState, Piece, PieceType};
2
3const ORTHOGONAL: [(i8, i8); 4] = [(1, 0), (-1, 0), (0, 1), (0, -1)];
4const DIAGONAL: [(i8, i8); 4] = [(1, 1), (-1, 1), (1, -1), (-1, -1)];
5const KNIGHT_JUMPS: [(i8, i8); 8] = [
6 (2, 1),
7 (1, 2),
8 (-1, 2),
9 (-2, 1),
10 (-2, -1),
11 (-1, -2),
12 (1, -2),
13 (2, -1),
14];
15
16pub fn get_pseudo_legal_attacks(state: &GameState, sq_idx: usize, piece: Piece) -> Vec<usize> {
18 let mut attacks = Vec::new();
19 let rank = (sq_idx / 8) as i8;
20 let file = (sq_idx % 8) as i8;
21
22 let mut slide = |directions: &[(i8, i8)], continuous: bool| {
23 for &(dr, df) in directions {
24 let mut r = rank;
25 let mut f = file;
26 loop {
27 r += dr;
28 f += df;
29
30 if !(0..=7).contains(&r) || !(0..=7).contains(&f) {
32 break;
33 }
34
35 let target_idx = (r * 8 + f) as usize;
36
37 attacks.push(target_idx);
39
40 if state.board[target_idx].is_some() {
42 break;
43 }
44
45 if !continuous {
46 break;
47 }
48 }
49 }
50 };
51
52 match piece.piece_type {
53 PieceType::Rook => slide(&ORTHOGONAL, true),
54 PieceType::Bishop => slide(&DIAGONAL, true),
55 PieceType::Queen => {
56 slide(&ORTHOGONAL, true);
57 slide(&DIAGONAL, true);
58 }
59 PieceType::Knight => {
60 for &(dr, df) in KNIGHT_JUMPS.iter() {
61 let r = rank + dr;
62 let f = file + df;
63 if (0..=7).contains(&r) && (0..=7).contains(&f) {
64 attacks.push((r * 8 + f) as usize);
65 }
66 }
67 }
68 PieceType::King => {
69 for &(dr, df) in ORTHOGONAL.iter().chain(DIAGONAL.iter()) {
70 let r = rank + dr;
71 let f = file + df;
72 if (0..=7).contains(&r) && (0..=7).contains(&f) {
73 attacks.push((r * 8 + f) as usize);
74 }
75 }
76
77 if false {
80 if state.castling_rights.contains('K') {
82 if state.board[61].is_none()
84 && state.board[62].is_none()
85 && !is_square_attacked(state, 60, Color::Black)
86 && !is_square_attacked(state, 61, Color::Black)
87 && !is_square_attacked(state, 62, Color::Black)
88 {
89 attacks.push(62);
90 }
91 }
92 if state.castling_rights.contains('Q') {
93 if state.board[59].is_none()
95 && state.board[58].is_none()
96 && state.board[57].is_none()
97 && !is_square_attacked(state, 60, Color::Black)
98 && !is_square_attacked(state, 59, Color::Black)
99 && !is_square_attacked(state, 58, Color::Black)
100 {
101 attacks.push(58);
102 }
103 }
104 } else if piece.color == Color::Black && sq_idx == 4 {
105 if state.castling_rights.contains('k') {
107 if state.board[5].is_none()
109 && state.board[6].is_none()
110 && !is_square_attacked(state, 4, Color::White)
111 && !is_square_attacked(state, 5, Color::White)
112 && !is_square_attacked(state, 6, Color::White)
113 {
114 attacks.push(6);
115 }
116 }
117 if state.castling_rights.contains('q') {
118 if state.board[3].is_none()
120 && state.board[2].is_none()
121 && state.board[1].is_none()
122 && !is_square_attacked(state, 4, Color::White)
123 && !is_square_attacked(state, 3, Color::White)
124 && !is_square_attacked(state, 2, Color::White)
125 {
126 attacks.push(2);
127 }
128 }
129 }
130 }
131 PieceType::Pawn => {
132 let dr = if piece.color == Color::White { -1 } else { 1 };
134 for df in [-1_i8, 1_i8] {
135 let r = rank + dr;
136 let f = file + df;
137 if (0..=7).contains(&r) && (0..=7).contains(&f) {
138 attacks.push((r * 8 + f) as usize);
139 }
140 }
141 }
142 }
143
144 attacks
145}
146
147pub fn get_legal_moves(state: &GameState, sq_idx: usize, piece: Piece) -> Vec<usize> {
149 let mut moves = Vec::new();
150 let rank = (sq_idx / 8) as i8;
151 let file = (sq_idx % 8) as i8;
152
153 if piece.piece_type != PieceType::Pawn {
154 let attacks = get_pseudo_legal_attacks(state, sq_idx, piece);
156 for target_idx in attacks {
157 if let Some(target_piece) = state.board[target_idx] {
158 if target_piece.color != piece.color {
160 moves.push(target_idx);
161 }
162 } else {
163 moves.push(target_idx);
164 }
165 }
166 } else {
167 let dir = if piece.color == Color::White { -1 } else { 1 };
169
170 let forward_r = rank + dir;
172 if (0..=7).contains(&forward_r) {
173 let forward_idx = (forward_r * 8 + file) as usize;
174 if state.board[forward_idx].is_none() {
175 moves.push(forward_idx);
176
177 let start_rank = if piece.color == Color::White { 6 } else { 1 };
179 if rank == start_rank {
180 let double_r = rank + 2 * dir;
181 let double_idx = (double_r * 8 + file) as usize;
182 if state.board[double_idx].is_none() {
183 moves.push(double_idx);
184 }
185 }
186 }
187 }
188
189 for df in [-1_i8, 1_i8] {
191 let cap_r = rank + dir;
192 let cap_f = file + df;
193 if (0..=7).contains(&cap_r) && (0..=7).contains(&cap_f) {
194 let cap_idx = (cap_r * 8 + cap_f) as usize;
195
196 if let Some(target_piece) = state.board[cap_idx] {
198 if target_piece.color != piece.color {
199 moves.push(cap_idx);
200 }
201 }
202
203 if let Some(ep_sq) = state.en_passant_target {
205 if cap_idx == ep_sq.index {
206 moves.push(cap_idx);
207 }
208 }
209 }
210 }
211 }
212
213 if piece.piece_type == PieceType::King {
215 if piece.color == Color::White && sq_idx == 60 {
216 if state.castling_rights.contains('K') {
218 if state.board[61].is_none()
220 && state.board[62].is_none()
221 && !is_square_attacked(state, 60, Color::Black)
222 && !is_square_attacked(state, 61, Color::Black)
223 && !is_square_attacked(state, 62, Color::Black)
224 {
225 moves.push(62); }
227 }
228 if state.castling_rights.contains('Q') {
229 if state.board[59].is_none()
231 && state.board[58].is_none()
232 && state.board[57].is_none()
233 && !is_square_attacked(state, 60, Color::Black)
234 && !is_square_attacked(state, 59, Color::Black)
235 && !is_square_attacked(state, 58, Color::Black)
236 {
237 moves.push(58); }
239 }
240 } else if piece.color == Color::Black && sq_idx == 4 {
241 if state.castling_rights.contains('k') {
243 if state.board[5].is_none()
245 && state.board[6].is_none()
246 && !is_square_attacked(state, 4, Color::White)
247 && !is_square_attacked(state, 5, Color::White)
248 && !is_square_attacked(state, 6, Color::White)
249 {
250 moves.push(6); }
252 }
253 if state.castling_rights.contains('q') {
254 if state.board[3].is_none()
256 && state.board[2].is_none()
257 && state.board[1].is_none()
258 && !is_square_attacked(state, 4, Color::White)
259 && !is_square_attacked(state, 3, Color::White)
260 && !is_square_attacked(state, 2, Color::White)
261 {
262 moves.push(2); }
264 }
265 }
266 }
267
268 let mut strictly_legal_moves = Vec::new();
270
271 for target_idx in moves {
272 let mut sim = state.clone();
273 sim.apply_move(sq_idx, target_idx, None);
274
275 if let Some(king_idx) = find_king(&sim, piece.color) {
276 if !is_square_attacked(&sim, king_idx, piece.color.opposite()) {
278 strictly_legal_moves.push(target_idx); }
280 } else {
281 strictly_legal_moves.push(target_idx);
283 }
284 }
285
286 strictly_legal_moves
287}
288
289pub fn is_square_attacked(state: &GameState, target_idx: usize, attacker_color: Color) -> bool {
291 for sq_idx in 0..64 {
292 if let Some(piece) = state.board[sq_idx] {
293 if piece.color == attacker_color {
294 let attacks = get_pseudo_legal_attacks(state, sq_idx, piece);
295 if attacks.contains(&target_idx) {
296 return true;
297 }
298 }
299 }
300 }
301 false
302}
303
304pub fn find_king(state: &GameState, color: Color) -> Option<usize> {
306 for sq_idx in 0..64 {
307 if let Some(piece) = state.board[sq_idx] {
308 if piece.color == color && piece.piece_type == PieceType::King {
309 return Some(sq_idx);
310 }
311 }
312 }
313 None
314}
315
316#[cfg(test)]
317mod tests {
318 use super::*;
319
320 #[test]
321 fn test_pinned_knight_cannot_move() {
322 let fen = "4r3/8/8/8/8/8/4N3/4K3 w - - 0 1";
324 let state = GameState::from_fen(fen).unwrap();
325
326 let knight_idx = 52; let piece = state.board[knight_idx].unwrap();
328
329 let legal_moves = get_legal_moves(&state, knight_idx, piece);
331 assert_eq!(
332 legal_moves.len(),
333 0,
334 "Pinned knight violently blocked from moving!"
335 );
336 }
337
338 #[test]
339 fn test_king_check_forces_responses() {
340 let fen = "8/8/8/8/8/8/P3r3/4K3 w - - 0 1";
342 let state = GameState::from_fen(fen).unwrap();
343
344 let pawn_idx = 48; let piece = state.board[pawn_idx].unwrap();
346 let legal_moves = get_legal_moves(&state, pawn_idx, piece);
347 assert_eq!(
348 legal_moves.len(),
349 0,
350 "Idle movement discarded when King is checked!"
351 );
352
353 let king_idx = 60; let piece = state.board[king_idx].unwrap();
355 let legal_moves = get_legal_moves(&state, king_idx, piece);
356
357 assert_eq!(
359 legal_moves.len(),
360 3,
361 "King physically forced to sidestep hostile checks or capture attackers!"
362 );
363 }
364}