rust-uttt-new.rs

pub const WHITE: usize = 0; // X
pub const BLACK: usize = 1; // O

pub const FIELD: u32 = 0x1FF << 16;
pub const SQUARE: u32 = 0xFFFF;
pub const ALL_FIELDS_LEGAL: u32 = 0x1 << 25;

pub const DIAGS: [u32; 2] = [0o421, 0o124];
pub const ROWS: [u32; 3] = [0o700, 0o070, 0o007];
pub const COLS: [u32; 3] = [0o111, 0o222, 0o444];
pub const ALL_FIELDS: u32 = 0o777;

#[macro_export]
macro_rules! field {
    ($val:expr) => {
        ($val & FIELD) >> 16
    };
}

#[macro_export]
macro_rules! square {
    ($val:expr) => {
        $val & SQUARE
    };
}

#[derive(Copy, Clone, Debug)]
pub struct Bitboard {
    pub valid_field: u32,
    pub board: [[u32; 9]; 2],
    pub turn: usize,
    cached_meta_field: Option<[u32; 2]>,
    cached_game_over: Option<bool>,
}

impl Default for Bitboard {
    fn default() -> Self {
        Bitboard {
            valid_field: ALL_FIELDS,
            board: [[0; 9]; 2],
            turn: 0,
            cached_meta_field: Option::None,
            cached_game_over: Option::None,
        }
    }
}

impl Bitboard {
    fn dirty(&mut self) {
        self.cached_game_over = Option::None;
        self.cached_meta_field = Option::None;
    }

    fn taken(&self, pos: u32) -> bool {
        (self.board[0][to_index(field!(pos))] | self.board[1][to_index(field!(pos))]) & square!(pos)
            != 0
    }

    pub fn make_move(&mut self, mov: u32) {
        let field = field!(mov);
        let square = square!(mov);
        self.board[self.turn][to_index(field)] |= square;
        if self.field_is_blocked(to_index(square) as u32) {
            self.valid_field = ALL_FIELDS;
        } else {
            self.valid_field = square;
        }
        self.turn = 1 - self.turn;
        self.dirty();
    }

    pub fn undo_move(&mut self, mov: u32) {
        let field = field!(mov);
        let square = square!(mov);
        self.board[1 - self.turn][to_index(field)] &= !square;
        if mov & ALL_FIELDS_LEGAL != 0 {
            self.valid_field = ALL_FIELDS;
        } else {
            self.valid_field = field;
        }
        self.turn = 1 - self.turn;
        self.dirty();
    }

    fn field_is_blocked(&self, field: u32) -> bool {
        let white_field = self.board[WHITE][field as usize];
        let black_field = self.board[BLACK][field as usize];
        is_won(white_field) || is_won(black_field) || is_tied(white_field | black_field)
    }

    pub fn get_meta_field(&mut self) -> [u32; 2] {
        if self.cached_meta_field.is_some() {
            return self.cached_meta_field.unwrap();
        }

        let mut field: [u32; 2] = [0; 2];
        for p in 0..2 {
            field[p] = (is_won(self.board[p][0]) as u32) << 8
                | (is_won(self.board[p][1]) as u32) << 7
                | (is_won(self.board[p][2]) as u32) << 6
                | (is_won(self.board[p][3]) as u32) << 5
                | (is_won(self.board[p][4]) as u32) << 4
                | (is_won(self.board[p][5]) as u32) << 3
                | (is_won(self.board[p][6]) as u32) << 2
                | (is_won(self.board[p][7]) as u32) << 1
                | (is_won(self.board[p][8]) as u32) << 0;
        }
        self.cached_meta_field = Option::Some(field);
        field
    }

    pub fn game_over(&mut self) -> bool {
        if self.cached_game_over.is_some() {
            return self.cached_game_over.unwrap();
        }

        let mut sum = 0;
        for i in 0..9 {
            sum += self.board[0][i].count_ones();
        }
        if sum < 9 {
            return false;
        }
        let meta_field = self.get_meta_field();
        let ret = is_won(meta_field[WHITE]) | is_won(meta_field[BLACK]) || self.game_tied();
        self.cached_game_over = Option::Some(ret);
        return ret;
    }

    pub fn game_tied(&self) -> bool {
        self.board[0]
            .iter()
            .zip(self.board[1].iter())
            .any(|(&f0, &f1)| is_won(f0) || is_won(f1) || is_tied(f0 | f1))
    }
}

pub fn to_index(no: u32) -> usize {
    (31 - no.leading_zeros()) as usize
}

pub fn is_tied(field: u32) -> bool {
    field == ALL_FIELDS
}

fn for_each_move<F>(board: &mut Bitboard, mut func: F)
where
    F: FnMut(&mut Bitboard, u32),
{
    for i in 0..9 {
        if (1 << i) & board.valid_field != 0 {
            for s in 0..9 {
                let mov = (1 << s) | ((1 << i) << 16);
                if !board.taken(mov) && !board.field_is_blocked(i) {
                    let tweak = if board.valid_field == 0o777 {
                        ALL_FIELDS_LEGAL
                    } else {
                        0
                    };
                    func(board, mov | tweak);
                }
            }
        }
    }
}

pub fn is_won(field: u32) -> bool {
    (field & DIAGS[0]) == DIAGS[0]
        || (field & DIAGS[1]) == DIAGS[1]
        || (field & ROWS[0]) == ROWS[0]
        || (field & ROWS[1]) == ROWS[1]
        || (field & ROWS[2]) == ROWS[2]
        || (field & COLS[0]) == COLS[0]
        || (field & COLS[1]) == COLS[1]
        || (field & COLS[2]) == COLS[2]
}

pub fn move_gen(board: &mut Bitboard, depth: u32) -> u32 {
    if board.game_over() {
        0
    } else if depth == 0 {
        let mut r = 0;
        for i in 0..9 {
            if (1 << i) & board.valid_field != 0 {
                for s in 0..9 {
                    if !board.taken((1 << s) | ((1 << i) << 16)) && !board.field_is_blocked(i) {
                        r += 1;
                    }
                }
            }
        }
        r
    } else {
        let mut r = 0;
        for_each_move(board, |board, mov| {
            board.make_move(mov);
            r += 1 + move_gen(board, depth - 1);
            board.undo_move(mov);
        });
        r
    }
}