SearchCore.bs

package SearchCore where

import ChessState
import MoveEval
import MoveUpdate
import FIFO
import FIFOF
import GetPut
import ClientServer
import Connectable
import qualified List
import Vector
import CShow

type Score = Int 11
type Heuristic c = c -> State -> Score

type RequestId = UInt 8
type Depth = UInt 8
type MoveCount = UInt 8

minScore :: Score
minScore = minBound + 1  -- We want -minScore = maxScore
maxScore :: Score
maxScore = maxBound

moveQueueSize :: Integer
moveQueueSize = 60

data Outcome = NoOutcome | Check | CheckMate | Draw
  deriving (Bits)

struct SearchQuery config =
  rid :: RequestId
  state :: State
  depth :: Depth
  getMoves :: Bool
  alpha :: Maybe Score
  beta :: Maybe Score
  config :: config
 deriving (Bits)

instance (DefaultValue config) => DefaultValue (SearchQuery config) where
  defaultValue =
    interface SearchQuery
      rid = 0
      state = initialState
      depth = 5
      getMoves = False
      alpha = Nothing
      beta = Nothing
      config = defaultValue

struct SearchResult =
  rid :: RequestId
  outcome :: Outcome
  bestMove :: Maybe Move
  forcedOutcome :: Bool  -- Can either player force a win
  score :: Score
  depth :: Depth
 deriving (Bits)

interface Frame =
  putState :: State -> Bool -> Score -> Score -> Action
  putScore :: Score -> Bool -> ActionValue Bool
  state :: State
  hasMove :: Bool
  forcedOutcome :: Bool
  score :: Score
  alpha :: Score
  beta :: Score
  moves :: FIFOF Move

mkFrame :: Module Frame
mkFrame = module
  state <- mkReg _
  inCheck <- mkReg False
  hasMove <- mkReg False
  forcedOutcome <- mkReg True
  bestScore <- mkReg 0
  alpha <- mkReg minScore
  beta <- mkReg maxScore
  moves <- mkSizedFIFOF moveQueueSize

  interface
    putState s c a b =
      do state := s
         inCheck := c
         hasMove := False
         forcedOutcome := True
         alpha := a
         beta := b
    putScore s f =
      do hasMove := True
         let newScore = if f then (if s > 0 then s - 1 else s + 1) else s
             isBestScore = newScore > bestScore || not hasMove
         if isBestScore
           then do bestScore := newScore
                   forcedOutcome := f
           else noAction
         if newScore > alpha then alpha := newScore else noAction
         return isBestScore
    state = state
    hasMove = hasMove
    forcedOutcome = forcedOutcome
    score =
      if hasMove then bestScore
      else if inCheck then minScore  -- Checkmate
      else 0  -- Draw
    alpha = alpha
    beta = beta
    moves = moves

