Skip to main content

hermes_simd_intrinsics/bitboard/
magic.rs

1//! Fancy Magic Bitboards sliding attack generation with lazy table initialization.
2
3use core::cell::UnsafeCell;
4use core::sync::atomic::{AtomicUsize, Ordering};
5
6/// Minimal `no_std`-compatible once-initialization cell used for the lazily
7/// built magic attack tables (spin-based; initialization races run `f` once).
8pub struct OnceLock<T> {
9    state: AtomicUsize, // 0 = uninitialized, 1 = initializing, 2 = initialized
10    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    /// Create an empty, uninitialized cell.
18    pub const fn new() -> Self {
19        Self {
20            state: AtomicUsize::new(0),
21            value: UnsafeCell::new(None),
22        }
23    }
24
25    /// Return the stored value, running `f` exactly once to initialize it;
26    /// concurrent callers spin until initialization completes.
27    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                    // Winner acquires no shared data on the 0->1 claim (it next
48                    // *writes* the value and publishes it via `store(2, Release)`),
49                    // so `Relaxed` success ordering is sufficient.
50                    .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/// ZST marker for Fancy Magic Bitboards backend.
74#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
75pub struct Magic;
76use super::swar::attack_ray;
77
78// Helper functions for directional shifting
79#[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/// Construct Rook occupancy mask (excluding edges).
113#[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/// Construct Bishop occupancy mask (excluding edges).
130#[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
160/// Generate occupancy pattern for a mask and index.
161fn 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
310// Flat lookup tables and offsets
311struct 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        // Populate tables using static precomputed magics
344        for sq in 0..64 {
345            // Rook
346            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            // Bishop
365            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}