1use core::cell::UnsafeCell;
4use core::sync::atomic::{AtomicUsize, Ordering};
5
6pub struct OnceLock<T> {
9 state: AtomicUsize, value: UnsafeCell<Option<T>>,
11}
12
13unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
14unsafe impl<T: Send> Send for OnceLock<T> {}
15
16impl<T> OnceLock<T> {
17 pub const fn new() -> Self {
19 Self {
20 state: AtomicUsize::new(0),
21 value: UnsafeCell::new(None),
22 }
23 }
24
25 pub fn get_or_init<F>(&self, f: F) -> &T
28 where
29 F: FnOnce() -> T,
30 {
31 if self.state.load(Ordering::Acquire) == 2 {
32 return unsafe {
33 (*self.value.get())
34 .as_ref()
35 .expect("OnceLock: state==2 implies value is initialized")
36 };
37 }
38
39 loop {
40 let current = self.state.load(Ordering::Acquire);
41 if current == 2 {
42 break;
43 }
44 if current == 0 {
45 if self
46 .state
47 .compare_exchange_weak(0, 1, Ordering::Relaxed, Ordering::Relaxed)
51 .is_ok()
52 {
53 unsafe {
54 *self.value.get() = Some(f());
55 }
56 self.state.store(2, Ordering::Release);
57 break;
58 }
59 } else {
60 core::hint::spin_loop();
61 }
62 }
63 unsafe {
64 (*self.value.get())
65 .as_ref()
66 .expect("OnceLock: state==2 implies value is initialized")
67 }
68 }
69}
70
71use hermes_simd_core::bitboard::BitBoardKernel;
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
75pub struct Magic;
76use super::swar::attack_ray;
77
78#[inline(always)]
80fn shift_n(v: u64) -> u64 {
81 v << 8
82}
83#[inline(always)]
84fn shift_s(v: u64) -> u64 {
85 v >> 8
86}
87#[inline(always)]
88fn shift_e(v: u64) -> u64 {
89 (v << 1) & 0xFEFEFEFEFEFEFEFE
90}
91#[inline(always)]
92fn shift_w(v: u64) -> u64 {
93 (v >> 1) & 0x7F7F7F7F7F7F7F7F
94}
95#[inline(always)]
96fn shift_ne(v: u64) -> u64 {
97 (v << 9) & 0xFEFEFEFEFEFEFEFE
98}
99#[inline(always)]
100fn shift_nw(v: u64) -> u64 {
101 (v << 7) & 0x7F7F7F7F7F7F7F7F
102}
103#[inline(always)]
104fn shift_se(v: u64) -> u64 {
105 (v >> 7) & 0xFEFEFEFEFEFEFEFE
106}
107#[inline(always)]
108fn shift_sw(v: u64) -> u64 {
109 (v >> 9) & 0x7F7F7F7F7F7F7F7F
110}
111
112#[inline]
114pub fn rook_mask(sq: u8) -> u64 {
115 let mut mask = 0u64;
116 let r = (sq / 8) as i32;
117 let f = (sq % 8) as i32;
118 for i in 1..7 {
119 if i != r {
120 mask |= 1u64 << (i * 8 + f);
121 }
122 if i != f {
123 mask |= 1u64 << (r * 8 + i);
124 }
125 }
126 mask
127}
128
129#[inline]
131pub fn bishop_mask(sq: u8) -> u64 {
132 let mut mask = 0u64;
133 let r = (sq / 8) as i32;
134 let f = (sq % 8) as i32;
135 for i in 1..7 {
136 let nr = r + i;
137 let nf = f + i;
138 if nr < 7 && nf < 7 {
139 mask |= 1u64 << (nr * 8 + nf);
140 }
141 let nr = r + i;
142 let nf = f - i;
143 if nr < 7 && nf > 0 {
144 mask |= 1u64 << (nr * 8 + nf);
145 }
146 let nr = r - i;
147 let nf = f + i;
148 if nr > 0 && nf < 7 {
149 mask |= 1u64 << (nr * 8 + nf);
150 }
151 let nr = r - i;
152 let nf = f - i;
153 if nr > 0 && nf > 0 {
154 mask |= 1u64 << (nr * 8 + nf);
155 }
156 }
157 mask
158}
159
160fn get_occupancy(index: usize, mask: u64) -> u64 {
162 let mut occupancy = 0u64;
163 let mut m = mask;
164 let mut i = index;
165 while m != 0 {
166 let lsb = m & m.wrapping_neg();
167 m ^= lsb;
168 if (i & 1) != 0 {
169 occupancy |= lsb;
170 }
171 i >>= 1;
172 }
173 occupancy
174}
175
176const ROOK_MAGICS: [u64; 64] = [
177 612507691268440096,
178 612507278951579776,
179 72092778695245832,
180 13907124514132861184,
181 36037595259731970,
182 4683752412853567490,
183 144116304901620224,
184 72064193267807488,
185 422213022908800,
186 2324420636851601920,
187 5066833117843584,
188 18437160864616448,
189 1191905823245468288,
190 73746496243761160,
191 217298750773725188,
192 1162491709658234966,
193 18014948810555456,
194 360571644730802304,
195 4503737068946432,
196 1154057300386783488,
197 2252351717508096,
198 2306969184032063552,
199 722409925791711888,
200 8070457129321923844,
201 140881369776129,
202 2305930988397535232,
203 727366526348300288,
204 2882338984544960769,
205 2314854612957921664,
206 1226109398745743488,
207 18016614716868624,
208 3940933141856388,
209 70438008914056,
210 1441292755694272521,
211 13935298335742103616,
212 650771245936148992,
213 8800404767760,
214 4611826764514067456,
215 18304738365800706,
216 2392641597604096,
217 18014675537002496,
218 72708505458589696,
219 2306124621734543376,
220 7318633399844880,
221 288239172530012164,
222 1125908513570880,
223 423793149607938,
224 72128522269097987,
225 5080060654553923648,
226 108368149504100864,
227 9228018641916657792,
228 621506644450222336,
229 1155173871364606080,
230 288934080773488768,
231 11530377299149456384,
232 6918392146485592576,
233 88510740570371,
234 5782639788609832065,
235 144396738350809281,
236 157913028229072929,
237 5765170767461884162,
238 562954517088258,
239 2307021690107888132,
240 2815026809283586,
241];
242
243const BISHOP_MAGICS: [u64; 64] = [
244 2900320361214181632,
245 4521240153325862,
246 307415763168788480,
247 577591189911371776,
248 9307640807030792,
249 299076021126434,
250 4648803916316672,
251 8368260145767256068,
252 297246444856279360,
253 9295447291819266081,
254 4416308641792,
255 81073727902580769,
256 720580351445188672,
257 2314992114330763778,
258 5718564289912896,
259 612666572850601984,
260 18159602805933060,
261 94575605329300608,
262 2778721004624613648,
263 2251853534609920,
264 9259682319860499464,
265 72198366053867524,
266 4612249038727946256,
267 2311648435029673024,
268 297246406417580304,
269 2253998907228672,
270 1127016866841633,
271 4644474689028384,
272 3207833970406400000,
273 142938675888640,
274 2265548827396352,
275 369437041106157696,
276 608023367451217920,
277 9224515546128126208,
278 126127195025834048,
279 70403640918528,
280 657526662304433160,
281 155391783633690752,
282 18586217575221400,
283 595040317506339856,
284 2451251326177329220,
285 2305915887359631364,
286 288511954241199106,
287 4611690700075958400,
288 316745583693888,
289 6926537335024648708,
290 5480885148919529568,
291 16141208930453029254,
292 2883465404411281498,
293 576976457634414720,
294 986410176759168,
295 702579169524056576,
296 4693455495865565184,
297 90297529081856,
298 4618466725024890896,
299 580965460101579777,
300 2305988728998461956,
301 565254270747148,
302 9223372451327787272,
303 72057733633020163,
304 1127034055164428,
305 5225583234828469520,
306 904169302867347712,
307 2612237321755557952,
308];
309
310struct MagicTable {
312 rook_table: alloc::vec::Vec<u64>,
313 bishop_table: alloc::vec::Vec<u64>,
314 rook_offsets: [usize; 64],
315 bishop_offsets: [usize; 64],
316 rook_magics: [u64; 64],
317 bishop_magics: [u64; 64],
318}
319
320static MAGIC_DATA: OnceLock<MagicTable> = OnceLock::new();
321
322fn get_magic_data() -> &'static MagicTable {
323 MAGIC_DATA.get_or_init(|| {
324 let mut rook_offsets = [0; 64];
325 let mut bishop_offsets = [0; 64];
326
327 let mut rook_total_size = 0;
328 let mut bishop_total_size = 0;
329
330 for sq in 0..64 {
331 rook_offsets[sq] = rook_total_size;
332 let rook_pop = rook_mask(sq as u8).count_ones() as usize;
333 rook_total_size += 1 << rook_pop;
334
335 bishop_offsets[sq] = bishop_total_size;
336 let bishop_pop = bishop_mask(sq as u8).count_ones() as usize;
337 bishop_total_size += 1 << bishop_pop;
338 }
339
340 let mut rook_table = alloc::vec![0u64; rook_total_size];
341 let mut bishop_table = alloc::vec![0u64; bishop_total_size];
342
343 for sq in 0..64 {
345 let r_mask = rook_mask(sq as u8);
347 let r_pop = r_mask.count_ones() as usize;
348 let r_num_patterns = 1 << r_pop;
349 let r_magic = ROOK_MAGICS[sq];
350 let r_shift = 64 - r_pop;
351 let r_offset = rook_offsets[sq];
352
353 for idx in 0..r_num_patterns {
354 let occ = get_occupancy(idx, r_mask);
355 let slider = 1u64 << sq;
356 let att = attack_ray(slider, occ, shift_n)
357 | attack_ray(slider, occ, shift_s)
358 | attack_ray(slider, occ, shift_e)
359 | attack_ray(slider, occ, shift_w);
360 let hash = ((occ.wrapping_mul(r_magic)) >> r_shift) as usize;
361 rook_table[r_offset + hash] = att;
362 }
363
364 let b_mask = bishop_mask(sq as u8);
366 let b_pop = b_mask.count_ones() as usize;
367 let b_num_patterns = 1 << b_pop;
368 let b_magic = BISHOP_MAGICS[sq];
369 let b_shift = 64 - b_pop;
370 let b_offset = bishop_offsets[sq];
371
372 for idx in 0..b_num_patterns {
373 let occ = get_occupancy(idx, b_mask);
374 let slider = 1u64 << sq;
375 let att = attack_ray(slider, occ, shift_ne)
376 | attack_ray(slider, occ, shift_nw)
377 | attack_ray(slider, occ, shift_se)
378 | attack_ray(slider, occ, shift_sw);
379 let hash = ((occ.wrapping_mul(b_magic)) >> b_shift) as usize;
380 bishop_table[b_offset + hash] = att;
381 }
382 }
383
384 MagicTable {
385 rook_table,
386 bishop_table,
387 rook_offsets,
388 bishop_offsets,
389 rook_magics: ROOK_MAGICS,
390 bishop_magics: BISHOP_MAGICS,
391 }
392 })
393}
394
395impl BitBoardKernel for Magic {
396 #[inline]
397 fn rook_attacks(square: u8, occupancy: u64) -> u64 {
398 let table = get_magic_data();
399 let mask = rook_mask(square);
400 let pop = mask.count_ones() as usize;
401 let magic = table.rook_magics[square as usize];
402 let shift = 64 - pop;
403 let offset = table.rook_offsets[square as usize];
404 let idx = (((occupancy & mask).wrapping_mul(magic)) >> shift) as usize;
405 table.rook_table[offset + idx]
406 }
407
408 #[inline]
409 fn bishop_attacks(square: u8, occupancy: u64) -> u64 {
410 let table = get_magic_data();
411 let mask = bishop_mask(square);
412 let pop = mask.count_ones() as usize;
413 let magic = table.bishop_magics[square as usize];
414 let shift = 64 - pop;
415 let offset = table.bishop_offsets[square as usize];
416 let idx = (((occupancy & mask).wrapping_mul(magic)) >> shift) as usize;
417 table.bishop_table[offset + idx]
418 }
419}