shared/bot/BitBoard.ts

import { isBoardCoordInBounds, serializeCoord } from '../board.ts'
import type { Cell, SerializedBitBoard } from '../types.ts'
import type { PlayerColor } from '../dataset/types.ts'
import { getWinMasks, type WinDirection } from './WinMasks.ts'

const PLAYERS = 2

export class BitBoard {
  readonly boardSize: number
  readonly cellCount: number
  readonly wordCount: number
  readonly words: Uint32Array
  occupiedCount: number

  constructor(boardSize: number, words?: Uint32Array | readonly number[], _occupiedCount?: number) {
    this.boardSize = boardSize
    this.cellCount = boardSize * boardSize
    this.wordCount = Math.ceil(this.cellCount / 32)
    const expectedWords = this.wordCount * PLAYERS
    this.words = words ? new Uint32Array(words) : new Uint32Array(expectedWords)
    if (this.words.length !== expectedWords) {
      throw new Error(`BitBoard expected ${expectedWords} words, got ${this.words.length}`)
    }
    sanitizePaddingBits(this.words, this.wordCount, this.cellCount)
    this.occupiedCount = countOccupied(this.words, this.wordCount)
  }

  static empty(boardSize: number): BitBoard {
    return new BitBoard(boardSize)
  }

  static deserialize(data: SerializedBitBoard): BitBoard {
    return new BitBoard(data.boardSize, data.words, data.occupiedCount)
  }

  static fromBoard(board: Map<string, Cell> | Record<string, Cell>, boardSize: number): BitBoard {
    const bitBoard = new BitBoard(boardSize)
    const entries = board instanceof Map ? board.entries() : Object.entries(board)
    for (const [key, cell] of entries) {
      const [x, y] = key.split(':').map(Number)
      const player = Number(cell.playerId)
      if ((player !== 0 && player !== 1) || !isBoardCoordInBounds(x, y, boardSize)) continue
      bitBoard.place(x, y, player)
    }
    return bitBoard
  }

  clone(): BitBoard {
    return new BitBoard(this.boardSize, this.words, this.occupiedCount)
  }

  serialize(): SerializedBitBoard {
    return {
      boardSize: this.boardSize,
      words: Array.from(this.words),
      occupiedCount: this.occupiedCount,
    }
  }

  has(x: number, y: number): boolean {
    if (!isBoardCoordInBounds(x, y, this.boardSize)) return false
    const { chunk, bit } = this.bitAt(x, y)
    return ((this.words[chunk] | this.words[this.wordCount + chunk]) & bit) !== 0
  }

  hasPlayer(x: number, y: number, player: PlayerColor): boolean {
    if (!isBoardCoordInBounds(x, y, this.boardSize)) return false
    const { chunk, bit } = this.bitAt(x, y)
    return (this.words[this.offset(player) + chunk] & bit) !== 0
  }

  playerAt(x: number, y: number): PlayerColor | null {
    if (this.hasPlayer(x, y, 0)) return 0
    if (this.hasPlayer(x, y, 1)) return 1
    return null
  }

  getCell(x: number, y: number, symbols: readonly [string, string], timestamp = 0): Cell | undefined {
    const player = this.playerAt(x, y)
    if (player === null) return undefined
    return { playerId: String(player), symbol: symbols[player], timestamp }
  }

  place(x: number, y: number, player: PlayerColor): void {
    if (!isBoardCoordInBounds(x, y, this.boardSize)) {
      throw new Error(`BitBoard move out of bounds: ${x}:${y}`)
    }
    if (this.has(x, y)) {
      throw new Error(`BitBoard move is occupied: ${x}:${y}`)
    }
    const { chunk, bit } = this.bitAt(x, y)
    this.words[this.offset(player) + chunk] |= bit
    this.occupiedCount++
  }

  remove(x: number, y: number): void {
    if (!isBoardCoordInBounds(x, y, this.boardSize) || !this.has(x, y)) return
    const { chunk, bit } = this.bitAt(x, y)
    this.words[chunk] &= ~bit
    this.words[this.wordCount + chunk] &= ~bit
    this.occupiedCount--
  }

  clear(): void {
    this.words.fill(0)
    this.occupiedCount = 0
  }

  withMove(x: number, y: number, player: PlayerColor): BitBoard {
    const next = this.clone()
    next.place(x, y, player)
    return next
  }

  swapPlayers(): void {
    for (let i = 0; i < this.wordCount; i++) {
      const tmp = this.words[i]
      this.words[i] = this.words[this.wordCount + i]
      this.words[this.wordCount + i] = tmp
    }
  }

