feat: migrate Stockfish analysis to use async generator instead of ca…

…llback methods

src/hooks/useAnalysisController/useAnalysisController.ts

@@ -1,5 +1,5 @@
 import { Chess } from 'chess.ts'
-import { useEffect, useMemo, useState } from 'react'
+import { useEffect, useMemo, useState, useCallback } from 'react'
 
 import {
   Color,
@@ -32,27 +32,14 @@ export const useAnalysisController = (
 ) => {
   const controller = useGameController(game, initialIndex, initialOrientation)
 
-  const parseStockfishEvaluation = (
-    message: StockfishEvaluation,
-    fen: string,
-  ) => {
-    setStockfishEvaluations((prev) => {
-      const newEvaluations = [...prev]
-      const moveIndex = game.moves.findIndex((move) => move.board === fen)
-      newEvaluations[moveIndex] = message
-
-      return newEvaluations
-    })
-  }
-
   const {
     maia,
     error: maiaError,
     status: maiaStatus,
     progress: maiaProgress,
     downloadModel: downloadMaia,
   } = useMaiaEngine()
-  const engine = useStockfishEngine(parseStockfishEvaluation)
+  const { streamEvaluations, stopEvaluation } = useStockfishEngine()
   const [currentMove, setCurrentMove] = useState<[string, string] | null>()
   const [stockfishEvaluations, setStockfishEvaluations] = useState<
     StockfishEvaluation[]
@@ -62,6 +49,19 @@ export const useAnalysisController = (
   >([])
   const [currentMaiaModel, setCurrentMaiaModel] = useState(MAIA_MODELS[0])
 
+  const parseStockfishEvaluation = useCallback(
+    (evaluation: StockfishEvaluation, fen: string) => {
+      setStockfishEvaluations((prev) => {
+        const newEvaluations = [...prev]
+        const moveIndex = game.moves.findIndex((move) => move.board === fen)
+        newEvaluations[moveIndex] = evaluation
+
+        return newEvaluations
+      })
+    },
+    [game.moves],
+  )
+
   useEffect(() => {
     const board = new Chess(game.moves[controller.currentIndex].board)
 
@@ -102,8 +102,30 @@ export const useAnalysisController = (
     const board = new Chess(game.moves[controller.currentIndex].board)
     if (stockfishEvaluations[controller.currentIndex]?.depth == 18) return
 
-    engine.evaluatePosition(board.fen(), board.moves().length)
-  }, [controller.currentIndex, game.moves, game.type, engine])
+    const evaluationStream = streamEvaluations(
+      board.fen(),
+      board.moves().length,
+    )
+
+    if (evaluationStream) {
+      ;(async () => {
+        for await (const evaluation of evaluationStream) {
+          parseStockfishEvaluation(evaluation, board.fen())
+        }
+      })()
+    }
+
+    return () => {
+      stopEvaluation()
+    }
+  }, [
+    controller.currentIndex,
+    game.moves,
+    game.type,
+    streamEvaluations,
+    stopEvaluation,
+    parseStockfishEvaluation,
+  ])
 
   const moves = useMemo(() => {
     const moveMap = new Map<string, string[]>()
@@ -227,7 +249,7 @@ export const useAnalysisController = (
     let okMoveChance = 0
     let goodMoveChance = 0
     if (!moveEvaluation || !moveEvaluation.maia || !moveEvaluation.stockfish) {
-      return { blunderMoveChance, okMoveChance, goodMoveChance }
+      return { blunderMoveChance: 0, okMoveChance: 0, goodMoveChance: 0 }
     }
 
     const { maia, stockfish } = moveEvaluation

src/hooks/useClientAnalysisController/useClientAnalysisController.ts

@@ -32,30 +32,15 @@ export const useClientAnalysisController = (game: AnalyzedGame) => {
 
   const [analysisState, setAnalysisState] = useState(0)
 
-  const parseStockfishEvaluation = (
-    message: StockfishEvaluation,
-    fen: string,
-  ) => {
-    console.log('Receiving Evaluation')
-    console.log(
-      fen === controller.currentNode?.fen,
-      fen,
-      controller.currentNode?.fen,
-    )
-    if (controller.currentNode && controller.currentNode.fen === fen) {
-      controller.currentNode.addStockfishAnalysis(message)
-      setAnalysisState((state) => state + 1)
-    }
-  }
-
   const {
     maia,
     error: maiaError,
     status: maiaStatus,
     progress: maiaProgress,
     downloadModel: downloadMaia,
   } = useMaiaEngine()
-  const engine = useStockfishEngine(parseStockfishEvaluation)
+
+  const { streamEvaluations, stopEvaluation } = useStockfishEngine()
   const [currentMove, setCurrentMove] = useState<[string, string] | null>()
   const [currentMaiaModel, setCurrentMaiaModel] = useState(MAIA_MODELS[0])
 
@@ -94,10 +79,28 @@ export const useClientAnalysisController = (game: AnalyzedGame) => {
     const board = new Chess(controller.currentNode.fen)
     if (controller.currentNode.analysis.stockfish?.depth == 18) return
 
-    console.log('Requesting evaluation')
+    const evaluationStream = streamEvaluations(
+      board.fen(),
+      board.moves().length,
+    )
 
-    engine.evaluatePosition(board.fen(), board.moves().length)
-  }, [controller.currentNode, game.type, engine])
+    if (evaluationStream) {
+      ;(async () => {
+        for await (const evaluation of evaluationStream) {
+          if (!controller.currentNode) {
+            stopEvaluation()
+            break
+          }
+          controller.currentNode.addStockfishAnalysis(evaluation)
+          setAnalysisState((state) => state + 1)
+        }
+      })()
+    }
+
+    return () => {
+      stopEvaluation()
+    }
+  }, [controller.currentNode, game.type, streamEvaluations, stopEvaluation])
 
   const moves = useMemo(() => {
     if (!controller.currentNode) return new Map<string, string[]>()

src/hooks/useMaiaEngine/model.ts

@@ -137,11 +137,6 @@ class Maia {
       legalMoves,
     )
 
-    console.log({
-      policy,
-      value,
-    })
-
     return {
       policy,
       value,

src/hooks/useStockfishEngine/engine.ts

@@ -7,21 +7,29 @@ class Engine {
   private fen: string
   private stockfish: StockfishWeb | null
   private moves: string[]
+  private isEvaluating: boolean
 
   private store: {
     [key: string]: StockfishEvaluation
   }
   private legalMoveCount: number
-  private callback: (data: StockfishEvaluation, fen: string) => void
+  private evaluationResolver: ((value: StockfishEvaluation) => void) | null
+  private evaluationRejecter: ((reason?: any) => void) | null
+  private evaluationPromise: Promise<StockfishEvaluation> | null
+  private evaluationGenerator: AsyncGenerator<StockfishEvaluation> | null
 
-  constructor(callback: (data: StockfishEvaluation, fen: string) => void) {
+  constructor() {
     this.fen = ''
     this.store = {}
     this.moves = []
+    this.isEvaluating = false
 
     this.legalMoveCount = 0
-    this.callback = callback
-    this.stockfish = null
+    this.evaluationResolver = null
+    this.evaluationRejecter = null
+    this.evaluationPromise = null
+    this.evaluationGenerator = null
+
     this.onMessage = this.onMessage.bind(this)
 
     setupStockfish().then((stockfish: StockfishWeb) => {
@@ -34,7 +42,10 @@ class Engine {
     })
   }
 
-  evaluatePosition(fen: string, legalMoveCount: number) {
+  async *streamEvaluations(
+    fen: string,
+    legalMoveCount: number,
+  ): AsyncGenerator<StockfishEvaluation> {
     if (this.stockfish) {
       if (typeof global.gc === 'function') {
         global.gc()
@@ -45,14 +56,41 @@ class Engine {
       const board = new Chess(fen)
       this.moves = board.moves({ verbose: true }).map((x) => x.from + x.to)
       this.fen = fen
+      this.isEvaluating = true
+      this.evaluationGenerator = this.createEvaluationGenerator()
 
       this.sendMessage('stop')
       this.sendMessage('ucinewgame')
       this.sendMessage(`position fen ${fen}`)
       this.sendMessage('go depth 18')
+
+      while (this.isEvaluating) {
+        try {
+          const evaluation = await this.getNextEvaluation()
+          if (evaluation) {
+            yield evaluation
+          } else {
+            break
+          }
+        } catch (error) {
+          console.error('Error in evaluation stream:', error)
+          break
+        }
+      }
     }
   }
 
+  private async getNextEvaluation(): Promise<StockfishEvaluation | null> {
+    return new Promise((resolve, reject) => {
+      this.evaluationResolver = resolve
+      this.evaluationRejecter = reject
+    })
+  }
+
+  private createEvaluationGenerator(): AsyncGenerator<StockfishEvaluation> | null {
+    return null
+  }
+
   private sendMessage(message: string) {
     if (this.stockfish) {
       this.stockfish.uci(message)
@@ -103,12 +141,26 @@ class Engine {
     if (!this.store[depth].sent && multipv === this.legalMoveCount) {
       this.store[depth].sent = true
 
-      this.callback(this.store[depth], this.fen)
+      if (this.evaluationResolver) {
+        this.evaluationResolver(this.store[depth])
+        this.evaluationResolver = null
+        this.evaluationRejecter = null
+      }
     }
   }
 
   private onError(msg: string) {
     console.error(msg)
+    if (this.evaluationRejecter) {
+      this.evaluationRejecter(msg)
+      this.evaluationResolver = null
+      this.evaluationRejecter = null
+    }
+    this.isEvaluating = false
+  }
+  stopEvaluation() {
+    this.isEvaluating = false
+    this.sendMessage('stop')
   }
 }
 

src/hooks/useStockfishEngine/useStockfishEngine.ts

@@ -1,11 +1,34 @@
-import { useMemo } from 'react'
-
+import { useMemo, useRef, useCallback } from 'react'
 import Engine from './engine'
 import { StockfishEvaluation } from 'src/types'
 
-export const useStockfishEngine = (
-  callback: (message: StockfishEvaluation, fen: string) => void,
-) => {
-  const engine = useMemo(() => new Engine(callback), [])
-  return engine
+export const useStockfishEngine = () => {
+  const engineRef = useRef<Engine | null>(null)
+
+  if (!engineRef.current) {
+    engineRef.current = new Engine()
+  }
+
+  const streamEvaluations = useCallback(
+    (fen: string, legalMoveCount: number) => {
+      if (!engineRef.current) {
+        console.error('Engine not initialized')
+        return null
+      }
+      return engineRef.current.streamEvaluations(fen, legalMoveCount)
+    },
+    [],
+  )
+
+  const stopEvaluation = useCallback(() => {
+    engineRef.current?.stopEvaluation()
+  }, [])
+
+  return useMemo(
+    () => ({
+      streamEvaluations,
+      stopEvaluation,
+    }),
+    [streamEvaluations, stopEvaluation],
+  )
 }

src/pages/analysis/[...id].tsx

@@ -96,7 +96,7 @@ const AnalysisPage: NextPage = () => {
         return
       }
       if (setCurrentMove) setCurrentMove(0)
-      console.log(game)
+
       setAnalyzedGame({ ...game, type: 'tournament' })
       setCurrentId(newId)
       router.push(`/analysis/${newId.join('/')}`, undefined, { shallow: true })
@@ -118,7 +118,7 @@ const AnalysisPage: NextPage = () => {
         return
       }
       if (setCurrentMove) setCurrentMove(0)
-      console.log(game)
+
       setAnalyzedGame({
         ...game,
         type: 'pgn',
@@ -144,7 +144,6 @@ const AnalysisPage: NextPage = () => {
       }
       if (setCurrentMove) setCurrentMove(0)
 
-      console.log(game)
       setAnalyzedGame({ ...game, type })
       setCurrentId([id, type])
       router.push(`/analysis/${id}/${type}`, undefined, { shallow: true })

src/types/base/tree.ts

@@ -198,7 +198,6 @@ export class GameNode {
   }
 
   addStockfishAnalysis(stockfishEval: StockfishEvaluation): void {
-    console.log('stockfish added')
     this._analysis.stockfish = stockfishEval
   }