interface (SearchCore :: * -> # -> *) config stackSize =
  server :: Server (SearchQuery config) SearchResult
  moves :: GetS MoveResponse
  clear :: Action

  status :: Bit 16 {-# always_ready, always_enabled #-}

mkSearchCore ::
  (Add dispatchDepth 1 maxDepth, Bits config cb, CShow config) =>
  Heuristic config -> Vector numDispatch (SearchCore config dispatchDepth) ->
  Module (SearchCore config maxDepth)
mkSearchCore heuristic dispatchCores = module
  queries :: FIFO (SearchQuery config) <- mkSizedFIFO 1
  results :: FIFO SearchResult <- mkFIFO
  cfg :: Reg config <- mkReg _

  let rid = queries.first.rid
  depth :: Reg Depth <- mkReg 0  -- This is a register in order to be always ready for the status

  stack :: Vector maxDepth Frame <- replicateM mkFrame
  stackSize :: Reg Depth <- mkReg 0
  let bottomFrame = stack !! 0

  state :: Wire State <- mkWire
  currentPlayerInCheck :: Wire Bool <- mkWire
  otherPlayerInCheck :: Wire Bool <- mkWire
  stateHeuristicScore :: Wire Score <- mkWire
  isDraw :: Wire Bool <- mkWire

  moveUpdate :: MoveUpdate <- mkMoveUpdate
  eval :: MoveEval <- mkMoveEval
  movesComplete :: Reg Bool <- mkReg True
  let pushState :: Frame -> Score -> Score -> Action
      pushState frame alpha beta = do
        frame.putState state otherPlayerInCheck alpha beta
        eval.state.put state
        movesComplete := False

  bestMove :: Reg Move <- mkReg _
  currentMove :: Reg Move <- mkReg _
  queryStateInCheck :: Reg Bool <- mkReg _
  initialMoves :: FIFO MoveResponse <- mkFIFO

  let enqInitialMove m = if queries.first.getMoves then initialMoves.enq $ NextMove m else noAction

  rules
    "get_query_state": when stackSize == 0 ==> state := queries.first.state
    "get_move_state": when stackSize > 0 ==> state := moveUpdate.nextState

    "eval_state": when True ==> do
      -- $display "eval_state " (cshow $ inCheck state.board (otherColor state.turn)) " " (cshow $ inCheck state.board state.turn) " " (cshow $ heuristic cfg state)
      -- This seems backwards but isn't: state is the state being considered next,
      -- while the "current" state is the top of the stack.
      currentPlayerInCheck := inCheck state.board (otherColor state.turn)
      otherPlayerInCheck := inCheck state.board state.turn
      stateHeuristicScore := heuristic cfg state
      isDraw := state.lastProgressMove >= 50

    "update_query_state_in_check": when stackSize == 0 ==> queryStateInCheck := otherPlayerInCheck

  let frameRules frameIndex =
        let nextFrame = if frameIndex < valueOf maxDepth - 1 then stack !! (frameIndex + 1) else _
            currentFrame = stack !! frameIndex
            prevFrame = if frameIndex > 0 then stack !! (frameIndex - 1) else _
            isTerminal = fromInteger frameIndex >= depth - 1 || isDraw
            depthStr = (List.foldr (+++) "" $ List.replicate (frameIndex + 1) "  ") +++ integerToString (frameIndex + 1)
            requestMoveUpdate = not moveUpdate.hasRequest || fromInteger frameIndex >= depth - 1
        in
          rules
            when stackSize > 0, stackSize - 1 == fromInteger frameIndex
              rules
                ("cutoff_" +++ integerToString frameIndex): when frameIndex > 0, currentFrame.alpha >= currentFrame.beta ==> do
                  $display depthStr " cutoff " currentFrame.score
                  prevFrame.putScore (negate currentFrame.score) currentFrame.forcedOutcome
                  stackSize := fromInteger frameIndex
                  movesComplete := True
                  currentFrame.moves.clear
                  moveUpdate.clear
                  eval.clear

                when currentFrame.alpha < currentFrame.beta
                  rules
                    ("request_move_update_" +++ integerToString frameIndex): when requestMoveUpdate ==> do
                      $display depthStr " request_move_update " (cshow currentFrame.moves.first)
                      moveUpdate.enq currentFrame.state currentFrame.moves.first
                      currentFrame.moves.deq

                    ("put_NextMove_" +++ integerToString frameIndex): when not movesComplete, NextMove m <- eval.move.first ==> do
                      $display depthStr " put_NextMove " (cshow m)
                      currentFrame.moves.enq m
                      eval.move.deq
                    ("put_NoMove_" +++ integerToString frameIndex): when not movesComplete, NoMove <- eval.move.first ==> do
                      $display depthStr " put_NoMove"
                      movesComplete := True
                      eval.move.deq

                    -- This exists to avoid deadlocking if a state somehow has > 60 valid moves
                    ("discard_overflow_state_" +++ integerToString frameIndex): when not movesComplete, not requestMoveUpdate, not currentFrame.moves.notFull ==> do
                      $display depthStr " discard_overflow_state " (cshow currentFrame.moves.first)
                      currentFrame.moves.deq

                    ("ignore_check_state_" +++ integerToString frameIndex): when currentPlayerInCheck, moveUpdate.ready ==> do
                      $display depthStr " ignore_check_state " (cshow moveUpdate.nextMove)
                      moveUpdate.deq

                    ("heuristic_state_" +++ integerToString frameIndex): when not currentPlayerInCheck, isTerminal, moveUpdate.ready ==> do
                      $display depthStr " heuristic_state " (cshow moveUpdate.nextMove) " " stateHeuristicScore
                      isBestScore <- currentFrame.putScore (if isDraw then 0 else negate stateHeuristicScore) False
                      if frameIndex == 0 && isBestScore then bestMove := moveUpdate.nextMove else noAction
                      if frameIndex == 0 then enqInitialMove moveUpdate.nextMove else noAction
                      moveUpdate.deq

                    when movesComplete
                      rules
                        ("push_state_" +++ integerToString frameIndex): when not currentPlayerInCheck, not isTerminal, moveUpdate.ready ==> do
                          $display depthStr " push_state " (cshow moveUpdate.nextMove) " " stateHeuristicScore
                          pushState nextFrame (negate currentFrame.beta) (negate currentFrame.alpha)
                          if frameIndex == 0
                            then do currentMove := moveUpdate.nextMove
                                    enqInitialMove moveUpdate.nextMove
                            else noAction
                          stackSize := fromInteger frameIndex + 2
                          moveUpdate.deq

                        ("pop_state_" +++ integerToString frameIndex): when frameIndex > 0, not currentFrame.moves.notEmpty, not moveUpdate.hasRequest, not moveUpdate.ready ==> do
                          $display depthStr " pop_state " currentFrame.score
                          isBestScore <- prevFrame.putScore (negate currentFrame.score) currentFrame.forcedOutcome
                          if frameIndex == 1 && isBestScore then bestMove := currentMove else noAction
                          stackSize := fromInteger frameIndex

  dispatchMoves :: Vector numDispatch (FIFOF Move) <- replicateM mkFIFOF
  dispatchInWaiting :: Reg (UInt 4) <- mkReg $ fromInteger $ valueOf numDispatch

  (dispatchRules, clearDispatchMoveUpdate) <-
    if valueOf numDispatch > 0
      then do
        dispatchMoveUpdate :: MoveUpdate <- mkMoveUpdate
        dispatchStateInCheck :: Wire Bool <- mkWire
        addRules $
          rules
            "eval_dispatch_state": when True ==>
              dispatchStateInCheck := inCheck dispatchMoveUpdate.nextState.board bottomFrame.state.turn

        let dispatchStateRules =
              rules
                "request_dispatch_move_update": when depth > 1, dispatchInWaiting > 0, bottomFrame.alpha < bottomFrame.beta ==> do
                  $display "request_dispatch_move_update " (cshow dispatchInWaiting) " " (cshow bottomFrame.moves.first)
                  dispatchMoveUpdate.enq bottomFrame.state bottomFrame.moves.first
                  bottomFrame.moves.deq
                  dispatchInWaiting := dispatchInWaiting - 1
              `rJoinDescendingUrgency`
              rules
                "ignore_dispatch_check_state": when dispatchStateInCheck, bottomFrame.alpha < bottomFrame.beta ==> do
                  $display "ignore_dispatch_check_state " (cshow dispatchInWaiting) " " (cshow dispatchMoveUpdate.nextMove)
                  dispatchInWaiting := dispatchInWaiting + 1
                  dispatchMoveUpdate.deq

            dispatchRules i =
              rules
                ("dispatch_query_" +++ integerToString i): when not dispatchStateInCheck, bottomFrame.alpha < bottomFrame.beta ==> do
                  let query = SearchQuery {
                    rid = rid;
                    state = dispatchMoveUpdate.nextState;
                    depth = depth - 1;
                    getMoves = False;
                    alpha = Just $ negate bottomFrame.beta;
                    beta = Just $ negate bottomFrame.alpha;
                    config = cfg;
                  }
                  $display ("dispatch_query_" +++ integerToString i +++ " ") (cshow dispatchMoveUpdate.nextMove) " " (cshow query)
                  (dispatchMoves !! i).enq dispatchMoveUpdate.nextMove
                  (dispatchCores !! i).server.request.put query
                  enqInitialMove dispatchMoveUpdate.nextMove
                  dispatchMoveUpdate.deq
                ("handle_dispatch_result_" +++ integerToString i): when stackSize > 0 ==> do
                  response <- (dispatchCores !! i).server.response.get
                  $display ("handle_dispatch_result_" +++ integerToString i +++ " ") (cshow response)
                  bestScore <- bottomFrame.putScore (negate response.score) response.forcedOutcome
                  if bestScore then bestMove := (dispatchMoves !! i).first else noAction
                  (dispatchMoves !! i).deq
                  dispatchInWaiting := dispatchInWaiting + 1
        return
          (foldr rJoinDescendingUrgency dispatchStateRules ((genWith dispatchRules) :: Vector numDispatch Rules),
           dispatchMoveUpdate.clear)
      else return (emptyRules, noAction)

  let reset = do
        joinActions $ map (\ frame -> (frame :: Frame).moves.clear) stack
        joinActions $ map (\ core -> (core :: SearchCore config dispatchDepth).clear) dispatchCores
        joinActions $ map (\ moves -> (moves :: (FIFOF Move)).clear) dispatchMoves
        stackSize := 0
        eval.clear
        moveUpdate.clear
        clearDispatchMoveUpdate
        dispatchInWaiting := fromInteger (valueOf numDispatch)
        movesComplete := True

      result =
        interface SearchResult
          rid = rid
          outcome =
            if queries.first.state.lastProgressMove >= 50 then Draw else
            case (queryStateInCheck, bottomFrame.hasMove) of
              (False, False) -> Draw
              (False, True) -> NoOutcome
              (True, False) -> CheckMate
              (True, True) -> Check
          bestMove = if bottomFrame.hasMove then Just bestMove else Nothing
          forcedOutcome = bottomFrame.forcedOutcome
          score = bottomFrame.score
          depth = depth

      searchRules =
        foldr1 rJoin ((genWith frameRules) :: Vector maxDepth Rules) <+>
        rules
          when movesComplete
            rules
              "invalid_query_depth": when stackSize == 0, queries.first.depth == 0 || queries.first.depth > fromInteger (valueOf maxDepth) ==> do
                $display "0 invalid_query_depth " (cshow queries.first)
                results.enq $ SearchResult {rid=rid; outcome=NoOutcome; bestMove=Nothing; forcedOutcome=False; score=0; depth=depth;}
                queries.deq

              "push_query_state": when stackSize == 0, queries.first.depth > 0, queries.first.depth <= fromInteger (valueOf maxDepth) ==> do
                $display "0 push_query_state " (cshow queries.first)
                depth := queries.first.depth
                cfg := queries.first.config
                pushState bottomFrame (fromMaybe minScore queries.first.alpha) (fromMaybe maxScore queries.first.beta)
                stackSize := 1

              "pop_result_state": when stackSize == 1, bottomFrame.alpha < bottomFrame.beta, not bottomFrame.moves.notEmpty, not moveUpdate.hasRequest, dispatchInWaiting == fromInteger (valueOf numDispatch), not moveUpdate.ready ==> do
                $display "  1 pop_result_state " (cshow result)
                if queries.first.getMoves then initialMoves.enq NoMove else noAction
                results.enq result
                queries.deq
                stackSize := 0

          "cutoff_result_state": when stackSize == 1, bottomFrame.alpha >= bottomFrame.beta, not queries.first.getMoves ==> do
            $display "  1 cutoff_result_state " (cshow result)
            results.enq result
            queries.deq
            reset

  clear :: PulseWire <- mkPulseWire
  let clearRule =
        rules
          "clear": when clear ==> do
            queries.clear
            results.clear
            initialMoves.clear
            reset

  addRules $ clearRule <+ (dispatchRules `rJoinDescendingUrgency` searchRules)

  interface
    server =
      interface Server
        request = toPut queries
        response = toGet results
    moves = fifoToGetS initialMoves
    clear = clear.send

    status = ((truncate $ pack depth) :: Bit 4) ++ ((truncate $ pack stackSize) :: Bit 4) ++ 0