chore: clean up tree logic

kevinjosethomas · kevinjosethomas · commit 31efe08fb851 · 2025-08-12T17:29:01.000-07:00
diff --git a/src/api/train.ts b/src/api/train.ts
@@ -61,12 +61,9 @@ export const fetchPuzzle = async () => {
   for (let i = 1; i < moves.length; i++) {
     const move = moves[i]
     if (move.uci && move.san) {
-      currentNode = gameTree.addMainMove(
-        currentNode,
-        move.board,
-        move.uci,
-        move.san,
-      )
+      currentNode = gameTree
+        .getLastMainlineNode()
+        .addChild(move.board, move.uci, move.san, true)
     }
   }
 
diff --git a/src/api/turing.ts b/src/api/turing.ts
@@ -46,12 +46,9 @@ export const fetchTuringGame = async () => {
   for (let i = 1; i < moves.length; i++) {
     const move = moves[i]
     if (move.uci && move.san) {
-      currentNode = gameTree.addMainMove(
-        currentNode,
-        move.board,
-        move.uci,
-        move.san,
-      )
+      currentNode = gameTree
+        .getLastMainlineNode()
+        .addChild(move.board, move.uci, move.san, true)
     }
   }
 
diff --git a/src/components/Analysis/StreamAnalysis.tsx b/src/components/Analysis/StreamAnalysis.tsx
@@ -92,7 +92,7 @@ export const StreamAnalysis: React.FC<Props> = ({
       } else {
         // For stream analysis, ALWAYS create variations for player moves
         // This preserves the live game mainline and allows exploration
-        const newVariation = game.tree.addVariation(
+        const newVariation = game.tree.addVariationNode(
           analysisController.currentNode,
           newFen,
           moveString,
diff --git a/src/components/Board/GameBoard.tsx b/src/components/Board/GameBoard.tsx
@@ -66,7 +66,7 @@ export const GameBoard: React.FC<Props> = ({
             if (currentNode.mainChild?.move === moveString) {
               goToNode(currentNode.mainChild)
             } else {
-              const newVariation = game.tree.addVariation(
+              const newVariation = game.tree.addVariationNode(
                 currentNode,
                 newFen,
                 moveString,
diff --git a/src/hooks/useAnalysisController/useAnalysisController.ts b/src/hooks/useAnalysisController/useAnalysisController.ts
@@ -475,7 +475,7 @@ export const useAnalysisController = (
     const moveResult = chess.move(currentMistake.bestMove, { sloppy: true })
 
     if (moveResult) {
-      const newVariation = game.tree.addVariation(
+      const newVariation = game.tree.addVariationNode(
         controller.currentNode,
         chess.fen(),
         currentMistake.bestMove,
diff --git a/src/hooks/useOpeningDrillController/useOpeningDrillController.ts b/src/hooks/useOpeningDrillController/useOpeningDrillController.ts
@@ -61,12 +61,9 @@ const parsePgnToTree = (pgn: string, gameTree: GameTree): GameNode | null => {
       if (existingChild) {
         currentNode = existingChild
       } else {
-        const newNode = gameTree.addMainMove(
-          currentNode,
-          chess.fen(),
-          moveUci,
-          moveObj.san,
-        )
+        const newNode = gameTree
+          .getLastMainlineNode()
+          .addChild(chess.fen(), moveUci, moveObj.san)
         if (newNode) {
           currentNode = newNode
         } else {
@@ -810,12 +807,9 @@ export const useOpeningDrillController = (
           try {
             const moveObj = chess.move(moveUci, { sloppy: true })
             if (moveObj) {
-              const newNode = gameTree.addMainMove(
-                currentNode,
-                chess.fen(),
-                moveUci,
-                moveObj.san,
-              )
+              const newNode = gameTree
+                .getLastMainlineNode()
+                .addChild(chess.fen(), moveUci, moveObj.san)
               if (newNode) {
                 currentNode = newNode
                 finalNode = newNode
@@ -834,12 +828,10 @@ export const useOpeningDrillController = (
           try {
             const moveObj = chess.move(moveUci, { sloppy: true })
             if (moveObj) {
-              const newNode = gameTree.addMainMove(
-                currentNode,
-                chess.fen(),
-                moveUci,
-                moveObj.san,
-              )
+              const newNode = gameTree
+                .getLastMainlineNode()
+                .addChild(currentNode, chess.fen(), moveUci, true, moveObj.san)
+
               if (newNode) {
                 currentNode = newNode
                 finalNode = newNode
diff --git a/src/hooks/usePlayController/usePlayController.ts b/src/hooks/usePlayController/usePlayController.ts
@@ -226,13 +226,27 @@ export const usePlayController = (id: string, config: PlayGameConfig) => {
 
   const addMoveWithTime = useCallback(
     (moveUci: string, moveTime: number) => {
-      const newNode = controller.tree.addMoveToMainLine(moveUci, moveTime)
-      if (newNode) {
-        console.log('old node', controller.currentNode.fen)
-        console.log('new move', moveUci)
-        console.log('new node', newNode.fen)
-        controller.setCurrentNode(newNode)
-        setTreeVersion((prev) => prev + 1)
+      const lastNode = controller.tree.getLastMainlineNode()
+      const board = new Chess(lastNode.fen)
+      const result = board.move(moveUci, { sloppy: true })
+
+      if (result) {
+        const newNode = lastNode.addChild(
+          board.fen(),
+          moveUci,
+          result.san,
+          true,
+          undefined,
+          moveTime,
+        )
+
+        if (newNode) {
+          console.log('old node', controller.currentNode.fen)
+          console.log('new move', moveUci)
+          console.log('new node', newNode.fen)
+          controller.setCurrentNode(newNode)
+          setTreeVersion((prev) => prev + 1)
+        }
       }
     },
     [controller.tree, controller],
diff --git a/src/hooks/useTrainingController/useTrainingController.ts b/src/hooks/useTrainingController/useTrainingController.ts
@@ -16,12 +16,9 @@ const buildTrainingGameTree = (game: PuzzleGame): GameTree => {
   for (let i = 1; i < game.moves.length; i++) {
     const move = game.moves[i]
     if (move.uci && move.san) {
-      currentNode = tree.addMainMove(
-        currentNode,
-        move.board,
-        move.uci,
-        move.san,
-      )
+      currentNode = tree
+        .getLastMainlineNode()
+        .addChild(move.board, move.uci, move.san, true),
     }
   }
 
diff --git a/src/hooks/useTuringController/useTuringController.ts b/src/hooks/useTuringController/useTuringController.ts
@@ -34,7 +34,9 @@ const buildTuringGameTree = (game: TuringGame): GameTree => {
     const move = game.moves[i]
     const uci = move.uci || (move.lastMove ? move.lastMove.join('') : undefined)
     if (uci && move.san) {
-      currentNode = tree.addMainMove(currentNode, move.board, uci, move.san)
+      currentNode = tree
+        .getLastMainlineNode()
+        .addChild(move.board, uci, move.san, true)
     }
   }
 
diff --git a/src/lib/common.ts b/src/lib/common.ts
@@ -8,13 +8,10 @@ export function buildGameTreeFromMoveList(moves: any[], initialFen: string) {
     const move = moves[i]
 
     if (move.lastMove) {
-      const [from, to] = move.lastMove
-      currentNode = tree.addMainMove(
-        currentNode,
-        move.board,
-        from + to,
-        move.san || '',
-      )
+      const [from, to, promotion] = move.lastMove
+      currentNode = tree
+        .getLastMainlineNode()
+        .addChild(move.board, from + to + promotion || '', move.san || '', true)
     }
   }
 
diff --git a/src/lib/lichess.ts b/src/lib/lichess.ts
@@ -119,7 +119,7 @@ export const handleLichessStreamMove = (
   }
 
   try {
-    currentGame.tree.addMainMove(currentNode, fen, uci, san)
+    currentGame.tree.getLastMainlineNode().addChild(fen, uci, san, true)
   } catch (error) {
     console.error('Error adding move to tree:', error)
     return currentGame
diff --git a/src/pages/analysis/[...id].tsx b/src/pages/analysis/[...id].tsx
@@ -438,17 +438,13 @@ const Analysis: React.FC<Props> = ({
         controller.currentNode.isMainline
       ) {
         // No main child exists AND we're on main line - create main line move
-        const newMainMove = analyzedGame.tree.addMainMove(
-          controller.currentNode,
-          newFen,
-          moveString,
-          san,
-          controller.currentMaiaModel,
-        )
+        const newMainMove = analyzedGame.tree
+          .getLastMainlineNode()
+          .addChild(newFen, moveString, san, true, controller.currentMaiaModel)
         controller.goToNode(newMainMove)
       } else {
         // Either main child exists but different move, OR we're in a variation - create variation
-        const newVariation = analyzedGame.tree.addVariation(
+        const newVariation = analyzedGame.tree.addVariationNode(
           controller.currentNode,
           newFen,
           moveString,
diff --git a/src/pages/puzzles.tsx b/src/pages/puzzles.tsx
@@ -435,7 +435,7 @@ const Train: React.FC<Props> = ({
               analysisController.currentNode.mainChild,
             )
           } else {
-            const newVariation = analyzedGame.tree.addVariation(
+            const newVariation = analyzedGame.tree.addVariationNode(
               analysisController.currentNode,
               newFen,
               moveString,
@@ -525,7 +525,7 @@ const Train: React.FC<Props> = ({
               analysisController.currentNode.mainChild,
             )
           } else {
-            const newVariation = analyzedGame.tree.addVariation(
+            const newVariation = analyzedGame.tree.addVariationNode(
               analysisController.currentNode,
               newFen,
               moveString,
@@ -636,7 +636,7 @@ const Train: React.FC<Props> = ({
         if (analysisController.currentNode.mainChild?.move === moveString) {
           analysisController.goToNode(analysisController.currentNode.mainChild)
         } else {
-          const newVariation = analyzedGame.tree.addVariation(
+          const newVariation = analyzedGame.tree.addVariationNode(
             analysisController.currentNode,
             newFen,
             moveString,
diff --git a/src/types/common.ts b/src/types/common.ts
@@ -1,8 +1,7 @@
-import { GameNode, GameTree } from './tree'
+import { GameNode } from './node'
+import { GameTree } from './tree'
 import { Dispatch, SetStateAction } from 'react'
 
-export * from './tree'
-
 export interface BaseGame {
   id: string
   tree: GameTree
diff --git a/src/types/index.ts b/src/types/index.ts
@@ -1,4 +1,6 @@
 export * from './play'
+export * from './tree'
+export * from './node'
 export * from './player'
 export * from './analysis'
 export * from './openings'
diff --git a/src/types/node.ts b/src/types/node.ts
diff --git a/src/types/tree.ts b/src/types/tree.ts

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,9 @@ const buildTuringGameTree = (game: TuringGame): GameTree => {`
`34`	`34`	`const move = game.moves[i]`
`35`	`35`	`const uci = move.uci \|\| (move.lastMove ? move.lastMove.join('') : undefined)`
`36`	`36`	`if (uci && move.san) {`
`37`		`- currentNode = tree.addMainMove(currentNode, move.board, uci, move.san)`
	`37`	`+ currentNode = tree`
	`38`	`+ .getLastMainlineNode()`
	`39`	`+ .addChild(move.board, uci, move.san, true)`
`38`	`40`	`}`
`39`	`41`	`}`
`40`	`42`