  forEachCell(fn: (x: number, y: number, player: PlayerColor) => void): void {
    for (let y = 0; y < this.boardSize; y++) {
      for (let x = 0; x < this.boardSize; x++) {
        const player = this.playerAt(x, y)
        if (player !== null) fn(x, y, player)
      }
    }
  }

  entries(symbols: readonly [string, string], timestamp = 0): [string, Cell][] {
    const entries: [string, Cell][] = []
    this.forEachCell((x, y, player) => {
      entries.push([serializeCoord(x, y), { playerId: String(player), symbol: symbols[player], timestamp }])
    })
    return entries
  }

  toMap(symbols: readonly [string, string], timestamp = 0): Map<string, Cell> {
    return new Map(this.entries(symbols, timestamp))
  }

  key(): string {
    return `${this.boardSize}:${Array.from(this.words).map(word => word.toString(16)).join(':')}`
  }

  private offset(player: PlayerColor): number {
    return player * this.wordCount
  }

  private bitAt(x: number, y: number): { chunk: number; bit: number } {
    const index = y * this.boardSize + x
    return { chunk: index >>> 5, bit: 1 << (index & 31) }
  }
}

export function bitBoardWinResult(
  board: BitBoard,
  x: number,
  y: number,
  player: PlayerColor,
  noOverlines?: boolean,
): { won: boolean; cells: { x: number; y: number }[] } {
  if (!isBoardCoordInBounds(x, y, board.boardSize)) return { won: false, cells: [] }
  const index = y * board.boardSize + x
  const offset = player * board.wordCount
  const masks = getWinMasks(board.boardSize).byCellIndex[index]
  for (const mask of masks) {
    if (!matchesMask(board.words, offset, mask.words)) continue
    const result = validateWinDirection(board, x, y, player, mask.direction, noOverlines)
    if (result.won) return result
  }
  return { won: false, cells: [] }
}

export function bitBoardCheckWin(
  board: BitBoard,
  x: number,
  y: number,
  player: PlayerColor,
  noOverlines?: boolean,
): boolean {
  return bitBoardWinResult(board, x, y, player, noOverlines).won
}

export function bitBoardToMap(board: BitBoard, symbols: [string, string]): Map<string, Cell> {
  return board.toMap(symbols)
}

function collectDirection(
  board: BitBoard,
  x: number,
  y: number,
  player: PlayerColor,
  dx: number,
  dy: number,
): { cells: { x: number; y: number }[]; end: { x: number; y: number } } {
  const cells: { x: number; y: number }[] = []
  let cx = x + dx
  let cy = y + dy
  while (board.hasPlayer(cx, cy, player)) {
    cells.push({ x: cx, y: cy })
    cx += dx
    cy += dy
  }
  return { cells, end: { x: cx, y: cy } }
}

function validateWinDirection(
  board: BitBoard,
  x: number,
  y: number,
  player: PlayerColor,
  direction: WinDirection,
  noOverlines?: boolean,
): { won: boolean; cells: { x: number; y: number }[] } {
  const [dx, dy] = direction
  const cells: { x: number; y: number }[] = [{ x, y }]
  const pos = collectDirection(board, x, y, player, dx, dy)
  const neg = collectDirection(board, x, y, player, -dx, -dy)
  const count = 1 + pos.cells.length + neg.cells.length
  if (count < 5) return { won: false, cells: [] }
  if (noOverlines && count > 5) return { won: false, cells: [] }
  const bothBlocked = board.has(pos.end.x, pos.end.y) && board.has(neg.end.x, neg.end.y)
  if (bothBlocked) return { won: false, cells: [] }
  return { won: true, cells: cells.concat(pos.cells, neg.cells) }
}

function matchesMask(words: Uint32Array, offset: number, maskWords: Uint32Array): boolean {
  for (let i = 0; i < maskWords.length; i++) {
    const mask = maskWords[i]
    if (mask === 0) continue
    if (((words[offset + i] & mask) >>> 0) !== mask) return false
  }
  return true
}

function sanitizePaddingBits(words: Uint32Array, wordCount: number, cellCount: number): void {
  const validBits = cellCount & 31
  if (validBits === 0) return
  const mask = 2 ** validBits - 1
  for (let player = 0; player < PLAYERS; player++) {
    const index = player * wordCount + wordCount - 1
    const next = words[index] & mask
    if (words[index] !== next) words[index] = next
  }
}

function countOccupied(words: Uint32Array, wordCount: number): number {
  let count = 0
  for (let i = 0; i < wordCount; i++) {
    count += popCount((words[i] | words[wordCount + i]) >>> 0)
  }
  return count
}

function popCount(value: number): number {
  let count = 0
  let word = value >>> 0
  while (word !== 0) {
    word &= word - 1
    count++
  }
  return count
}